[tuukkjs]

Hybrid attention+mamba models can expose non-contiguous KV cache views via as_strided_. The shuffled cache writer assumed a contiguous layout and could corrupt shuffle=1 decode inputs, so compute shuffled write offsets from the actual cache view strides.

Follow-up to #32914. After the fix, we get correct results for Qwen 3.5 for the paged_attention_rocm path with shuffle=1, i.e., for low batch sizes. At batch_size=512 the issue persists because paged_attention_common routes to the ASM paged-attention kernel (pa_fwd_asm), which only supports head_dim == 128 . Qwen 3.5 uses head_dim=256. A fix in aiter (ROCm/aiter@b4885d2) restricts the ASM kernel to head_dim == 128 and falls back to paged_attention_rocm for other head sizes. Once vLLM upgrades aiter to v0.1.12 or later, shuffle=1 should work correctly for Qwen 3.5 at all batch sizes.

Adding head_dim != 128 support to the ASM kernel is out of scope of this PR.

AI assistance was used for this PR.

Test Plan

Test with Qwen 3.5 (hybrid model) and Qwen 3 (attention-only model).

Test Result

lm_eval (gsm8k, 5-shot) — Qwen3.5-397B (hybrid) MI355X

Variant	batch_size	flexible-extract	strict-match
Before fix (shuffle=1)	4	0.0038 (±0.0017)	0.0000 (±0.0000)
Before fix (shuffle=1)	512	0.0030 (±0.0015)	0.0000 (±0.0000)
After fix (shuffle=1)	4	0.8613 (±0.0095)	0.8431 (±0.0100)
After fix (shuffle=1)	512	0.0068 (±0.0023)	0.0000 (±0.0000)
pa_v1 (shuffle=0)	4	0.8590 (±0.0096)	0.8370 (±0.0102)
pa_v1 (shuffle=0)	512	0.8582 (±0.0096)	0.8332 (±0.0103)

The fix results in correct results for paged_attention_rocm path of shuffle=1, i.e., for low batch sizes. At batch_size=512 the issue persists because asm-kernel only supports head_dim == 128.

lm_eval (gsm8k, 5-shot) — Qwen3-235B (attention-only) MI355X

Variant	batch_size	flexible-extract	strict-match
After fix (shuffle=1)	4	0.9113 (±0.0078)	0.9014 (±0.0082)
After fix (shuffle=1)	512	0.9060 (±0.0080)	0.8999 (±0.0083)

The standard attention-only model works correctly at both batch sizes with the fix applied, confirming no regression for standard attention-only models.

Summary

Model	Type	batch_size=4	batch_size=512
Qwen3.5-397B	Hybrid	Fixed	Fix out of scope for this PR
Qwen3-235B	Attention-only	Pass	Pass

[gemini-code-assist]

Code Review

This pull request refactors the reshape_and_cache_shuffle_kernel in the ROCm Aiter Flash Attention backend to use explicit tensor strides for offset calculations. By passing kc_stride and vc_stride parameters directly to the kernel, the implementation replaces manual offset arithmetic, improving the robustness of the cache indexing logic. I have no feedback to provide as there were no review comments.

[tjtanaa]

@tuukkjs From your description, this PR requires a new aiter version v0.1.12, am I right?

[tuukkjs]

@tuukkjs From your description, this PR requires a new aiter version v0.1.12, am I right?

No, this fix will apply with the current aiter version as well. Aiter v0.1.12 will just change the routing of paged_attention_common such that calls relating to models with head_dim != 128 are no longer sent to pa_fwd_asm even for high batch sizes. Routing issue and issue with kv cache write (this PR) are unrelated. However, the both issues affect Qwen 3.5.

[tuukkjs]

To clarify the situation further, in my understanding, there are three separate (though somewhat related) issues here:

1. shuffled KV-cache writes for hybrid models is not working, see this note in the current codebase. This PR intends to solve that issue.
2. paged_attention_common is routing calls relating to models with head_dim != 128 to pa_fwd_asm which doesn't support that. This is solved when aiter is upgraded.
3. pa_fwd_asm doesn't support models with head_dim != 128. This is a separate issue that could be, if it were a priority, adressed by implementing new kernels or extending the old ones.

Solving each of the above brings some value own its own irrespective of what is done to others (with the exception of 3 and 2 as solving 3 would void the need for 2).

[tjtanaa]

2. paged_attention_common is routing calls relating to models with head_dim != 128 to pa_fwd_asm which doesn't support that. This is solved when aiter is upgraded.

do you have the commit to the necessary upgrade?

And we just landed the aiter v0.1.12 upgrade, could you help me to find out if it contains the fix? Thanks a lot.

[tuukkjs]

2. paged_attention_common is routing calls relating to models with head_dim != 128 to pa_fwd_asm which doesn't support that. This is solved when aiter is upgraded.

do you have the commit to the necessary upgrade?

And we just landed the aiter v0.1.12 upgrade, could you help me to find out if it contains the fix? Thanks a lot.

The aiter commit is b4885d2. It is included in aiter v0.1.12. It seems vllm is still using aiter v0.1.10.post3?

[tjtanaa]

2. paged_attention_common is routing calls relating to models with head_dim != 128 to pa_fwd_asm which doesn't support that. This is solved when aiter is upgraded.

do you have the commit to the necessary upgrade?
And we just landed the aiter v0.1.12 upgrade, could you help me to find out if it contains the fix? Thanks a lot.

The aiter commit is b4885d2. It is included in aiter v0.1.12. It seems vllm is still using aiter v0.1.10.post3?

Yes. We have to revert to v0.1.10.post3 as over last week, the v0.1.12 tag on AITER is a moving tag.

[tuukkjs]

I ran the results with b4885d2 aiter commit applied on top of aiter v0.1.10.post3 which showed that then all batch sizes give correct results.

Config	batch_size	kernel	lm_eval flexible-extract	lm_eval strict-match
vllm + aiter-v0.1.10.post3+b4885d2 + this PR	4	paged_attention_rocm	0.8704 (±0.0093)	0.8491 (±0.0099)
vllm + aiter-v0.1.10.post3+b4885d2 + this PR	512	paged_attention_rocm	0.8469 (±0.0099)	0.8332 (±0.0103)

I also ran lm_eval with aiter v0.1.12. Batch size 4 gave score about 0.2 and batch size 512 gave score near 0. This makes me think there is something wrong with that tag (or at least was few days ago).

[tuukkjs]

@tjtanaa should we merge this? As far as I understand, without aiter upgrade, this fixes low batch cases. After aiter upgrade (see previous message) also high batch cases are fixed. As far as I understand this affects only ROCM_AITER_FA backend with shuffle (so there should be no regressions elsewhere in my understanding).

[mergify]

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

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

[AndreasKaratzas]

I think that hybrid attention is good now. Could you please test the most recent vLLM release and see if the issue you are trying to resolve with this PR has already been resolved?

[tuukkjs]

Yes, it is now working on main already, so this PR can be closed it seems to me. However, I think there might be a stale TODO now in code. This PR proposes to remove it: #44178.

---

Image:
vllm/vllm-openai-rocm:nightly-3fd9d2d35714e80b4cb3fcd3c408a0398fa2525f

Qwen3.5 hybrid lm eval gsm8k, 5-shot:

Variant	batch_size	flexible-extract	strict-match
Plain 3fd9d2d35714 image	4	0.8628	0.8491
Plain 3fd9d2d35714 image	512	0.8537	0.8362