diff --git a/docs/configuration/env_variables.md b/docs/configuration/env_variables.md index 0794f74b9d..b90fd52023 100644 --- a/docs/configuration/env_variables.md +++ b/docs/configuration/env_variables.md @@ -33,6 +33,31 @@ This document lists the supported diagnostic and profiling, as well as performan | `VLLM_MINIMAX_M3_MOE_DECODE_GATHER` | Enables the MiniMax-M3 routed-expert gather path for low-token decode. Set to `0` or `false` to use the dense expert path. | `true` | | `VLLM_MINIMAX_M3_MOE_GATHER_MAX_TOKENS` | Maximum token count for the MiniMax-M3 routed-expert gather path. Larger batches use the dense expert path. | `16` | +## Experimental: Custom FP8 MoE Gather Combine + +These variables control an **experimental** pure-PyTorch gathered-expert MoE +combine for silu + FP8-per-channel weights, an alternative to the Habana +`mixture_of_experts` op. It is off by default and intended for low-token +(small batch / decode) workloads. Verification runs both the custom and stock +paths and reduces their maximum FP8-ULP over the expert-parallel group in-memory +without writing model-derived tensors to disk. Note that verify mode adds a +**per-layer host sync** (``max_ulp.item()`` on CPU during every forward pass), +so it must not be enabled on performance runs. + +The default `VLLM_HPU_MOE_GATHER_RATIO` of `0.4` is based on a crossover sweep +across the Qwen 3.5 MoE family (35B / 122B / 397B) at several expert-parallel +levels; this optimization has only been observed to help that family. The win/loss +cutoff most closely tracks the gathered-to-local-experts ratio and lands around +this value, so raise it only if you have measured the gather path to still win at +higher ratios on your model/config, and lower it for configs where it loses sooner +(e.g. high-EP deployments with wide experts). + +| Parameter name | Description | Default value | +| ---------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------- | +| `VLLM_HPU_MOE_GATHER` | Enables the custom gathered-expert FP8 MoE combine (silu only). Falls back to the stock fused op when disabled or when the gather ratio is exceeded. | `false` | +| `VLLM_HPU_MOE_GATHER_RATIO` | Fraction of this rank's `local_experts` up to which the custom gather path is used. The gathered count is `min(local_experts, tokens * top_k)`; above `ratio * local_experts` the stock fused op is used. | `0.4` | +| `VLLM_HPU_MOE_GATHER_VERIFY` | Runs both the custom and stock paths and reduces their maximum FP8-ULP across the EP group in-memory. Logs an info line at startup when enabled and warns if any element exceeds 2 ULP. Requires `VLLM_HPU_MOE_GATHER`. | `false` | + Use `VLLM_BUCKETING_STRATEGY=exp` for the default exponential warm-up, `VLLM_BUCKETING_STRATEGY=lin` for explicitly configured linear ranges, or `VLLM_BUCKETING_STRATEGY=pad` for padding-aware ranges with absolute and relative padding limits. Leave `VLLM_EXPONENTIAL_BUCKETING` unset when using `VLLM_BUCKETING_STRATEGY`. The legacy flag is checked for backward compatibility and still overrides the selected strategy when present. diff --git a/vllm_gaudi/envs.py b/vllm_gaudi/envs.py index 5fe7c07c08..63cd0d6bca 100644 --- a/vllm_gaudi/envs.py +++ b/vllm_gaudi/envs.py @@ -17,6 +17,9 @@ VLLM_MINIMAX_M3_MOE_DECODE_GATHER: bool = True VLLM_MINIMAX_M3_MOE_GATHER_MAX_TOKENS: int = 16 VLLM_MM_WARMUP_OUTSIDE_COMPILE_ONLY: bool = False + VLLM_HPU_MOE_GATHER: bool = False + VLLM_HPU_MOE_GATHER_RATIO: float = 0.4 + VLLM_HPU_MOE_GATHER_VERIFY: bool = False VLLM_COMPACT_GDN: bool = False # The begin-* and end* here are used by the documentation generator @@ -89,6 +92,24 @@ "VLLM_MM_WARMUP_OUTSIDE_COMPILE_ONLY": lambda: os.environ.get("VLLM_MM_WARMUP_OUTSIDE_COMPILE_ONLY", "false").strip().lower() in ("1", "true"), + # EXPERIMENTAL custom gathered-expert FP8 MoE combine (silu only). + # Off by default; falls back to the stock Habana fused op when disabled or + # when the number of gathered experts (min(local_experts, tokens*top_k)) + # would exceed this fraction of this rank's local_experts. Parameterizing on + # the ratio to local_experts (rather than an absolute count) keeps the gate + # sane across EP configs, which change local_experts. + # + # Default rationale (0.4): this optimization has only been observed to help + # the Qwen 3.5 MoE family. A crossover sweep across Qwen3.5 35B/122B/397B at + # several EP levels shows the custom-gather win/loss cutoff tracks + # g/local_experts most closely, and it lands around this ratio. + "VLLM_HPU_MOE_GATHER": + lambda: os.environ.get("VLLM_HPU_MOE_GATHER", "0").lower() in ("1", "true"), + "VLLM_HPU_MOE_GATHER_RATIO": + lambda: float(os.environ.get("VLLM_HPU_MOE_GATHER_RATIO", "0.4")), + "VLLM_HPU_MOE_GATHER_VERIFY": + lambda: os.environ.get("VLLM_HPU_MOE_GATHER_VERIFY", "0").lower() in ("1", "true"), + # Use the compact recurrent-state (conv/ssm) layout for gated delta net # models. The model runner auto-detects and sets this during init, so read # it lazily rather than caching it at import time. diff --git a/vllm_gaudi/ops/hpu_fp8.py b/vllm_gaudi/ops/hpu_fp8.py index 7e7f0397c8..4d500466d3 100644 --- a/vllm_gaudi/ops/hpu_fp8.py +++ b/vllm_gaudi/ops/hpu_fp8.py @@ -2,6 +2,8 @@ from typing import Optional import torch +from vllm.distributed import get_ep_group +from vllm.logger import init_logger from vllm_gaudi import envs from torch.nn.parameter import Parameter from vllm.model_executor.layers.fused_moe.layer import FusedMoEFactory as FusedMoE @@ -23,6 +25,99 @@ ChannelWiseTorchFP8ScaledMMLinearKernel, ) +logger = init_logger(__name__) + +# EXPERIMENTAL custom MoE combine: replace the Habana mixture_of_experts op (a +# fixed per-layer stage pipeline) with a pure-PyTorch gathered-expert path. +# Default stock. `VLLM_HPU_MOE_GATHER_VERIFY=1` (with VLLM_HPU_MOE_GATHER=1) runs +# BOTH the custom path and the Habana op on the same inputs and reduces their +# maximum FP8-ULP over the expert-parallel group in-memory (no files). +_HPU_MOE_GATHER = envs.VLLM_HPU_MOE_GATHER +# Only meaningful together with the gather path. +_HPU_MOE_GATHER_VERIFY = envs.VLLM_HPU_MOE_GATHER_VERIFY and _HPU_MOE_GATHER +# Correctness bar for the verify path: any element exceeding this many FP8-ULP +# is an error (matches the archived offline analysis' "no element > 2 ULP" bar). +_HPU_MOE_GATHER_VERIFY_MAX_ULP = 2 +# Fraction of this rank's local_experts up to which the custom gather path is +# used. The gathered pure-PyTorch path reads only the routed experts, so it wins +# while the gathered count g stays well below local_experts (stock must touch all +# of them); it loses once g approaches local_experts (the dense gather + fp32 bmm +# path is slower than the Habana op). See envs.py for the default rationale. +_HPU_MOE_GATHER_RATIO = envs.VLLM_HPU_MOE_GATHER_RATIO +if _HPU_MOE_GATHER: + from vllm_gaudi.ops.hpu_moe_combine import gather_silu_fp8_moe # noqa: E402 +else: + gather_silu_fp8_moe = None # type: ignore[assignment] + +if _HPU_MOE_GATHER_VERIFY: + logger.info("MoE gather combine VERIFY mode enabled: comparing custom vs stock " + "per layer (FP8-ULP bar = %d)", _HPU_MOE_GATHER_VERIFY_MAX_ULP) + + +def _fp8_ulp(a: torch.Tensor, b: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Sign-aware FP8-ULP distance + out-of-range mask. + + Same-sign pairs use real E4M3 ulps: |uint8(quantize(a)) - uint8(quantize(b))|. + In E4M3FN adjacent same-sign representable values differ by exactly +/-1 in + uint8, so this IS the representable-step count. Cross-sign pairs use + universal subnormal units |a - b| / 2^-9. Returns elementwise Float32 ulps + (set to 0.0 wherever either input is out of range) plus an ``out_of_range`` + bool mask flagging those positions. + + Elements above the E4M3 maximum representable value (448) saturate to + ``0x7F`` / ``0xFF`` (NaN/inf codes) during the uint8 cast and would falsely + compare as 0 ULP, so the caller must not fold them into the finite-ULP bar; + they are handled separately via ``out_of_range``. + """ + a_fp32 = a.float() + b_fp32 = b.float() + e4m3_max = 448.0 + out_of_range = (a_fp32.abs() > e4m3_max) | (b_fp32.abs() > e4m3_max) + a_bits = a_fp32.to(torch.float8_e4m3fn).view(torch.uint8).int() + b_bits = b_fp32.to(torch.float8_e4m3fn).view(torch.uint8).int() + same_sign = (a_bits >= 128) == (b_bits >= 128) + same_sign_ulp = (a_bits - b_bits).abs().float() + cross_ulp = (a_fp32 - b_fp32).abs() / (2.0**-9) + ulp = torch.where(same_sign, same_sign_ulp, cross_ulp) + ulp = torch.where(out_of_range, torch.zeros_like(ulp), ulp) + return ulp, out_of_range + + +# Relative-error bar for out-of-range elements (|a-b| / (|a|+|b|)). The +# denominator is at least the E4M3 max (448) out of range, so this cannot +# false-positive on small in-range values. +_HPU_MOE_GATHER_VERIFY_REL_ERR = 0.05 + + +def _verify_moe_combine(stock: torch.Tensor, custom: torch.Tensor) -> None: + ulp, out_of_range = _fp8_ulp(stock, custom) + max_in_range_ulp = ulp.amax() + # Out-of-range positions cannot be compared in ULP (the fp8 cast saturates + # them), so measure their divergence as a relative error instead. Both paths + # agreeing out of range (rel_err 0) stays quiet. + a_fp32 = stock.float() + b_fp32 = custom.float() + denom = a_fp32.abs() + b_fp32.abs() + rel_err = ((a_fp32 - b_fp32).abs() / denom.clamp(min=1e-9)).where(out_of_range, torch.zeros_like(a_fp32)) + max_out_of_range_rel = rel_err.amax() + reduced = torch.stack([max_in_range_ulp, max_out_of_range_rel]) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + ep_group = get_ep_group() + if ep_group.world_size > 1: + torch.distributed.all_reduce( + reduced, + op=torch.distributed.ReduceOp.MAX, + group=ep_group.device_group, + ) + max_in_range_ulp, max_out_of_range_rel = reduced[0], reduced[1] + if max_out_of_range_rel.item() > _HPU_MOE_GATHER_VERIFY_REL_ERR: + logger.warning( + "MoE gather combine: out-of-range element(s) (value > 448, " + "outside E4M3 finite range) diverge by relative error %s", max_out_of_range_rel.item()) + if max_in_range_ulp.item() > _HPU_MOE_GATHER_VERIFY_MAX_ULP: + logger.warning("MoE gather combine mismatch: max FP8-ULP = %s exceeds %d", max_in_range_ulp.item(), + _HPU_MOE_GATHER_VERIFY_MAX_ULP) + class HPUPerTensorTorchFP8ScaledMMLinearKernel(PerTensorTorchFP8ScaledMMLinearKernel): @@ -164,6 +259,14 @@ def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module): # Snapshot the (static) quant-config flag while the vLLM config context # is set; the forward hot path reads this cached value instead. self.has_moe_quant_config = model_has_quant_config() + # The custom gathered-expert path reads per-channel scales + # (w13_weight_scale_inv / w2_weight_scale_inv, shape [E, 2I] / [E, H]), + # which exist only when block_quant + force_channel_fp8 are true + # (fp8_block_moe_prepare_weights branch in process_weights_after_loading). + # Non-block FP8 checkpoints use a per-tensor scale (w13_weight_scale, + # shape [E, 2]) and block FP8 without force_channel_fp8 uses block-shaped + # scales; both would crash or silently mis-scale in gather_silu_fp8_moe. + self._moe_gather_ok = self.block_quant and envs.VLLM_HPU_FORCE_CHANNEL_FP8 @property def is_monolithic(self) -> bool: @@ -268,13 +371,35 @@ def apply_monolithic( topk_ids = topk_ids.view(-1, topk_ids.shape[-1]) topk_weights = topk_weights.view(-1, topk_weights.shape[-1]) - output = layer.moe_op( - x, - topk_ids, - topk_weights, - permuted_weights=True, - activation=_normalize_moe_activation(layer.activation), - ) + activation = _normalize_moe_activation(layer.activation) + # Use the custom gathered-expert combine only when it wins: the number of + # distinct routed experts g = min(local_experts, tokens*K) must stay below + # a fraction of this rank's local_experts. Beyond that (large batch / long + # prefill, or few local experts) fall back to the stock fused op, which is + # faster and keeps the graph shapes fixed. `tokens`/`K` are static (T, K + # from x/topk_ids); local_num_experts is fixed per layer. + use_gather = (_HPU_MOE_GATHER and activation == "silu" and self.quant_config.activation_scheme != "static" + and self._moe_gather_ok + and x.shape[0] * topk_ids.shape[-1] <= layer.local_num_experts * _HPU_MOE_GATHER_RATIO) + if use_gather: + # EXPERIMENTAL custom combine: gather only the routed experts + # (bypasses the Habana op's fixed per-layer stage pipeline). + if _HPU_MOE_GATHER_VERIFY: + stock = layer.moe_op(x, topk_ids, topk_weights, permuted_weights=True, activation=activation) + custom = gather_silu_fp8_moe(layer, x, topk_ids, topk_weights) + _verify_moe_combine(stock, custom) + del stock + output = custom + else: + output = gather_silu_fp8_moe(layer, x, topk_ids, topk_weights) + else: + output = layer.moe_op( + x, + topk_ids, + topk_weights, + permuted_weights=True, + activation=activation, + ) return output.view(*(output.size(0), *input_shape[1:])) diff --git a/vllm_gaudi/ops/hpu_moe_combine.py b/vllm_gaudi/ops/hpu_moe_combine.py new file mode 100644 index 0000000000..d2f88e768d --- /dev/null +++ b/vllm_gaudi/ops/hpu_moe_combine.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pure-PyTorch gathered-expert MoE combine for silu + FP8-per-channel weights. + +Replaces the Habana ``mixture_of_experts`` combine (a fixed per-layer launch +pipeline) with a leaner active-expert gather + GEMM + weighted-reduce path, +mirroring ``vllm_gaudi.ops.hpu_fused_moe._gather_swigluoai_moe`` but for a silu +gated activation and the FP8 per-channel weight layout +(``extension/ops.py:fp8_channel_moe_prepare_weights``). + +Only ``g = min(E_local, tokens * K)`` distinct routed experts are read from HBM +and computed (<= 8 at BS=1 for K=8), instead of the Habana op's fixed stage +pipeline. The gathered count is static (no ``torch.nonzero``/host branch), so the +path captures cleanly into a compiled HPU graph. + +The UP/GATE projection runs in native FP8 via ``torch.ops.hpu.fp8_gemm_v2`` +(fp32 internal accumulation): ``x`` is shared across all routed experts, so the G +per-expert GEMMs collapse into ONE wide fp8 GEMM over concatenated expert weights +``[T,H] x [H, G*2I] -> [T, G*2I]`` (expert ``g`` at columns +``[g*2I:(g+1)*2I]`` via ``permute(2,0,1)``), and the gathered ``w13`` weights are +NEVER dequantized to fp32. + +The DOWN projection keeps ``act`` high-precision (fp32): ``act`` is the silu +output with a wide dynamic range, so fp8-quantizing it would lose ~2-6% precision +(3 mantissa bits) and blow the 2-ULP bar. HPU's fp8 GEMM ops accept only fp8 +activations, so ``w2`` is dequantized to fp32 for a plain fp32 bmm (as the +baseline did). Only the up/gate GEMM benefits from native fp8. + +Numeric fidelity: ``x``/``w13`` are FP8 with per-channel scales folded into +``fp8_gemm_v2`` (fp32 internal accumulation); the silu and weighted sum run in +fp32, rounding to the output dtype only at the end. Exact bit-equality with the +black-box op is NOT expected; correctness is assessed via an FP8-ULP bar. +""" +from __future__ import annotations + +import torch + + +def _dynamic_quant(data: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # import lazily to avoid pulling heavy deps at module import + from vllm_gaudi.extension.ops import dynamic_quant + return dynamic_quant(data) + + +def gather_silu_fp8_moe( + layer, + x: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, +) -> torch.Tensor: + """Compute the rank-local FP8 silu MoE partial for routed experts only. + + topk_ids [T,K] int64 (global expert ids), topk_weights [T,K] bf16. + Mirrors the Habana op's data flow: fp8-quantize x, per-token weighted sum + over the routed experts, each expert computed as silu(w13 x) w2. + """ + tokens, hidden_size = x.shape + num_topk = topk_ids.shape[-1] + w13 = layer.w13_weight # [E, 2I, H] fp8 + w2 = layer.w2_weight # [E, H, I] fp8 + scale13 = layer.w13_weight_scale_inv # [E, 2I] + scale2 = layer.w2_weight_scale_inv # [E, H] + local_experts = w13.shape[0] + intermediate_size = w13.shape[1] // 2 + + # ---- FP8-quantize x per token (match the op's input quantization) ---- + # x_scale is kept as [T,1] (NOT squeezed): fp8_gemm_v2 requires the row + # scale in 2-D form; squeezing only works at T=1 and silently mis-scales at + # larger T. + x_fp8, x_scale = _dynamic_quant(x) # x_fp8 [T,H], x_scale [T,1] f32 + + # ---- per-token expert combine weights (scatter over E) ---- + # topk_ids are GLOBAL expert ids; remap to this rank's local ids and mask + # experts owned by other EP ranks (they contribute +0.0 to this rank's + # partial, which is then reduce-scattered by the runner). At TP=1 + # ep_rank=0 -> experts_min=0, identical to the unremapped path. + experts_min = int(layer.moe_config.ep_rank * layer.local_num_experts) + local_ids = topk_ids - experts_min # [T,K] + in_range = (local_ids >= 0) & (local_ids < local_experts) + safe_ids = torch.where(in_range, local_ids, torch.zeros_like(local_ids)) + safe_weights = torch.where(in_range, topk_weights, torch.zeros_like(topk_weights)).float() + gate_weights = x.new_zeros(tokens, local_experts, dtype=torch.float32) + gate_weights.scatter_add_(1, safe_ids, safe_weights) # [T,E] f32 + + # ---- static gathered-expert count (>= # distinct hit experts) ---- + # The number of distinct routed experts is provably <= tokens * K, so with + # g = min(E, tokens*K) every real hit is included and the extra (zero-weight) + # padding experts contribute exactly +0.0 to the weighted sum. Keeping `g` + # static (no torch.nonzero, no `if G == 0`) lets this capture cleanly into a + # compiled HPU graph (mirrors _gather_swigluoai_moe). + gathered_experts = min(local_experts, tokens * num_topk) + hit_counts = (gate_weights != 0).float().sum(0) # [E] + gather_ids = torch.topk(hit_counts, gathered_experts, sorted=False).indices # [G] + gather_ids, _ = torch.sort(gather_ids) # ascending ids + + # ---- gather only the active experts' weights + per-channel scales ---- + w13_gathered = w13.index_select(0, gather_ids) # [G, 2I, H] fp8 + w2_gathered = w2.index_select(0, gather_ids) # [G, H, I] fp8 + scale13_gathered = scale13.index_select(0, gather_ids) # [G, 2I] + scale2_gathered = scale2.index_select(0, gather_ids) # [G, H] + + # ---- per-expert MLP, all in FP8 (no dequant to fp32) ---- + # + # UP/GATE projection: x is shared across all routed experts, so the G + # per-expert GEMMs `x @ w13[g]` collapse into ONE wide fp8 GEMM over the + # concatenated expert weights [T,H] x [H, G*2I] -> [T, G*2I]. Expert g's + # block lands on columns [g*2I:(g+1)*2I] via the permute(2,0,1) ordering + # (NOT permute(1,0,2), which interleaves experts and is wrong). The + # per-channel B_scale_inv is concatenated to match, with B_scale_shape + # declaring per-channel (not per-block) scaling. fp8_gemm_v2 accumulates in + # fp32 internally and returns out_dtype; A_scale_inv MUST be passed as [T,1]. + w13_columns = w13_gathered.permute(2, 0, 1).reshape(hidden_size, -1) # [H, G*2I] fp8 + scale13_columns = scale13_gathered.reshape(-1) # [G*2I] + projected = torch.ops.hpu.fp8_gemm_v2( + A=x_fp8, + trans_A=False, + B=w13_columns, + trans_B=False, + D=None, + out_dtype=torch.float32, + A_scale_inv=x_scale, + B_scale_inv=scale13_columns, + B_scale_shape=[gathered_experts * 2 * intermediate_size], + bias=None, + accumulate=False, + ) # [T, G*2I] f32 + projected = projected.reshape(tokens, gathered_experts, 2 * intermediate_size).permute(1, 0, 2) # [G, T, 2I] f32 + gate, up = projected[..., :intermediate_size], projected[..., intermediate_size:] + activations = gate * torch.sigmoid(gate) * up # silu(gate)*up, [G,T,I] f32 + + # DOWN projection: act has a wide dynamic range (silu output), so it CANNOT + # be fp8-quantized without losing ~2-6% precision (fp8 has only 3 mantissa + # bits), which would blow the 2-ULP bar. The stock op keeps act high-precision + # here, and HPU's fp8 GEMM ops (fp8_gemm_v2 / fp8_gemm) only accept fp8 + # activations. So the down projection dequantizes w2 to fp32 and runs a plain + # fp32 bmm (same as the baseline); only the UP/GATE projection runs in native + # fp8. + w2_float = w2_gathered.float() * scale2_gathered.unsqueeze(-1) # [G, H, I] f32 + expert_outputs = torch.bmm(activations, w2_float.transpose(1, 2)) # [G, T, H] f32 + + # ---- weighted sum over experts -> [T, H] ---- + gathered_weights = gate_weights.index_select(1, gather_ids).t() # [G, T] + out = (expert_outputs * gathered_weights.unsqueeze(-1)).sum(0) # [T, H] f32 + return out.to(x.dtype)