[K3 Perf] Optimize k3 mamba metadata preparation, 6.6~7.6x kernel performance improvement (#52388)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 25, 2026 at 11:00 UTC
41729fc53b02cb09beda2a1ced690d021443be43
4 files changed
+187
-6
tests/models/kimi_k3/test_kda_metadata.py
+37
@@ -330,6 +330,13 @@ def test_recoverssm_spec_uses_one_state_slot_and_current_window(
330
use_recoverssm=True,
331
)
332
assert isinstance(builder, KimiK3KDAMetadataBuilder)
333
+ if mamba_cache_mode == "align":
334
+ builder.mamba_aligned_state_indices = mamba_get_block_table_tensor(
335
+ common_attn_metadata.block_table_tensor,
336
+ common_attn_metadata.seq_lens,
337
+ builder.kv_cache_spec,
338
+ mamba_cache_mode,
339
+ )
340
context = builder.recoverssm_context
341
assert context is not None
342
actual = builder.build(
@@ -506,6 +513,36 @@ def test_kimi_k3_kda_backend_uses_private_metadata_builder():
513
assert issubclass(KimiK3KDAMetadataBuilder, GDNAttentionMetadataBuilder)
514
515
516
+def test_kimi_k3_metadata_uses_precomputed_aligned_state_indices():
517
+ batch = BatchSpec(seq_lens=[40, 30], query_lens=[1, 1])
518
+ common_attn_metadata = create_common_attn_metadata(
519
+ batch, BLOCK_SIZE, DEVICE
520
+ ).replace(is_prefilling=torch.tensor([False, False]))
521
+ builder = _make_builder(
522
+ KimiK3KDAMetadataBuilder,
523
+ num_speculative_tokens=2,
524
+ full_cuda_graph=False,
525
+ mamba_cache_mode="align",
526
+ )
527
+ assert isinstance(builder, KimiK3KDAMetadataBuilder)
528
+ precomputed_indices = torch.tensor(
529
+ [
530
+ [101, 102, 103],
531
+ [201, 202, 203],
532
+ [301, 302, 303],
533
+ ],
534
+ dtype=torch.int32,
535
+ )
536
+ builder.mamba_aligned_state_indices = precomputed_indices
537
+
538
+ metadata = builder.build(0, common_attn_metadata)
539
+
540
+ torch.testing.assert_close(
541
+ metadata.non_spec_state_indices_tensor,
542
+ precomputed_indices[: batch.batch_size, 0],
543
+ )
544
+
545
+
546
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
547
def test_stage_spec_decode_metadata_matches_pytorch():
548
device = torch.device("cuda")
vllm/models/kimi_k3/nvidia/kda_metadata.py
+18
-6
@@ -306,6 +306,8 @@ class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
306
307
308
class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
309
+ mamba_aligned_state_indices: torch.Tensor | None = None
310
+
311
def __init__(
312
self,
313
kv_cache_spec: MambaSpec,
@@ -367,12 +369,22 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
369
# offsets = torch.arange(1 + num_speculative_blocks, dtype=torch.int32)
370
# indices = (start[:, None] + offsets).to(torch.int64)
371
# block_table_tensor = torch.gather(block_table, 1, indices)
370
- block_table_tensor = _mamba_get_block_table_tensor(
371
- m.block_table_tensor,
372
- m.seq_lens,
373
- self.kv_cache_spec,
374
- self.vllm_config.cache_config.mamba_cache_mode,
375
- )
372
+ if self.vllm_config.cache_config.mamba_cache_mode == "align":
373
+ if self.mamba_aligned_state_indices is not None:
374
+ block_table_tensor = self.mamba_aligned_state_indices[: m.num_reqs]
375
+ else:
376
+ assert not self.vllm_config.use_v2_model_runner, (
377
+ "Aligned Mamba state indices must be precomputed"
378
+ )
379
+ # TODO: remove this MRV1 fallback once MRV2 is the default runner.
380
+ block_table_tensor = _mamba_get_block_table_tensor(
381
+ m.block_table_tensor,
382
+ m.seq_lens,
383
+ self.kv_cache_spec,
384
+ self.vllm_config.cache_config.mamba_cache_mode,
385
+ )
386
+ else:
387
+ block_table_tensor = m.block_table_tensor
388
389
if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
390
spec_sequence_masks_cpu = None
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+18
@@ -271,6 +271,24 @@ class MambaHybridModelState(DefaultModelState):
271
)
272
num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np)
273
274
+ if self._align_mode:
275
+ mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config)
276
+ aligned_index_builders = []
277
+ for group_idx, group_id in enumerate(mamba_group_ids):
278
+ for group in attn_groups[group_id]:
279
+ builder = group.get_metadata_builder(0)
280
+ if hasattr(builder, "mamba_aligned_state_indices"):
281
+ aligned_index_builders.append((group_idx, builder))
282
+ if aligned_index_builders:
283
+ ctx = self._ensure_align_ctx(
284
+ kv_cache_config, mamba_group_ids, block_tables
285
+ )
286
+ all_group_indices = ctx.compute_aligned_state_indices(
287
+ input_batch.seq_lens, num_reqs
288
+ )
289
+ for group_idx, builder in aligned_index_builders:
290
+ builder.mamba_aligned_state_indices = all_group_indices[group_idx]
291
+
292
mamba_attn_metadata = MambaHybridAttnMetadata(
293
is_prefilling=is_prefilling,
294
num_accepted_tokens=num_accepted_tokens,
vllm/v1/worker/mamba_utils.py
+114
@@ -31,6 +31,70 @@ logger = init_logger(__name__)
31
_TEMPORAL_TILES = 16
32
33
34
+@triton.jit(do_not_specialize=["num_requests"])
35
+def get_aligned_state_indices_multi_group_kernel(
36
+ block_table_ptrs_ptr,
37
+ seq_lens_ptr,
38
+ state_indices_ptr,
39
+ block_table_stride_req: tl.int64,
40
+ seq_lens_stride: tl.constexpr,
41
+ state_indices_stride_0: tl.constexpr,
42
+ state_indices_stride_1: tl.constexpr,
43
+ state_indices_stride_2: tl.constexpr,
44
+ num_requests,
45
+ CACHE_BLOCK_SIZE: tl.constexpr,
46
+ NUM_GROUPS: tl.constexpr,
47
+ BLOCK_GROUPS: tl.constexpr,
48
+ NUM_STATE_SLOTS: tl.constexpr,
49
+ BLOCK_STATE_SLOTS: tl.constexpr,
50
+ BLOCK_ROWS: tl.constexpr,
51
+):
52
+ rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
53
+ valid_row = rows < num_requests
54
+
55
+ seq_lens = tl.load(
56
+ seq_lens_ptr + rows * seq_lens_stride,
57
+ mask=valid_row,
58
+ other=1,
59
+ )
60
+ first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0)
61
+
62
+ # load multiple block table for each group
63
+ groups = tl.arange(0, BLOCK_GROUPS)
64
+ valid_group = groups < NUM_GROUPS
65
+ group_base_addrs = tl.load(
66
+ block_table_ptrs_ptr + groups,
67
+ mask=valid_group,
68
+ other=0,
69
+ )
70
+ block_tables = group_base_addrs.to(tl.pointer_type(tl.int32))
71
+ state_slots = tl.arange(0, BLOCK_STATE_SLOTS)
72
+ valid_state_slot = state_slots < NUM_STATE_SLOTS
73
+ state_indices = tl.load(
74
+ block_tables[:, None, None]
75
+ + rows[None, :, None] * block_table_stride_req
76
+ + first_state_slot[None, :, None]
77
+ + state_slots[None, None, :],
78
+ mask=(
79
+ valid_group[:, None, None]
80
+ & valid_row[None, :, None]
81
+ & valid_state_slot[None, None, :]
82
+ ),
83
+ )
84
+ tl.store(
85
+ state_indices_ptr
86
+ + groups[:, None, None] * state_indices_stride_0
87
+ + rows[None, :, None] * state_indices_stride_1
88
+ + state_slots[None, None, :] * state_indices_stride_2,
89
+ state_indices,
90
+ mask=(
91
+ valid_group[:, None, None]
92
+ & valid_row[None, :, None]
93
+ & valid_state_slot[None, None, :]
94
+ ),
95
+ )
96
+
97
+
98
@triton.jit
99
def _memcpy_u64_tiled(
100
src_addr,
@@ -689,6 +753,10 @@ class MambaSpecDecodeGPUContext:
753
block_table_ptrs: torch.Tensor
754
block_table_stride_req: int = 0
755
756
+ # persistent output for the once-per-step, all-group aligned-index launch.
757
+ # shape: [num_groups, max_num_reqs, 1 + num_speculative_blocks].
758
+ aligned_state_indices: torch.Tensor | None = None
759
+
760
# Per-request staging buffers (CPU+GPU mirrors). The runner stages
761
# values into the CPU view in ``_prepare_inputs`` and the fused kernel
762
# reads the GPU side. These only exist when the postprocess kernel is
@@ -758,6 +826,15 @@ class MambaSpecDecodeGPUContext:
826
block_table_ptrs=torch.zeros(
827
len(mamba_group_ids), dtype=torch.int64, device=device
828
),
829
+ aligned_state_indices=torch.empty(
830
+ (
831
+ len(mamba_group_ids),
832
+ max_num_reqs,
833
+ 1 + mamba_spec.num_speculative_blocks,
834
+ ),
835
+ dtype=torch.int32,
836
+ device=device,
837
+ ),
838
mamba_state_idx_buf=make_buffer(max_num_reqs, dtype=torch.int32),
839
num_scheduled_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
840
num_computed_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32),
@@ -934,6 +1011,43 @@ class MambaSpecDecodeGPUContext:
1011
1012
self.is_initialized = True
1013
1014
+ def compute_aligned_state_indices(
1015
+ self,
1016
+ seq_lens: torch.Tensor,
1017
+ num_reqs: int,
1018
+ ) -> torch.Tensor:
1019
+ """compute every Mamba group's aligned physical state IDs in one launch."""
1020
+ assert self.is_initialized
1021
+ assert seq_lens.is_cuda
1022
+ assert 0 <= num_reqs <= seq_lens.shape[0]
1023
+ assert self.aligned_state_indices is not None
1024
+ assert num_reqs <= self.aligned_state_indices.shape[1]
1025
+ if num_reqs == 0:
1026
+ return self.aligned_state_indices[:, :0]
1027
+
1028
+ num_state_slots = self.aligned_state_indices.shape[2]
1029
+ block_rows = 32
1030
+ grid = (triton.cdiv(num_reqs, block_rows),)
1031
+ get_aligned_state_indices_multi_group_kernel[grid](
1032
+ self.block_table_ptrs,
1033
+ seq_lens,
1034
+ self.aligned_state_indices,
1035
+ self.block_table_stride_req,
1036
+ seq_lens.stride(0),
1037
+ self.aligned_state_indices.stride(0),
1038
+ self.aligned_state_indices.stride(1),
1039
+ self.aligned_state_indices.stride(2),
1040
+ num_reqs,
1041
+ CACHE_BLOCK_SIZE=self.block_size,
1042
+ NUM_GROUPS=self.num_groups,
1043
+ BLOCK_GROUPS=triton.next_power_of_2(self.num_groups),
1044
+ NUM_STATE_SLOTS=num_state_slots,
1045
+ BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots),
1046
+ BLOCK_ROWS=block_rows,
1047
+ num_warps=1,
1048
+ )
1049
+ return self.aligned_state_indices[:, :num_reqs]
1050
+
1051
def run_fused_postprocess(
1052
self,
1053
num_reqs: int,