[Perf] Support minimum KV splits configuration for TokenSpeed MLA decode (#60345)

Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Co-authored-by: Codex <noreply@openai.com>

tests/engine/test_arg_utils.py

@@ -503,6 +503,12 @@ def test_attention_config():
     engine_args = EngineArgs.from_cli_args(args)
     assert engine_args.attention_config == AttentionConfig()
 
+    args = parser.parse_args(
+        ["--attention-config", '{"tokenspeed_mla_min_split_kv": 8}']
+    )
+    engine_args = EngineArgs.from_cli_args(args)
+    assert engine_args.attention_config.tokenspeed_mla_min_split_kv == 8
+
     # set backend via dot notation
     args = parser.parse_args(["--attention-config.backend", "FLASH_ATTN"])
     assert args is not None

tests/v1/attention/test_mla_backends.py

@@ -1098,6 +1098,7 @@ def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch):
     ],
 )
 @pytest.mark.parametrize("cp_interleave_size", [1, 16, 896])
+@pytest.mark.parametrize("min_split_kv", [1, 8])
 def test_tokenspeed_mla_decode_contract(
     monkeypatch,
     causal,
@@ -1105,6 +1106,7 @@ def test_tokenspeed_mla_decode_contract(
     dcp_world_size,
     dcp_rank,
     cp_interleave_size,
+    min_split_kv,
 ):
     decode_call = None
     num_decodes = 2
@@ -1148,7 +1150,12 @@ def test_tokenspeed_mla_decode_contract(
     impl.scale = 1.0
     impl.softmax_scale = 1.0
     impl.output_scale = 1.0
-    impl._workspace_buffer = torch.empty(1, dtype=torch.int8)
+    impl._min_split_kv = min_split_kv
+    impl._max_decode_tokens = num_decode_tokens
+    workspace = torch.empty(1, dtype=torch.int8)
+    monkeypatch.setattr(
+        tokenspeed_mla_module, "_get_workspace", lambda *args: workspace
+    )
 
     metadata = SimpleNamespace(
         num_decodes=num_decodes,
@@ -1183,6 +1190,8 @@ def test_tokenspeed_mla_decode_contract(
     assert lse.shape == (num_decode_tokens, num_heads)
 
     assert decode_call is not None
+    assert decode_call["workspace_buffer"] is workspace
+    assert decode_call["min_split_kv"] == min_split_kv
     assert decode_call["query"].shape == (
         num_decodes,
         tokens_per_decode,
@@ -1217,16 +1226,23 @@ def test_tokenspeed_mla_decode_contract(
         pytest.param(896, 8 * 896 + 17, id="block-interleave-full-cycle"),
     ],
 )
-def test_tokenspeed_mla_dcp_matches_unsharded_decode(cp_interleave_size, seq_len):
+def test_tokenspeed_mla_dcp_matches_unsharded_decode(
+    monkeypatch, cp_interleave_size, seq_len
+):
     from tokenspeed_mla import tokenspeed_mla_decode
 
     from vllm.v1.attention.ops.dcp import (
         _lse_weighted_combine,
         mask_dcp_empty_shards_,
     )
+    from vllm.v1.worker.workspace import WorkspaceManager
 
     torch.manual_seed(7)
     device = torch.device("cuda")
+    workspace_manager = WorkspaceManager(device)
+    monkeypatch.setattr(
+        tokenspeed_mla_module, "current_workspace_manager", lambda: workspace_manager
+    )
     dcp_world_size = 8
     kernel_block_size = 64
     num_heads = 128
@@ -1257,7 +1273,9 @@ def test_tokenspeed_mla_dcp_matches_unsharded_decode(cp_interleave_size, seq_len
     global_tokens = (torch.randn(seq_len, head_size, device=device) * 0.1).to(
         torch.float8_e4m3fn
     )
-    workspace = tokenspeed_mla_module._get_workspace(device, num_heads, kv_lora_rank)
+    workspace = tokenspeed_mla_module._get_workspace(
+        device, num_heads, kv_lora_rank, 1, query.shape[1]
+    )
     global_cache, global_block_table = make_paged_cache(global_tokens)
     global_seq_len = torch.tensor([seq_len], dtype=torch.int32, device=device)
     reference, _ = tokenspeed_mla_decode(
@@ -1301,7 +1319,8 @@ def test_tokenspeed_mla_dcp_matches_unsharded_decode(cp_interleave_size, seq_len
         impl.scale = head_size**-0.5
         impl.softmax_scale = None
         impl.output_scale = None
-        impl._workspace_buffer = workspace
+        impl._min_split_kv = 1
+        impl._max_decode_tokens = query.shape[1]
         metadata = SimpleNamespace(
             num_decodes=1,
             num_decode_tokens=query.shape[1],

vllm/config/attention.py

@@ -61,6 +61,11 @@ class AttentionConfig:
     flash_attn_max_num_splits_for_cuda_graph: int = 32
     """Flash Attention max number splits for cuda graph decode."""
 
+    tokenspeed_mla_min_split_kv: int = Field(default=1, ge=1, le=256)
+    """Minimum KV splits for TokenSpeed MLA decode. The default of 1 preserves
+    automatic scheduling. Larger values request a split floor and may increase
+    workspace memory. Only applies to the TokenSpeed MLA backend."""
+
     tq_max_kv_splits_for_cuda_graph: int = 32
     """TurboQuant max NUM_KV_SPLITS for cuda graph decode.
     Fixes the split count so grid dimensions are constant across captures,

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -24,6 +24,7 @@ from vllm.v1.attention.backend import (
     AttentionType,
     MultipleOf,
 )
+from vllm.v1.worker.workspace import current_workspace_manager
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -38,21 +39,38 @@ logger = init_logger(__name__)
 # own buffer check.
 _TOKENSPEED_MAX_Q_LEN = 8
 
-_g_workspace: dict[torch.device, torch.Tensor] = {}
-
 
 def _get_workspace(
-    device: torch.device, num_heads: int, kv_lora_rank: int
+    device: torch.device,
+    num_heads: int,
+    kv_lora_rank: int,
+    min_split_kv: int,
+    max_decode_tokens: int,
 ) -> torch.Tensor:
     from tokenspeed_mla import get_num_sm
 
+    # TokenSpeed reduction scratch: B * H * Q * split_kv * (D + 1) * 4 bytes,
+    # using the effective query layout and FP32 partial outputs plus LSE.
+    # Automatic splitting lowers split_kv as the batch grows; split_kv=1
+    # needs no reduction scratch, so an SM-based bound covers all batch sizes.
     needed = (
         get_num_sm(device) * num_heads * _TOKENSPEED_MAX_Q_LEN * (kv_lora_rank + 1) * 4
     )
-    existing = _g_workspace.get(device)
-    if existing is None or existing.numel() < needed:
-        _g_workspace[device] = torch.empty(needed, dtype=torch.int8, device=device)
-    return _g_workspace[device]
+    if min_split_kv > 1:
+        # A forced floor breaks that bound: reserve for the maximum batch.
+        # Packed queries pad head rows to 128; reserve one tile per query token.
+        needed = max(
+            needed,
+            max_decode_tokens
+            * max(num_heads, 128)
+            * min_split_kv
+            * (kv_lora_rank + 1)
+            * 4,
+        )
+        logger.info_once(
+            "TokenSpeed MLA decode workspace: %.2f GiB", needed / (1 << 30)
+        )
+    return current_workspace_manager().get_simultaneous(((needed,), torch.int8))[0]
 
 
 class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
@@ -212,15 +230,25 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
                 f"got kv_cache_dtype={self.kv_cache_dtype!r}."
             )
 
-        # Allocate (or fetch the cached) workspace lazily on first forward —
-        # __init__ runs before the device is necessarily set on the worker;
-        # we know it for sure at forward time when we see the input tensor.
-        self._workspace_buffer: torch.Tensor | None = None
+        config = get_current_vllm_config()
+        self._min_split_kv = config.attention_config.tokenspeed_mla_min_split_kv
+        self._max_decode_tokens = 0
+        if self._min_split_kv > 1:
+            spec = config.speculative_config
+            query_len = 1 + (spec.num_speculative_tokens if spec else 0)
+            self._max_decode_tokens = max(
+                min(
+                    config.scheduler_config.max_num_batched_tokens,
+                    config.scheduler_config.max_num_seqs * query_len,
+                ),
+                config.compilation_config.max_cudagraph_capture_size or 0,
+            )
+
         self.softmax_scale: float | None = None
         self.output_scale: float | None = None
         # NIXL resolves interleaving after model construction; retain the config
         # rather than caching its initial interleave size.
-        self._parallel_config = get_current_vllm_config().parallel_config
+        self._parallel_config = config.parallel_config
 
         # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always
         # runs when tokenspeed is selected, prefill backend may not (user can
@@ -290,11 +318,13 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
             )
             self.output_scale = layer._k_scale_float
 
-        if self._workspace_buffer is None:
-            # Parallelism can change the runtime query head count.
-            self._workspace_buffer = _get_workspace(
-                q.device, q.shape[-2], self.kv_lora_rank
-            )
+        workspace = _get_workspace(
+            q.device,
+            q.shape[-2],
+            self.kv_lora_rank,
+            self._min_split_kv,
+            self._max_decode_tokens,
+        )
 
         # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size).
         # tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm).
@@ -302,7 +332,7 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
         kernel_out = tokenspeed_mla_decode(
             query=q,
             kv_cache=kv_c_and_k_pe_cache,
-            workspace_buffer=self._workspace_buffer,
+            workspace_buffer=workspace,
             kv_lora_rank=self.kv_lora_rank,
             qk_rope_head_dim=self.qk_rope_head_dim,
             block_tables=block_tables,
@@ -317,6 +347,7 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
             cp_world=self.dcp_world_size,
             cp_rank=self.dcp_rank,
             cp_interleave_size=self._parallel_config.cp_kv_cache_interleave_size,
+            min_split_kv=self._min_split_kv,
         )
         if return_lse:
             o, lse = kernel_out