[Bugfix] Disable cross-layer KV cache for MLA attention backends

[haosdent]

Purpose

Fixes #37032

MLA models (e.g., GLM-4.7-Flash) produce garbage output with --kv-offloading-size because cross-layer KV cache allocation creates non-contiguous per-layer views. MLA decode kernels assume contiguous block layout, so they read wrong memory for block_id > 0.

Fix (3 parts)

1. mla_attention.py, indexer.py: MLA backends return identity permutation (0, 1, 2, 3) from get_kv_cache_stride_order(include_num_layers_dimension=True), keeping num_layers first in physical layout to signal cross-layer is unsupported.

2. kv_connector_model_runner_mixin.py: use_uniform_kv_cache() checks stride_order[0] == 0 (layers dim first) and skips cross-layer allocation, falling back to per-layer contiguous KV caches.

3. offloading_connector.py: register_kv_caches() uses AttentionLayerBase instead of Attention to match both Attention and MLAAttention layers.

Test Plan

pytest tests/v1/kv_offload/test_cpu_gpu.py::test_mla_backend_rejects_cross_layer_kv_cache tests/v1/kv_offload/test_cpu_gpu.py::test_deepseek_v32_indexer_rejects_cross_layer_kv_cache -v

End-to-end with deepseek-ai/DeepSeek-V2-Lite-Chat + KV offloading on NVIDIA GB10.

Test Result

Unit tests: 2/2 passed. End-to-end:

• Without fix: ' 夹缝缝缝缝缝缝缝缝缝缝缝缝缝缝缝缝缝缝' (garbage)
• With fix: ' 5, 6, 7, 8, 9, 10.\n' (correct)

[gemini-code-assist]

Code Review

This pull request addresses a bug causing garbage output with MLA models when KV offloading is enabled. The root cause is the incorrect handling of non-contiguous memory layouts by MLA attention kernels when cross-layer KV cache is used. The fix correctly disables this feature for MLA backends by raising a NotImplementedError, which leverages an existing fallback mechanism to use per-layer contiguous KV caches. The changes are well-targeted, clearly explained, and are accompanied by new tests that verify the corrected behavior. The solution appears robust and effectively resolves the issue.

[mergify]

Hi @haosdent, 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

[orozery]

@haosdent thanks for this fix!
I was able to reproduce the issue.
I was not aware that MLA kernels did not support cross layers.

However, instead of returning NotImplementedError, I think we should return the identity permutation (e.g. (0, 1, 2, ...)).
This will yield contiguous per-layer tensors.
For DeepseekV32IndexerBackend this is already the case, so no change is needed there.

[haosdent]

However, instead of returning NotImplementedError, I think we should return the identity permutation (e.g. (0, 1, 2, ...)).
This will yield contiguous per-layer tensors.

Thanks @orozery ! Let me update the pull request

[haosdent]

However, instead of returning NotImplementedError, I think we should return the identity permutation (e.g. (0, 1, 2, ...)).
This will yield contiguous per-layer tensors.

Hi, @orozery, I checked again, and this approach does not seem to work. The identity permutation still allocates a single cross-layer tensor. Slicing tensor[i] gives a non-contiguous view where stride(0) = num_layers * block_size * head_size. MLA kernels assume stride(0) = block_size * head_size which no permutation fixes this. We need to opt out of cross-layer allocation entirely, which is what NotImplementedError does.

[orozery]

Hi, @orozery, I checked again, and this approach does not seem to work. The identity permutation still allocates a single cross-layer tensor. Slicing tensor[i] gives a non-contiguous view where stride(0) = num_layers * block_size * head_size. MLA kernels assume stride(0) = block_size * head_size which no permutation fixes this. We need to opt out of cross-layer allocation entirely, which is what NotImplementedError does.

It will allocate a single cross-layer tensor, but will create per-layer views kv_caches: dict[str, torch.tensor] which are contiguous.
I tested it now but it actually fails since the offloading connector currently assumes that when using cross layers the first dimension is num_blocks.

I think that if get_kv_cache_stride_order starts with (0, ...), e.g. the layers dimension is first we can abort creating a cross-layers tensor, here:

vllm/vllm/v1/worker/kv_connector_model_runner_mixin.py

Line 195 in 2754231

return len(kv_cache_stride_order) == len(kv_cache_shape) + 1

This also aligns with #34742, aiming to remove all exception throwing by get_kv_cache_stride_order.

[haosdent]

Thanks @orozery , I referred to #34742 and updated the PR, can you help review again? Many thanks!

[orozery]

@haosdent LGTM!
Can we just move the new unit tests to a new e.g. tests/v1/kv_connector/unit/test_kv_cache_layout.py ?

[haosdent]

Can we just move the new unit tests to a new e.g. tests/v1/kv_connector/unit/test_kv_cache_layout.py ?

Done, thanks @orozery

[orozery]

Thanks again @haosdent !