[DeepEPv2] Support MXFp8 Activation Scale Dispatch (#51398)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Aug 26, 2026 at 08:26 UTC
6a5e8f5979376e666c930f668d0f58d78a9933c6
1 file changed
+63
-2
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
+63
-2
@@ -13,6 +13,10 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
13
TopKWeightAndReduceDelegate,
14
)
15
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
16
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
17
+ MXFP8_BLOCK_SIZE,
18
+ swizzle_mxfp8_scale,
19
+)
20
from vllm.triton_utils import tl, triton
21
from vllm.utils.math_utils import round_up
22
from vllm.v1.worker.ubatching import (
@@ -20,6 +24,49 @@ from vllm.v1.worker.ubatching import (
24
)
25
26
27
+def _quantize_before_dispatch(
28
+ quant_config: FusedMoEQuantConfig, defer_input_quant: bool
29
+) -> bool:
30
+ """
31
+ Do quantized dispatch for blockfp8 and mxfp8, unless the
32
+ subsequent moe kernel requires bf16 inputs.
33
+ """
34
+ if defer_input_quant:
35
+ return False
36
+ return quant_config.is_block_quantized or quant_config.quant_dtype == "mxfp8"
37
+
38
+
39
+def _pack_mxfp8_scale(scale: torch.Tensor) -> torch.Tensor:
40
+ """Pack row-major [M, K/32] UE8M0 scales into [M, K/128] int32.
41
+
42
+ DeepEP moves scale factors as opaque 4-byte packs (`sf_pack_t` is a
43
+ float/UE8M0x4 union), so 1-byte UE8M0 scales must be packed 4-per-int32.
44
+ """
45
+ assert scale.dtype == torch.uint8 and scale.ndim == 2, (
46
+ f"expected 2D uint8 mxfp8 scales, got {scale.shape} {scale.dtype}"
47
+ )
48
+ assert scale.size(1) % 4 == 0, (
49
+ f"mxfp8 dispatch needs hidden_size % {MXFP8_BLOCK_SIZE * 4} == 0, "
50
+ f"got {scale.size(1)} scale columns"
51
+ )
52
+ return scale.contiguous().view(torch.int32)
53
+
54
+
55
+def _unpack_mxfp8_scale(
56
+ scale: torch.Tensor, hidden_size: int, is_scale_swizzled: bool
57
+) -> torch.Tensor:
58
+ """Inverse of `_pack_mxfp8_scale`, restoring the expert kernel's layout.
59
+
60
+ TRTLLM consumes the row-major [M, K/32] scales as-is; CUTLASS wants them
61
+ swizzled into F8_128x4, which can only happen here because the swizzle
62
+ interleaves scales across a 128-row tile (i.e. across tokens).
63
+ """
64
+ scale = scale.contiguous().view(torch.uint8)
65
+ if is_scale_swizzled:
66
+ scale = swizzle_mxfp8_scale(scale, M=scale.size(0), K=hidden_size)
67
+ return scale
68
+
69
+
70
class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
71
"""
72
Prepare/Finalize using DeepEP v2 ElasticBuffer (unified API).
@@ -278,7 +325,14 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
325
)
326
expert_tokens_meta.psum_recv_per_rank = psum_recv_per_rank
327
281
- if not quant_config.is_block_quantized and not defer_input_quant:
328
+ if _quantize_before_dispatch(quant_config, defer_input_quant):
329
+ if quant_config.quant_dtype == "mxfp8" and expert_x_scale is not None:
330
+ expert_x_scale = _unpack_mxfp8_scale(
331
+ expert_x_scale,
332
+ hidden_size=expert_x.size(-1),
333
+ is_scale_swizzled=quant_config.is_scale_swizzled,
334
+ )
335
+ elif not defer_input_quant:
336
expert_x_scale = None
337
if expert_x.numel() != 0:
338
expert_x, expert_x_scale = moe_kernel_quantize_input(
@@ -319,16 +373,23 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
373
)
374
a1 = a1 * topk_weights.to(a1.dtype)
375
322
- if quant_config.is_block_quantized and not defer_input_quant:
376
+ if _quantize_before_dispatch(quant_config, defer_input_quant):
377
+ # Scales must be row-major [num_tokens, ...] here so each token's
378
+ # scales can be shuffled with it; any swizzling the expert kernel
379
+ # wants is applied post-dispatch in `_receiver`.
380
a1q, a1q_scale = moe_kernel_quantize_input(
381
a1,
382
quant_config.a1_scale,
383
quant_dtype=quant_config.quant_dtype,
384
per_act_token_quant=quant_config.per_act_token_quant,
385
block_shape=quant_config.block_shape,
386
+ is_scale_swizzled=False,
387
+ mx_alignment=quant_config.mx_alignment,
388
)
389
if a1q_scale is not None and a1q_scale.numel() == 1:
390
a1q_scale = a1q_scale.view(1, 1)
391
+ if quant_config.quant_dtype == "mxfp8":
392
+ a1q_scale = _pack_mxfp8_scale(a1q_scale)
393
a1_post_scale = None
394
else:
395
a1q = a1