[Bugfix][ROCm] Give KV-first attention blocks their own page in hybrid models - #51837
Conversation
…d models Hybrid models share one page pool between attention and Mamba layers, which relies on block id b addressing page b. ROCm's (2, num_blocks, ...) layout splits a block into two far-apart halves, so attention blocks and Mamba state pages resolve to overlapping bytes and silently corrupt each other. This shows up as NaN inside the attention KV cache and garbage output: Qwen3.6-35B DSpark speculative decoding on MI300 scores 0.20 on GSM8K instead of 0.95. Lay KV-first views out page-first when the allocation is shared with Mamba state, so every block owns exactly its own page. Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
mergify
Bot
added
rocm
|
/ci run |
|
✅ Triggered Buildkite CI #83379 for commit |
AndreasKaratzas
added
the
ready-run-all-tests
|
✅ @stefankoncarevic, CI is now available for this PR.
|
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
Thanks for the work
Please provide with full reproduce command lm_eval .. line with full logs, also please add metrics for e2e performance using vllm bench ...
|
I ran 40 GSM8K questions with 5-shot prompting and greedy decoding. LMEVAL_LOG_LEVEL=DEBUG .venv/bin/lm_eval run \
--model vllm \
--model_args \
pretrained=RedHatAI/Qwen3.6-35B-A3B-NVFP4 \
trust_remote_code=True \
max_model_len=4096 \
max_num_seqs=32 \
gpu_memory_utilization=0.85 \
language_model_only=True \
enable_prefix_caching=False \
disable_log_stats=False \
spec_method=dspark \
spec_model=RedHatAI/Qwen3.6-35B-A3B-speculator.dspark \
spec_tokens=8 \
enable_thinking=False \
seed=0 \
--tasks gsm8k \
--num_fewshot 5 \
--limit 40 \
--batch_size auto \
--apply_chat_template \
--fewshot_as_multiturn false \
--gen_kwargs temperature=0 max_gen_toks=256 \
--log_samples \
--output_path ./lm-eval-results \
--confirm_run_unsafe_code \
--write_out \
--show_config \
2>&1 | tee lm-eval.log
Serving benchmarkFor the E2E benchmark, I used 512 random requests with 512 input tokens and 128 output tokens, unlimited request rate, and a maximum concurrency of 32. I ran the same workload three times on each revision. .venv/bin/python -m vllm.entrypoints.cli.main serve \
RedHatAI/Qwen3.6-35B-A3B-NVFP4 \
--trust-remote-code \
--max-model-len 4096 \
--max-num-seqs 32 \
--gpu-memory-utilization 0.85 \
--language-model-only \
--no-enable-prefix-caching \
--speculative-config '{"method":"dspark","model":"RedHatAI/Qwen3.6-35B-A3B-speculator.dspark","num_speculative_tokens":8,"draft_sample_method":"probabilistic"}' \
--default-chat-template-kwargs '{"enable_thinking":false}' \
--port 8000.venv/bin/python -m vllm.entrypoints.cli.main bench serve \
--backend vllm \
--base-url http://127.0.0.1:8000 \
--endpoint /v1/completions \
--model RedHatAI/Qwen3.6-35B-A3B-NVFP4 \
--dataset-name random \
--random-input-len 512 \
--random-output-len 128 \
--random-range-ratio 0 \
--num-prompts 512 \
--num-warmups 32 \
--request-rate inf \
--max-concurrency 32 \
--ignore-eos \
--temperature 0 \
--seed 42 \
--percentile-metrics ttft,tpot,itl,e2el \
--metric-percentiles 50,90,99 \
--save-resultThese are the averages across the three runs:
All requests completed successfully. Throughput and median latency improved, although P99 E2E latency was worse on this particular random workload. |
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
LGTM, thanks for the work!
|
We are trying to get #51718 merged asap, and I think this one will need to be reconsidered with that. Can |
|
Probably TRITON_ATTN is fine. I'd like to get a second opinion on that by @Rohan138 |
|
Yeah TRITON_ATTN is fine for CI, for perf for e.g. upcoming Qwen 3.8 we'll probably use e.g. ROCM_AITER_FA or ROCM_AITER_UNIFIED_ATTN. We will eventually work on consolidating these, in the meantime I'm fine marking off these paths in the attention backend selector so that e.g. ROCM_ATTN is never available for hybrid models after the layout refactor. |
|
Let's go with ROCM_AITER_FA or ROCM_AITER_UNIFIED_ATTN. I want to stay as close as possible to perf pipeline. |
…d models (vllm-project#51837) Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
5 participants
Purpose
On MI300, DSpark speculative decoding on
RedHatAI/Qwen3.6-35B-A3B-NVFP4produces garbage for most requests in a batch. GSM8K accuracy drops to 0.20 against 0.95 for the same target model without speculation, andtest_dspark_correctness_and_acceptance_rate[qwen3.6-speculators]fails. The same test passes on H200.Hybrid models share a single page pool between the full-attention layers, the Mamba/GDN state layers, and the draft model's layers. That sharing is only safe because every group agrees that block id
boccupies pageb, i.e. the bytes[b * page_size, (b + 1) * page_size).ROCm's KV cache layout breaks that agreement. Its shape is
(2, num_blocks, block_size, num_kv_heads, head_size), so K and V live in two far-apart halves of the allocation and blockbactually covers half of pageb / 2plus half of pagenum_blocks / 2 + b / 2. Attention blocks and Mamba state pages therefore resolve to overlapping bytes and scribble over each other. Backends whose block dimension comes first, such asTRITON_ATTNand the NVIDIA paths, keep the agreement and are unaffected, which is why this is ROCm-only.The corruption is silent. Instrumenting the first full-attention layer showed clean inputs and
NaNoutputs for two of four sequences, withNaNvalues sitting inside the valid context range of those sequences' KV cache blocks. The affected sequences were exactly the ones holding an even block id, matching the factor of two between the layout and the page size. A singleNaNkey poisons the whole softmax row, so the target returns degenerate logits and the decoded text collapses.The fix builds the view page-first, then swaps the dimensions back, so each block owns exactly its own page. It applies only to KV-first attention layers whose allocation is also used by Mamba state; blocks-first backends and attention-only allocations keep their current layout. The logical shape is unchanged, so kernels see exactly what they saw before.
Test Plan
Unless noted, everything below ran on a single MI300 (gfx942) with
RedHatAI/Qwen3.6-35B-A3B-NVFP4plusRedHatAI/Qwen3.6-35B-A3B-speculator.dspark.pytest "tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py::test_dspark_correctness_and_acceptance_rate[qwen3.6-speculators]"pytest tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.pypytest tests/v1/worker/test_attn_utils.py tests/v1/core/test_kv_cache_utils.py tests/v1/core/test_contiguous_kv_packing.pytests/v1/worker/test_attn_utils.py: one asserts that a KV-first layer sharing its allocation with Mamba gives every block its own page, the other asserts that an attention-only allocation keeps the existing layout.Test Result
TRITON_ATTNon both target and draft scored 0.925.Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.