[ROCm] Reduce AITER MLA FP8 BMM warmup sizes

tests/kernels/attention/test_rocm_aiter_mla_fp8_bmm_precompile.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+from vllm.model_executor.layers.attention.mla_attention import (
+    _get_aiter_fp8_bmm_precompile_batch_sizes,
+)
+
+
+def _vllm_config(
+    max_num_seqs: int,
+    max_num_batched_tokens: int = 2048,
+    max_num_scheduled_tokens: int | None = None,
+    num_speculative_tokens: int = 0,
+    max_cudagraph_capture_size: int | None = 512,
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        scheduler_config=SimpleNamespace(
+            max_num_seqs=max_num_seqs,
+            max_num_batched_tokens=max_num_batched_tokens,
+            max_num_scheduled_tokens=max_num_scheduled_tokens,
+        ),
+        compilation_config=SimpleNamespace(
+            max_cudagraph_capture_size=max_cudagraph_capture_size,
+        ),
+        num_speculative_tokens=num_speculative_tokens,
+    )
+
+
+def test_aiter_fp8_bmm_precompile_uses_graph_limit() -> None:
+    vllm_config = _vllm_config(
+        max_num_seqs=1024,
+        max_cudagraph_capture_size=512,
+    )
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(range(1, 513))
+
+
+def test_aiter_fp8_bmm_precompile_uses_scheduler_batch_size_below_graph_limit() -> None:
+    vllm_config = _vllm_config(max_num_seqs=128)
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(range(1, 129))
+
+
+def test_aiter_fp8_bmm_precompile_uses_scheduled_token_limit() -> None:
+    vllm_config = _vllm_config(
+        max_num_seqs=128,
+        max_num_scheduled_tokens=64,
+    )
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(range(1, 65))
+
+
+def test_aiter_fp8_bmm_precompile_uses_batched_token_limit() -> None:
+    vllm_config = _vllm_config(
+        max_num_seqs=128,
+        max_num_batched_tokens=32,
+    )
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(range(1, 33))
+
+
+def test_aiter_fp8_bmm_precompile_accounts_for_spec_decode_tokens() -> None:
+    vllm_config = _vllm_config(
+        max_num_seqs=16,
+        num_speculative_tokens=3,
+    )
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(range(1, 65))
+
+
+def test_aiter_fp8_bmm_precompile_caps_spec_decode_at_graph_limit() -> None:
+    vllm_config = _vllm_config(
+        max_num_seqs=256,
+        max_num_batched_tokens=2048,
+        max_cudagraph_capture_size=512,
+        num_speculative_tokens=3,
+    )
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(range(1, 513))
+
+
+def test_aiter_fp8_bmm_precompile_caps_spec_decode_at_token_limit() -> None:
+    vllm_config = _vllm_config(
+        max_num_seqs=128,
+        max_num_batched_tokens=256,
+        num_speculative_tokens=3,
+    )
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(range(1, 257))
+
+
+def test_aiter_fp8_bmm_precompile_caps_at_old_limit() -> None:
+    vllm_config = _vllm_config(
+        max_num_seqs=2048,
+        max_cudagraph_capture_size=None,
+    )
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(
+        range(1, 1025)
+    )
+
+
+def test_aiter_fp8_bmm_precompile_defaults_to_old_limit() -> None:
+    vllm_config = _vllm_config(
+        max_num_seqs=0,
+        max_cudagraph_capture_size=None,
+    )
+
+    assert _get_aiter_fp8_bmm_precompile_batch_sizes(vllm_config) == list(
+        range(1, 1025)
+    )

vllm/model_executor/layers/attention/mla_attention.py

@@ -306,6 +306,43 @@
 logger = init_logger(__name__)
 
 _FP8_DTYPE = current_platform.fp8_dtype()
+_AITER_FP8_BMM_PRECOMPILE_MAX_BATCH_SIZE = 1024
+
+
+def _get_aiter_fp8_bmm_precompile_batch_sizes(
+    vllm_config: VllmConfig,
+) -> list[int]:
+    # Keep the old upper bound as a safety cap. The AITER Triton BMM kernels
+    # can still JIT-compile an unseen size later; this only controls startup
+    # warming for the decode-token counts vLLM can actually schedule.
+    scheduler_config = vllm_config.scheduler_config
+    max_num_seqs = scheduler_config.max_num_seqs or 0
+    decode_query_len = 1 + max(getattr(vllm_config, "num_speculative_tokens", 0), 0)
+    max_decode_tokens = max_num_seqs * decode_query_len
+
+    token_limit = (
+        scheduler_config.max_num_scheduled_tokens
+        or scheduler_config.max_num_batched_tokens
+        or 0
+    )
+    if token_limit > 0 and max_decode_tokens > 0:
+        max_decode_tokens = min(max_decode_tokens, token_limit)
+
+    # MLAAttention.forward_impl strips CUDA-graph padding before calling this
+    # BMM path, so warm a contiguous range of real sizes. Limit that range to
+    # graph-covered decode sizes when CUDA graphs are enabled; larger batches
+    # can still JIT later, but they already leave the captured fast path.
+    graph_limit = (
+        getattr(vllm_config.compilation_config, "max_cudagraph_capture_size", None) or 0
+    )
+    if graph_limit > 0 and max_decode_tokens > 0:
+        max_decode_tokens = min(max_decode_tokens, graph_limit)
+
+    max_size = min(
+        _AITER_FP8_BMM_PRECOMPILE_MAX_BATCH_SIZE,
+        max_decode_tokens or _AITER_FP8_BMM_PRECOMPILE_MAX_BATCH_SIZE,
+    )
+    return list(range(1, max_size + 1))
 
 
 def _detect_output_quant_key(
@@ -1086,17 +1123,18 @@ def process_weights_after_loading(self, act_dtype: torch.dtype):
                 W_V, dtype=current_platform.fp8_dtype()
             )
 
-            # The kernel operates on non-padded inputs. Hence, pre-compiling
-            # triton kernel to avoid runtime compilation for unseen batch sizes
-            # Pre-compile for batch sizes 1 to 1024 to cover most use-cases.
-            # On DS-R1, this step adds roughly 50s to the model loading time.
-            max_batch_size = 1024  # [ToDo] Find the optimal upper limit
-            pre_compilation_list = list(range(1, max_batch_size + 1))
+            # The kernel operates on non-padded decode tokens. Pre-compile the
+            # real decode-token counts the scheduler can emit; unseen sizes can
+            # still be JIT-compiled later, but skipping unreachable warmup
+            # shapes reduces model loading time.
+            pre_compilation_list = _get_aiter_fp8_bmm_precompile_batch_sizes(
+                self._vllm_config
+            )
             if is_global_first_rank():
                 pre_compilation_list = tqdm(
                     pre_compilation_list,
                     desc="[Aiter Triton] Pre-compiling fp8 BMM kernel",
-                    total=max_batch_size,
+                    total=len(pre_compilation_list),
                 )
 
             for m in pre_compilation_list: