feat(moe): add unified block-scale FP8 support

flashinfer/fused_moe/__init__.py

@@ -42,6 +42,7 @@
     B12xW4A16Runner,
     CuteDslNvfp4Runner,
     TrtllmFp4RoutedRunner,
+    TrtllmFp8BlockRunner,
 )
 
 # Legacy flat-argument APIs (unchanged, not deprecated)
@@ -138,6 +139,7 @@
     "MoELayer",
     "MoEWeightPack",
     "TrtllmFp4RoutedRunner",
+    "TrtllmFp8BlockRunner",
     "QuantConfig",
     "QuantVariant",
     "RoutingConfig",

flashinfer/fused_moe/api.py

@@ -274,7 +274,46 @@ class TrtllmFp8BlockConfig:
 
     @classmethod
     def supported(cls, arch: int) -> bool:
-        return arch >= 80
+        # The available TRTLLM block-FP8 BMM cubins are validated only on the
+        # SM100 family. The outer JIT can compile for major 12, but its FP8
+        # kernels currently fail at runtime on SM120/121.
+        return arch in (100, 103)
+
+    @staticmethod
+    def prepare_weights(
+        w1_bf16,
+        w2_bf16,
+        *,
+        variant: QuantVariant,
+        num_local_experts: int,
+        hidden_size: int,
+        intermediate_size: int,
+        device=None,
+    ):
+        """Build the ``trtllm_fp8_block`` weight view from canonical BF16.
+
+        ``variant`` must be :attr:`QuantVariant.DeepSeekFp8` or
+        :attr:`QuantVariant.MxFp8`; their scale formats are intentionally
+        prepared by separate paths.
+        """
+        from .prepare import prepare_trtllm_fp8_block_weights
+
+        return prepare_trtllm_fp8_block_weights(
+            w1_bf16,
+            w2_bf16,
+            variant=variant,
+            num_local_experts=num_local_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            device=device,
+        )
+
+    @staticmethod
+    def prepare_activations(hidden_states_bf16, *, variant: QuantVariant):
+        """Quantize BF16 activations for the selected block-FP8 convention."""
+        from .prepare import prepare_trtllm_fp8_block_activations
+
+        return prepare_trtllm_fp8_block_activations(hidden_states_bf16, variant=variant)
 
     def __repr__(self) -> str:
         return "TrtllmFp8BlockConfig()"
@@ -617,7 +656,17 @@ def __getitem__(self, key: str):
 
 @dataclass
 class MoEActivationPack:
-    """Per-call transient data — pre-quantized NVFP4 activations plus routing inputs.
+    """Per-call backend-native activations plus routing inputs.
+
+    Activation encoding depends on ``QuantConfig.variant``:
+
+    * NVFP4: packed ``uint8 [M, H/2]`` values with
+      ``float8_e4m3fn [M, H/16]`` block scales.
+    * BF16: raw ``bfloat16 [M, H]`` values with no scale tensor.
+    * DeepSeek FP8: ``float8_e4m3fn [M, H]`` values with transposed
+      ``float32 [H/128, M]`` block scales.
+    * MXFP8: ``float8_e4m3fn [M, H]`` values with token-major
+      ``uint8 [M, H/32]`` UE8M0 scales.
 
     ``routing_input_mode`` selects how routing reaches the kernel (the runner reads it directly):
 
@@ -629,20 +678,20 @@ class MoEActivationPack:
       methods like DeepSeekV3/MiniMax2, ``routing_bias``); the kernel computes the top-k selection
       itself per ``RoutingConfig.method``.  ``topk_ids`` / ``topk_weights`` stay ``None`` — the
       runner allocates internal kernel-filled buffers, and the routing result is not surfaced
-      back through the pack (routing replay is a separate, future capability).  Currently only
-      the TRTLLM FP4 runner supports this mode; ``MoELayer`` dispatches a logits pack only to
-      capable backends (see each runner's ``supported_routing_modes``).
+      back through the pack (routing replay is a separate, future capability). TRTLLM FP4 and
+      block-FP8 runners support this mode; ``MoELayer`` dispatches a logits pack only to capable
+      backends (see each runner's ``supported_routing_modes``).
 
     ``topk_ids`` / ``topk_weights`` follow the routed-MoE naming convention (gh #2425); they
     keep the field positions of the former ``selected_experts`` / ``final_scales``, so
     positional construction of pre-routed packs is unchanged.  The in-kernel routing fields
     are keyword-only.
     """
 
-    hidden_states_q: Tensor  # [M, H//2] uint8 (packed NVFP4) or [M, H] bf16
-    hidden_states_scale: Optional[
-        Tensor
-    ]  # [M, H//16] float8_e4m3fn block scales; None for BF16
+    # Backend-native activation payload; layouts documented above.
+    hidden_states_q: Tensor
+    # Variant-specific scales documented above; None for BF16.
+    hidden_states_scale: Optional[Tensor]
     # Pre-routed top-k selection (Packed/Unpacked modes); None under FromLogits.
     topk_ids: Optional[Tensor] = None  # [M, top_k] int32 (expert indices)
     topk_weights: Optional[Tensor] = None  # [M, top_k] float32 (routing weights)

flashinfer/fused_moe/core.py

@@ -1363,9 +1363,16 @@ def forward(
             expert_weights = moe_inputs.expert_weights
             topk_weights = expert_weights
             hidden_states = moe_inputs.hidden_states
+            # The generic helper identifies TRTLLM dtypes whose ABI normally
+            # consumes an auxiliary scale tensor (FP4 and MX formats). Plain
+            # E4m3 returns false, but DeepSeek block-FP8 is an exception: it
+            # requires the real per-1x128-block scales from the activation pack.
             hidden_states_scale = (
                 moe_inputs.hidden_states_scale
-                if trtllm_gen_dtype_has_scale(self.dtype_act)
+                if (
+                    trtllm_gen_dtype_has_scale(self.dtype_act)
+                    or self.fp8_quantization_type == Fp8QuantizationType.DeepSeekFp8
+                )
                 else None
             )
 
@@ -1452,30 +1459,13 @@ def forward(
                     or self.fp8_quantization_type == Fp8QuantizationType.MxFp8
                 ):
                     # FP8 block scale
-                    current_num_tokens = hidden_states.shape[0]
-                    current_hidden_size = hidden_states.shape[1]
-                    if self.fp8_quantization_type == Fp8QuantizationType.DeepSeekFp8:
-                        current_hidden_states_scale = torch.full(
-                            (current_hidden_size // 128, current_num_tokens),
-                            2.0,
-                            dtype=torch.float,
-                            device=hidden_states.device,
-                        )
-                    elif self.fp8_quantization_type == Fp8QuantizationType.MxFp8:
-                        current_hidden_states_scale = hidden_states_scale
-
-                    else:
-                        raise ValueError(
-                            f"Unsupported FP8 quantization type: {self.fp8_quantization_type}"
-                        )
-
                     moe_op.trtllm_fp8_block_scale_moe(
                         routing_logits,
                         topk_ids,
                         topk_weights,
                         kwargs["routing_bias"],
                         hidden_states,
-                        current_hidden_states_scale,
+                        hidden_states_scale,
                         kwargs["gemm1_weights"],
                         kwargs["gemm1_weights_scale"],
                         moe_inputs.gemm1_lora_delta,

flashinfer/fused_moe/layer.py

@@ -38,13 +38,15 @@
     MoEWeightPack,
     TrtllmBf16Config,
     TrtllmFp4Config,
+    TrtllmFp8BlockConfig,
 )
 from .runners import (
     B12xNvfp4Runner,
     B12xW4A16Runner,
     CuteDslNvfp4Runner,
     TrtllmBf16RoutedRunner,
     TrtllmFp4RoutedRunner,
+    TrtllmFp8BlockRunner,
 )
 from .utils import map_to_hybrid_bucket
 
@@ -56,6 +58,7 @@
     CuteDslNvfp4Runner,
     TrtllmFp4RoutedRunner,
     TrtllmBf16RoutedRunner,
+    TrtllmFp8BlockRunner,
     B12xNvfp4Runner,
     B12xW4A16Runner,
 ]
@@ -65,6 +68,7 @@
     CuteDslConfig: CuteDslNvfp4Runner,
     TrtllmFp4Config: TrtllmFp4RoutedRunner,
     TrtllmBf16Config: TrtllmBf16RoutedRunner,
+    TrtllmFp8BlockConfig: TrtllmFp8BlockRunner,
     B12xNvfp4Config: B12xNvfp4Runner,
     B12xW4A16Config: B12xW4A16Runner,
 }

flashinfer/fused_moe/prepare.py

@@ -47,6 +47,7 @@
 # Module-level permute-index cache.  Permute indices depend only on weight
 # dims, so the cache is safe to reuse across shapes and calls.
 _TRTLLM_PERMUTE_CACHE: dict = {}
+_TRTLLM_FP8_PERMUTE_CACHE: dict = {}
 
 
 # The E8M0 range clamp and residual-scale factorization are adapted from
@@ -538,6 +539,196 @@ def prepare_trtllm_fp4_weights(
     }
 
 
+def _deepseek_fp8_quantize_activations(
+    x: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize ``[M, K]`` BF16 per 1x128 block.
+
+    TRTLLM's DeepSeek path consumes scales transposed as ``[K // 128, M]``.
+    """
+    block = 128
+    m, k = x.shape
+    blocks = x.float().reshape(m, k // block, block)
+    fp8_max = torch.finfo(torch.float8_e4m3fn).max
+    scales = (blocks.abs().amax(dim=-1, keepdim=True) / fp8_max).clamp(min=1e-12)
+    quantized = (blocks / scales).clamp(-fp8_max, fp8_max)
+    return (
+        quantized.reshape(m, k).to(torch.float8_e4m3fn),
+        scales.squeeze(-1).transpose(0, 1).contiguous(),
+    )
+
+
+def _deepseek_fp8_quantize_weights(
+    weights: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize ``[E, N, K]`` BF16 per 128x128 block."""
+    block = 128
+    e, n, k = weights.shape
+    blocks = (
+        weights.float()
+        .reshape(e, n // block, block, k // block, block)
+        .permute(0, 1, 3, 2, 4)
+    )
+    fp8_max = torch.finfo(torch.float8_e4m3fn).max
+    scales = (blocks.abs().amax(dim=(-1, -2), keepdim=True) / fp8_max).clamp(min=1e-12)
+    quantized = (blocks / scales).clamp(-fp8_max, fp8_max)
+    quantized = quantized.permute(0, 1, 3, 2, 4).reshape(e, n, k)
+    return quantized.to(torch.float8_e4m3fn), scales[..., 0, 0].contiguous()
+
+
+def _validate_trtllm_fp8_block_inputs(
+    w1_bf16: torch.Tensor,
+    w2_bf16: torch.Tensor,
+    *,
+    num_local_experts: int,
+    hidden_size: int,
+    intermediate_size: int,
+) -> None:
+    if w1_bf16.dtype != torch.bfloat16 or w2_bf16.dtype != torch.bfloat16:
+        raise ValueError(
+            "prepare_trtllm_fp8_block_weights expects BF16 weights, got "
+            f"w1={w1_bf16.dtype}, w2={w2_bf16.dtype}."
+        )
+    expected_w1 = (num_local_experts, 2 * intermediate_size, hidden_size)
+    expected_w2 = (num_local_experts, hidden_size, intermediate_size)
+    if tuple(w1_bf16.shape) != expected_w1 or tuple(w2_bf16.shape) != expected_w2:
+        raise ValueError(
+            f"weight shapes {tuple(w1_bf16.shape)}/{tuple(w2_bf16.shape)} != "
+            f"expected {expected_w1}/{expected_w2}."
+        )
+
+
+def prepare_trtllm_fp8_block_weights(
+    w1_bf16: torch.Tensor,
+    w2_bf16: torch.Tensor,
+    *,
+    variant,
+    num_local_experts: int,
+    hidden_size: int,
+    intermediate_size: int,
+    device: Optional[torch.device] = None,
+) -> Dict[str, torch.Tensor]:
+    """Prepare canonical BF16 expert weights for TRTLLM block-FP8 MoE.
+
+    DeepSeek FP8 uses E4M3 payloads with FP32 128x128 block scales. MXFP8
+    uses E4M3 payloads with linear UE8M0 scales over 32-element K blocks.
+    Both native views remain in ``MajorK`` layout; the unified runner records
+    the exact variant and passes the corresponding kernel enum.
+    """
+    from .api import QuantVariant
+
+    if variant not in (QuantVariant.DeepSeekFp8, QuantVariant.MxFp8):
+        raise ValueError(
+            "variant must be QuantVariant.DeepSeekFp8 or QuantVariant.MxFp8, "
+            f"got {variant!r}."
+        )
+    _validate_trtllm_fp8_block_inputs(
+        w1_bf16,
+        w2_bf16,
+        num_local_experts=num_local_experts,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+    )
+    if device is None:
+        device = w1_bf16.device
+    w1_bf16 = w1_bf16.to(device).contiguous()
+    w2_bf16 = w2_bf16.to(device).contiguous()
+
+    if variant is QuantVariant.DeepSeekFp8:
+        for name, dim in (
+            ("hidden_size", hidden_size),
+            ("intermediate_size", intermediate_size),
+            ("2 * intermediate_size", 2 * intermediate_size),
+        ):
+            if dim % 128 != 0:
+                raise ValueError(f"DeepSeek FP8 requires {name} divisible by 128.")
+        w1_q, w1_sf = _deepseek_fp8_quantize_weights(w1_bf16)
+        w2_q, w2_sf = _deepseek_fp8_quantize_weights(w2_bf16)
+    else:
+        if hidden_size % 32 != 0 or intermediate_size % 32 != 0:
+            raise ValueError(
+                "MXFP8 requires hidden_size and intermediate_size divisible by 32."
+            )
+        from ..quantization.fp8_quantization import mxfp8_quantize
+        from .core import (
+            _maybe_get_cached_w3_w1_permute_indices,
+            get_w2_permute_indices_with_cache,
+        )
+
+        w1_q, w1_sf, w2_q, w2_sf = [], [], [], []
+        for expert in range(num_local_experts):
+            q, sf = mxfp8_quantize(w1_bf16[expert], is_sf_swizzled_layout=False)
+            sf = sf.view(torch.uint8).reshape(2 * intermediate_size, hidden_size // 32)
+            permute = _maybe_get_cached_w3_w1_permute_indices(
+                _TRTLLM_FP8_PERMUTE_CACHE,
+                q.view(torch.uint8),
+                128,
+                is_gated_act_gemm=True,
+            )
+            permute_sf = _maybe_get_cached_w3_w1_permute_indices(
+                _TRTLLM_FP8_PERMUTE_CACHE,
+                sf,
+                128,
+                num_elts_per_sf=32,
+                is_gated_act_gemm=True,
+            )
+            w1_q.append(q.view(torch.uint8)[permute.to(device)].view(q.dtype))
+            w1_sf.append(sf[permute_sf.to(device)])
+
+            q, sf = mxfp8_quantize(w2_bf16[expert], is_sf_swizzled_layout=False)
+            sf = sf.view(torch.uint8).reshape(hidden_size, intermediate_size // 32)
+            permute = get_w2_permute_indices_with_cache(
+                _TRTLLM_FP8_PERMUTE_CACHE, q.view(torch.uint8), 128
+            )
+            permute_sf = get_w2_permute_indices_with_cache(
+                _TRTLLM_FP8_PERMUTE_CACHE,
+                sf,
+                128,
+                num_elts_per_sf=32,
+            )
+            w2_q.append(q.view(torch.uint8)[permute.to(device)].view(q.dtype))
+            w2_sf.append(sf[permute_sf.to(device)])
+        w1_q, w1_sf = torch.stack(w1_q), torch.stack(w1_sf)
+        w2_q, w2_sf = torch.stack(w2_q), torch.stack(w2_sf)
+
+    return {
+        "gemm1_weights": w1_q,
+        "gemm1_weights_scale": w1_sf,
+        "gemm2_weights": w2_q,
+        "gemm2_weights_scale": w2_sf,
+    }
+
+
+def prepare_trtllm_fp8_block_activations(
+    hidden_states_bf16: torch.Tensor,
+    *,
+    variant,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize ``[M, H]`` BF16 activations for TRTLLM block-FP8 MoE."""
+    from .api import QuantVariant
+
+    if hidden_states_bf16.dtype != torch.bfloat16 or hidden_states_bf16.dim() != 2:
+        raise ValueError(
+            "prepare_trtllm_fp8_block_activations expects a 2D BF16 tensor, "
+            f"got shape={tuple(hidden_states_bf16.shape)}, "
+            f"dtype={hidden_states_bf16.dtype}."
+        )
+    hidden_states_bf16 = hidden_states_bf16.contiguous()
+    if variant is QuantVariant.DeepSeekFp8:
+        if hidden_states_bf16.shape[1] % 128 != 0:
+            raise ValueError("DeepSeek FP8 hidden_size must be divisible by 128.")
+        return _deepseek_fp8_quantize_activations(hidden_states_bf16)
+    if variant is QuantVariant.MxFp8:
+        from ..quantization.fp8_quantization import mxfp8_quantize
+
+        q, sf = mxfp8_quantize(hidden_states_bf16, is_sf_swizzled_layout=False)
+        return q, sf.view(torch.uint8).reshape(hidden_states_bf16.shape[0], -1)
+    raise ValueError(
+        "variant must be QuantVariant.DeepSeekFp8 or QuantVariant.MxFp8, "
+        f"got {variant!r}."
+    )
+
+
 def prepare_trtllm_bf16_weights(
     w1_bf16: torch.Tensor,
     w2_bf16: torch.Tensor,