[GLM5.2 Perf] `fused_indexer_q_rope_quant` triton kernel, 1.9% ~ 3.3% E2E Throughput improvement. (#46862)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 27, 2026 at 01:16 UTC
b588f66dc2982fe3228e0aee55b80387d31db2e4
2 files changed
+167
vllm/model_executor/layers/sparse_attn_indexer.py
+131
@@ -11,7 +11,11 @@ from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11
from vllm.forward_context import get_forward_context
12
from vllm.logger import init_logger
13
from vllm.model_executor.custom_op import CustomOp
14
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
+ get_fp8_min_max,
16
+)
17
from vllm.platforms import current_platform
18
+from vllm.triton_utils import tl, triton
19
from vllm.utils.deep_gemm import (
20
fp8_fp4_mqa_logits,
21
fp8_fp4_paged_mqa_logits,
@@ -37,6 +41,133 @@ RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
41
MXFP4_BLOCK_SIZE = 32
42
43
44
+@triton.jit
45
+def _fused_indexer_q_rope_quant_kernel(
46
+ positions,
47
+ q,
48
+ q_s0,
49
+ q_s1,
50
+ cos_sin_cache,
51
+ cos_sin_s0,
52
+ q_fp8,
53
+ q_fp8_s0,
54
+ q_fp8_s1,
55
+ weights,
56
+ weights_s0,
57
+ weights_s1,
58
+ weights_out,
59
+ weights_out_s0,
60
+ weights_out_s1,
61
+ softmax_scale,
62
+ head_scale,
63
+ fp8_min: tl.constexpr,
64
+ fp8_max: tl.constexpr,
65
+ is_neox: tl.constexpr,
66
+):
67
+ token = tl.program_id(0)
68
+ head = tl.program_id(1)
69
+ offs32 = tl.arange(0, 32)
70
+ offs64 = tl.arange(0, 64)
71
+
72
+ pos = tl.load(positions + token)
73
+ cos = tl.load(cos_sin_cache + pos * cos_sin_s0 + offs32).to(tl.float32)
74
+ sin = tl.load(cos_sin_cache + pos * cos_sin_s0 + 32 + offs32).to(tl.float32)
75
+ q_base = q + token * q_s0 + head * q_s1
76
+ out_base = q_fp8 + token * q_fp8_s0 + head * q_fp8_s1
77
+
78
+ if is_neox:
79
+ # NeoX layout, x0 = q[0:32], x1 = q[32:64]
80
+ x0 = tl.load(q_base + offs32).to(tl.float32)
81
+ x1 = tl.load(q_base + 32 + offs32).to(tl.float32)
82
+ else:
83
+ # interleaved layout
84
+ # x0 = q[0, 2, 4, ...], x1 = q[1, 3, 5, ...]
85
+ x0 = tl.load(q_base + offs32 * 2).to(tl.float32)
86
+ x1 = tl.load(q_base + offs32 * 2 + 1).to(tl.float32)
87
+ r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32)
88
+ r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32)
89
+ amax = tl.maximum(tl.max(tl.abs(r0)), tl.max(tl.abs(r1)))
90
+
91
+ q_nope = tl.load(q_base + 64 + offs64).to(tl.float32)
92
+ amax = tl.maximum(amax, tl.max(tl.abs(q_nope)))
93
+ scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max)
94
+ # e8m0 format
95
+ q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
96
+
97
+ if is_neox:
98
+ tl.store(
99
+ out_base + offs32,
100
+ tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
101
+ )
102
+ tl.store(
103
+ out_base + 32 + offs32,
104
+ tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
105
+ )
106
+ else:
107
+ tl.store(
108
+ out_base + offs32 * 2,
109
+ tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
110
+ )
111
+ tl.store(
112
+ out_base + offs32 * 2 + 1,
113
+ tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
114
+ )
115
+ tl.store(
116
+ out_base + 64 + offs64,
117
+ tl.clamp(q_nope / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty),
118
+ )
119
+
120
+ weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32)
121
+ tl.store(
122
+ weights_out + token * weights_out_s0 + head * weights_out_s1,
123
+ weight * q_scale * softmax_scale * head_scale,
124
+ )
125
+
126
+
127
+def fused_indexer_q_rope_quant(
128
+ positions: torch.Tensor,
129
+ q: torch.Tensor,
130
+ cos_sin_cache: torch.Tensor,
131
+ weights: torch.Tensor,
132
+ softmax_scale: float,
133
+ head_scale: float,
134
+ is_neox: bool,
135
+) -> tuple[torch.Tensor, torch.Tensor]:
136
+ assert current_platform.is_cuda()
137
+ assert q.dtype == torch.bfloat16
138
+ assert q.shape[-1] == 128
139
+ assert cos_sin_cache.shape[-1] == 64
140
+ assert weights.shape == q.shape[:2]
141
+
142
+ q_fp8 = torch.empty_like(q, dtype=current_platform.fp8_dtype())
143
+ weights_out = torch.empty_like(weights, dtype=torch.float32)
144
+ fp8_min, fp8_max = get_fp8_min_max()
145
+ _fused_indexer_q_rope_quant_kernel[(q.shape[0], q.shape[1])](
146
+ positions,
147
+ q,
148
+ q.stride(0),
149
+ q.stride(1),
150
+ cos_sin_cache,
151
+ cos_sin_cache.stride(0),
152
+ q_fp8,
153
+ q_fp8.stride(0),
154
+ q_fp8.stride(1),
155
+ weights,
156
+ weights.stride(0),
157
+ weights.stride(1),
158
+ weights_out,
159
+ weights_out.stride(0),
160
+ weights_out.stride(1),
161
+ softmax_scale,
162
+ head_scale,
163
+ fp8_min=fp8_min,
164
+ fp8_max=fp8_max,
165
+ is_neox=is_neox,
166
+ num_warps=1,
167
+ )
168
+ return q_fp8, weights_out
169
+
170
+
171
def _gather_workspace_shapes(
172
total_seq_lens: int,
173
head_dim: int,
vllm/model_executor/models/deepseek_v2.py
+36
@@ -73,6 +73,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
73
from vllm.model_executor.layers.rotary_embedding import get_rope
74
from vllm.model_executor.layers.sparse_attn_indexer import (
75
SparseAttnIndexer,
76
+ fused_indexer_q_rope_quant,
77
)
78
from vllm.model_executor.layers.vocab_parallel_embedding import (
79
ParallelLMHead,
@@ -674,6 +675,13 @@ class Indexer(nn.Module):
675
)
676
677
self.is_inplace_rope = is_inplace_rope
678
+ self.use_fused_indexer_q = (
679
+ current_platform.is_cuda()
680
+ and self.quant_block_size == self.head_dim
681
+ and self.head_dim == 128
682
+ and self.rope_dim == 64
683
+ and self.scale_fmt is not None
684
+ )
685
686
def forward(
687
self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
@@ -698,6 +706,34 @@ class Indexer(nn.Module):
706
rotary_emb(
707
positions, q[..., : self.rope_dim], k[..., : self.rope_dim].unsqueeze(1)
708
)
709
+ elif self.use_fused_indexer_q and q.dtype == torch.bfloat16:
710
+ # fused wk + weights_proj: one GEMM, then split
711
+ kw, _ = self.wk_weights_proj(hidden_states)
712
+ k = kw[:, : self.head_dim]
713
+ weights = kw[:, self.head_dim :]
714
+
715
+ k = self.k_norm(k)
716
+ k_pe, k_nope = torch.split(
717
+ k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
718
+ )
719
+
720
+ q_fp8, weights = fused_indexer_q_rope_quant(
721
+ positions,
722
+ q,
723
+ rotary_emb.cos_sin_cache,
724
+ weights,
725
+ self.softmax_scale,
726
+ self.n_head**-0.5,
727
+ rotary_emb.is_neox_style,
728
+ )
729
+
730
+ # rotate only the MQA K
731
+ q_dummy = torch.empty_like(k_pe.unsqueeze(1))
732
+ _, k_pe = rotary_emb(positions, q_dummy, k_pe.unsqueeze(1))
733
+ k_pe = k_pe.reshape(-1, 1, self.rope_dim)
734
+ k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
735
+
736
+ return self.indexer_op(hidden_states, q_fp8, k, weights)
737
else:
738
q_pe, q_nope = torch.split(
739
q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1