Support non-contiguous KV cache in TRTLLM fp8 dequant kernel

[vadiklyutiy]

Summary

Fix the trtllm_prefill_attn_kvfp8_dequant Triton kernel to support non-contiguous KV cache tensors (e.g. cross-layer unified allocation used by KV offloading).

The kernel previously computed flat pointer offsets from tensor shape, assuming contiguity. With cross-layer KV caches the per-layer view is non-contiguous (strides skip over other layers), causing incorrect memory reads. Now the kernel uses actual tensor strides for the page, K/V, and head dimensions.

Related: vllm-project/vllm#34158 — relaxes the same assertion and moves the contiguity check closer to the dequant kernel. Our change goes further by fixing the dequant kernel itself to work with non-contiguous inputs.

Test

New standalone test: tests/kernels/attention/test_trtllm_kvfp8_dequant.py

python -m pytest tests/kernels/attention/test_trtllm_kvfp8_dequant.py -v

71 tests total (64 parametrized + 7 corner cases), all passing:

• Parametrized matrix: num_kv_heads={1,8}, head_size={64,128}, block_size={16,32}, batch_size={1,4}, num_pages_per_seq={3,8}, contiguous={True,False}
• Corner cases: zero/padding pages in block tables, all-zero block tables, different K/V scales, single page per sequence, large page indices (32K), large block_size (64), cross-layer with 36 layers (real gpt-oss-120b pattern)

[gemini-code-assist]

Code Review

This pull request introduces support for non-contiguous KV cache tensors in the trtllm_prefill_attn_kvfp8_dequant Triton kernel. The kernel is updated to use tensor strides for memory access, correctly handling layouts like those from cross-layer unified allocation. A comprehensive new test suite is also added, which thoroughly validates the changes against a reference implementation for both contiguous and non-contiguous cases, including numerous corner cases. The implementation appears correct and robust, and the tests provide strong confidence in the fix.

[pavanimajety]

@Etelis What do you think of this fix? Seems reasonable to me to remove the stride check for kv_cache throughout

[Etelis]

Looks good to me — using actual strides in the dequant kernel is the right fix. With this merged I can drop the is_strictly_contiguous assertion I kept for the FP8 path in #34158 and replace it with the same minimal stride check you have here.

[pavanimajety]

LGTM, thanks!

[mergify]

Hi @vadiklyutiy, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[vadiklyutiy]

if there are no further comments could we merge this PR?