[GLM 5.3 Perf] Skip qlnorm calculation for MHA, 4.4~7.7% E2E TTFT Improvement (#58845)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Oct 1, 2026 at 03:13 UTC
4d6874d13756868453f78d5e717b2fbe9f8281d0
5 files changed
+99
-26
tests/kernels/test_fused_deepseek_v32_norm_rope.py
+36
-10
@@ -24,6 +24,7 @@ outputs are checked within 1 representable-step (ULP); bf16 norm/RoPE outputs us
24
rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses).
25
"""
26
27
+from functools import partial
28
from typing import NamedTuple
29
30
import pytest
@@ -880,22 +881,36 @@ def test_fused_q_no_indexer(num_tokens: int, cfg: ModelConfig):
881
@pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS)
882
@pytest.mark.parametrize("num_tokens", [1, 17, 512])
883
@pytest.mark.parametrize("has_indexer", [True, False])
883
-def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool, cfg: ModelConfig):
884
+@pytest.mark.parametrize("use_mha", [True, False])
885
+@pytest.mark.parametrize(
886
+ "index_interleave,cache_dtype", [(False, torch.float32), (True, torch.bfloat16)]
887
+)
888
+@pytest.mark.parametrize("capture", [False, True])
889
+def test_fused_q_bf16_query(
890
+ num_tokens: int,
891
+ has_indexer: bool,
892
+ use_mha: bool,
893
+ index_interleave: bool,
894
+ cache_dtype: torch.dtype,
895
+ capture: bool,
896
+ cfg: ModelConfig,
897
+):
898
"""bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is
885
- produced (bf16, unquantized); ql_nope is consumed directly by the caller."""
899
+ produced; MHA does not need an absorbed ql_nope projection."""
900
torch.manual_seed(6)
901
dev = "cuda"
902
max_pos = 8192
903
pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos
904
905
+ # Attention takes q_pe as a strided slice of the full Q projection.
906
q_pe = torch.randn(
892
- num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16
893
- )
907
+ num_tokens, cfg.num_heads, 4 * cfg.rope_dim, device=dev, dtype=torch.bfloat16
908
+ )[..., -cfg.rope_dim :]
909
ql_nope = torch.randn(
910
num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16
911
)
912
q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32)
898
- q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
913
+ q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev).to(cache_dtype)
914
915
index_q = index_w = idx_cos_sin = None
916
if has_indexer:
@@ -909,24 +924,35 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool, cfg: ModelConfig
924
index_w = torch.randn(
925
num_tokens, cfg.index_heads, device=dev, dtype=torch.float32
926
)
912
- idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev)
927
+ idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev).to(cache_dtype)
928
914
- iq_fp8, iw_out, q_pe_out = K.fused_q(
929
+ run = partial(
930
+ K.fused_q,
931
pos,
932
q_pe,
933
q_cos_sin,
934
index_q,
935
idx_cos_sin,
920
- ql_nope,
936
+ q_pe if use_mha else ql_nope,
937
q_scale,
938
index_w,
939
cfg.index_head_dim**-0.5,
940
cfg.index_heads**-0.5,
941
has_indexer=has_indexer,
926
- index_rope_interleave=False,
942
+ index_rope_interleave=index_interleave,
943
quantize_mqa=False,
944
)
945
946
+ iq_fp8, iw_out, q_pe_out = run()
947
+ if capture:
948
+ graph = torch.cuda.CUDAGraph()
949
+ with torch.cuda.graph(graph):
950
+ iq_fp8, iw_out, q_pe_out = run()
951
+ q_pe.mul_(0.5)
952
+ if index_q is not None:
953
+ index_q.mul_(1.5)
954
+ graph.replay()
955
+
956
# MQA query: only the RoPE'd q_pe, bf16, unquantized.
957
assert q_pe_out.dtype == torch.bfloat16
958
assert q_pe_out.shape == (num_tokens, cfg.num_heads, cfg.rope_dim)
@@ -945,7 +971,7 @@ def test_fused_q_bf16_query(num_tokens: int, has_indexer: bool, cfg: ModelConfig
971
index_q.float(),
972
pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads),
973
idx_cos_sin,
948
- interleave=False,
974
+ interleave=index_interleave,
975
)
976
q_ref, scale_ref = ue8m0_quant(iq_ref)
977
assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)")
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+7
-1
@@ -267,9 +267,11 @@ def test_skipped_k_cache_insert_accepts_no_k(
267
268
269
@pytest.mark.parametrize("fp8_query", [False, True])
270
+@pytest.mark.parametrize("prepared_mha", [False, True])
271
def test_deepseek_v32_dispatches_selected_mha(
272
monkeypatch: pytest.MonkeyPatch,
273
fp8_query: bool,
274
+ prepared_mha: bool,
275
) -> None:
276
attn_metadata = SimpleNamespace(num_actual_tokens=2)
277
kv_cache = torch.empty(1)
@@ -301,7 +303,11 @@ def test_deepseek_v32_dispatches_selected_mha(
303
q_nope = torch.randn(2, 1, 2)
304
q_pe = torch.randn(2, 1, 2)
305
mqa_q = torch.randn(2, 1, 2)
304
- mha_q = torch.cat((q_nope, q_pe + 1 if fp8_query else mqa_q), dim=-1)
306
+ if fp8_query and not prepared_mha:
307
+ mqa_q = mqa_q.to(torch.float8_e4m3fn)
308
+ mha_q = torch.cat(
309
+ (q_nope, q_pe + 1 if fp8_query and not prepared_mha else mqa_q), dim=-1
310
+ )
311
kv_c = torch.empty(2, 2)
312
k_pe = torch.empty(2, 2)
313
output = torch.empty(2, 2)
vllm/models/deepseek_v32/attention.py
+26
-4
@@ -7,7 +7,7 @@ import torch.nn as nn
7
from transformers import DeepseekV2Config, DeepseekV3Config
8
9
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
10
-from vllm.config import CacheConfig, VllmConfig
10
+from vllm.config import CacheConfig, CUDAGraphMode, VllmConfig
11
from vllm.distributed import get_tensor_model_parallel_world_size
12
from vllm.distributed.parallel_state import get_tp_group
13
from vllm.forward_context import get_forward_context
@@ -283,6 +283,18 @@ class DeepseekV32Attention(MLAAttention):
283
is_neox_style=not getattr(config, "indexer_rope_interleave", False),
284
)
285
286
+ def _should_prepare_mqa_query(
287
+ self,
288
+ attn_metadata: "MLACommonMetadata | None",
289
+ cudagraph_runtime_mode: CUDAGraphMode,
290
+ ) -> bool:
291
+ return (
292
+ attn_metadata is None
293
+ or not self._use_sparse_mha(attn_metadata)
294
+ or cudagraph_runtime_mode == CUDAGraphMode.FULL
295
+ or torch.cuda.is_current_stream_capturing()
296
+ )
297
+
298
def forward( # type: ignore[override]
299
self,
300
positions: torch.Tensor,
@@ -321,6 +333,9 @@ class DeepseekV32Attention(MLAAttention):
333
layer_attn_metadata, _, _, _ = get_attention_context(self.layer_name)
334
self.impl.prepare_for_batch(layer_attn_metadata)
335
336
+ prepare_mqa_query = self._should_prepare_mqa_query(
337
+ layer_attn_metadata, forward_context.cudagraph_runtime_mode
338
+ )
339
if self.indexer is not None and not self.skip_topk:
340
has_indexer = True
341
indexer_k_norm_w = self.indexer.k_norm.weight
@@ -404,7 +419,10 @@ class DeepseekV32Attention(MLAAttention):
419
420
q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
421
q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
407
- ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1)
422
+ if prepare_mqa_query:
423
+ ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1)
424
+ else:
425
+ ql_nope = q_nope
426
427
if self.indexer is not None and not self.skip_topk:
428
index_q = self.indexer.wq_b(q_c)[0]
@@ -425,7 +443,7 @@ class DeepseekV32Attention(MLAAttention):
443
indexer_n_head_scale,
444
has_indexer=has_indexer,
445
index_rope_interleave=self._index_rope_interleave,
428
- quantize_mqa=self._fp8_query,
446
+ quantize_mqa=self._fp8_query and prepare_mqa_query,
447
)
448
449
self._sparse_indexer_and_attn(
@@ -535,7 +553,11 @@ class DeepseekV32Attention(MLAAttention):
553
554
if self._use_sparse_mha(attn_metadata):
555
assert kv_c is not None and k_pe is not None
538
- mha_q_pe = self.rotary_emb(positions, q_pe)[0] if self._fp8_query else mqa_q
556
+ mha_q_pe = (
557
+ self.rotary_emb(positions, q_pe)[0]
558
+ if mqa_q.dtype != q_pe.dtype
559
+ else mqa_q
560
+ )
561
mha_q = torch.cat((q_nope, mha_q_pe), dim=-1)
562
self.forward_impl(
563
mha_q,
vllm/models/deepseek_v32/common/kernels.py
+1
@@ -1052,6 +1052,7 @@ def fused_q(
1052
index_weights_out,
1053
has_indexer=has_indexer,
1054
index_rope_interleave=index_rope_interleave,
1055
+ quantize_mqa=quantize_mqa,
1056
)
1057
return index_q_fp8, index_weights_out, mqa_q
1058
vllm/models/deepseek_v32/nvidia/ops/fused_q_cutedsl.py
+29
-11
@@ -35,10 +35,12 @@ def is_fused_q_cutedsl_supported(
35
) -> bool:
36
if not (
37
current_platform.has_device_capability(100)
38
- and quantize_mqa
39
- and q_pe.dtype == ql_nope.dtype == torch.bfloat16
38
+ and q_pe.dtype == torch.bfloat16
39
and q_pe.shape[-1] == 64
41
- and ql_nope.shape[-1] == 512
40
+ and (
41
+ not quantize_mqa
42
+ or (ql_nope.dtype == torch.bfloat16 and ql_nope.shape[-1] == 512)
43
+ )
44
# One warp per head group; a non-multiple would trip the kernel's own
45
# assert instead of falling back to Triton.
46
and q_pe.shape[1] % 4 == 0
@@ -68,9 +70,10 @@ def fused_q_cutedsl(
70
idx_weights_out: torch.Tensor,
71
has_indexer: bool = True,
72
index_rope_interleave: bool = True,
73
+ quantize_mqa: bool = True,
74
) -> None:
75
_, num_heads, rope_dim = q_pe.shape
73
- _, _, nope_dim = ql_nope.shape
76
+ nope_dim = ql_nope.shape[-1] if quantize_mqa else 0
77
_, num_idx_heads, idx_dim = idx_q.shape
78
79
if has_indexer:
@@ -98,8 +101,8 @@ def fused_q_cutedsl(
101
positions,
102
q_pe,
103
rope_cache,
101
- ql_nope,
102
- q_scale.view(1),
104
+ ql_nope if quantize_mqa else None,
105
+ q_scale.view(1) if quantize_mqa else None,
106
mqa_output,
107
idx_q,
108
idx_rope_cache,
@@ -121,7 +124,7 @@ class FusedQKernel:
124
index_rope_interleave: bool,
125
) -> None:
126
assert rope_dim == 64
124
- assert nope_dim == 512
127
+ assert nope_dim in (0, 512)
128
assert idx_dim in (128, 0)
129
130
self.rope_dim = rope_dim
@@ -262,6 +265,17 @@ class FusedQKernel:
265
cute.arch.griddepcontrol_wait()
266
267
pos = positions[token_id]
268
+ # no need to do Q calculation for MHA
269
+ if cutlass.const_expr(self.nope_dim == 0):
270
+ q0, q1 = cvt.bf16x2_to_fp32x2(
271
+ cute.recast_tensor(q_pe, Uint32)[token_id, head_id, lane_id]
272
+ )
273
+ cos = q_pe_rope_cache[pos, lane_id].to(Float32)
274
+ sin = q_pe_rope_cache[pos, 32 + lane_id].to(Float32)
275
+ result = cvt.fp32x2_to_bf16x2(q0 * cos - q1 * sin, q1 * cos + q0 * sin)
276
+ cute.arch.griddepcontrol_launch_dependents()
277
+ cute.recast_tensor(mqa_output, Uint32)[token_id, head_id, lane_id] = result
278
+ return
279
inv_scale = 1.0 / q_scale[0]
280
281
cp_op = cute.nvgpu.CopyUniversalOp()
@@ -479,12 +493,16 @@ class FusedQKernel:
493
BFloat16, (num_tokens, num_heads, rope_dim), divisibility=16
494
)
495
rope_cache = _make_fake_tensor(rope_type, (max_pos, rope_dim), divisibility=8)
482
- ql_nope = _make_fake_tensor(
483
- BFloat16, (num_tokens, num_heads, nope_dim), divisibility=16
496
+ ql_nope = (
497
+ _make_fake_tensor(
498
+ BFloat16, (num_tokens, num_heads, nope_dim), divisibility=16
499
+ )
500
+ if nope_dim
501
+ else None
502
)
485
- q_scale = _make_fake_tensor(Float32, (1,), divisibility=4)
503
+ q_scale = _make_fake_tensor(Float32, (1,), divisibility=4) if nope_dim else None
504
mqa_output = _make_fake_tensor(
487
- Float8E4M3FN,
505
+ Float8E4M3FN if nope_dim else BFloat16,
506
(num_tokens, num_heads, nope_dim + rope_dim),
507
divisibility=16,
508
)