From 34d73a3375bd1f7510db2d60a31a4bbbcb59af48 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 30 Mar 2026 22:28:33 +0000 Subject: [PATCH] Fuse indexer k cache update Signed-off-by: Woosuk Kwon --- .../attention/test_triton_cache_ops.py | 412 ++++++++++++++++++ .../deepseek_v3_2_monolithic/decoder_layer.py | 29 +- .../models/deepseek_v3_2_monolithic/ops.py | 231 +++++++--- .../sparse_indexer.py | 19 - 4 files changed, 602 insertions(+), 89 deletions(-) create mode 100644 tests/kernels/attention/test_triton_cache_ops.py diff --git a/tests/kernels/attention/test_triton_cache_ops.py b/tests/kernels/attention/test_triton_cache_ops.py new file mode 100644 index 00000000000..b763be60b03 --- /dev/null +++ b/tests/kernels/attention/test_triton_cache_ops.py @@ -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 diff --git a/vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py b/vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py index 19974faf6a5..a2176c2ff93 100644 --- a/vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py +++ b/vllm/model_executor/models/deepseek_v3_2_monolithic/decoder_layer.py @@ -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, diff --git a/vllm/model_executor/models/deepseek_v3_2_monolithic/ops.py b/vllm/model_executor/models/deepseek_v3_2_monolithic/ops.py index 7cd50e4e8a9..ef01e75c1c7 100644 --- a/vllm/model_executor/models/deepseek_v3_2_monolithic/ops.py +++ b/vllm/model_executor/models/deepseek_v3_2_monolithic/ops.py @@ -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), diff --git a/vllm/model_executor/models/deepseek_v3_2_monolithic/sparse_indexer.py b/vllm/model_executor/models/deepseek_v3_2_monolithic/sparse_indexer.py index d4d33600c9b..490074ef120 100644 --- a/vllm/model_executor/models/deepseek_v3_2_monolithic/sparse_indexer.py +++ b/vllm/model_executor/models/deepseek_v3_2_monolithic/sparse_indexer.py @@ -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