feat: custom gathered-expert MoE combine for silu/FP8 (specifically Qwen 3.5 family)

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.

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.

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:]))