@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dd3bc9c72f

[Perf][MRV2] Reuse Mamba/GDN metadata across KV cache groups (#58762)

Signed-off-by: Luca Motz <luca.motz@icloud.com> Co-authored-by: Xin Yang <xyangx@amazon.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Luca Motz committed Sep 28, 2026 at 06:36 UTC dd3bc9c72faf410043b453e18e58553a1ad4d388

5 files changed +189 -21

tests/models/kimi_k3/test_kda_metadata.py

@@ -49,6 +49,7 @@ PRUNED_METADATA_FIELDS = {

49 "prefill_state_indices",

50 "prefill_has_initial_state",

51 "spec_sequence_masks",

52 + "spec_sequence_masks_cpu",

53 "uniform_spec_sequence_length",

54 "flashinfer_prefill_query_start_loc",

55 "flashinfer_prefill_seq_order",

tests/v1/attention/test_gdn_metadata_builder.py

+63

@@ -21,6 +21,7 @@ from vllm.v1.attention.backends.gdn_attn import (

21 GDNAttentionMetadata,

22 GDNAttentionMetadataBuilder,

23 )

24 +from vllm.v1.attention.backends.utils import mamba_get_block_table_tensor

25 from vllm.v1.kv_cache_interface import MambaSpec

27 BLOCK_SIZE = 16

@@ -75,6 +76,17 @@ GDN_BUILD_TEST_CASES = {

76 expected_num_prefill_tokens=0,

77 expected_num_spec_decodes=0,

78 ),

79 + # No speculative config, decode alongside prefill

80 + "regular_decode_with_prefill": GDNBuildTestCase(

81 + seq_lens=[40, 100],

82 + query_lens=[1, 50],

83 + num_decode_draft_tokens=None,

84 + num_speculative_tokens=0,

85 + expected_num_decodes=1,

86 + expected_num_prefills=1,

87 + expected_num_prefill_tokens=50,

88 + expected_num_spec_decodes=0,

89 + ),

90 # Multi-token prefill alongside spec decode — no decode to reclassify

91 "spec_decode_with_real_prefill": GDNBuildTestCase(

92 seq_lens=[100, 20],

@@ -155,9 +167,12 @@ def _build(

167 builder: GDNAttentionMetadataBuilder,

168 batch_spec: BatchSpec,

169 num_decode_draft_tokens: list[int] | None = None,

170 + block_table: torch.Tensor | None = None,

171 ) -> GDNAttentionMetadata:

172 """Build GDN attention metadata, optionally with spec-decode kwargs."""

173 common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE)

174 + if block_table is not None:

175 + common = common.replace(block_table_tensor=block_table)

176 kwargs: dict = {}

177 if num_decode_draft_tokens is not None:

178 kwargs["num_decode_draft_tokens_cpu"] = torch.tensor(

@@ -184,6 +199,54 @@ def test_gdn_build_classification(test_case: GDNBuildTestCase):

199 assert meta.num_spec_decodes == test_case.expected_num_spec_decodes

200

201

202 +@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])

203 +@pytest.mark.parametrize("full_cuda_graph", [False, True])

204 +@pytest.mark.parametrize(

205 + "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys()

206 +)

207 +def test_update_block_table_matches_build(

208 + test_case: GDNBuildTestCase, full_cuda_graph: bool, mamba_cache_mode: str

209 +):

210 + """update_block_table() on another group's metadata matches build()."""

211 + batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens)

212 + src, dst, ref = (

213 + _create_gdn_builder(test_case.num_speculative_tokens, full_cuda_graph)

214 + for _ in range(3)

215 + )

216 + for builder in (src, dst, ref):

217 + builder.vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode

218 + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE)

219 + if mamba_cache_mode == "align":

220 + # MRV2 precomputes these for every Mamba group each step.

221 + dst.mamba_aligned_state_indices = ref.mamba_aligned_state_indices = (

222 + mamba_get_block_table_tensor(

223 + common.block_table_tensor, common.seq_lens, ref.kv_cache_spec, "align"

224 + )

225 + )

226 + draft_tokens = test_case.num_decode_draft_tokens

227 + expected = _build(ref, batch, draft_tokens, common.block_table_tensor)

228 + source = _build(src, batch, draft_tokens)

229 + fields = (

230 + "spec_state_indices_tensor",

231 + "non_spec_state_indices_tensor",

232 + "prefill_state_indices",

233 + )

234 + source_indices = [getattr(source, f) for f in fields]

235 + source_indices = [t if t is None else t.clone() for t in source_indices]

236 + meta = dst.update_block_table(

237 + source, common.block_table_tensor, common.slot_mapping

238 + )

239 +

240 + for field, source_index in zip(fields, source_indices):

241 + actual = getattr(meta, field)

242 + torch.testing.assert_close(actual, getattr(expected, field))

243 + # The source group's indices are untouched.

244 + torch.testing.assert_close(getattr(source, field), source_index)

245 + # FULL graph state indices land in this group's own buffers.

246 + if full_cuda_graph and meta.num_prefills == 0 and actual is not None:

247 + assert actual.data_ptr() == getattr(dst, field).data_ptr()

248 +

249 +

250 def test_has_initial_state_after_reclassification():

251 """After reclassification, num_prefills > 0 so the prefill kernel path

252 should compute has_initial_state. For the reclassified request with

tests/v1/worker/test_gpu_autoregressive_speculator.py

+3 -1

@@ -98,7 +98,9 @@ def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode):

98 speculator.attn_groups = [

99 [

100 SimpleNamespace(

101 - get_metadata_builder=lambda _: SimpleNamespace(build=build),

101 + get_metadata_builder=lambda _: SimpleNamespace(

102 + build=build, supports_update_block_table=False

103 + ),

104 layer_names=["draft"],

105 )

106 ]

vllm/v1/attention/backends/gdn_attn.py

+105 -19

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Backend for GatedDeltaNet attention."""

5 -from dataclasses import dataclass

5 +from dataclasses import dataclass, replace

6 from typing import Literal

8 import torch

@@ -60,6 +60,7 @@ class GDNAttentionMetadata:

60 None # shape: [batch - num_spec_decodes,]

61 )

62 spec_sequence_masks: torch.Tensor | None = None # shape: [batch,]

63 + spec_sequence_masks_cpu: torch.Tensor | None = None # shape: [batch,]

64 spec_token_indx: torch.Tensor | None = None

65 non_spec_token_indx: torch.Tensor | None = None

@@ -114,6 +115,18 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

115 self.use_full_cuda_graph: bool = (

116 self.compilation_config.cudagraph_mode.has_full_cudagraphs()

117 )

118 + # update_block_table() keeps the source group's batch-level FULL graph

119 + # buffers, so only MRV2, which also reuses metadata at capture, may use

120 + # it.

121 + self.supports_update_block_table = (

122 + vllm_config.use_v2_model_runner

123 + and device.type == "cuda"

124 + # Not isinstance: KDA's RecoverSSM/checkpoint metadata is per group.

125 + and type(self) is GDNAttentionMetadataBuilder

126 + )

127 + if self.supports_update_block_table:

128 + # Opts into MRV2's CUDA-only aligned-index precompute.

129 + self.mamba_aligned_state_indices: torch.Tensor | None = None

130

131 self.decode_cudagraph_max_bs: int = (

132 self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1)

@@ -220,12 +233,16 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

233 query_start_loc = m.query_start_loc

234 query_start_loc_cpu = m.query_start_loc_cpu

235 nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None

223 - block_table_tensor = mamba_get_block_table_tensor(

224 - m.block_table_tensor,

225 - m.seq_lens,

226 - self.kv_cache_spec,

227 - self.vllm_config.cache_config.mamba_cache_mode,

228 - )

236 + aligned_state_indices = getattr(self, "mamba_aligned_state_indices", None)

237 + if aligned_state_indices is not None:

238 + block_table_tensor = aligned_state_indices[: m.num_reqs]

239 + else:

240 + block_table_tensor = mamba_get_block_table_tensor(

241 + m.block_table_tensor,

242 + m.seq_lens,

243 + self.kv_cache_spec,

244 + self.vllm_config.cache_config.mamba_cache_mode,

245 + )

246

247 uniform_spec_sequence_length = None

248 spec_sequence_masks_cpu: torch.Tensor | None = None

@@ -460,12 +477,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

477 # metadata below is indexed by request.

478 batch_size = m.num_reqs

479

463 - if (

464 - self.use_full_cuda_graph

465 - and num_prefills == 0

466 - and num_decodes == 0

467 - and num_spec_decodes <= self.decode_cudagraph_max_bs

468 - and num_spec_decode_tokens <= self.decode_cudagraph_max_bs

480 + if self._stage_spec_decode(

481 + num_prefills, num_decodes, num_spec_decodes, num_spec_decode_tokens

482 ):

483 assert spec_sequence_masks is not None

484 self.spec_state_indices_tensor[:num_spec_decodes].copy_(

@@ -506,12 +519,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

519 num_accepted_tokens = self.num_accepted_tokens[:batch_size]

520 num_accepted_tokens[num_spec_decodes:].fill_(1)

521

509 - if (

510 - self.use_full_cuda_graph

511 - and num_prefills == 0

512 - and num_spec_decodes == 0

513 - and num_decodes <= self.decode_cudagraph_max_bs

514 - ):

522 + if self._stage_decode(num_prefills, num_decodes, num_spec_decodes):

523 self.non_spec_state_indices_tensor[:num_decodes].copy_(

524 non_spec_state_indices_tensor, non_blocking=True

525 )

@@ -546,6 +554,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

554 spec_state_indices_tensor=spec_state_indices_tensor,

555 non_spec_state_indices_tensor=non_spec_state_indices_tensor,

556 spec_sequence_masks=spec_sequence_masks,

557 + spec_sequence_masks_cpu=spec_sequence_masks_cpu,

558 spec_token_indx=spec_token_indx,

559 non_spec_token_indx=non_spec_token_indx,

560 num_accepted_tokens=num_accepted_tokens,

@@ -556,6 +565,83 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

565 )

566 return attn_metadata

567

568 + def _stage_spec_decode(

569 + self,

570 + num_prefills: int,

571 + num_decodes: int,

572 + num_spec_decodes: int,

573 + num_spec_decode_tokens: int,

574 + ) -> bool:

575 + """Whether spec-decode metadata goes into the FULL cudagraph buffers."""

576 + return (

577 + self.use_full_cuda_graph

578 + and num_prefills == 0

579 + and num_decodes == 0

580 + and num_spec_decodes <= self.decode_cudagraph_max_bs

581 + and num_spec_decode_tokens <= self.decode_cudagraph_max_bs

582 + )

583 +

584 + def _stage_decode(

585 + self, num_prefills: int, num_decodes: int, num_spec_decodes: int

586 + ) -> bool:

587 + """Whether decode metadata goes into the FULL cudagraph buffers."""

588 + return (

589 + self.use_full_cuda_graph

590 + and num_prefills == 0

591 + and num_spec_decodes == 0

592 + and num_decodes <= self.decode_cudagraph_max_bs

593 + )

594 +

595 + def update_block_table(

596 + self,

597 + metadata: GDNAttentionMetadata,

598 + blk_table: torch.Tensor,

599 + slot_mapping: torch.Tensor,

600 + ) -> GDNAttentionMetadata:

601 + """Re-gather this group's state indices. The other fields are

602 + batch-level and stay shared with ``metadata``."""

603 + m = metadata

604 + if self.vllm_config.cache_config.mamba_cache_mode == "align":

605 + assert self.mamba_aligned_state_indices is not None

606 + blk_table = self.mamba_aligned_state_indices

607 + masks = m.spec_sequence_masks_cpu

608 + spec_indices = non_spec_indices = prefill_indices = None

609 + if masks is None:

610 + non_spec_indices = blk_table[:, 0]

611 + if m.num_prefills > 0:

612 + prefill_indices = non_spec_indices[m.num_decodes :]

613 + else:

614 + spec_indices = blk_table[masks, : self.num_spec + 1]

615 + if m.num_prefills > 0:

616 + non_spec_indices = prefill_indices = blk_table[~masks, 0]

617 +

618 + if self._stage_spec_decode(

619 + m.num_prefills, m.num_decodes, m.num_spec_decodes, m.num_spec_decode_tokens

620 + ):

621 + assert m.spec_state_indices_tensor is not None

622 + batch_size = m.spec_state_indices_tensor.shape[0]

623 + self.spec_state_indices_tensor[: m.num_spec_decodes].copy_(

624 + spec_indices, non_blocking=True

625 + )

626 + spec_indices = self.spec_state_indices_tensor[:batch_size]

627 + spec_indices[m.num_spec_decodes :].fill_(NULL_BLOCK_ID)

628 +

629 + if self._stage_decode(m.num_prefills, m.num_decodes, m.num_spec_decodes):

630 + assert m.non_spec_state_indices_tensor is not None

631 + batch_size = m.non_spec_state_indices_tensor.shape[0]

632 + self.non_spec_state_indices_tensor[: m.num_decodes].copy_(

633 + non_spec_indices, non_blocking=True

634 + )

635 + non_spec_indices = self.non_spec_state_indices_tensor[:batch_size]

636 + non_spec_indices[m.num_decodes :].fill_(NULL_BLOCK_ID)

637 +

638 + return replace(

639 + m,

640 + spec_state_indices_tensor=spec_indices,

641 + non_spec_state_indices_tensor=non_spec_indices,

642 + prefill_state_indices=prefill_indices,

643 + )

644 +

645 def build_for_cudagraph_capture(

646 self, common_attn_metadata: CommonAttentionMetadata

647 ):

vllm/v1/worker/gpu/attn_utils.py

+17 -1

@@ -29,6 +29,7 @@ from vllm.v1.kv_cache_interface import (

29 AttentionSpec,

30 KVCacheConfig,

31 KVCacheSpec,

32 + MambaSpec,

33 UniformTypeKVCacheSpecs,

34 )

35 from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata

@@ -462,6 +463,10 @@ def build_attn_metadata(

463

464 attn_metadata: dict[str, Any] = {}

465 token_to_req_indices: torch.Tensor | None = None

466 + # Mamba groups with the same spec and builder differ only in their state

467 + # indices, so later groups re-gather those from the first group's metadata.

468 + # Also at capture, so FULL graphs share the batch-level buffers.

469 + cached_metadata: dict[tuple[KVCacheSpec, type], Any] = {}

470 num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)

471 for i in range(num_kv_cache_groups):

472 if not attn_groups[i]:

@@ -514,7 +519,16 @@ def build_attn_metadata(

519

520 for attn_group in attn_groups[i]:

521 attn_metadata_builder = attn_group.get_metadata_builder(ubatch_idx)

517 - if for_cudagraph_capture:

522 + reuse_key = None

523 + if attn_metadata_builder.supports_update_block_table and isinstance(

524 + attn_group.kv_cache_spec, MambaSpec

525 + ):

526 + reuse_key = (attn_group.kv_cache_spec, type(attn_metadata_builder))

527 + if reuse_key in cached_metadata:

528 + metadata = attn_metadata_builder.update_block_table(

529 + cached_metadata[reuse_key], block_table, slot_mapping

530 + )

531 + elif for_cudagraph_capture:

532 metadata = attn_metadata_builder.build_for_cudagraph_capture(

533 common_attn_metadata

534 )

@@ -532,6 +546,8 @@ def build_attn_metadata(

546 common_attn_metadata=common_attn_metadata,

547 **attn_metadata_extra_kwargs,

548 )

549 + if reuse_key is not None and reuse_key not in cached_metadata:

550 + cached_metadata[reuse_key] = metadata

551 for layer_name in attn_group.layer_names:

552 attn_metadata[layer_name] = metadata

553 token_to_req_indices = common_attn_metadata._token_to_req_indices_cache