[ROCm][Bugfix] Pass num_kv_splits to aiter mla_reduce_v1 (#46760)
Signed-off-by: Rohan Potdar <rohanpotdar138@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rohan Potdar committed
Jun 26, 2026 at 16:58 UTC
2ff76a5e856e385f72aa49cadc4d0a724d1f7da8
1 file changed
+3
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+3
@@ -828,6 +828,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
828
attn_metadata.fp8_prefill_reduce_final_map,
829
attn_metadata.fp8_prefill_reduce_partial_map,
830
tile_q,
831
+ # num_kv_splits added by ROCm/aiter#3391; 0 selects the kernel
832
+ # default max(cu_num, 0) == cu_num, matching pre-#3391 behavior.
833
+ 0,
834
out_3d,
835
final_lse,
836
)