[DeepEPv2] Support MXFp8 Activation Scale Dispatch (#51398)

Signed-off-by: Robert Shaw <robshaw@redhat.com>
Co-authored-by: Robert Shaw <robshaw@redhat.com>

vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py

@@ -13,6 +13,10 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceDelegate,
 )
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    swizzle_mxfp8_scale,
+)
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import round_up
 from vllm.v1.worker.ubatching import (
@@ -20,6 +24,49 @@ from vllm.v1.worker.ubatching import (
 )
 
 
+def _quantize_before_dispatch(
+    quant_config: FusedMoEQuantConfig, defer_input_quant: bool
+) -> bool:
+    """
+    Do quantized dispatch for blockfp8 and mxfp8, unless the
+    subsequent moe kernel requires bf16 inputs.
+    """
+    if defer_input_quant:
+        return False
+    return quant_config.is_block_quantized or quant_config.quant_dtype == "mxfp8"
+
+
+def _pack_mxfp8_scale(scale: torch.Tensor) -> torch.Tensor:
+    """Pack row-major [M, K/32] UE8M0 scales into [M, K/128] int32.
+
+    DeepEP moves scale factors as opaque 4-byte packs (`sf_pack_t` is a
+    float/UE8M0x4 union), so 1-byte UE8M0 scales must be packed 4-per-int32.
+    """
+    assert scale.dtype == torch.uint8 and scale.ndim == 2, (
+        f"expected 2D uint8 mxfp8 scales, got {scale.shape} {scale.dtype}"
+    )
+    assert scale.size(1) % 4 == 0, (
+        f"mxfp8 dispatch needs hidden_size % {MXFP8_BLOCK_SIZE * 4} == 0, "
+        f"got {scale.size(1)} scale columns"
+    )
+    return scale.contiguous().view(torch.int32)
+
+
+def _unpack_mxfp8_scale(
+    scale: torch.Tensor, hidden_size: int, is_scale_swizzled: bool
+) -> torch.Tensor:
+    """Inverse of `_pack_mxfp8_scale`, restoring the expert kernel's layout.
+
+    TRTLLM consumes the row-major [M, K/32] scales as-is; CUTLASS wants them
+    swizzled into F8_128x4, which can only happen here because the swizzle
+    interleaves scales across a 128-row tile (i.e. across tokens).
+    """
+    scale = scale.contiguous().view(torch.uint8)
+    if is_scale_swizzled:
+        scale = swizzle_mxfp8_scale(scale, M=scale.size(0), K=hidden_size)
+    return scale
+
+
 class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
     """
     Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API).
@@ -278,7 +325,14 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
                 )
             expert_tokens_meta.psum_recv_per_rank = psum_recv_per_rank
 
-        if not quant_config.is_block_quantized and not defer_input_quant:
+        if _quantize_before_dispatch(quant_config, defer_input_quant):
+            if quant_config.quant_dtype == "mxfp8" and expert_x_scale is not None:
+                expert_x_scale = _unpack_mxfp8_scale(
+                    expert_x_scale,
+                    hidden_size=expert_x.size(-1),
+                    is_scale_swizzled=quant_config.is_scale_swizzled,
+                )
+        elif not defer_input_quant:
             expert_x_scale = None
             if expert_x.numel() != 0:
                 expert_x, expert_x_scale = moe_kernel_quantize_input(
@@ -319,16 +373,23 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
             )
             a1 = a1 * topk_weights.to(a1.dtype)
 
-        if quant_config.is_block_quantized and not defer_input_quant:
+        if _quantize_before_dispatch(quant_config, defer_input_quant):
+            # Scales must be row-major [num_tokens, ...] here so each token's
+            # scales can be shuffled with it; any swizzling the expert kernel
+            # wants is applied post-dispatch in `_receiver`.
             a1q, a1q_scale = moe_kernel_quantize_input(
                 a1,
                 quant_config.a1_scale,
                 quant_dtype=quant_config.quant_dtype,
                 per_act_token_quant=quant_config.per_act_token_quant,
                 block_shape=quant_config.block_shape,
+                is_scale_swizzled=False,
+                mx_alignment=quant_config.mx_alignment,
             )
             if a1q_scale is not None and a1q_scale.numel() == 1:
                 a1q_scale = a1q_scale.view(1, 1)
+            if quant_config.quant_dtype == "mxfp8":
+                a1q_scale = _pack_mxfp8_scale(a1q_scale)
             a1_post_scale = None
         else:
             a1q = a1