From ad21c0000535afe1e6b1f61af92772900fa4eed4 Mon Sep 17 00:00:00 2001 From: Michael Gschwind Date: Thu, 3 Sep 2026 01:08:56 -0700 Subject: [PATCH 1/3] [Bugfix] Fix DeepSeek V4 K-cache gather for strided block tables `_dequantize_and_gather_k_kernel` accepts `block_table` tensor views but advances between rows using `block_table.shape[-1]`. A narrowed view of a fixed-width backing table retains the backing table's larger `stride(0)`, so requests after the first read incorrect physical block IDs and can access memory outside the K-cache allocation. Two fixes were considered: 1. Materialize the view with `block_table.contiguous()` before launching the kernel. 2. Pass `block_table.stride(0)` to the kernel and use the physical row stride for address calculation. We chose the stride argument because the kernel already accepts tensor views and already uses physical strides for `out` and `k_cache`. Calling `contiguous()` would repair the addressing indirectly by allocating temporary storage and copying the active block table on every invocation of this inference-path operation. Passing one scalar stride fixes the incorrect assumption at its source without adding GPU work, memory traffic, or allocator pressure. For contiguous block tables, `stride(0) == shape[-1]`, so the generated addresses remain unchanged. For narrowed views, the kernel now advances by the backing table's actual row layout. The production failure occurred on four 8xA100 Slurm nodes with TP16, PP2, EP16, block size 256, and compiled CUDA graphs: ```bash #!/bin/bash #SBATCH --nodes=4 #SBATCH --ntasks=4 #SBATCH --ntasks-per-node=1 #SBATCH --gpus-per-node=8 #SBATCH --time=02:00:00 set -euo pipefail export MODEL=nvidia/DeepSeek-V4-Pro-0813-NVFP4 export MASTER_ADDR MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1) export MASTER_PORT=29501 export COMPILATION_CONFIG='{"cudagraph_capture_sizes":[1,2,4,8,16,32,64],"inductor_compile_config":{"enable_auto_functionalized_v2":true}}' srun --nodes=4 --ntasks=4 --ntasks-per-node=1 bash -lc ' headless=() if (( SLURM_NODEID > 0 )); then headless=(--headless) fi exec vllm serve "$MODEL" \ --revision 2ff4ec53ee664e54571f670276d2c89d0fcc7b82 \ --served-model-name deepseek-v4-pro \ --host 0.0.0.0 \ --port 8000 \ --trust-remote-code \ --tokenizer-mode deepseek_v4 \ --load-format fastsafetensors \ --safetensors-load-strategy lazy \ --distributed-executor-backend mp \ --master-addr "$MASTER_ADDR" \ --master-port "$MASTER_PORT" \ --nnodes "$SLURM_NNODES" \ --node-rank "$SLURM_NODEID" \ --tensor-parallel-size 16 \ --pipeline-parallel-size 2 \ --enable-expert-parallel \ --kv-cache-dtype fp8 \ --block-size 256 \ --max-model-len 10240 \ --max-num-seqs 64 \ --max-num-batched-tokens 16384 \ --gpu-memory-utilization 0.9 \ --linear-backend auto \ --moe-backend auto \ --no-enable-prefix-caching \ --distributed-timeout-seconds 7200 \ --cpu-distributed-timeout-seconds 7200 \ --compilation-config "$COMPILATION_CONFIG" \ "${headless[@]}" ' ``` The first 64 concurrent requests compiled `_dequantize_and_gather_k_kernel` and then failed with: ```text Triton kernel JIT compilation during inference: _dequantize_and_gather_k_kernel CUDA error: an illegal memory access was encountered ``` Add `test_dequantize_and_gather_k_cache_triton` alongside the existing `test_dequantize_and_gather_k_cache`. The new test mirrors the existing test's cache layout and PyTorch-reference comparison but calls the Triton implementation directly with a narrowed block-table view at request counts 1, 2, 4, 8, 16, 32, 64, 128, and 256. Direct coverage is necessary because the platform dispatcher can select another implementation and allow the dispatcher test to pass without executing the faulty Triton kernel. This is not testing a different caller interface: `dequantize_and_gather_k_cache` and `dequantize_and_gather_k_cache_triton` remain plug-compatible. The added stride is private to the Triton wrapper-to-kernel launch. Before the fix: ```text sm80: c1 passed; c2-c256 failed sm86: c1 passed; c2-c256 failed sm89: c1 passed; c2-c256 failed A100 Compute Sanitizer: 27,983 errors ``` After the fix: ```text sm80: c1-c256 passed sm86: c1-c256 passed sm89: c1-c256 passed A100 Compute Sanitizer: 0 errors ``` Signed-off-by: Michael Gschwind --- tests/kernels/test_compressor_kv_cache.py | 42 +++++++++++++++++++ .../deepseek_v4/common/ops/cache_utils.py | 16 +++---- 2 files changed, 51 insertions(+), 7 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index cc4f4c72a931..89a3cb815d14 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -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, @@ -1417,6 +1420,45 @@ 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): + 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 ──────────────────────────────────────────────────── diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index cea755bd1a90..916d10729b2c 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -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 @@ -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 @@ -289,7 +289,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 @@ -378,8 +378,10 @@ def dispatch( # type: ignore[override] 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), @@ -455,7 +457,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, @@ -479,7 +481,7 @@ def __call__( 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, From 68229d591e1be5df34241475db594d37e771ce25 Mon Sep 17 00:00:00 2001 From: Michael Gschwind Date: Thu, 8 Oct 2026 16:35:05 -0700 Subject: [PATCH 2/3] Fix repository pre-commit checks after rebase Co-authored-by: Codex Signed-off-by: Michael Gschwind --- vllm/models/deepseek_v4/common/ops/cache_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 916d10729b2c..594034b53f37 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -379,8 +379,7 @@ def dispatch( # type: ignore[override] return self.CompileKey( **compile_key_fields, block_table_stride=triton_scalar_specialization_rep( - (max_model_len + block_table_block_size - 1) - // block_table_block_size + (max_model_len + block_table_block_size - 1) // block_table_block_size ), cache_block_size=cache_block_size, block_stride=block_stride, From f4a087a1d9827279a8877b3afffc8698ddc2401c Mon Sep 17 00:00:00 2001 From: Michael Gschwind Date: Thu, 8 Oct 2026 17:38:13 -0700 Subject: [PATCH 3/3] Document strided DeepSeek cache gathering Co-authored-by: Codex Signed-off-by: Michael Gschwind --- tests/kernels/test_compressor_kv_cache.py | 1 + vllm/models/deepseek_v4/common/ops/cache_utils.py | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 89a3cb815d14..e63d4577bc2a 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -1422,6 +1422,7 @@ def test_dequantize_and_gather_k_cache( @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 diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 594034b53f37..b0f0823317b6 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -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) @@ -372,6 +373,7 @@ 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) @@ -440,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( @@ -476,6 +479,7 @@ 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),