Skip to content

[Bugfix] Fix DeepSeek V4 K-cache gather for strided block tables #55109

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mikekg wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[Bugfix] Fix DeepSeek V4 K-cache gather for strided block tables #55109

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ad21c00
[Bugfix] Fix DeepSeek V4 K-cache gather for strided block tables
mikekg Sep 3, 2026
68229d5
Fix repository pre-commit checks after rebase
mikekg Oct 8, 2026
f4a087a
Document strided DeepSeek cache gathering
mikekg Oct 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions tests/kernels/test_compressor_kv_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,9 @@
dequantize_and_gather_k_cache,
quantize_and_insert_k_cache,
)
from vllm.models.deepseek_v4.common.ops.cache_utils import (
_DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL,
)
from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
_fused_kv_compress_norm_rope_insert_indexer_attn,
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
Expand Down Expand Up @@ -1417,6 +1420,46 @@ def test_dequantize_and_gather_k_cache(
torch.testing.assert_close(actual, expected, rtol=0, atol=0)


@pytest.mark.parametrize("num_reqs", [1, 2, 4, 8, 16, 32, 64, 128, 256])
def test_dequantize_and_gather_k_cache_triton(num_reqs: int):
"""Verify cache gathering respects a noncontiguous block-table view."""
block_size = 256
head_dim = 512
nope_dim = 448
scale_dim = 8
head_bytes = nope_dim + (head_dim - nope_dim) * 2 + scale_dim
device = "cuda"

compressed_kv = torch.randn(num_reqs, head_dim, dtype=torch.bfloat16, device=device)
physical_blocks = torch.randperm(num_reqs, device=device)
block_table_storage = torch.full(
(num_reqs, 4), int(-1e6), dtype=torch.int32, device=device
)
block_table = block_table_storage[:, :1]
block_table[:, 0] = physical_blocks
assert block_table.stride(0) != block_table.shape[-1]

slot_mapping = physical_blocks.to(torch.int64) * block_size
k_cache = torch.empty(
num_reqs, block_size, head_bytes, dtype=torch.uint8, device=device
)
quantize_and_insert_k_cache(
compressed_kv, k_cache.view(num_reqs, -1), slot_mapping, block_size
)

seq_lens = torch.ones(num_reqs, dtype=torch.int32, device=device)
expected = torch.empty(num_reqs, 1, head_dim, dtype=torch.bfloat16, device=device)
actual = torch.empty_like(expected)
_dequantize_and_gather_k_cache_reference(
expected, k_cache, seq_lens, None, block_table, block_size, 0
)
_DEQUANTIZE_AND_GATHER_K_CACHE_KERNEL(
actual, k_cache, seq_lens, None, block_table, block_size, 0
)
torch.accelerator.synchronize()
torch.testing.assert_close(actual, expected, rtol=0, atol=0)


# ── Test C: Indexer path ────────────────────────────────────────────────────


Expand Down
19 changes: 12 additions & 7 deletions vllm/models/deepseek_v4/common/ops/cache_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ class DequantizeAndGatherKCacheKernel(

@dataclass(frozen=True)
class CompileKey:
max_blocks_per_seq: int
block_table_stride: int
cache_block_size: int
block_stride: int
use_fnuz: bool
Expand All @@ -252,10 +252,10 @@ def kernel(
k_cache_ptr,
seq_lens_ptr,
block_table_ptr,
block_table_stride,
offset,
gather_lens_ptr,
# Constants
max_blocks_per_seq: tl.constexpr,
fp8_dim: tl.constexpr, # 448
bf16_dim: tl.constexpr, # 64
scale_dim: tl.constexpr, # 8
Expand All @@ -268,6 +268,7 @@ def kernel(
n_quant_blocks: tl.constexpr, # 7 real blocks
use_fnuz: tl.constexpr = False,
):
"""Gather paged cache rows and dequantize their FP8 portions into BF16."""
batch_idx = tl.program_id(0)
worker_id = tl.program_id(1)
num_workers = tl.num_programs(1)
Expand All @@ -289,7 +290,7 @@ def kernel(
pos_in_block = pos % cache_block_size

# Get physical block index from block table
block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq
block_table_row_ptr = block_table_ptr + batch_idx * block_table_stride
physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32

# int64: physical_block_idx * block_stride can exceed 2^31 with many
Expand Down Expand Up @@ -372,14 +373,16 @@ def dispatch( # type: ignore[override]
offset: int,
**compile_key_fields: bool,
) -> CompileKey:
"""Specialize the gather kernel for the padded cache-block layout."""
token_stride = 576
scale_stride = 8
unpadded = cache_block_size * (token_stride + scale_stride)
block_stride = ((unpadded + token_stride - 1) // token_stride) * token_stride
return self.CompileKey(
**compile_key_fields,
max_blocks_per_seq=(max_model_len + block_table_block_size - 1)
// block_table_block_size,
block_table_stride=triton_scalar_specialization_rep(
(max_model_len + block_table_block_size - 1) // block_table_block_size
),
cache_block_size=cache_block_size,
block_stride=block_stride,
offset=triton_scalar_specialization_rep(offset),
Expand Down Expand Up @@ -439,6 +442,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
)

def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
"""Construct synthetic cache and block-table inputs for kernel warmup."""
int32_ptr = TritonWarmupTensor(torch.int32)
return dict(
out=TritonWarmupTensor(
Expand All @@ -455,7 +459,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
gather_lens=(int32_ptr if compile_key.has_gather_lens else None),
block_table=TritonWarmupTensor(
torch.int32,
shape=(1, compile_key.max_blocks_per_seq),
shape=(1, compile_key.block_table_stride),
),
block_size=compile_key.cache_block_size,
offset=compile_key.offset,
Expand All @@ -475,11 +479,12 @@ def __call__(
*,
use_fnuz: bool = False,
) -> LaunchSpec:
"""Build the gather launch using the actual cache and block-table strides."""
num_reqs = seq_lens.shape[0]
return (num_reqs, self.NUM_WORKERS), dict(
out_stride0=out.stride(0),
out_stride1=out.stride(1),
max_blocks_per_seq=block_table.shape[-1],
block_table_stride=block_table.stride(0),
fp8_dim=448,
bf16_dim=64,
scale_dim=8,
Expand Down
Loading