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

[Perf][GDN] Internal prefill checkpoints for the FlashInfer GDN prefill backend (#60050)

Browse Source 
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Jiangyun Zhu
2026-10-14 14:49:46 +08:00
committed by GitHub
co-authored by Claude Opus 5.5
parent a2fb482732
commit badfcdbf91
12 changed files with 794 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/engine.yaml
+1 -1
View File
@@ -141,7 +141,7 @@ steps:
- tests/v1/e2e/general/
commands:
- pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
- pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async
- pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async_flashinfer
mirror:
amd:
label: ":amd: (MI355 DPX) E2E Core Large Memory"
.buildkite/test_areas/kernels.yaml
+2
View File
@@ -513,8 +513,10 @@ steps:
source_file_dependencies:
- csrc/mamba/
- tests/kernels/mamba
- vllm/model_executor/layers/mamba/checkpoint.py
- vllm/model_executor/layers/mamba/gdn
- vllm/model_executor/layers/mamba/ops
- vllm/v1/attention/backends/gdn_attn.py
commands:
- pytest -v -s kernels/mamba
mirror:
tests/kernels/mamba/test_gdn_prefill_flashinfer.py
+136
View File
@@ -13,10 +13,19 @@ if current_platform.is_rocm():
)
import flashinfer.gdn_prefill # noqa: E402
import torch.nn.functional as F # noqa: E402
from vllm.model_executor.layers.mamba.checkpoint import ( # noqa: E402
MambaPrefillCheckpointExporter,
MambaPrefillCheckpointMetadata,
)
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
fi_chunk_gated_delta_rule,
) # noqa: E402
from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402
plan_flashinfer_checkpoints,
)
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID # noqa: E402
def test_flashinfer_gdn_prefill_uses_int64_cu_seqlens(monkeypatch):
@@ -51,3 +60,130 @@ def test_flashinfer_gdn_prefill_uses_int64_cu_seqlens(monkeypatch):
assert captured_cu_seqlens.dtype == torch.int64
assert output.shape == q.shape
assert final_state is None
@pytest.mark.skipif(
not (
current_platform.is_device_capability(90)
or current_platform.is_device_capability_family(100)
or current_platform.is_device_capability_family(120)
),
reason="FlashInfer GDN prefill checkpoints need SM90, SM100 or SM120.",
)
@pytest.mark.parametrize("use_cp", [True, False])
@pytest.mark.parametrize("ssm_dtype", [torch.float32, torch.bfloat16])
def test_flashinfer_checkpoints_export_the_state_at_each_offset(use_cp, ssm_dtype):
"""Each exported checkpoint is the state of the prefill stopped there.
The leading decode row is outside the prefill kernel. The two checkpointed
prefills put FlashInfer's stride (their offsets' gcd) below both offsets,
so the exporter must pick each state out of the shared emitted buffer;
the third prefill takes no checkpoint. Both FlashInfer routes are covered
because ``use_cp="auto"`` may take either. The conv cache has MTP's spec
slots after the window, which the conv kernels leave alone.
"""
torch.manual_seed(0)
device = "cuda"
num_qk_heads, num_v_heads, head_dim, conv_dim, state_len = 2, 4, 128, 16, 3
num_spec = 2
offsets = [0, 128, 0, 192]
query_lens = [200, 50, 300]
state_indices = [NULL_BLOCK_ID, 3, NULL_BLOCK_ID, 5]
num_blocks = 8
num_tokens = sum(query_lens)
q = F.normalize(
torch.randn(1, num_tokens, num_qk_heads, head_dim, device=device), dim=-1
).bfloat16()
k = F.normalize(
torch.randn(1, num_tokens, num_qk_heads, head_dim, device=device), dim=-1
).bfloat16()
v = torch.randn(1, num_tokens, num_v_heads, head_dim, device=device).bfloat16()
g = F.logsigmoid(torch.randn(1, num_tokens, num_v_heads, device=device))
beta = torch.sigmoid(torch.randn(1, num_tokens, num_v_heads, device=device))
initial_state = 0.1 * torch.randn(
len(query_lens), num_v_heads, head_dim, head_dim, device=device
)
prefill_cu = [0]
for query_len in query_lens:
prefill_cu.append(prefill_cu[-1] + query_len)
every_n_tokens, cu_starts, rows = plan_flashinfer_checkpoints(offsets, query_lens)
state_checkpoints = torch.empty(
cu_starts[-1], num_v_heads, head_dim, head_dim, device=device
)
fi_chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state=initial_state,
output_final_state=True,
cu_seqlens=torch.tensor(prefill_cu, device=device),
use_qk_l2norm_in_kernel=False,
state_checkpoints=state_checkpoints,
checkpoint_cu_starts=torch.tensor(cu_starts, device=device),
checkpoint_every_n_tokens=every_n_tokens,
use_cp=use_cp,
)
# The exporter sees the decode row's token in front of the prefills.
raw_qkv = torch.randn(1 + num_tokens, conv_dim, device=device).bfloat16()
cu_seqlens = torch.tensor([0] + [1 + n for n in prefill_cu], device=device)
conv_state = torch.zeros(
num_blocks,
state_len + num_spec,
conv_dim,
device=device,
dtype=torch.bfloat16,
).transpose(-1, -2)
ssm_state = torch.zeros(
num_blocks, num_v_heads, head_dim, head_dim, device=device, dtype=ssm_dtype
)
MambaPrefillCheckpointExporter(state_len=state_len).export(
MambaPrefillCheckpointMetadata(
checkpoint_offsets=torch.tensor(offsets, device=device, dtype=torch.int32),
state_indices=torch.tensor(state_indices, device=device),
),
raw_qkv=raw_qkv,
conv_state=conv_state,
recurrent_checkpoint=state_checkpoints,
recurrent_state=ssm_state,
cu_seqlens=cu_seqlens,
recurrent_rows=torch.tensor(rows, device=device),
)
written = []
for row, (offset, block) in enumerate(zip(offsets, state_indices)):
if not offset:
continue
written.append(block)
seq = row - 1
start = prefill_cu[seq]
_, expected_state = fi_chunk_gated_delta_rule(
q=q[:, start : start + offset],
k=k[:, start : start + offset],
v=v[:, start : start + offset],
g=g[:, start : start + offset],
beta=beta[:, start : start + offset],
initial_state=initial_state[seq : seq + 1],
output_final_state=True,
cu_seqlens=torch.tensor([0, offset], device=device),
use_qk_l2norm_in_kernel=False,
)
torch.testing.assert_close(
ssm_state[block].float(), expected_state[0], atol=2e-2, rtol=2e-2
)
end = cu_seqlens[row].item() + offset
torch.testing.assert_close(
conv_state[block, :, :state_len],
raw_qkv[end - state_len : end].T,
atol=0,
rtol=0,
)
assert not conv_state[block, :, state_len:].any()
untouched = [i for i in range(num_blocks) if i not in written]
assert not ssm_state[untouched].any()
assert not conv_state[untouched].any()
tests/v1/attention/test_gdn_metadata_builder.py
+18
View File
@@ -21,6 +21,7 @@ from vllm.v1.attention.backend import CommonAttentionMetadata
from vllm.v1.attention.backends.gdn_attn import (
GDNAttentionMetadata,
GDNAttentionMetadataBuilder,
plan_flashinfer_checkpoints,
)
from vllm.v1.attention.backends.utils import mamba_get_block_table_tensor
from vllm.v1.kv_cache_interface import MambaSpec
@@ -488,3 +489,20 @@ def test_update_block_table_regathers_checkpoint(num_spec):
assert not torch.equal(
actual.checkpoint.state_indices, source.checkpoint.state_indices
)
def test_plan_flashinfer_checkpoints_selects_the_state_at_each_offset():
"""Each checkpoint row maps to the state FlashInfer emits at its offset.
The leading decode row is not in the prefill kernel and never checkpoints;
the 50-token prefill is shorter than the stride and emits no states.
"""
every_n_tokens, cu_starts, rows = plan_flashinfer_checkpoints(
offsets=[0, 128, 0, 192], query_lens=[200, 50, 300]
)
assert every_n_tokens == 64
assert cu_starts == [0, 3, 3, 7]
# Row 1: state after 128 tokens of the first prefill. Row 3: after 192
# tokens of the third prefill, whose states start at index 3.
assert rows == [0, 1, 0, 5]
tests/v1/e2e/general/test_mamba_prefix_cache.py
+154 -3
View File
@@ -16,7 +16,10 @@ from tests.utils import create_new_process_for_each_test
from vllm import LLM, SamplingParams, TokensPrompt
from vllm.config import CacheConfig
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
from vllm.model_executor.layers.mamba.mamba_utils import (
MambaStateCopyFunc,
is_conv_state_dim_first,
)
from vllm.platforms import current_platform
from vllm.sequence import IntermediateTensors
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
@@ -59,6 +62,10 @@ MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
# where block boundaries fall. XPU rounds GDN block sizes up to a multiple of 64
# (XPUPlatform.update_block_size_for_backend), so use one.
BLOCK_SIZE = 576
# The step actions also encode the split at the prompt's last block boundary,
# which a GDN prefill backend with internal checkpoints (FlashInfer) skips; see
# get_internal_checkpoint_step_configs.
GDN_PREFILL_CONFIG = {"gdn_prefill_backend": "triton"}
DEVICE_TYPE = current_platform.device_type
NUM_HIDDEN_LAYERS = 1
cur_step_action_idx = 0
@@ -882,6 +889,73 @@ def get_mamba_prefix_cache_step_configs(
return tests
def get_internal_checkpoint_step_configs() -> dict[str, TestConfig]:
"""Async-scheduling steps when the GDN prefill exports the checkpoint at
the prompt's last block boundary itself, so the final chunk runs whole.
The checkpoint goes in the second-to-last block the chunk reaches, keyed to
the boundary: column 1 holds the state at B for a 2B + 10 prompt.
"""
B = BLOCK_SIZE
return {
"prompt_2_block_size": TestConfig(
num_prompt_tokens=B * 2,
num_generated_tokens=10,
num_accepted_tokens=4,
step_actions=[
StepAction(0, B * 2, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
StepAction(B * 2, 4, [1, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
],
),
"prompt_2_block_size_10": TestConfig(
num_prompt_tokens=B * 2 + 10,
num_generated_tokens=10,
num_accepted_tokens=4,
step_actions=[
StepAction(0, B * 2 + 10, [0, 1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
StepAction(B * 2 + 10, 4, [0, 1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
],
),
"prompt_10_block_size": TestConfig(
num_prompt_tokens=B * 10,
num_generated_tokens=10,
num_accepted_tokens=4,
step_actions=[
StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
StepAction(
B * 5,
B * 5,
[0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1],
(4, 9),
(-1, -1),
),
StepAction(
B * 10,
4,
[0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1],
(9, 10),
(-1, -1),
),
],
),
"prompt_10_block_size_10": TestConfig(
num_prompt_tokens=B * 10 + 10,
num_generated_tokens=10,
num_accepted_tokens=4,
step_actions=[
StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
StepAction(
B * 5,
B * 5 + 10,
[0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1],
(4, 10),
(-1, -1),
),
],
),
}
def _run_mamba_prefix_cache_mrv1(
monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
):
@@ -902,6 +976,7 @@ def _run_mamba_prefix_cache_mrv1(
load_format="dummy",
enable_prefix_caching=True,
block_size=BLOCK_SIZE,
additional_config=GDN_PREFILL_CONFIG,
speculative_config={
"method": "qwen3_next_mtp",
"num_speculative_tokens": num_speculative_tokens,
@@ -957,8 +1032,11 @@ def test_mamba_prefix_cache_mrv1_async(monkeypatch: pytest.MonkeyPatch):
def _run_mamba_prefix_cache_mrv2(
monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
monkeypatch: pytest.MonkeyPatch,
async_scheduling: bool,
additional_config: dict[str, Any] = GDN_PREFILL_CONFIG,
):
internal_checkpoints = additional_config["gdn_prefill_backend"] == "flashinfer"
global async_scheduling_mode
async_scheduling_mode = async_scheduling
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
@@ -1098,8 +1176,24 @@ def _run_mamba_prefix_cache_mrv2(
with gpu_sync_allowed():
start_pos = input_batch.positions[input_batch.query_start_loc[0]].item()
assert cur_step_action.num_computed_tokens_start == start_pos
if scheduler_output.scheduled_new_reqs and "snapshot_col" in captured:
captured["snapshot"] = block_states(self, captured.pop("snapshot_col"))
return ret
def block_states(runner: MRV2GPUModelRunner, col: int) -> list[list[torch.Tensor]]:
# Each Mamba layer's (conv, temporal) states in block column col.
kv_cache_config = captured["kv_cache_config"]
forward_context = runner.vllm_config.compilation_config.static_forward_context
return [
[
temporal_block(state, captured["block_tables"][group_id], col).clone()
for state in forward_context[layer_name].kv_cache
]
for group_ids in get_mamba_groups(kv_cache_config).values()
for group_id in group_ids
for layer_name in kv_cache_config.kv_cache_groups[group_id].layer_names
]
def fake_sample(
self: MRV2GPUModelRunner,
hidden_states: torch.Tensor,
@@ -1161,6 +1255,7 @@ def _run_mamba_prefix_cache_mrv2(
skip_tokenizer_init=True,
enable_prefix_caching=True,
block_size=BLOCK_SIZE,
additional_config=additional_config,
speculative_config={
"method": "qwen3_next_mtp",
"num_speculative_tokens": num_speculative_tokens,
@@ -1173,7 +1268,11 @@ def _run_mamba_prefix_cache_mrv2(
)
try:
tests = get_mamba_prefix_cache_step_configs(async_scheduling)
tests = (
get_internal_checkpoint_step_configs()
if internal_checkpoints
else get_mamba_prefix_cache_step_configs(async_scheduling)
)
global step_actions
global cur_step_action_idx
@@ -1219,6 +1318,35 @@ def _run_mamba_prefix_cache_mrv2(
f"got events={second_events!r}"
)
assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
if internal_checkpoints:
# The checkpoint a 2B + 10 prefill exports at B (column 1) is the
# final state of a prefill that stops at B (column 0). Distinct
# tokens make a shifted conv window visible.
tokens = [(i * 7919) % 30000 + 1 for i in range(BLOCK_SIZE * 2 + 10)]
scheduler = engine.llm_engine.engine_core.engine_core.scheduler
states = []
for num_tokens, col in ((BLOCK_SIZE, 0), (len(tokens), 1)):
captured["snapshot_col"] = col
_ = engine.generate(
[TokensPrompt(prompt_token_ids=tokens[:num_tokens])],
sampling_params=sampling_params,
)
states.append(captured.pop("snapshot"))
assert scheduler.reset_prefix_cache()
hf_config = engine.llm_engine.model_config.hf_text_config
state_len = hf_config.linear_conv_kernel_dim - 1
for (ref_conv, ref_ssm), (conv, ssm) in zip(*states):
if not is_conv_state_dim_first():
ref_conv, conv = ref_conv.mT, conv.mT
# Under MTP the conv state has num_spec scratch columns after
# the window.
for new, ref in (
(conv[..., :state_len], ref_conv[..., :state_len]),
(ssm, ref_ssm),
):
rel_err = (new.float() - ref.float()).norm() / ref.float().norm()
assert rel_err < 1e-2, rel_err
finally:
del engine
torch.accelerator.empty_cache()
@@ -1233,3 +1361,26 @@ def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch):
@create_new_process_for_each_test()
def test_mamba_prefix_cache_mrv2_async(monkeypatch: pytest.MonkeyPatch):
_run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=True)
@pytest.mark.skipif(
not current_platform.is_cuda()
or not (
current_platform.is_device_capability(90)
or (
(
current_platform.is_device_capability_family(100)
or current_platform.is_device_capability_family(120)
)
and current_platform.get_cuda_runtime_major() >= 13
)
),
reason="FlashInfer GDN prefill needs SM90, or SM100/SM120 with CUDA 13.",
)
@create_new_process_for_each_test()
def test_mamba_prefix_cache_mrv2_async_flashinfer(monkeypatch: pytest.MonkeyPatch):
_run_mamba_prefix_cache_mrv2(
monkeypatch,
async_scheduling=True,
additional_config={"gdn_prefill_backend": "flashinfer"},
)
tests/v1/worker/test_gpu_model_runner.py
+52 -1
View File
@@ -29,6 +29,7 @@ from vllm.distributed.parallel_state import (
from vllm.lora.layers import LoRAMappingType
from vllm.lora.request import LoRARequest
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.mamba.abstract import MambaBase
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
from vllm.platforms import current_platform
@@ -42,7 +43,10 @@ from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
ROCMAiterMLASparseBackend,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.attention.backends.registry import (
AttentionBackendEnum,
MambaAttentionBackendEnum,
)
from vllm.v1.core.kv_cache_planning import (
estimate_max_model_len,
get_kv_cache_configs,
@@ -512,6 +516,53 @@ def test_alignment_rejected_by_a_sibling_backend_raises(monkeypatch, cpu):
platform.update_block_size_for_backend(vllm_config)
class _CheckpointingMambaLayer(MambaBase):
mamba_type = MambaAttentionBackendEnum.GDN_ATTN
def __init__(self, prefill_checkpoint_alignment: int | None):
self._prefill_checkpoint_alignment = prefill_checkpoint_alignment
def get_prefill_checkpoint_alignment(self) -> int | None:
return self._prefill_checkpoint_alignment
def get_state_shape(self):
return ()
def get_state_dtype(self):
return ()
def _align_qwen3_5_blocks(*prefill_checkpoint_alignments: int | None):
from tests.v1.attention.utils import create_vllm_config
vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B")
cache_config = vllm_config.cache_config
cache_config.mamba_cache_mode = "align"
forward_context = vllm_config.compilation_config.static_forward_context
for i, alignment in enumerate(prefill_checkpoint_alignments):
forward_context[f"mamba.{i}"] = _CheckpointingMambaLayer(alignment)
Platform._align_hybrid_block_size(vllm_config, _mock_backend([MultipleOf(16)]))
return cache_config
@pytest.mark.parametrize(
("prefill_checkpoint_alignments", "block_size"),
[((64,), 576), ((None,), 544), ((64, None), 544)],
)
def test_blocks_land_on_the_mamba_checkpoint_grid(
prefill_checkpoint_alignments, block_size
):
"""A layer that checkpoints every 64 tokens needs 64-aligned blocks.
Unaligned, Qwen3.5-0.8B's block (544) puts every other checkpoint off the
grid. A Mamba layer that cannot checkpoint (like Qwen4Exp's PLE short conv)
stops the scheduler from taking any, so the block stays.
"""
cache_config = _align_qwen3_5_blocks(*prefill_checkpoint_alignments)
assert cache_config.block_size == block_size
def test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises(monkeypatch):
# XPU's GDN rounding runs after super()'s check, so it must check again.
try:
vllm/model_executor/layers/mamba/abstract.py
+27 -2
View File
@@ -2,11 +2,11 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import abstractmethod
from collections.abc import Iterable
from math import prod
from math import lcm, prod
import torch
from vllm.config import VllmConfig
from vllm.config import VllmConfig, get_layers_from_vllm_config
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.utils.torch_utils import get_dtype_size
from vllm.v1.attention.backend import AttentionBackend
@@ -62,6 +62,13 @@ class MambaBase(AttentionLayerBase):
def get_state_dtype(self) -> tuple[torch.dtype, ...]:
pass
def get_prefill_checkpoint_alignment(self) -> int | None:
"""Token alignment of the checkpoint offsets this layer can export from
inside a prefill, if any; in align mode, the platform makes Mamba blocks
multiples of it.
"""
return None
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
mamba_block_size = vllm_config.cache_config.mamba_block_size
assert mamba_block_size is not None
@@ -86,3 +93,21 @@ class MambaBase(AttentionLayerBase):
def get_attn_backend(self) -> type[AttentionBackend]:
"""Get the attention backend class for this Mamba layer."""
return get_mamba_attn_backend(self.mamba_type)
def get_mamba_prefill_checkpoint_alignment(vllm_config: VllmConfig) -> int | None:
"""LCM of the Mamba layers' prefill checkpoint alignments.
None unless every Mamba layer declares one: the scheduler takes internal
checkpoints only when every Mamba group can export them.
"""
alignment = 1
for layer in get_layers_from_vllm_config(
vllm_config,
MambaBase, # type: ignore[type-abstract]
).values():
layer_alignment = layer.get_prefill_checkpoint_alignment()
if layer_alignment is None:
return None
alignment = lcm(alignment, layer_alignment)
return alignment
vllm/model_executor/layers/mamba/checkpoint.py
+126 -10
View File
@@ -1,12 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import ABC, abstractmethod
from dataclasses import dataclass, replace
import torch
from vllm.config import VllmConfig
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import async_tensor_h2d
from vllm.v1.attention.backend import CommonAttentionMetadata
@@ -71,9 +71,11 @@ def _gather_checkpoint_state_indices(
class MambaPrefillCheckpointMetadata:
checkpoint_offsets: torch.Tensor
state_indices: torch.Tensor
# Host copy of ``checkpoint_offsets``, for backends that lay out the
# checkpoint before any tensor exists (Mamba2's SSD chunk split).
# Host copies of ``checkpoint_offsets`` and of each row's query length, for
# backends that lay out the checkpoint before any tensor exists (Mamba2's
# SSD chunk split, FlashInfer GDN's checkpoint plan).
offsets: list[int] | None = None
query_lens: list[int] | None = None
# Gather indices into the block table, kept so that another KV cache group
# with the same spec can re-derive its own ``state_indices``.
request_rows: torch.Tensor | None = None
@@ -150,20 +152,134 @@ class MambaPrefillCheckpointBuilder:
m.block_table_tensor, request_rows_tensor, checkpoint_cols_tensor
),
offsets=checkpoint_offsets,
query_lens=query_lens,
request_rows=request_rows_tensor,
block_cols=checkpoint_cols_tensor,
)
class MambaPrefillCheckpointExporter(ABC):
"""Export a mid-prefill checkpoint into backend-specific paged states."""
@dataclass(frozen=True)
class MambaPrefillCheckpointExporter:
"""Store each checkpoint's pre-convolution window and recurrent state.
The recurrent states are rows of a buffer the prefill kernel emitted;
``recurrent_rows`` picks each checkpoint row's state out of it, and by
default row ``i`` holds checkpoint ``i``'s.
"""
state_len: int | None = None
@abstractmethod
def export(
self,
checkpoint: MambaPrefillCheckpointMetadata,
*args,
**kwargs,
*,
raw_qkv: torch.Tensor,
conv_state: torch.Tensor,
recurrent_checkpoint: torch.Tensor,
recurrent_state: torch.Tensor,
cu_seqlens: torch.Tensor,
recurrent_rows: torch.Tensor | None = None,
) -> None:
"""Write checkpoint state into the paged cache."""
raise NotImplementedError
state_len = (
self.state_len if self.state_len is not None else conv_state.shape[-1]
)
width = raw_qkv.shape[-1]
recurrent_row_size = recurrent_checkpoint[0].numel()
block_size = 256
store_cache_checkpoints_kernel[
(
checkpoint.checkpoint_offsets.numel(),
triton.cdiv(max(width * state_len, recurrent_row_size), block_size),
)
](
raw_qkv,
conv_state,
recurrent_checkpoint,
recurrent_state,
cu_seqlens,
checkpoint.checkpoint_offsets,
checkpoint.state_indices,
recurrent_rows if recurrent_rows is not None else checkpoint.state_indices,
raw_qkv.stride(0),
raw_qkv.stride(1),
conv_state.stride(0),
conv_state.stride(1),
conv_state.stride(2),
recurrent_checkpoint.stride(0),
recurrent_state.stride(0),
checkpoint.checkpoint_offsets.stride(0),
state_len,
width,
recurrent_row_size,
NULL_BLOCK_ID,
recurrent_rows is not None,
block_size,
)
@triton.jit
def store_cache_checkpoints_kernel(
x_ptr,
conv_state_ptr,
recurrent_checkpoint_ptr,
recurrent_state_ptr,
query_start_loc_ptr,
checkpoint_offsets_ptr,
checkpoint_state_indices_ptr,
recurrent_rows_ptr,
x_stride_0: tl.constexpr,
x_stride_1: tl.constexpr,
state_stride_0: tl.constexpr,
state_stride_1: tl.constexpr,
state_stride_2: tl.constexpr,
checkpoint_stride_0: tl.constexpr,
recurrent_state_stride_0: tl.constexpr,
checkpoint_offset_stride: tl.constexpr,
STATE_LEN: tl.constexpr,
WIDTH: tl.constexpr,
RECURRENT_ROW_SIZE: tl.constexpr,
NULL_STATE_IDX: tl.constexpr,
HAS_RECURRENT_ROWS: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
seq_idx = tl.program_id(0)
cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64)
checkpoint_offset = tl.load(
checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
)
valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)
valid_conv = (
(cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)
)
width_idx = cols // STATE_LEN
history_idx = cols % STATE_LEN
checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset
token_idx = checkpoint_end - STATE_LEN + history_idx
values = tl.load(
x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
mask=valid_conv,
)
tl.store(
conv_state_ptr
+ state_idx * state_stride_0
+ width_idx * state_stride_1
+ history_idx * state_stride_2,
values,
mask=valid_conv,
)
if HAS_RECURRENT_ROWS:
recurrent_row = tl.load(recurrent_rows_ptr + seq_idx).to(tl.int64)
else:
recurrent_row = seq_idx.to(tl.int64)
valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint
recurrent = tl.load(
recurrent_checkpoint_ptr + recurrent_row * checkpoint_stride_0 + cols,
mask=valid_recurrent,
)
tl.store(
recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,
recurrent,
mask=valid_recurrent,
)
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+198 -18
View File
@@ -2,7 +2,11 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Inference-only Qwen3-Next/Qwen3.5 model."""
import functools
import itertools
import math
import os
from dataclasses import replace
from typing import Literal
import torch
@@ -15,6 +19,7 @@ from vllm import envs
from vllm._aiter_ops import rocm_aiter_ops
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
from vllm.config import (
CacheConfig,
VllmConfig,
get_current_vllm_config,
)
@@ -30,6 +35,12 @@ from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
RowParallelLinear,
)
from vllm.model_executor.layers.mamba.abstract import (
get_mamba_prefill_checkpoint_alignment,
)
from vllm.model_executor.layers.mamba.checkpoint import (
MambaPrefillCheckpointExporter,
)
from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weight_loader
from vllm.model_executor.layers.mamba.mamba_utils import (
@@ -59,13 +70,19 @@ from vllm.third_party.flash_linear_attention.ops import (
)
from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd
from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import (
LayerNameType,
_encode_layer_name,
_resolve_layer_name,
direct_register_custom_op,
)
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
from vllm.v1.attention.backends.gdn_attn import (
FLASHINFER_GDN_CHECKPOINT_ALIGNMENT,
GDNAttentionMetadata,
)
from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
from vllm.v1.worker.workspace import current_workspace_manager
# Optional ROCm AITER Triton kernels for the GDN decode path.
# Availability is checked centrally via rocm_aiter_ops; the actual function
@@ -241,6 +258,10 @@ def fi_chunk_gated_delta_rule(
output_final_state: bool,
cu_seqlens: torch.Tensor | None = None,
use_qk_l2norm_in_kernel: bool = True,
state_checkpoints: torch.Tensor | None = None,
checkpoint_cu_starts: torch.Tensor | None = None,
checkpoint_every_n_tokens: int = 0,
use_cp: Literal["auto"] | bool = "auto",
):
from flashinfer.gdn_prefill import (
chunk_gated_delta_rule as chunk_gated_delta_rule_fi,
@@ -271,6 +292,10 @@ def fi_chunk_gated_delta_rule(
initial_state=fi_state,
output_final_state=output_final_state,
cu_seqlens=cu_seqlens,
state_checkpoints=state_checkpoints,
checkpoint_cu_starts=checkpoint_cu_starts,
checkpoint_every_n_tokens=checkpoint_every_n_tokens,
use_cp=use_cp,
backend="flashinfer",
)
# FlashInfer returns (output, state) when output_final_state=True,
@@ -283,6 +308,20 @@ def fi_chunk_gated_delta_rule(
return result.unsqueeze(0), None
def _checkpoint_grid(cache_config: CacheConfig) -> int:
"""Token grid that FlashInfer GDN checkpoint offsets are kept on.
Prefix-cache boundaries, made multiples of FlashInfer's checkpoint
alignment; the platform already keeps blocks on it.
"""
if cache_config.prefix_match_unit is not None:
return math.lcm(
cache_config.prefix_match_unit, FLASHINFER_GDN_CHECKPOINT_ALIGNMENT
)
assert cache_config.mamba_block_size is not None
return cache_config.mamba_block_size
@CustomOp.register("chunk_gated_delta_rule")
class ChunkGatedDeltaRule(CustomOp):
def __init__(self) -> None:
@@ -591,6 +630,21 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
self.gdn_prefill_backend = self.chunk_gated_delta_rule.gdn_prefill_backend
self._prefill_kernels_warmed_up = False
# FlashInfer's prefill emits mid-prefill states, so an align-mode
# checkpoint needs no split prefill.
self._checkpoint_alignment: int | None = None
if (
vllm_config.cache_config.mamba_cache_mode == "align"
and self.gdn_prefill_backend == "flashinfer"
):
self._checkpoint_alignment = FLASHINFER_GDN_CHECKPOINT_ALIGNMENT
self._vllm_config = vllm_config
self._max_num_batched_tokens = (
vllm_config.scheduler_config.max_num_batched_tokens
)
self._checkpoint_exporter = MambaPrefillCheckpointExporter(
state_len=self.conv_kernel_size - 1
)
self.enable_packed_recurrent_decode = (
envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
)
@@ -624,6 +678,46 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
raise ValueError(f"Duplicate layer name: {prefix}")
compilation_config.static_forward_context[prefix] = self
def get_prefill_checkpoint_alignment(self) -> int | None:
return self._checkpoint_alignment
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
spec = super().get_kv_cache_spec(vllm_config)
if not self._exports_checkpoints or not isinstance(spec, MambaSpec):
return spec
# FlashInfer emits states at one stride for the whole batch, the gcd of
# the offsets; keeping every offset on the prefix-cache grid keeps that
# stride a multiple of the grid, which bounds the states emitted.
return replace(
spec,
num_prefill_checkpoint_blocks=1,
prefill_checkpoint_alignment=_checkpoint_grid(vllm_config.cache_config),
)
@functools.cached_property
def _exports_checkpoints(self) -> bool:
# Known only once every layer is built: a Mamba layer that cannot
# checkpoint (e.g. Qwen4Exp's PLE short conv) keeps the scheduler
# splitting the prefill instead.
return (
self._checkpoint_alignment is not None
and get_mamba_prefill_checkpoint_alignment(self._vllm_config) is not None
)
@functools.cached_property
def _checkpoint_buffer_spec(self) -> tuple[tuple[int, ...], torch.dtype]:
# Offsets on the grid emit at most one state per grid step of a batch.
# Read lazily: the block size is final only after the model is built.
return (
(
cdiv(self._max_num_batched_tokens, _checkpoint_grid(self.cache_config)),
self.num_v_heads // self.tp_size,
self.head_v_dim,
self.head_k_dim,
),
torch.float32,
)
def _fused_gdn_decode_unsupported_reason(
self, vllm_config: VllmConfig
) -> str | None:
@@ -1167,6 +1261,9 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
if self._prefill_kernels_warmed_up:
return
self._prefill_kernels_warmed_up = True
if self._exports_checkpoints:
# Size the workspace while profiling; it is locked afterwards.
current_workspace_manager().get_simultaneous(self._checkpoint_buffer_spec)
device = qkv_or_qkvz.device
dtype = qkv_or_qkvz.dtype
@@ -1242,6 +1339,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
use_qk_l2norm_in_kernel=False,
aiter_prefill_metadata=aiter_prefill_metadata,
)
if self.gdn_prefill_backend == "flashinfer":
self._warmup_flashinfer_prefill_variants(q, k, v, g, beta, state)
except Exception:
logger.warning(
"GDN prefill kernel warmup (T=%d) failed for "
@@ -1276,6 +1375,50 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
torch.accelerator.empty_cache()
def _warmup_flashinfer_prefill_variants(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
state: torch.Tensor,
) -> None:
"""Compile every FlashInfer GDN prefill kernel a batch can dispatch to.
FlashInfer builds separate kernels for its context-parallel route (few
sequences) and its regular route, each with and without checkpoint
output. A single sequence under ``use_cp="auto"`` takes the CP route
wherever the heuristic can pick it; ``use_cp=False`` forces the other.
"""
checkpoint_buffers: list[torch.Tensor | None] = [None]
if self._exports_checkpoints:
(state_checkpoints,) = current_workspace_manager().get_simultaneous(
self._checkpoint_buffer_spec
)
checkpoint_buffers.append(state_checkpoints[:1])
cu_seqlens = torch.tensor([0, q.shape[1]], device=q.device, dtype=torch.int32)
checkpoint_cu_starts = torch.tensor([0, 1], device=q.device)
for use_cp, state_checkpoints in itertools.product(
("auto", False), checkpoint_buffers
):
has_checkpoint = state_checkpoints is not None
fi_chunk_gated_delta_rule(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state=state,
output_final_state=True,
cu_seqlens=cu_seqlens,
use_qk_l2norm_in_kernel=False,
use_cp=use_cp,
state_checkpoints=state_checkpoints,
checkpoint_cu_starts=checkpoint_cu_starts if has_checkpoint else None,
checkpoint_every_n_tokens=q.shape[1] if has_checkpoint else 0,
)
def _forward_core_rocm(
self,
qkvz: torch.Tensor,
@@ -1443,6 +1586,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
)
# 1.2: Process the remaining part
conv_input_non_spec = mixed_qkv_non_spec
if attn_metadata.num_prefills > 0:
assert mixed_qkv_non_spec is not None
mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)
@@ -1598,23 +1742,59 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
assert prefill_has_initial_state is not None
initial_state = ssm_state[prefill_state_indices]
initial_state[~prefill_has_initial_state, ...] = 0
(
core_attn_out_non_spec,
last_recurrent_state,
) = self.chunk_gated_delta_rule(
q=query_non_spec,
k=key_non_spec,
v=value_non_spec,
g=g_non_spec,
beta=beta_non_spec,
initial_state=initial_state,
output_final_state=True,
cu_seqlens=attn_metadata.prefill_query_start_loc,
chunk_indices=attn_metadata.chunk_indices,
chunk_offsets=attn_metadata.chunk_offsets,
use_qk_l2norm_in_kernel=False,
aiter_prefill_metadata=attn_metadata.aiter_prefill_metadata,
)
plan = attn_metadata.flashinfer_checkpoint_plan
if plan is None:
(
core_attn_out_non_spec,
last_recurrent_state,
) = self.chunk_gated_delta_rule(
q=query_non_spec,
k=key_non_spec,
v=value_non_spec,
g=g_non_spec,
beta=beta_non_spec,
initial_state=initial_state,
output_final_state=True,
cu_seqlens=attn_metadata.prefill_query_start_loc,
chunk_indices=attn_metadata.chunk_indices,
chunk_offsets=attn_metadata.chunk_offsets,
use_qk_l2norm_in_kernel=False,
aiter_prefill_metadata=attn_metadata.aiter_prefill_metadata,
)
else:
# Only FlashInfer's prefill emits checkpoints.
(state_checkpoints,) = current_workspace_manager().get_simultaneous(
self._checkpoint_buffer_spec
)
state_checkpoints = state_checkpoints[: plan.num_checkpoints]
(
core_attn_out_non_spec,
last_recurrent_state,
) = fi_chunk_gated_delta_rule(
q=query_non_spec,
k=key_non_spec,
v=value_non_spec,
g=g_non_spec,
beta=beta_non_spec,
initial_state=initial_state,
output_final_state=True,
cu_seqlens=attn_metadata.prefill_query_start_loc,
use_qk_l2norm_in_kernel=False,
state_checkpoints=state_checkpoints,
checkpoint_cu_starts=plan.cu_starts,
checkpoint_every_n_tokens=plan.every_n_tokens,
)
assert attn_metadata.checkpoint is not None
assert non_spec_query_start_loc is not None
self._checkpoint_exporter.export(
attn_metadata.checkpoint,
raw_qkv=conv_input_non_spec,
conv_state=conv_state,
recurrent_checkpoint=state_checkpoints,
recurrent_state=ssm_state,
cu_seqlens=non_spec_query_start_loc,
recurrent_rows=plan.rows,
)
# Init cache
ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype)
vllm/model_executor/layers/mamba/kda_checkpoint.py
-116
View File
@@ -1,130 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
import torch
from vllm.model_executor.layers.mamba.checkpoint import (
MambaPrefillCheckpointExporter,
MambaPrefillCheckpointMetadata,
)
from vllm.triton_utils import tl, triton
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
def kda_prefill_checkpoint_alignment(backend: str) -> int | None:
return 16 if backend == "flashkda" else None
@dataclass(frozen=True)
class FlashKDAPrefillCheckpointExporter(MambaPrefillCheckpointExporter):
"""Store FlashKDA recurrent and convolution checkpoint states."""
state_len: int | None = None
def export(
self,
checkpoint: MambaPrefillCheckpointMetadata,
*,
raw_qkv: torch.Tensor,
conv_state: torch.Tensor,
recurrent_checkpoint: torch.Tensor,
recurrent_state: torch.Tensor,
cu_seqlens: torch.Tensor,
) -> None:
state_len = (
self.state_len if self.state_len is not None else conv_state.shape[-1]
)
width = raw_qkv.shape[-1]
recurrent_row_size = recurrent_checkpoint[0].numel()
block_size = 256
store_cache_checkpoints_kernel[
(
checkpoint.checkpoint_offsets.numel(),
triton.cdiv(max(width * state_len, recurrent_row_size), block_size),
)
](
raw_qkv,
conv_state,
recurrent_checkpoint,
recurrent_state,
cu_seqlens,
checkpoint.checkpoint_offsets,
checkpoint.state_indices,
raw_qkv.stride(0),
raw_qkv.stride(1),
conv_state.stride(0),
conv_state.stride(1),
conv_state.stride(2),
recurrent_checkpoint.stride(0),
recurrent_state.stride(0),
checkpoint.checkpoint_offsets.stride(0),
state_len,
width,
recurrent_row_size,
NULL_BLOCK_ID,
block_size,
)
@triton.jit
def store_cache_checkpoints_kernel(
x_ptr,
conv_state_ptr,
recurrent_checkpoint_ptr,
recurrent_state_ptr,
query_start_loc_ptr,
checkpoint_offsets_ptr,
checkpoint_state_indices_ptr,
x_stride_0: tl.constexpr,
x_stride_1: tl.constexpr,
state_stride_0: tl.constexpr,
state_stride_1: tl.constexpr,
state_stride_2: tl.constexpr,
checkpoint_stride_0: tl.constexpr,
recurrent_state_stride_0: tl.constexpr,
checkpoint_offset_stride: tl.constexpr,
STATE_LEN: tl.constexpr,
WIDTH: tl.constexpr,
RECURRENT_ROW_SIZE: tl.constexpr,
NULL_STATE_IDX: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
seq_idx = tl.program_id(0)
cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64)
checkpoint_offset = tl.load(
checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
)
valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0)
valid_conv = (
(cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN)
)
width_idx = cols // STATE_LEN
history_idx = cols % STATE_LEN
checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset
token_idx = checkpoint_end - STATE_LEN + history_idx
values = tl.load(
x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1,
mask=valid_conv,
)
tl.store(
conv_state_ptr
+ state_idx * state_stride_0
+ width_idx * state_stride_1
+ history_idx * state_stride_2,
values,
mask=valid_conv,
)
valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint
recurrent = tl.load(
recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols,
mask=valid_recurrent,
)
tl.store(
recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols,
recurrent,
mask=valid_recurrent,
)
vllm/platforms/interface.py
+9
View File
@@ -947,6 +947,9 @@ class Platform:
from math import lcm
from vllm.config.vllm import set_current_vllm_config
from vllm.model_executor.layers.mamba.abstract import (
get_mamba_prefill_checkpoint_alignment,
)
from vllm.model_executor.models import ModelRegistry
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
@@ -1080,6 +1083,12 @@ class Platform:
# For hybrid MLA/Mamba models, make the manager block size a
# multiple of 128 so split kernel blocks keep that invariant.
kernel_block_alignment_size = max(kernel_block_alignment_size, 128)
if cache_config.mamba_cache_mode == "align":
# Internal prefill checkpoints land on block boundaries.
kernel_block_alignment_size = lcm(
kernel_block_alignment_size,
get_mamba_prefill_checkpoint_alignment(vllm_config) or 1,
)
# Use minimum block size that satisfies both backend alignment and
# mamba page size compatibility
vllm/v1/attention/backends/gdn_attn.py
+71
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Backend for GatedDeltaNet attention."""
import math
from dataclasses import dataclass, replace
from typing import Literal
@@ -28,6 +29,47 @@ from vllm.v1.attention.backends.utils import (
)
from vllm.v1.kv_cache_interface import MambaSpec
# FlashInfer's GDN prefill emits states every N tokens, N a multiple of its
# 64-token chunk, so checkpoint offsets must be multiples of 64.
FLASHINFER_GDN_CHECKPOINT_ALIGNMENT = 64
def plan_flashinfer_checkpoints(
offsets: list[int], query_lens: list[int]
) -> tuple[int, list[int], list[int]]:
"""Lay out internal checkpoints as FlashInfer GDN prefill emits them.
Args:
offsets: Checkpoint offset into each checkpoint row's query, 0 for
none. The last ``len(query_lens)`` rows are the prefill kernel's.
query_lens: Query length of each sequence the prefill kernel runs.
Returns:
``(every_n_tokens, cu_starts, rows)``: the stride shared by every
checkpoint in the batch, cumulative per-sequence checkpoint counts,
and the index of each checkpoint row's state among the emitted ones.
"""
every_n_tokens = math.gcd(*offsets)
cu_starts = [0]
for query_len in query_lens:
cu_starts.append(cu_starts[-1] + query_len // every_n_tokens)
first = len(offsets) - len(query_lens)
rows = [0] * len(offsets)
for i, offset in enumerate(offsets[first:]):
if offset:
rows[first + i] = cu_starts[i] + offset // every_n_tokens - 1
return every_n_tokens, cu_starts, rows
@dataclass
class GDNFlashInferCheckpointPlan:
every_n_tokens: int
num_checkpoints: int
cu_starts: torch.Tensor
# Per checkpoint row, its state among the ``num_checkpoints`` emitted.
rows: torch.Tensor
class GDNAttentionBackend(AttentionBackend):
@staticmethod
@@ -54,6 +96,7 @@ class GDNAttentionMetadata:
num_actual_tokens: int
checkpoint: MambaPrefillCheckpointMetadata | None = None
flashinfer_checkpoint_plan: GDNFlashInferCheckpointPlan | None = None
has_initial_state: torch.Tensor | None = None
spec_query_start_loc: torch.Tensor | None = None # shape: [num_spec_decodes + 1,]
@@ -116,6 +159,14 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
]
_, self.gdn_prefill_backend = _resolve_gdn_prefill_backend(vllm_config)
self._check_chunk_metadata_override(type(self), self.gdn_prefill_backend)
# Only the Qwen GDN layer opts in on FlashInfer; KDA layers sharing this
# backend export checkpoints their own way.
checkpoint_alignment = kv_cache_spec.prefill_checkpoint_alignment
self.plan_flashinfer_checkpoints = (
self.gdn_prefill_backend == "flashinfer"
and checkpoint_alignment is not None
and checkpoint_alignment % FLASHINFER_GDN_CHECKPOINT_ALIGNMENT == 0
)
if self.speculative_config:
assert self.speculative_config.num_speculative_tokens is not None
@@ -528,6 +579,25 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
request_rows = (~spec_sequence_masks_cpu).nonzero().flatten().tolist()
checkpoint = self.checkpoint_builder.build(m, request_rows)
flashinfer_checkpoint_plan: GDNFlashInferCheckpointPlan | None = None
if checkpoint is not None and self.plan_flashinfer_checkpoints:
assert checkpoint.offsets is not None
assert checkpoint.query_lens is not None
# The prefill kernel runs the trailing checkpoint rows.
num_prefill_rows = prefill_query_start_loc_cpu.shape[0] - 1
every_n_tokens, cu_starts, rows = plan_flashinfer_checkpoints(
checkpoint.offsets, checkpoint.query_lens[-num_prefill_rows:]
)
plan = async_tensor_h2d(
cu_starts + rows, query_start_loc.device, torch.int64
)
flashinfer_checkpoint_plan = GDNFlashInferCheckpointPlan(
every_n_tokens=every_n_tokens,
num_checkpoints=cu_starts[-1],
cu_starts=plan[: len(cu_starts)],
rows=plan[len(cu_starts) :],
)
# Function code counted on either presency non-spec decode or spec decode,
# but not both.
assert not (num_decodes > 0 and num_spec_decodes > 0), (
@@ -606,6 +676,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
num_spec_decode_tokens=num_spec_decode_tokens,
num_actual_tokens=m.num_actual_tokens,
checkpoint=checkpoint,
flashinfer_checkpoint_plan=flashinfer_checkpoint_plan,
has_initial_state=has_initial_state,
chunk_indices=chunk_indices,
chunk_offsets=chunk_offsets,