From b588f66dc2982fe3228e0aee55b80387d31db2e4 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Sat, 27 Jun 2026 01:16:20 -0400 Subject: [PATCH] [GLM5.2 Perf] `fused_indexer_q_rope_quant` triton kernel, 1.9% ~ 3.3% E2E Throughput improvement. (#46862) Signed-off-by: yewentao256 --- .../layers/sparse_attn_indexer.py | 131 ++++++++++++++++++ vllm/model_executor/models/deepseek_v2.py | 36 +++++ 2 files changed, 167 insertions(+) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index fe2b268cde..c1bc731ee6 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -11,7 +11,11 @@ from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.utils.deep_gemm import ( fp8_fp4_mqa_logits, fp8_fp4_paged_mqa_logits, @@ -37,6 +41,133 @@ RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 MXFP4_BLOCK_SIZE = 32 +@triton.jit +def _fused_indexer_q_rope_quant_kernel( + positions, + q, + q_s0, + q_s1, + cos_sin_cache, + cos_sin_s0, + q_fp8, + q_fp8_s0, + q_fp8_s1, + weights, + weights_s0, + weights_s1, + weights_out, + weights_out_s0, + weights_out_s1, + softmax_scale, + head_scale, + fp8_min: tl.constexpr, + fp8_max: tl.constexpr, + is_neox: tl.constexpr, +): + token = tl.program_id(0) + head = tl.program_id(1) + offs32 = tl.arange(0, 32) + offs64 = tl.arange(0, 64) + + pos = tl.load(positions + token) + cos = tl.load(cos_sin_cache + pos * cos_sin_s0 + offs32).to(tl.float32) + sin = tl.load(cos_sin_cache + pos * cos_sin_s0 + 32 + offs32).to(tl.float32) + q_base = q + token * q_s0 + head * q_s1 + out_base = q_fp8 + token * q_fp8_s0 + head * q_fp8_s1 + + if is_neox: + # NeoX layout, x0 = q[0:32], x1 = q[32:64] + x0 = tl.load(q_base + offs32).to(tl.float32) + x1 = tl.load(q_base + 32 + offs32).to(tl.float32) + else: + # interleaved layout + # x0 = q[0, 2, 4, ...], x1 = q[1, 3, 5, ...] + x0 = tl.load(q_base + offs32 * 2).to(tl.float32) + x1 = tl.load(q_base + offs32 * 2 + 1).to(tl.float32) + r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32) + r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32) + amax = tl.maximum(tl.max(tl.abs(r0)), tl.max(tl.abs(r1))) + + q_nope = tl.load(q_base + 64 + offs64).to(tl.float32) + amax = tl.maximum(amax, tl.max(tl.abs(q_nope))) + scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max) + # e8m0 format + q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) + + if is_neox: + tl.store( + out_base + offs32, + tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + 32 + offs32, + tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + else: + tl.store( + out_base + offs32 * 2, + tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + offs32 * 2 + 1, + tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + tl.store( + out_base + 64 + offs64, + tl.clamp(q_nope / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), + ) + + weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32) + tl.store( + weights_out + token * weights_out_s0 + head * weights_out_s1, + weight * q_scale * softmax_scale * head_scale, + ) + + +def fused_indexer_q_rope_quant( + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + softmax_scale: float, + head_scale: float, + is_neox: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + assert current_platform.is_cuda() + assert q.dtype == torch.bfloat16 + assert q.shape[-1] == 128 + assert cos_sin_cache.shape[-1] == 64 + assert weights.shape == q.shape[:2] + + q_fp8 = torch.empty_like(q, dtype=current_platform.fp8_dtype()) + weights_out = torch.empty_like(weights, dtype=torch.float32) + fp8_min, fp8_max = get_fp8_min_max() + _fused_indexer_q_rope_quant_kernel[(q.shape[0], q.shape[1])]( + positions, + q, + q.stride(0), + q.stride(1), + cos_sin_cache, + cos_sin_cache.stride(0), + q_fp8, + q_fp8.stride(0), + q_fp8.stride(1), + weights, + weights.stride(0), + weights.stride(1), + weights_out, + weights_out.stride(0), + weights_out.stride(1), + softmax_scale, + head_scale, + fp8_min=fp8_min, + fp8_max=fp8_max, + is_neox=is_neox, + num_warps=1, + ) + return q_fp8, weights_out + + def _gather_workspace_shapes( total_seq_lens: int, head_dim: int, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 814118f8a7..9b08ca9825 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -73,6 +73,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.sparse_attn_indexer import ( SparseAttnIndexer, + fused_indexer_q_rope_quant, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -674,6 +675,13 @@ class Indexer(nn.Module): ) self.is_inplace_rope = is_inplace_rope + self.use_fused_indexer_q = ( + current_platform.is_cuda() + and self.quant_block_size == self.head_dim + and self.head_dim == 128 + and self.rope_dim == 64 + and self.scale_fmt is not None + ) def forward( self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb @@ -698,6 +706,34 @@ class Indexer(nn.Module): rotary_emb( positions, q[..., : self.rope_dim], k[..., : self.rope_dim].unsqueeze(1) ) + elif self.use_fused_indexer_q and q.dtype == torch.bfloat16: + # fused wk + weights_proj: one GEMM, then split + kw, _ = self.wk_weights_proj(hidden_states) + k = kw[:, : self.head_dim] + weights = kw[:, self.head_dim :] + + k = self.k_norm(k) + k_pe, k_nope = torch.split( + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + + q_fp8, weights = fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + weights, + self.softmax_scale, + self.n_head**-0.5, + rotary_emb.is_neox_style, + ) + + # rotate only the MQA K + q_dummy = torch.empty_like(k_pe.unsqueeze(1)) + _, k_pe = rotary_emb(positions, q_dummy, k_pe.unsqueeze(1)) + k_pe = k_pe.reshape(-1, 1, self.rope_dim) + k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) + + return self.indexer_op(hidden_states, q_fp8, k, weights) else: q_pe, q_nope = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1