iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Luca Motz
2026-09-28 04:36:58 +00:00
committed by GitHub
co-authored by Xin Yang Claude Opus 5.5 Codex Lucas Wilkinson
parent 55de40a2fc
commit dd3bc9c72f
5 changed files with 189 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/kimi_k3/test_kda_metadata.py
+1
View File
@@ -49,6 +49,7 @@ PRUNED_METADATA_FIELDS = {
"prefill_state_indices",
"prefill_has_initial_state",
"spec_sequence_masks",
"spec_sequence_masks_cpu",
"uniform_spec_sequence_length",
"flashinfer_prefill_query_start_loc",
"flashinfer_prefill_seq_order",
tests/v1/attention/test_gdn_metadata_builder.py
+63
View File
@@ -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
tests/v1/worker/test_gpu_autoregressive_speculator.py
+3 -1
View File
@@ -98,7 +98,9 @@ def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode):
speculator.attn_groups = [
[
SimpleNamespace(
get_metadata_builder=lambda _: SimpleNamespace(build=build),
get_metadata_builder=lambda _: SimpleNamespace(
build=build, supports_update_block_table=False
),
layer_names=["draft"],
)
]
vllm/v1/attention/backends/gdn_attn.py
+105 -19
View File
@@ -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
@@ -114,6 +115,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)
@@ -220,12 +233,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,
)
uniform_spec_sequence_length = None
spec_sequence_masks_cpu: torch.Tensor | None = None
@@ -460,12 +477,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_(
@@ -506,12 +519,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
)
@@ -546,6 +554,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,
@@ -556,6 +565,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
):
vllm/v1/worker/gpu/attn_utils.py
+17 -1
View File
@@ -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
@@ -462,6 +463,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]:
@@ -514,7 +519,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
)
@@ -532,6 +546,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