[Bugfix] Relax TRTLLM KV cache contiguity assertion for cross-layer layout

[Etelis]

Fixes #33572

The is_strictly_contiguous assertion on kv_cache_permute rejects the non-canonical strides views, and #33192 worked around it by just disabling TRTLLM entirely when KV transfer is enabled.

I dug into the FlashInfer kernel source to check if full-tensor contiguity is actually needed. Turns out the TRTLLM kernels in csrc/trtllm_fmha_kernel_launcher.cu read actual tensor strides:

int kv_stride_keys_values = key_cache.stride(-2);
int kv_stride_heads       = key_cache.stride(-3);
int kv_stride_batch       = key_cache.stride(0);

So the kernel handles non-contiguous layouts fine — the assertion was a vLLM-side guard,. The one place that needs contiguity is the Triton FP8 dequant kernel (_trtllm_prefill_attn_kvfp8_dequant), which computes strides from shape. Moved the assertion there.

Testing

Ran on B200 (Blackwell, SM100) with vLLM 0.15.1, FlashInfer 0.6.1, Llama-3.2-1B-Instruct.

test_cpu_offloading[FLASHINFER-48]: PASSED
test_flashinfer_trtllm_attention: 224 passed, 16 skipped

Using TRTLLM prefill attention (auto-detected).
Using HND KV cache layout for FLASHINFER backend.
Allocating a cross layer KV cache of shape (167623, 2, 8, 16, 16, 64)

Running tests: 100% 10/10 [00:02, 3.88it/s]
Average times:
    Cold: 69.95ms
    GPU hit: 19.26ms
    CPU hit: 40.50ms
PASSED (83.24s)

[gemini-code-assist]

Code Review

This pull request relaxes the contiguity assertion for the TRTLLM KV cache to allow its use with cross-layer layouts, which can result from KV cache transfer. The change correctly removes the assertion from the general prefill and decode paths and moves it to the specific FP8 dequantization kernel path where contiguity is strictly required.

My review identifies one potential issue: while the logic to disable TRTLLM for decode with KV transfer is removed, a similar logic for prefill seems to have been missed, which would prevent the fix from being fully effective. I've added a comment with a suggestion to address this.

[mgoin]

This seems reasonable to me given your investigation cc @vadiklyutiy @pavanimajety to confirm

[mergify]

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

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

[vadiklyutiy]

I introduced is_strictly_contiguous check in #32008 and #32417.
There was a problem in degenerate strides

Shape: torch.Size([1, 32, 128])
Stride: (4608, 128, 1)

4608 != 32*128

So, the quote from PR description

I dug into the FlashInfer kernel source to check if full-tensor contiguity is actually needed. Turns out the TRTLLM kernels in csrc/trtllm_fmha_kernel_launcher.cu read actual tensor strides:

int kv_stride_keys_values = key_cache.stride(-2);
int kv_stride_heads       = key_cache.stride(-3);
int kv_stride_batch       = key_cache.stride(0);

cause a fail with this degenerate strides - kv_stride_batch will be incorrect.

Want to note that is_strictly_contiguous actually check 2 things:

• is_contiguous
• strides are canonical

If we are talking only about remove is_contiguous, I'd propose to break is_strictly_contiguous by 2 part and left "strides are canonical" check in the same place.

[vadiklyutiy]

I did dig/think a bit more about the topic.

Fully disabling TRTLLM attn when using kv-cache offloading (implemented in #33192) is really expensive workaround - TRTLLM attn is fastest.

This PR is right direction but let fully solve the problem - modify trtllm_prefill_attn_kvfp8_dequant to work with non-contiguous kv-cache.

[vadiklyutiy]

I did dig/think a bit more about the topic.

Fully disabling TRTLLM attn when using kv-cache offloading (implemented in #33192) is really expensive workaround - TRTLLM attn is fastest.

This PR is right direction but let fully solve the problem - modify trtllm_prefill_attn_kvfp8_dequant to work with non-contiguous kv-cache.

I created #36867 that add support of non-contiguous kv-cache in trtllm_prefill_attn_kvfp8_dequant. If it ok, we can fully remove is_strictly_contiguous assert for kv-cache

[pavanimajety]

@Etelis Just FYI on what happens if the strides are not canonical flashinfer-ai/flashinfer#2232 (comment)

I vote to keep the stride check like Vadim has in #36867

[Etelis]

Thanks @vadiklyutiy — your analysis of the degenerate stride issue and the kernel fix in #36867 are great. I've updated this PR to use a targeted canonical stride check (matching your approach) instead of the blanket contiguity assertion.

[pavanimajety]

Thanks @Etelis!

[vadiklyutiy]

@Etelis Could you clarify why did you keep checks for -1 and -2 dims? Shouldn't we check -3 as well?
This flashinfer-ai/flashinfer#2232 (comment) was about incorrectness of the stride[-3].