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

[mikekg]

Purpose

_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.

Test Plan

Run the existing platform-dispatch test and the new direct Triton test:

python3 -m pytest -v \
  tests/kernels/test_compressor_kv_cache.py::test_dequantize_and_gather_k_cache \
  tests/kernels/test_compressor_kv_cache.py::test_dequantize_and_gather_k_cache_triton

Run the direct Triton test under Compute Sanitizer:

compute-sanitizer --tool memcheck \
  python3 -m pytest -v \
  tests/kernels/test_compressor_kv_cache.py::test_dequantize_and_gather_k_cache_triton

Reproduce the production configuration on four 8xA100 Slurm nodes with TP16, PP2, EP16, block size 256, and compiled CUDA graphs:

#!/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 existing test_dequantize_and_gather_k_cache continues to test the platform dispatcher. The new test_dequantize_and_gather_k_cache_triton mirrors its 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. Both Python functions retain the same caller interface; the added stride is private to the Triton wrapper-to-kernel launch.

Test Result

The production run reached the first 64 concurrent requests, compiled _dequantize_and_gather_k_kernel, and then failed with:

Triton kernel JIT compilation during inference: _dequantize_and_gather_k_kernel
CUDA error: an illegal memory access was encountered

Before the fix, the direct test passed at request count 1 and failed at every request count from 2 through 256 on sm80, sm86, and sm89. Compute Sanitizer on an A100 reported 27,983 invalid memory accesses.

After the fix, the direct test passed at request counts 1, 2, 4, 8, 16, 32, 64, 128, and 256 on sm80, sm86, and sm89. Compute Sanitizer on an A100 reported:

9 passed
ERROR SUMMARY: 0 errors

---

Essential Elements of an Effective PR Description Checklist

• The purpose and root cause are described.
• The unit, sanitizer, and end-to-end test commands are provided.
• Before-and-after results are provided.
• Documentation impact was considered; this change does not alter a user-facing interface.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: d196b337-1829-401b-bff8-61384cc9b526

📥 Commits

Reviewing files that changed from the base of the PR and between 6cbb3c1 and 6d0e0616838f992ed1bd2ec8f5500cb443040ec3.

📒 Files selected for processing (2)

• tests/kernels/test_compressor_kv_cache.py
• vllm/models/deepseek_v4/common/ops/cache_utils.py

Included review availability: Your plan provides up to 10 included reviews per hour; 6 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Bug Fixes

  • Improved cache dequantization for block tables with non-contiguous memory layouts.
  • Ensured correct block-table addressing across varying request counts and physical page mappings.

• Tests

  • Added coverage validating Triton cache dequantization against the PyTorch reference implementation across multiple request sizes.

Walkthrough

The Triton dequantization kernel now uses the block table’s runtime row stride. A parameterized test validates bit-exact output across request counts and non-contiguous block-table layouts.

Changes

K Cache Dequantization

Layer / File(s)	Summary
Runtime block-table stride vllm/models/deepseek_v4/common/ops/cache_utils.py	The kernel receives block_table.stride(0) at runtime and uses it for block-table row addressing.
Non-contiguous block-table validation tests/kernels/test_compressor_kv_cache.py	A parameterized test compares Triton output with the PyTorch reference for request counts from 1 to 256, including a non-contiguous block-table slice.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 6d0e0

K-cache gathering now correctly addresses strided block-table views while preserving contiguous-table behavior. Regression coverage validates the corrected path, with no remaining merge-blocking risk.

Suggested reviewers: chaunceyjiang

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the DeepSeek V4 K-cache gather bug and the fix for strided block tables.
Description check	✅ Passed	The description directly explains the root cause, chosen fix, test coverage, reproduction steps, and before-and-after results for the changeset.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mikekg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork