iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[GLM5.2 Perf] fused_indexer_q_rope_quant triton kernel, 1.9% ~ 3.3% E2E Throughput improvement. (#46862)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-06-26 22:16:20 -07:00
committed by GitHub
parent 455f25aa13
commit b588f66dc2
2 changed files with 167 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/sparse_attn_indexer.py
+131
View File
@@ -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,
vllm/model_executor/models/deepseek_v2.py
+36
View File
@@ -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