karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

Fuse indexer k cache update

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
This commit is contained in:
Woosuk Kwon
2026-03-30 22:28:33 +00:00
parent 9d7beab915
commit 34d73a3375
4 changed files with 602 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_triton_cache_ops.py
+412
View File
@@ -0,0 +1,412 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Bitwise equivalence tests: Triton cache kernels vs CUDA cache kernels."""
import random
import pytest
import torch
from vllm import _custom_ops as ops
from vllm.utils.torch_utils import set_random_seed
# ---------- indexer_k_quant_and_cache parameters ----------
# head_dim == quant_block_size == 128 always for the indexer kernel.
HEAD_DIMS = [128]
QUANT_BLOCK_SIZES = [128]
SCALE_FMTS = ["ue8m0", "other"]
CACHE_BLOCK_SIZES = [16]
NUM_BLOCKS = [32]
NUM_TOKENS = [1, 42]
DTYPES = [torch.bfloat16]
SEEDS = [0]
# ---------- concat_and_cache_mla parameters ----------
KV_LORA_RANKS = [256, 512]
PE_DIMS = [64]
MLA_BLOCK_SIZES = [16]
MLA_NUM_BLOCKS = [8]
MLA_NUM_TOKENS = [1, 42]
MLA_KV_CACHE_DTYPES = ["auto", "fp8_e4m3"]
# =====================================================================
# indexer_k_quant_and_cache
# =====================================================================
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
@pytest.mark.parametrize("quant_block_size", QUANT_BLOCK_SIZES)
@pytest.mark.parametrize("scale_fmt", SCALE_FMTS)
@pytest.mark.parametrize("cache_block_size", CACHE_BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@torch.inference_mode()
def test_indexer_k_quant_and_cache_equivalence(
head_dim: int,
quant_block_size: int,
scale_fmt: str,
cache_block_size: int,
num_blocks: int,
num_tokens: int,
dtype: torch.dtype,
seed: int,
):
set_random_seed(seed)
device = "cuda"
# cache_stride = fp8 data + float32 scales (in fp8-element units)
cache_stride = head_dim + head_dim * 4 // quant_block_size
total_slots = num_blocks * cache_block_size
slot_mapping_lst = random.sample(range(total_slots), num_tokens)
slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
k = torch.randn(num_tokens, head_dim, dtype=dtype, device=device)
# Create two identical caches
kv_cache_cuda = torch.zeros(
num_blocks,
cache_block_size,
cache_stride,
dtype=torch.float8_e4m3fn,
device=device,
)
kv_cache_triton = kv_cache_cuda.clone()
# Run CUDA kernel
ops.indexer_k_quant_and_cache(
k,
kv_cache_cuda,
slot_mapping,
quant_block_size,
scale_fmt,
)
# Run Triton kernel
from vllm.v1.attention.ops.indexer_k_quant_and_cache import (
indexer_k_quant_and_cache as triton_indexer_k_quant_and_cache,
)
triton_indexer_k_quant_and_cache(
k,
kv_cache_triton,
slot_mapping,
quant_block_size,
scale_fmt,
)
# Bitwise comparison (view as uint8 to catch any bit-level differences)
cuda_bytes = kv_cache_cuda.view(torch.uint8)
triton_bytes = kv_cache_triton.view(torch.uint8)
torch.testing.assert_close(cuda_bytes, triton_bytes, atol=0, rtol=0)
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
@pytest.mark.parametrize("quant_block_size", QUANT_BLOCK_SIZES)
@pytest.mark.parametrize("cache_block_size", CACHE_BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("seed", SEEDS)
@torch.inference_mode()
def test_indexer_k_quant_and_cache_padding(
head_dim: int,
quant_block_size: int,
cache_block_size: int,
num_blocks: int,
seed: int,
):
"""Verify that padded tokens (slot_mapping == -1) leave the cache untouched."""
set_random_seed(seed)
device = "cuda"
cache_stride = head_dim + head_dim * 4 // quant_block_size
num_tokens = 4
# All slots are -1 (padding)
slot_mapping = torch.full((num_tokens,), -1, dtype=torch.long, device=device)
k = torch.randn(num_tokens, head_dim, dtype=torch.bfloat16, device=device)
kv_cache = torch.zeros(
num_blocks,
cache_block_size,
cache_stride,
dtype=torch.float8_e4m3fn,
device=device,
)
snapshot = kv_cache.clone()
from vllm.v1.attention.ops.indexer_k_quant_and_cache import (
indexer_k_quant_and_cache as triton_indexer_k_quant_and_cache,
)
triton_indexer_k_quant_and_cache(
k,
kv_cache,
slot_mapping,
quant_block_size,
"ue8m0",
)
torch.testing.assert_close(
kv_cache.view(torch.uint8),
snapshot.view(torch.uint8),
atol=0,
rtol=0,
)
# =====================================================================
# concat_and_cache_mla
# =====================================================================
@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
@pytest.mark.parametrize("pe_dim", PE_DIMS)
@pytest.mark.parametrize("block_size", MLA_BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", MLA_NUM_BLOCKS)
@pytest.mark.parametrize("num_tokens", MLA_NUM_TOKENS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kv_cache_dtype", MLA_KV_CACHE_DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@torch.inference_mode()
def test_concat_and_cache_mla_equivalence(
kv_lora_rank: int,
pe_dim: int,
block_size: int,
num_blocks: int,
num_tokens: int,
dtype: torch.dtype,
kv_cache_dtype: str,
seed: int,
):
set_random_seed(seed)
device = "cuda"
total_slots = num_blocks * block_size
slot_mapping_lst = random.sample(range(total_slots), num_tokens)
slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
k_pe = torch.randn(num_tokens, pe_dim, dtype=dtype, device=device)
entry_size = kv_lora_rank + pe_dim
scale = torch.tensor(0.1, dtype=torch.float32, device=device)
cache_elem_dtype = torch.uint8 if kv_cache_dtype == "fp8_e4m3" else dtype
kv_cache_cuda = torch.zeros(
num_blocks,
block_size,
entry_size,
dtype=cache_elem_dtype,
device=device,
)
kv_cache_triton = kv_cache_cuda.clone()
# Run CUDA kernel
ops.concat_and_cache_mla(
kv_c,
k_pe,
kv_cache_cuda,
slot_mapping,
kv_cache_dtype,
scale,
)
# Run Triton kernel
from vllm.v1.attention.ops.concat_and_cache_mla import (
concat_and_cache_mla as triton_concat_and_cache_mla,
)
triton_concat_and_cache_mla(
kv_c,
k_pe,
kv_cache_triton,
slot_mapping,
kv_cache_dtype,
scale,
)
# Bitwise comparison
cuda_bytes = kv_cache_cuda.view(torch.uint8)
triton_bytes = kv_cache_triton.view(torch.uint8)
torch.testing.assert_close(cuda_bytes, triton_bytes, atol=0, rtol=0)
@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
@pytest.mark.parametrize("pe_dim", PE_DIMS)
@pytest.mark.parametrize("block_size", MLA_BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", MLA_NUM_BLOCKS)
@pytest.mark.parametrize("seed", SEEDS)
@torch.inference_mode()
def test_concat_and_cache_mla_padding(
kv_lora_rank: int,
pe_dim: int,
block_size: int,
num_blocks: int,
seed: int,
):
"""Verify that padded tokens (slot_mapping == -1) leave the cache untouched."""
set_random_seed(seed)
device = "cuda"
num_tokens = 4
slot_mapping = torch.full((num_tokens,), -1, dtype=torch.long, device=device)
kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=torch.bfloat16, device=device)
k_pe = torch.randn(num_tokens, pe_dim, dtype=torch.bfloat16, device=device)
entry_size = kv_lora_rank + pe_dim
scale = torch.tensor(0.1, dtype=torch.float32, device=device)
kv_cache = torch.zeros(
num_blocks,
block_size,
entry_size,
dtype=torch.bfloat16,
device=device,
)
snapshot = kv_cache.clone()
from vllm.v1.attention.ops.concat_and_cache_mla import (
concat_and_cache_mla as triton_concat_and_cache_mla,
)
triton_concat_and_cache_mla(
kv_c,
k_pe,
kv_cache,
slot_mapping,
"auto",
scale,
)
torch.testing.assert_close(
kv_cache.view(torch.uint8),
snapshot.view(torch.uint8),
atol=0,
rtol=0,
)
# =====================================================================
# fused_norm_rope + indexer_k_quant_and_cache
# =====================================================================
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("seed", SEEDS)
@torch.inference_mode()
def test_fused_norm_rope_indexer_equivalence(
num_tokens: int,
seed: int,
):
"""Verify the fused LayerNorm+RoPE+FP8-quant path in fused_norm_rope
produces a bitwise-identical kv_cache to the unfused path
(layer_norm -> rope -> indexer_k_quant_and_cache)."""
from vllm.model_executor.models.deepseek_v3_2_monolithic.ops import (
fused_norm_rope,
)
from vllm.v1.attention.ops.indexer_k_quant_and_cache import (
indexer_k_quant_and_cache,
)
set_random_seed(seed)
device = "cuda"
dtype = torch.bfloat16
# Dimensions matching DeepSeek V3
q_dim = 1536
kv_dim = 512
kpe_dim = 64
index_k_dim = 128
rot_dim = 64 # qk_rope_head_dim for indexer
topk = 8
cache_block_size = 16
num_blocks = 32
cache_stride = index_k_dim + 4 # 128 fp8 bytes + 4 scale bytes
# Inputs
positions = torch.randint(0, 1024, (num_tokens,), device=device)
q_c = torch.randn(num_tokens, q_dim, dtype=dtype, device=device)
q_rms_w = torch.randn(q_dim, dtype=dtype, device=device)
kv_c = torch.randn(num_tokens, kv_dim, dtype=dtype, device=device)
kv_rms_w = torch.randn(kv_dim, dtype=dtype, device=device)
k_pe = torch.randn(num_tokens, kpe_dim, dtype=dtype, device=device)
kpe_cos_sin = torch.randn(4096, kpe_dim, dtype=dtype, device=device)
index_k = torch.randn(num_tokens, index_k_dim, dtype=dtype, device=device)
index_k_w = torch.randn(index_k_dim, dtype=dtype, device=device)
index_k_b = torch.randn(index_k_dim, dtype=dtype, device=device)
ik_cos_sin = torch.randn(4096, rot_dim, dtype=dtype, device=device)
topk_buf = torch.zeros(num_tokens, topk, dtype=torch.int32, device=device)
eps = 1e-6
total_slots = num_blocks * cache_block_size
slot_mapping_lst = random.sample(range(total_slots), num_tokens)
slot_mapping = torch.tensor(
slot_mapping_lst,
dtype=torch.long,
device=device,
)
kv_cache_fused = torch.zeros(
num_blocks,
cache_block_size,
cache_stride,
dtype=torch.float8_e4m3fn,
device=device,
)
kv_cache_ref = kv_cache_fused.clone()
# --- Reference path: standalone LayerNorm + RoPE + FP8 quant ---
from vllm.model_executor.models.deepseek_v3_2_monolithic.ops import (
layer_norm,
qk_rope,
)
index_k_ref = layer_norm(index_k, index_k_w, index_k_b, eps)
# RoPE in-place (1 head, non-interleaved, no start offset).
# qk_rope applies to both Q and K; pass a dummy for Q.
dummy_q = torch.empty_like(index_k_ref.unsqueeze(1))
qk_rope(
positions,
dummy_q,
index_k_ref.unsqueeze(1),
ik_cos_sin,
q_start_offset=0,
interleaved=False,
)
indexer_k_quant_and_cache(
index_k_ref,
kv_cache_ref,
slot_mapping,
128,
"ue8m0",
)
# --- Fused path: all in one kernel ---
index_k_fused = index_k.clone()
fused_norm_rope(
positions,
q_c.clone(),
q_rms_w,
eps,
kv_c.clone(),
kv_rms_w,
eps,
k_pe.clone(),
kpe_cos_sin,
index_k_fused,
index_k_w,
index_k_b,
eps,
ik_cos_sin,
topk_buf.clone(),
slot_mapping=slot_mapping,
indexer_k_cache=kv_cache_fused,
)
# The fused path keeps full fp32 precision through LayerNorm →
# RoPE → FP8 quant (no intermediate bf16 truncation), so the
# FP8 values may differ from the bf16 standalone path. Verify
# via dequantization that the results are close.
ref_bytes = kv_cache_ref.view(torch.uint8)
fused_bytes = kv_cache_fused.view(torch.uint8)
assert (ref_bytes.int() - fused_bytes.int()).abs().float().mean() < 0.5
vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py
+22 -7
View File
@@ -165,8 +165,22 @@ class MonolithicDecoderLayer(nn.Module):
index_k, _ = self.attn.indexer_wk(hidden_states)
index_weights, _ = self.attn.indexer_weights_proj(hidden_states)
# Step 2. Q RMS norm + KV RMS norm + KV RoPE + Index K layer norm + RoPE
# Step 2. Q RMS norm
# + KV RMS norm + KV RoPE
# + Index K layer norm + RoPE + FP8 quant + cache write
# + Init topk indices
#
# Fetch slot_mapping early so fused_norm_rope can write FP8 data
# directly into the indexer KV cache (saves a separate kernel).
from vllm.forward_context import get_forward_context
attn_metadata = get_forward_context().attn_metadata
if isinstance(attn_metadata, dict):
idx_meta = attn_metadata[self.attn.indexer_k_cache.prefix]
slot_mapping = idx_meta.slot_mapping
else:
slot_mapping = None
q_c, kv_c = fused_norm_rope(
positions,
# Q RMS norm
@@ -188,6 +202,11 @@ class MonolithicDecoderLayer(nn.Module):
self.attn.indexer_rope_emb.cos_sin_cache,
# Top k indices
self.attn.topk_indices_buffer,
# Fused FP8 quant + cache write
slot_mapping=slot_mapping,
indexer_k_cache=self.attn.indexer_k_cache.kv_cache
if slot_mapping is not None
else None,
)
# Step 3. q_c -> q
@@ -208,8 +227,6 @@ class MonolithicDecoderLayer(nn.Module):
# Index Q RoPE
index_q,
self.attn.indexer_rope_emb.cos_sin_cache,
# Index Q Quantize
1e-10, # quant_eps
# Index weights
index_weights,
self.attn.indexer_softmax_scale,
@@ -217,15 +234,13 @@ class MonolithicDecoderLayer(nn.Module):
)
# Step 5. Sparse indexer.
# The FP8 quant + cache write for index_k is already done in
# fused_norm_rope (step 2) when slot_mapping is available.
sparse_attn_indexer(
hidden_states,
self.attn.indexer_k_cache.prefix,
self.attn.indexer_k_cache.kv_cache,
index_q_fp8,
index_k,
index_weights,
self.attn.indexer_quant_block_size, # 128
"ue8m0", # scale_fmt
self.attn.topk_tokens,
self.attn.index_head_dim,
self.max_model_len,
vllm/model_executor/models/deepseek_v3_2_monolithic/ops.py
+168 -63
View File
@@ -391,6 +391,47 @@ def qk_rope(
)
@triton.jit
def _fp8_ue8m0_quantize(vals):
"""Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
"""
vals = vals.to(tl.float32)
amax = tl.max(tl.abs(vals))
scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv)
return fp8_vals, scale
@triton.jit
def _fp8_quant_and_cache_write(
vals,
mask,
slot_idx,
kv_cache_ptr,
kv_cache_scale_ptr,
cache_block_size,
cache_stride,
offsets,
HEAD_DIM: tl.constexpr,
):
k_fp8, scale = _fp8_ue8m0_quantize(vals)
block_idx = slot_idx // cache_block_size
block_offset = slot_idx % cache_block_size
block_start = block_idx * cache_block_size * cache_stride
tl.store(
kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
k_fp8,
mask=mask,
)
scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
@triton.jit
def _fused_norm_rope_kernel(
pos_ptr,
@@ -429,14 +470,43 @@ def _fused_norm_rope_kernel(
index_k_rope_cos_sin_cache_ptr,
index_k_rope_cos_sin_cache_stride,
INDEX_K_HALF_ROT_DIM: tl.constexpr,
# Index K fp32 scratch buffer for layernorm → RoPE handoff
index_k_normed_ptr,
# Index K FP8 quant + cache write
slot_mapping_ptr,
kv_cache_ptr,
kv_cache_scale_ptr,
cache_block_size,
cache_stride,
# Top k indices
topk_indices_ptr,
topk_indices_stride,
TOPK: tl.constexpr,
TOPK_BLOCK_SIZE: tl.constexpr,
):
pid = tl.program_id(0)
tok_idx = tl.program_id(1)
if tl.program_id(0) == 0:
if pid == 4:
# Fill top k indices buffer with -1
for i in range(0, TOPK, TOPK_BLOCK_SIZE):
offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
mask = offset < TOPK
tl.store(
topk_indices_ptr + tok_idx * topk_indices_stride + offset,
-1,
mask=mask,
)
return
if slot_mapping_ptr is None:
# Memory profiling run.
return
slot_idx = tl.load(slot_mapping_ptr + tok_idx)
if slot_idx < 0:
# Padding
return
if pid == 1:
# Q RMS norm
q_block = tl.arange(0, Q_BLOCK_SIZE)
q_mask = q_block < Q_DIM
@@ -444,16 +514,14 @@ def _fused_norm_rope_kernel(
q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
return
elif tl.program_id(0) == 1:
elif pid == 3:
# KV RMS Norm
kv_block = tl.arange(0, KV_DIM)
kv_c = tl.load(kv_ptr + tok_idx * kv_stride + kv_block)
kv_c_rms_w = tl.load(kv_rms_norm_w_ptr + kv_block)
kv_c = _rms_norm(kv_c, kv_c_rms_w, kv_rms_eps, KV_DIM)
tl.store(kv_c_out_ptr + tok_idx * kv_c_out_stride + kv_block, kv_c)
return
elif tl.program_id(0) == 2:
elif pid == 2:
# KV RoPE
pos = tl.load(pos_ptr + tok_idx)
cos, sin = _cos_sin_cache_kernel(
@@ -472,13 +540,13 @@ def _fused_norm_rope_kernel(
0,
True,
)
return
elif tl.program_id(0) == 3:
# Index K layer norm + RoPE
elif pid == 0:
# Fused: Index K LayerNorm + RoPE + FP8 quant + cache write.
# Eliminates the separate indexer_k_quant_and_cache kernel launch.
# 1. LayerNorm → fp32 temp buffer
index_k_block = tl.arange(0, INDEX_K_BLOCK_SIZE)
index_k_mask = index_k_block < INDEX_K_DIM
# Layer Norm
index_k = tl.load(
index_k_ptr + tok_idx * index_k_stride + index_k_block,
mask=index_k_mask,
@@ -488,7 +556,7 @@ def _fused_norm_rope_kernel(
index_k_b = tl.load(
index_k_layer_norm_bias_ptr + index_k_block, mask=index_k_mask
)
index_k = _layer_norm(
normed = _layer_norm(
index_k,
index_k_w,
index_k_b,
@@ -496,44 +564,59 @@ def _fused_norm_rope_kernel(
index_k_mask,
INDEX_K_DIM,
)
# Write to a fp32 scratch buffer so RoPE can read the two
# halves without Triton pointer-aliasing issues.
scratch = index_k_normed_ptr + tok_idx * INDEX_K_DIM
tl.store(scratch + index_k_block, normed, mask=index_k_mask)
# Save to the original buffer
tl.store(
index_k_ptr + tok_idx * index_k_stride + index_k_block,
index_k,
mask=index_k_mask,
)
# RoPE
# 2. RoPE (neox / non-interleaved) on the full vector.
pos = tl.load(pos_ptr + tok_idx)
cos, sin = _cos_sin_cache_kernel(
index_k_rope_cos_sin_cache_ptr,
index_k_rope_cos_sin_cache_stride,
pos,
INDEX_K_HALF_ROT_DIM,
cos_full = tl.load(
index_k_rope_cos_sin_cache_ptr
+ pos * index_k_rope_cos_sin_cache_stride
+ index_k_block % INDEX_K_HALF_ROT_DIM,
mask=index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
other=1.0,
).to(tl.float32)
sin_full = tl.load(
index_k_rope_cos_sin_cache_ptr
+ pos * index_k_rope_cos_sin_cache_stride
+ INDEX_K_HALF_ROT_DIM
+ index_k_block % INDEX_K_HALF_ROT_DIM,
mask=index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
other=0.0,
).to(tl.float32)
# XOR with HALF swaps the first/second half of the rotation
# region to get each element's partner.
partner_offs = tl.where(
index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
index_k_block ^ INDEX_K_HALF_ROT_DIM,
index_k_block,
)
_rope_kernel(
index_k_ptr + tok_idx * index_k_stride,
0,
cos,
sin,
1,
INDEX_K_HALF_ROT_DIM,
0,
False,
full = tl.load(scratch + index_k_block, mask=index_k_mask)
# Atomic read for the partner: tl.atomic_add(ptr, 0) returns the
# current value with guaranteed store visibility, avoiding the
# Triton compiler's aliasing issue with different offset expressions.
zeros = tl.zeros([INDEX_K_BLOCK_SIZE], dtype=tl.float32)
partner = tl.atomic_add(scratch + partner_offs, zeros, mask=index_k_mask)
sign = tl.where(index_k_block < INDEX_K_HALF_ROT_DIM, -1.0, 1.0)
roped = full * cos_full + sign * partner * sin_full
result = tl.where(index_k_block < 2 * INDEX_K_HALF_ROT_DIM, roped, full)
# 3. FP8 quantize + cache write from registers.
# No need to write back to index_k_ptr — the only consumer
# (sparse_attn_indexer) reads from the cache, not index_k.
_fp8_quant_and_cache_write(
result,
index_k_mask,
slot_idx,
kv_cache_ptr,
kv_cache_scale_ptr,
cache_block_size,
cache_stride,
index_k_block,
INDEX_K_DIM,
)
return
elif tl.program_id(0) == 4:
# Fill top k indices buffer with -1
for i in range(0, TOPK, TOPK_BLOCK_SIZE):
offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
mask = offset < TOPK
tl.store(
topk_indices_ptr + tok_idx * topk_indices_stride + offset,
-1,
mask=mask,
)
return
def fused_norm_rope(
@@ -552,6 +635,9 @@ def fused_norm_rope(
index_k_layer_norm_eps: float,
index_k_rope_cos_sin_cache: torch.Tensor,
topk_indices_buffer: torch.Tensor,
# Cache params for fused index-k FP8 quant + write
slot_mapping: torch.Tensor | None = None,
indexer_k_cache: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
assert positions.ndim == 1
assert q_c.ndim == 2
@@ -566,6 +652,33 @@ def fused_norm_rope(
index_k_dim = index_k.shape[-1]
topk = topk_indices_buffer.shape[-1]
# When indexer_k_cache is provided, program 0 writes FP8 data + scale
# directly into the cache, eliminating a separate
# indexer_k_quant_and_cache call.
if indexer_k_cache is not None:
assert slot_mapping is not None
cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
cache_block_size = indexer_k_cache.shape[1]
cache_stride = indexer_k_cache.shape[2]
# Ensure the pointer is fp8-typed so tl.store accepts fp8 values.
if indexer_k_cache.dtype == torch.uint8:
indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
else:
# Dummy values — program 0 will still do LayerNorm + RoPE but
# skip the FP8 cache write (slot_idx will be < 0 for all tokens).
cache_scale_view = torch.empty(0, dtype=torch.float32, device=positions.device)
indexer_k_cache = torch.empty(
0, dtype=torch.float8_e4m3fn, device=positions.device
)
slot_mapping = torch.full(
(num_tokens,), -1, dtype=torch.int64, device=positions.device
)
cache_block_size = 1
cache_stride = 1
# fp32 scratch buffer for layernorm output → RoPE handoff.
index_k_normed = torch.empty_like(index_k, dtype=torch.float32)
q_c_out = torch.empty_like(q_c)
kv_c_out = torch.empty_like(kv_c)
_fused_norm_rope_kernel[(5, num_tokens)](
@@ -593,7 +706,7 @@ def fused_norm_rope(
k_rope_cos_sin_cache,
k_rope_cos_sin_cache.stride(0),
k_rope_cos_sin_cache.shape[-1] // 2,
# Index K layer norm + RoPE
# Index K layer norm + RoPE + FP8 quant
index_k,
index_k.stride(0),
index_k_layer_norm_w,
@@ -604,6 +717,13 @@ def fused_norm_rope(
index_k_rope_cos_sin_cache,
index_k_rope_cos_sin_cache.stride(0),
index_k_rope_cos_sin_cache.shape[-1] // 2,
index_k_normed,
# FP8 cache write
slot_mapping,
indexer_k_cache,
cache_scale_view,
cache_block_size,
cache_stride,
# Top k indices buffer
topk_indices_buffer,
topk_indices_buffer.stride(0),
@@ -635,9 +755,6 @@ def _fused_q_kernel(
INDEX_Q_HALF_ROT_DIM: tl.constexpr,
# Index Q Quantize
index_q_fp8_ptr,
index_q_fp8_eps,
FP8_MIN: tl.constexpr,
FP8_MAX: tl.constexpr,
INDEX_Q_HEAD_DIM: tl.constexpr,
# Index weights
index_weights_ptr,
@@ -705,13 +822,8 @@ def _fused_q_kernel(
+ head_idx * index_q_stride1
+ index_q_block
)
index_q = index_q.to(tl.float32)
index_q_abs_max = tl.maximum(tl.max(tl.abs(index_q)), index_q_fp8_eps)
s = index_q_abs_max * (1.0 / FP8_MAX)
index_q_scale = tl.exp2(tl.ceil(tl.log2(s)))
index_q_fp8 = tl.clamp(index_q / index_q_scale, FP8_MIN, FP8_MAX)
index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q)
tl.store(
index_q_fp8_ptr
+ tok_idx * index_q_stride0
@@ -741,8 +853,6 @@ def fused_q(
q_start_offset: int,
index_q: torch.Tensor,
index_q_cos_sin_cache: torch.Tensor,
# Index Q Quantize
quant_eps: float,
# Index weights
index_weights: torch.Tensor,
index_weights_softmax_scale: float,
@@ -759,9 +869,7 @@ def fused_q(
num_index_q_heads = index_q.shape[1]
index_q_head_dim = index_q.shape[2]
FP8_DTYPE = torch.float8_e4m3fn
FP8_FINFO = torch.finfo(FP8_DTYPE)
index_q_fp8 = torch.empty_like(index_q, dtype=FP8_DTYPE)
index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
_fused_q_kernel[(2, num_tokens, num_index_q_heads)](
positions,
@@ -781,9 +889,6 @@ def fused_q(
index_q_cos_sin_cache.stride(0),
index_q_cos_sin_cache.shape[-1] // 2,
index_q_fp8,
quant_eps,
FP8_FINFO.min,
FP8_FINFO.max,
index_q_head_dim,
index_weights,
index_weights.stride(0),
vllm/model_executor/models/deepseek_v3_2_monolithic/sparse_indexer.py
-19
View File
@@ -14,14 +14,10 @@ logger = init_logger(__name__)
def sparse_attn_indexer(
hidden_states: torch.Tensor,
k_cache_prefix: str,
kv_cache: torch.Tensor,
q_fp8: torch.Tensor,
k: torch.Tensor,
weights: torch.Tensor,
quant_block_size: int,
scale_fmt: str | None,
topk_tokens: int,
head_dim: int,
max_model_len: int,
@@ -43,25 +39,10 @@ def sparse_attn_indexer(
attn_metadata = attn_metadata[k_cache_prefix]
assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
slot_mapping = attn_metadata.slot_mapping
has_decode = attn_metadata.num_decodes > 0
has_prefill = attn_metadata.num_prefills > 0
num_decode_tokens = attn_metadata.num_decode_tokens
# During speculative decoding, k may be padded to the CUDA graph batch
# size while slot_mapping only covers actual tokens. Truncate k to avoid
# out-of-bounds reads in the kernel.
num_tokens = slot_mapping.shape[0]
k = k[:num_tokens]
ops.indexer_k_quant_and_cache(
k,
kv_cache,
slot_mapping,
quant_block_size,
scale_fmt,
)
if has_prefill:
prefill_metadata = attn_metadata.prefill
assert prefill_metadata is not None