From 03a7d3776fc2f520a70c8d2da1724eab6df6d0e1 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:33:47 -0400 Subject: [PATCH] [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 Co-authored-by: Codex --- tests/engine/test_arg_utils.py | 6 ++ tests/v1/attention/test_mla_backends.py | 27 ++++++-- vllm/config/attention.py | 5 ++ .../attention/backends/mla/tokenspeed_mla.py | 67 ++++++++++++++----- 4 files changed, 83 insertions(+), 22 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index fc9f0229c4..3e3a82f80f 100644 --- a/tests/engine/test_arg_utils.py +++ b/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 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9c504f31ef..91bd9008f0 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/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], diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 8a9be6336c..80548f58f2 100644 --- a/vllm/config/attention.py +++ b/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, diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 96859d7f98..8d41cd178e 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/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