diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 21a65bb5ed..bddee97a3a 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -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" diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 0a9bc886c3..a1224f6231 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -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: diff --git a/tests/kernels/mamba/test_gdn_prefill_flashinfer.py b/tests/kernels/mamba/test_gdn_prefill_flashinfer.py index 157e589eca..be6ab89b3c 100644 --- a/tests/kernels/mamba/test_gdn_prefill_flashinfer.py +++ b/tests/kernels/mamba/test_gdn_prefill_flashinfer.py @@ -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() diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index 4b52eb9c09..937bf2f30d 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.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] diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index cf8aee141b..629450dba1 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -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"}, + ) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 534d813124..4bce1e16ee 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -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: diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 5322f899dd..7f1ad6207d 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -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 diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py index 2f3122f870..57635a7005 100644 --- a/vllm/model_executor/layers/mamba/checkpoint.py +++ b/vllm/model_executor/layers/mamba/checkpoint.py @@ -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, + ) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 516ef13236..1a3f8183c4 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -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) diff --git a/vllm/model_executor/layers/mamba/kda_checkpoint.py b/vllm/model_executor/layers/mamba/kda_checkpoint.py index 20741623d1..35acf936d9 100644 --- a/vllm/model_executor/layers/mamba/kda_checkpoint.py +++ b/vllm/model_executor/layers/mamba/kda_checkpoint.py @@ -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, - ) diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 3df44a9275..3034d7bc1a 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -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 diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 08369adcb1..16a2c39176 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -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,