A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 22:35 UTC
[Bugfix] [Performance] DeepEPHighThroughput + DeepSeek : Quant before Dispatch (#21837)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Aug 1, 2025 at 22:44 UTC
ac45c44d98e77f30e47b8fb69134f4635183070d
1 file changed
+8
-5
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+8
-5
@@ -144,12 +144,13 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
144
"apply_router_weight_on_input is only implemented for topk=1")
145
a1 = a1 * topk_weights.to(a1.dtype)
146
147
- if quant_config.per_act_token_quant:
147
+ if quant_config.is_block_quantized:
148
+ # Quant and Dispatch
149
a1q, a1q_scale = moe_kernel_quantize_input(
150
a1,
151
a1_scale,
152
quant_dtype=quant_config.quant_dtype,
152
- per_act_token_quant=True,
153
+ per_act_token_quant=quant_config.per_act_token_quant,
154
block_shape=quant_config.block_shape,
155
)
156
if a1q_scale is not None and a1q_scale.numel() == 1:
@@ -162,8 +163,10 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
163
rank_topk_weights=topk_weights,
164
num_experts=num_experts)
165
else:
165
- # DeepEP kernels only support dispatching per-token-quant
166
- # quantization. dispatch in bfloat16.
166
+ # Dispatch and Quant
167
+ # DeepEP kernels only support dispatching block-quantized
168
+ # activation scales.
169
+ # Dispatch in bfloat16
170
(expert_x, _, expert_tokens_meta, expert_topk_ids,
171
expert_topk_weights) = self._do_dispatch(
172
tokens=a1,
@@ -171,7 +174,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
174
rank_topk_ids=topk_ids,
175
rank_topk_weights=topk_weights,
176
num_experts=num_experts)
174
- # quantize now
177
+ # Quantize after dispatch.
178
expert_x_scale = None
179
if expert_x.numel() != 0:
180
expert_x, expert_x_scale = moe_kernel_quantize_input(