[ROCm] Validate block_size for explicitly selected attention backends

[AndreasKaratzas]

#36274 stripped block_size from attn_selector_config before backend validation in get_attn_backend_cls, which was correct for auto-selection (block_size may not be finalized at that point). However, this also bypassed block_size validation for explicitly user-selected backends, breaking the contract established in #36292.

• Add an explicit supports_block_size check for the selected-backend path, before the strip
• Auto-selection path is unaffected and block_size is still stripped there
• Fixes: test_mla_backend_selection[env_vars1-TRITON_MLA-1-None-True]

cc @kenroche

[gemini-code-assist]

Code Review

This pull request modifies the get_attn_backend_cls method in vllm/platforms/rocm.py to validate the block_size for explicitly selected attention backends. A new check is added to verify that the selected backend supports the specified block_size before this parameter is stripped from the configuration. This fixes a bug where this validation was being bypassed. The change is targeted and does not affect the automatic backend selection path. I have reviewed the changes and found no issues.

[tjtanaa]

@AndreasKaratzas We should also guard this block size. Because there are also a use case where users do not specify a specific backend, but he specified the --block-size explicitly.

This #36274 looks more like a hotpatch.

We should be looking into solving supporting the correct size through get_supported_kernel_block_sizes of attention backend class.

[AndreasKaratzas]

@tjtanaa Indeed. I removed both my patch and patch from #36274

[tjtanaa]

@AndreasKaratzas Let's try to fix the get_supported_kernel_block_sizes of attention backend in this PR. Else Meta will encounter issue running Qwen3.5 after this PR.

[AndreasKaratzas]

@tjtanaa But Qwen works fine. Are you referring to:

vllm bench throughput --model Qwen/Qwen3-Next-80B-A3B-Instruct --kv-cache-dtype auto --load-format dummy --input-len 1024 --output-len 1024 --num-prompts 128 --tensor-parallel-size 8 --dtype float16

Throughput: 3.61 requests/s, 4157.86 total tokens/s, 461.98 output tokens/s
Total num prompt tokens:  131072
Total num output tokens:  16384

vllm bench throughput --model Qwen/Qwen3-Next-80B-A3B-Instruct --kv-cache-dtype auto --load-format dummy --input-len 1024 --output-len 1024 --num-prompts 128 --tensor-parallel-size 8 --dtype float16 --attention-backend ROCM_ATTN

Throughput: 8.01 requests/s, 9228.46 total tokens/s, 1025.38 output tokens/s
Total num prompt tokens:  131072
Total num output tokens:  16384

The above are with this PR.

[AndreasKaratzas]

Tried Qwen 3.5 today as well.

vllm bench throughput --model Qwen/Qwen3.5-35B-A3B --load-format dummy --input-len 1024 --output-len 1024 --num-prompts 128 --tensor-parallel-size 8 --dtype float16

Throughput: 5.53 requests/s, 6367.64 total tokens/s, 707.52 output tokens/s
Total num prompt tokens:  131072
Total num output tokens:  16384

vllm bench throughput --model Qwen/Qwen3.5-35B-A3B --load-format dummy --input-len 1024 --output-len 1024 --num-prompts 128 --tensor-parallel-size 8 --dtype float16 --attention-backend ROCM_ATTN

Throughput: 8.34 requests/s, 9602.49 total tokens/s, 1066.94 output tokens/s
Total num prompt tokens:  131072
Total num output tokens:  16384

cc @jennyyyyzhen @Rohan138 @tjtanaa

[jennyyyyzhen]

There is a code refactor effort on cuda side #35122 and probably fixed the previous issue.

[gshtras]

cc @Rohan138

[AndreasKaratzas]

AMD CI is red due to new regression from: #36204
I will be addressing it in a different PR.

EDIT: Fixed here: #37219