iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
vLLM-4x-RTX-3090/patches/0064-mamba-reuse-Mamba-GDN-metadata-across-KV-cache-group.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

363 lines
16 KiB
Diff
Raw Permalink Blame History

From 26f4ff1480ab7307aaad9d27309fb5f5a85061e3 Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Mon, 28 Sep 2026 06:36:58 +0200
Subject: [PATCH 064/103] mamba: reuse Mamba/GDN metadata across KV cache
groups (port of upstream #58762)
Port of the upstream vLLM change #58762 (Apache-2.0). Mamba/GDN attention metadata is built once and reused by groups that share it.
Test-only conflicts with upstream changes not in v0.30.0 were resolved by taking only this change's tests.
---
tests/models/kimi_k3/test_kda_metadata.py | 1 +
.../v1/attention/test_gdn_metadata_builder.py | 63 +++++++++
vllm/v1/attention/backends/gdn_attn.py | 124 +++++++++++++++---
vllm/v1/worker/gpu/attn_utils.py | 18 ++-
4 files changed, 186 insertions(+), 20 deletions(-)
diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py
index c3f4238..8ff7aff 100644
--- a/tests/models/kimi_k3/test_kda_metadata.py
+++ b/tests/models/kimi_k3/test_kda_metadata.py
@@ -49,6 +49,7 @@ PRUNED_METADATA_FIELDS = {
"prefill_state_indices",
"prefill_has_initial_state",
"spec_sequence_masks",
+ "spec_sequence_masks_cpu",
"flashinfer_prefill_query_start_loc",
"flashinfer_prefill_seq_order",
}
diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py
index b4d3acf..44dee49 100644
--- a/tests/v1/attention/test_gdn_metadata_builder.py
+++ b/tests/v1/attention/test_gdn_metadata_builder.py
@@ -21,6 +21,7 @@ from vllm.v1.attention.backends.gdn_attn import (
GDNAttentionMetadata,
GDNAttentionMetadataBuilder,
)
+from vllm.v1.attention.backends.utils import mamba_get_block_table_tensor
from vllm.v1.kv_cache_interface import MambaSpec
BLOCK_SIZE = 16
@@ -75,6 +76,17 @@ GDN_BUILD_TEST_CASES = {
expected_num_prefill_tokens=0,
expected_num_spec_decodes=0,
),
+ # No speculative config, decode alongside prefill
+ "regular_decode_with_prefill": GDNBuildTestCase(
+ seq_lens=[40, 100],
+ query_lens=[1, 50],
+ num_decode_draft_tokens=None,
+ num_speculative_tokens=0,
+ expected_num_decodes=1,
+ expected_num_prefills=1,
+ expected_num_prefill_tokens=50,
+ expected_num_spec_decodes=0,
+ ),
# Multi-token prefill alongside spec decode — no decode to reclassify
"spec_decode_with_real_prefill": GDNBuildTestCase(
seq_lens=[100, 20],
@@ -155,9 +167,12 @@ def _build(
builder: GDNAttentionMetadataBuilder,
batch_spec: BatchSpec,
num_decode_draft_tokens: list[int] | None = None,
+ block_table: torch.Tensor | None = None,
) -> GDNAttentionMetadata:
"""Build GDN attention metadata, optionally with spec-decode kwargs."""
common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE)
+ if block_table is not None:
+ common = common.replace(block_table_tensor=block_table)
kwargs: dict = {}
if num_decode_draft_tokens is not None:
kwargs["num_decode_draft_tokens_cpu"] = torch.tensor(
@@ -184,6 +199,54 @@ def test_gdn_build_classification(test_case: GDNBuildTestCase):
assert meta.num_spec_decodes == test_case.expected_num_spec_decodes
+@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])
+@pytest.mark.parametrize("full_cuda_graph", [False, True])
+@pytest.mark.parametrize(
+ "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys()
+)
+def test_update_block_table_matches_build(
+ test_case: GDNBuildTestCase, full_cuda_graph: bool, mamba_cache_mode: str
+):
+ """update_block_table() on another group's metadata matches build()."""
+ batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens)
+ src, dst, ref = (
+ _create_gdn_builder(test_case.num_speculative_tokens, full_cuda_graph)
+ for _ in range(3)
+ )
+ for builder in (src, dst, ref):
+ builder.vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
+ common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE)
+ if mamba_cache_mode == "align":
+ # MRV2 precomputes these for every Mamba group each step.
+ dst.mamba_aligned_state_indices = ref.mamba_aligned_state_indices = (
+ mamba_get_block_table_tensor(
+ common.block_table_tensor, common.seq_lens, ref.kv_cache_spec, "align"
+ )
+ )
+ draft_tokens = test_case.num_decode_draft_tokens
+ expected = _build(ref, batch, draft_tokens, common.block_table_tensor)
+ source = _build(src, batch, draft_tokens)
+ fields = (
+ "spec_state_indices_tensor",
+ "non_spec_state_indices_tensor",
+ "prefill_state_indices",
+ )
+ source_indices = [getattr(source, f) for f in fields]
+ source_indices = [t if t is None else t.clone() for t in source_indices]
+ meta = dst.update_block_table(
+ source, common.block_table_tensor, common.slot_mapping
+ )
+
+ for field, source_index in zip(fields, source_indices):
+ actual = getattr(meta, field)
+ torch.testing.assert_close(actual, getattr(expected, field))
+ # The source group's indices are untouched.
+ torch.testing.assert_close(getattr(source, field), source_index)
+ # FULL graph state indices land in this group's own buffers.
+ if full_cuda_graph and meta.num_prefills == 0 and actual is not None:
+ assert actual.data_ptr() == getattr(dst, field).data_ptr()
+
+
def test_has_initial_state_after_reclassification():
"""After reclassification, num_prefills > 0 so the prefill kernel path
should compute has_initial_state. For the reclassified request with
diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py
index dfc04a8..f11e2a9 100644
--- a/vllm/v1/attention/backends/gdn_attn.py
+++ b/vllm/v1/attention/backends/gdn_attn.py
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Backend for GatedDeltaNet attention."""
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
from typing import Literal
import torch
@@ -60,6 +60,7 @@ class GDNAttentionMetadata:
None # shape: [batch - num_spec_decodes,]
)
spec_sequence_masks: torch.Tensor | None = None # shape: [batch,]
+ spec_sequence_masks_cpu: torch.Tensor | None = None # shape: [batch,]
spec_token_indx: torch.Tensor | None = None
non_spec_token_indx: torch.Tensor | None = None
@@ -113,6 +114,18 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
self.use_full_cuda_graph: bool = (
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
)
+ # update_block_table() keeps the source group's batch-level FULL graph
+ # buffers, so only MRV2, which also reuses metadata at capture, may use
+ # it.
+ self.supports_update_block_table = (
+ vllm_config.use_v2_model_runner
+ and device.type == "cuda"
+ # Not isinstance: KDA's RecoverSSM/checkpoint metadata is per group.
+ and type(self) is GDNAttentionMetadataBuilder
+ )
+ if self.supports_update_block_table:
+ # Opts into MRV2's CUDA-only aligned-index precompute.
+ self.mamba_aligned_state_indices: torch.Tensor | None = None
self.decode_cudagraph_max_bs: int = (
self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1)
@@ -219,12 +232,16 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
query_start_loc = m.query_start_loc
query_start_loc_cpu = m.query_start_loc_cpu
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
- block_table_tensor = mamba_get_block_table_tensor(
- m.block_table_tensor,
- m.seq_lens,
- self.kv_cache_spec,
- self.vllm_config.cache_config.mamba_cache_mode,
- )
+ aligned_state_indices = getattr(self, "mamba_aligned_state_indices", None)
+ if aligned_state_indices is not None:
+ block_table_tensor = aligned_state_indices[: m.num_reqs]
+ else:
+ block_table_tensor = mamba_get_block_table_tensor(
+ m.block_table_tensor,
+ m.seq_lens,
+ self.kv_cache_spec,
+ self.vllm_config.cache_config.mamba_cache_mode,
+ )
spec_sequence_masks_cpu: torch.Tensor | None = None
if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
@@ -451,12 +468,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
# metadata below is indexed by request.
batch_size = m.num_reqs
- if (
- self.use_full_cuda_graph
- and num_prefills == 0
- and num_decodes == 0
- and num_spec_decodes <= self.decode_cudagraph_max_bs
- and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
+ if self._stage_spec_decode(
+ num_prefills, num_decodes, num_spec_decodes, num_spec_decode_tokens
):
assert spec_sequence_masks is not None
self.spec_state_indices_tensor[:num_spec_decodes].copy_(
@@ -497,12 +510,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
num_accepted_tokens = self.num_accepted_tokens[:batch_size]
num_accepted_tokens[num_spec_decodes:].fill_(1)
- if (
- self.use_full_cuda_graph
- and num_prefills == 0
- and num_spec_decodes == 0
- and num_decodes <= self.decode_cudagraph_max_bs
- ):
+ if self._stage_decode(num_prefills, num_decodes, num_spec_decodes):
self.non_spec_state_indices_tensor[:num_decodes].copy_(
non_spec_state_indices_tensor, non_blocking=True
)
@@ -537,6 +545,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
spec_state_indices_tensor=spec_state_indices_tensor,
non_spec_state_indices_tensor=non_spec_state_indices_tensor,
spec_sequence_masks=spec_sequence_masks,
+ spec_sequence_masks_cpu=spec_sequence_masks_cpu,
spec_token_indx=spec_token_indx,
non_spec_token_indx=non_spec_token_indx,
num_accepted_tokens=num_accepted_tokens,
@@ -546,6 +555,83 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
)
return attn_metadata
+ def _stage_spec_decode(
+ self,
+ num_prefills: int,
+ num_decodes: int,
+ num_spec_decodes: int,
+ num_spec_decode_tokens: int,
+ ) -> bool:
+ """Whether spec-decode metadata goes into the FULL cudagraph buffers."""
+ return (
+ self.use_full_cuda_graph
+ and num_prefills == 0
+ and num_decodes == 0
+ and num_spec_decodes <= self.decode_cudagraph_max_bs
+ and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
+ )
+
+ def _stage_decode(
+ self, num_prefills: int, num_decodes: int, num_spec_decodes: int
+ ) -> bool:
+ """Whether decode metadata goes into the FULL cudagraph buffers."""
+ return (
+ self.use_full_cuda_graph
+ and num_prefills == 0
+ and num_spec_decodes == 0
+ and num_decodes <= self.decode_cudagraph_max_bs
+ )
+
+ def update_block_table(
+ self,
+ metadata: GDNAttentionMetadata,
+ blk_table: torch.Tensor,
+ slot_mapping: torch.Tensor,
+ ) -> GDNAttentionMetadata:
+ """Re-gather this group's state indices. The other fields are
+ batch-level and stay shared with ``metadata``."""
+ m = metadata
+ if self.vllm_config.cache_config.mamba_cache_mode == "align":
+ assert self.mamba_aligned_state_indices is not None
+ blk_table = self.mamba_aligned_state_indices
+ masks = m.spec_sequence_masks_cpu
+ spec_indices = non_spec_indices = prefill_indices = None
+ if masks is None:
+ non_spec_indices = blk_table[:, 0]
+ if m.num_prefills > 0:
+ prefill_indices = non_spec_indices[m.num_decodes :]
+ else:
+ spec_indices = blk_table[masks, : self.num_spec + 1]
+ if m.num_prefills > 0:
+ non_spec_indices = prefill_indices = blk_table[~masks, 0]
+
+ if self._stage_spec_decode(
+ m.num_prefills, m.num_decodes, m.num_spec_decodes, m.num_spec_decode_tokens
+ ):
+ assert m.spec_state_indices_tensor is not None
+ batch_size = m.spec_state_indices_tensor.shape[0]
+ self.spec_state_indices_tensor[: m.num_spec_decodes].copy_(
+ spec_indices, non_blocking=True
+ )
+ spec_indices = self.spec_state_indices_tensor[:batch_size]
+ spec_indices[m.num_spec_decodes :].fill_(NULL_BLOCK_ID)
+
+ if self._stage_decode(m.num_prefills, m.num_decodes, m.num_spec_decodes):
+ assert m.non_spec_state_indices_tensor is not None
+ batch_size = m.non_spec_state_indices_tensor.shape[0]
+ self.non_spec_state_indices_tensor[: m.num_decodes].copy_(
+ non_spec_indices, non_blocking=True
+ )
+ non_spec_indices = self.non_spec_state_indices_tensor[:batch_size]
+ non_spec_indices[m.num_decodes :].fill_(NULL_BLOCK_ID)
+
+ return replace(
+ m,
+ spec_state_indices_tensor=spec_indices,
+ non_spec_state_indices_tensor=non_spec_indices,
+ prefill_state_indices=prefill_indices,
+ )
+
def build_for_cudagraph_capture(
self, common_attn_metadata: CommonAttentionMetadata
):
diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py
index 89f54d4..cdf0eb8 100644
--- a/vllm/v1/worker/gpu/attn_utils.py
+++ b/vllm/v1/worker/gpu/attn_utils.py
@@ -29,6 +29,7 @@ from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheConfig,
KVCacheSpec,
+ MambaSpec,
UniformTypeKVCacheSpecs,
)
from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
@@ -427,6 +428,10 @@ def build_attn_metadata(
attn_metadata: dict[str, Any] = {}
token_to_req_indices: torch.Tensor | None = None
+ # Mamba groups with the same spec and builder differ only in their state
+ # indices, so later groups re-gather those from the first group's metadata.
+ # Also at capture, so FULL graphs share the batch-level buffers.
+ cached_metadata: dict[tuple[KVCacheSpec, type], Any] = {}
num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
for i in range(num_kv_cache_groups):
if not attn_groups[i]:
@@ -479,7 +484,16 @@ def build_attn_metadata(
for attn_group in attn_groups[i]:
attn_metadata_builder = attn_group.get_metadata_builder(ubatch_idx)
- if for_cudagraph_capture:
+ reuse_key = None
+ if attn_metadata_builder.supports_update_block_table and isinstance(
+ attn_group.kv_cache_spec, MambaSpec
+ ):
+ reuse_key = (attn_group.kv_cache_spec, type(attn_metadata_builder))
+ if reuse_key in cached_metadata:
+ metadata = attn_metadata_builder.update_block_table(
+ cached_metadata[reuse_key], block_table, slot_mapping
+ )
+ elif for_cudagraph_capture:
metadata = attn_metadata_builder.build_for_cudagraph_capture(
common_attn_metadata
)
@@ -497,6 +511,8 @@ def build_attn_metadata(
common_attn_metadata=common_attn_metadata,
**attn_metadata_extra_kwargs,
)
+ if reuse_key is not None and reuse_key not in cached_metadata:
+ cached_metadata[reuse_key] = metadata
for layer_name in attn_group.layer_names:
attn_metadata[layer_name] = metadata
token_to_req_indices = common_attn_metadata._token_to_req_indices_cache
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink