[Perf][DSv4.1] Restore the fused query RMSNorm + MXFP8 quantization path (#57679)
Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
Juntian Liu committed
Sep 24, 2026 at 18:39 UTC
609731e8f56c465c4afa2435298fcedcf8301496
1 file changed
+8
-3
vllm/models/deepseek_v41/common/ops/query_quant.py
+8
-3
@@ -3,7 +3,10 @@
3
4
import torch
5
6
-from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
6
+from vllm.model_executor.layers.fusion.quant_activation import (
7
+ QuantizedActivation,
8
+ get_input_quant_key,
9
+)
10
from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
11
from vllm.platforms import current_platform
12
from vllm.triton_utils import tl, triton
@@ -102,6 +105,7 @@ def fused_q_kv_rmsnorm_quant(
105
)
106
if tokens:
107
block = triton.next_power_of_2(max(q_size, kv_size))
108
+ launch_pdl = current_platform.is_arch_support_pdl()
109
_q_kv_norm_quant_kernel[(padded_tokens, 2)](
110
qr,
111
kv,
@@ -117,8 +121,9 @@ def fused_q_kv_rmsnorm_quant(
121
q_size,
122
kv_size,
123
block,
120
- current_platform.is_arch_support_pdl(),
124
+ launch_pdl,
125
num_warps=8 if block >= 2048 else 4,
126
+ launch_pdl=launch_pdl,
127
)
128
return QuantizedActivation(qo, scales, qr.dtype, qr.shape, kMxfp8Dynamic), kvo
129
@@ -135,7 +140,7 @@ def can_fuse_query_quant(linears: list[torch.nn.Module]) -> bool:
140
# QuantKey does not encode scale layout; restrict this producer to the
141
# consumers that accept F8_128x4 swizzled scales.
142
return all(
138
- getattr(linear, "input_quant_key", None) == kMxfp8Dynamic
143
+ get_input_quant_key(linear) == kMxfp8Dynamic
144
and type(getattr(getattr(linear, "quant_method", None), "kernel", None))
145
in (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel)
146
for linear in linears