From a4439adb5fa76df3445fd9cac53d9ee5fac64ae8 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 22 Oct 2026 21:33:29 -0400 Subject: [PATCH 1/2] [Perf][GLM-5.3-Flash] KDA decode through FlashInfer's fused kernels GLM-5.3-Flash decodes each KDA layer with three Triton kernels (conv update, recurrence, gated RMSNorm). FlashInfer's fused_kda_decode does all three in one CuTe DSL launch per layer (packed_fused_kda_decode for draft-verify rows), as Kimi K3 already uses. - Plain decode steps, the decode rows of mixed steps and draft-verify rows run the FlashInfer kernels. In a mixed step the conv skips the decode rows (null slots) and the prefill kernel sees only the builder's prefill tail, as the GDN layer peels its decode rows; FlashKDA's prepare kernel no longer scans the decode rows. - Selection follows Kimi K3: additional_config kda_decode_backend / kda_spec_decode_backend (--kda-decode-backend) in auto / flashinfer / triton. auto uses FlashInfer on SM10x when a static check accepts the layer: 8/12/16/24/32/48/64/96 local heads; draft-verify rows 12-96 heads and the SD conv-state layout. - Requires FlashInfer 0.7.3: flashinfer-ai/flashinfer#7132 adds 16/64 heads (TP4, TP1/DP attention) and the DS conv-state layout. - The fused kernel applies o_norm, so with it _forward returns normed output for every step type (a breakable graph replays the code after _forward as captured, whatever the live step). - FlashInfer's [3, width, P] conv weight and fp32 norm weight are built from the fused conv1d parameter and o_norm in process_weights_after_loading and refilled in place on reload. Pure decode steps run the fused kernel in one FULL CUDA graph; eager (PIECEWISE) mixed steps also enqueue fewer kernels, and the decode-row peel keeps one-token rows out of the chunked prefill kernel. Error vs an fp64 reference equals the Triton path's. End-to-end serving results (GB300 and GB200) are in the PR description. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ao Tang --- .../glm5next/test_kda_conv_weight_refit.py | 3 + tests/models/glm5next/test_kda_flashinfer.py | 413 ++++++++++++++++++ tests/models/glm5next/test_kda_recurrent.py | 4 +- vllm/models/glm5next/common/kda.py | 307 ++++++++++++- 4 files changed, 715 insertions(+), 12 deletions(-) create mode 100644 tests/models/glm5next/test_kda_flashinfer.py diff --git a/tests/models/glm5next/test_kda_conv_weight_refit.py b/tests/models/glm5next/test_kda_conv_weight_refit.py index 0cd6c11203df..dcd0de687ec5 100644 --- a/tests/models/glm5next/test_kda_conv_weight_refit.py +++ b/tests/models/glm5next/test_kda_conv_weight_refit.py @@ -80,6 +80,8 @@ def _decode_layer(conv1d: nn.Module) -> kda.Glm5NextLinearAttention: layer.kda_safe_gate, layer.kda_lower_bound = True, -5.0 layer.A_log = layer.dt_bias = None layer._conv_state_dim_first = True + layer.kda_decode_backend = layer.kda_spec_decode_backend = "triton" + layer._norm_in_forward = False layer.kv_cache = (torch.zeros(1, 3 * CHANNELS, KERNEL - 1), None) return layer @@ -122,6 +124,7 @@ def decode() -> torch.Tensor: torch.zeros(1, 1, 1, CHANNELS), torch.zeros(1, 1, 1), torch.zeros(1, 1, 1, CHANNELS), + torch.zeros(1, 1, CHANNELS), ) return seen.pop() diff --git a/tests/models/glm5next/test_kda_flashinfer.py b/tests/models/glm5next/test_kda_flashinfer.py new file mode 100644 index 000000000000..edfabcaa2192 --- /dev/null +++ b/tests/models/glm5next/test_kda_flashinfer.py @@ -0,0 +1,413 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash KDA decode through FlashInfer's fused kernels (conv + SiLU, +the recurrence and the gated RMSNorm in one launch). + +Each test runs the same layer, weights, caches and step once on the Triton +path FlashInfer replaces and once on FlashInfer: conv states must match +exactly, recurrent states and outputs within kernel rounding, and null or +unused slots must stay untouched. +""" + +from math import prod +from typing import Any + +import pytest +import torch +from transformers import Glm5NextTextConfig + +from tests.v1.attention.utils import create_vllm_config +from vllm.config import SpeculativeConfig, VllmConfig, set_current_vllm_config +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.mamba.mamba_utils import get_conv_state_layout +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.models.glm5next.common.kda import Glm5NextLinearAttention +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_default_torch_dtype +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata + +if not current_platform.is_cuda(): + pytest.skip("FlashInfer KDA kernels are CUDA-only", allow_module_level=True) + +pytestmark = pytest.mark.usefixtures("dist_init") + +HIDDEN = 256 +D, W = 128, 4 +NUM_SLOTS = 24 +PREFIX = "model.layers.0.self_attn" + + +@pytest.fixture(scope="module") +def model_dir(tmp_path_factory) -> str: + """A one-layer GLM-5.3-Flash config; the layer needs no checkpoint.""" + path = tmp_path_factory.mktemp("glm5next") + config = Glm5NextTextConfig( + hidden_size=HIDDEN, num_hidden_layers=1, layer_types=["linear_attention"] + ) + config.architectures = ["Glm5NextForCausalLM"] + config.save_pretrained(path) + return str(path) + + +@pytest.fixture +def conv_state_layout(request, monkeypatch): + monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", request.param) + get_conv_state_layout.cache_clear() + yield request.param + get_conv_state_layout.cache_clear() + + +def _make_layer( + model_dir: str, + num_heads: int, + decode: str = "triton", + prefill: str = "triton", + num_spec: int = 0, + state_dtype: torch.dtype = torch.float32, +) -> tuple[Glm5NextLinearAttention, VllmConfig]: + """A layer with seeded weights and caches, post-processed as after a load.""" + vllm_config = create_vllm_config( + model_name=model_dir, + dtype=torch.bfloat16, + hf_config_override={"linear_num_heads": num_heads}, + ) + vllm_config.cache_config.mamba_ssm_cache_dtype = ( + "float32" if state_dtype == torch.float32 else "bfloat16" + ) + vllm_config.additional_config = { + "kda_decode_backend": decode, + "kda_spec_decode_backend": decode, + "kda_prefill_backend": prefill, + } + if num_spec: + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=num_spec + ) + with ( + set_current_vllm_config(vllm_config), + set_default_torch_dtype(torch.bfloat16), + torch.device("cuda"), + ): + layer = Glm5NextLinearAttention( + vllm_config.model_config.hf_text_config, vllm_config, PREFIX + ) + + g = torch.Generator(device="cuda").manual_seed(0) + for param in layer.parameters(): + weight = torch.randn(param.shape, device="cuda", generator=g) + default_weight_loader(param, weight / param.shape[-1] ** 0.5) + heads = torch.linspace(1, 16, layer.local_num_heads, device="cuda") + default_weight_loader(layer.A_log, heads.log().view_as(layer.A_log)) + layer.o_norm.weight.data.add_(1) + layer.process_weights_after_loading() + + shapes, dtypes = layer.get_state_shape(), layer.get_state_dtype() + page_bytes = sum(prod(s) * d.itemsize for s, d in zip(shapes, dtypes)) + layer.bind_kv_cache( + torch.zeros(NUM_SLOTS, 1, 1, page_bytes, dtype=torch.int8, device="cuda") + ) + for cache, scale in zip(layer.kv_cache, (1.0, 0.2)): + cache.copy_(scale * torch.randn(cache.shape, device="cuda", generator=g)) + return layer, vllm_config + + +def _require_flashinfer_decode(layer: Glm5NextLinearAttention) -> None: + backend = ( + layer.kda_spec_decode_backend if layer.num_spec else layer.kda_decode_backend + ) + if backend != "flashinfer": + pytest.skip("FlashInfer fused KDA decode does not support this config") + + +def _hidden(num_tokens: int, seed: int = 1) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(num_tokens, HIDDEN, device="cuda", generator=g).bfloat16() + + +def _forward( + layer: Glm5NextLinearAttention, + vllm_config: VllmConfig, + meta: GDNAttentionMetadata, + hidden: torch.Tensor, +) -> torch.Tensor: + positions = torch.arange(hidden.size(0), device="cuda") + with ( + set_current_vllm_config(vllm_config), + set_forward_context({PREFIX: meta}, vllm_config), + ): + return layer(hidden, positions) + + +def _non_spec_metadata( + query_lens: list[int], + num_decodes: int, + slots: list[int], + has_state: list[bool], + block_table_width: int = 1, +) -> GDNAttentionMetadata: + """Non-spec metadata as the builder makes it: decodes first, then the + prefill tail rebased for the peeled path. The slot indices are column 0 + of a [N, block_table_width] block table (1 + num_spec columns with + speculative decoding, so a strided view).""" + qsl_cpu = torch.tensor([0, *torch.tensor(query_lens).cumsum(0).tolist()]).int() + qsl = qsl_cpu.cuda() + block_table = torch.zeros( + len(slots), block_table_width, dtype=torch.int32, device="cuda" + ) + block_table[:, 0] = torch.tensor(slots, dtype=torch.int32) + indices = block_table[:, 0] + has_initial_state = torch.tensor(has_state, device="cuda") + num_tokens = int(qsl_cpu[-1]) + prefill: dict[str, Any] = {} + if num_decodes < len(query_lens): + nums_dict, batch_ptr, token_chunk_offset_ptr = compute_causal_conv1d_metadata( + qsl_cpu, device=qsl.device + ) + prefill = dict( + has_initial_state=has_initial_state, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + prefill_query_start_loc=qsl[num_decodes:] - num_decodes, + prefill_state_indices=indices[num_decodes:], + prefill_has_initial_state=has_initial_state[num_decodes:], + ) + return GDNAttentionMetadata( + num_prefills=len(query_lens) - num_decodes, + num_prefill_tokens=num_tokens - num_decodes, + num_decodes=num_decodes, + num_decode_tokens=num_decodes, + num_spec_decodes=0, + num_spec_decode_tokens=0, + num_actual_tokens=num_tokens, + non_spec_query_start_loc=qsl, + non_spec_state_indices_tensor=indices, + **prefill, + ) + + +def _spec_metadata( + num_spec: int, + accepted: list[int], + prefill_lens: list[int], + slots: list[list[int]], +) -> GDNAttentionMetadata: + """Draft-verify rows first (1 + num_spec tokens each), then prefills.""" + T = 1 + num_spec + num_spec_rows = len(accepted) + lens = [T] * num_spec_rows + prefill_lens + is_spec = torch.tensor([1] * num_spec_rows + [0] * len(prefill_lens)).bool() + token_is_spec = torch.repeat_interleave(is_spec, torch.tensor(lens)) + prefill: dict[str, Any] = {} + if prefill_lens: + qsl_cpu = torch.tensor( + [0, *torch.tensor(prefill_lens).cumsum(0).tolist()] + ).int() + qsl = qsl_cpu.cuda() + indices = torch.tensor( + [s[0] for s in slots[num_spec_rows:]], dtype=torch.int32, device="cuda" + ) + has_initial_state = torch.tensor([False, True], device="cuda")[ + : len(prefill_lens) + ] + nums_dict, batch_ptr, token_chunk_offset_ptr = compute_causal_conv1d_metadata( + qsl_cpu, device=qsl.device + ) + prefill = dict( + non_spec_query_start_loc=qsl, + non_spec_state_indices_tensor=indices, + has_initial_state=has_initial_state, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + prefill_query_start_loc=qsl, + prefill_state_indices=indices, + prefill_has_initial_state=has_initial_state, + ) + return GDNAttentionMetadata( + num_prefills=len(prefill_lens), + num_prefill_tokens=sum(prefill_lens), + num_decodes=0, + num_decode_tokens=0, + num_spec_decodes=num_spec_rows, + num_spec_decode_tokens=T * num_spec_rows, + num_actual_tokens=sum(lens), + spec_query_start_loc=torch.arange( + 0, T * num_spec_rows + 1, T, dtype=torch.int32, device="cuda" + ), + spec_state_indices_tensor=torch.tensor( + slots[:num_spec_rows], dtype=torch.int32, device="cuda" + ), + spec_sequence_masks=is_spec.cuda(), + spec_token_indx=token_is_spec.nonzero().flatten().cuda(), + non_spec_token_indx=(~token_is_spec).nonzero().flatten().cuda(), + num_accepted_tokens=torch.tensor(accepted, dtype=torch.int32, device="cuda"), + **prefill, + ) + + +@torch.inference_mode() +def test_backend_selection(model_dir): + """``auto`` runs GLM-5.3's decode shapes (64 heads at TP1, 16 at TP4) on + FlashInfer on SM10x, keeps Triton for 8-head draft-verify rows, and an + explicit ``flashinfer`` fails at startup where it is unsupported.""" + if not current_platform.is_device_capability_family(100): + with pytest.raises(RuntimeError, match="SM10x"): + _make_layer(model_dir, 64, decode="flashinfer") + return + for num_heads in (16, 64): + layer, _ = _make_layer(model_dir, num_heads, decode="auto") + assert layer.kda_decode_backend == "flashinfer" + layer, _ = _make_layer(model_dir, num_heads, decode="auto", num_spec=2) + assert layer.kda_spec_decode_backend == "flashinfer" + layer, _ = _make_layer(model_dir, 8, decode="auto", num_spec=2) + assert layer.kda_spec_decode_backend == "triton" + with pytest.raises(RuntimeError, match="SM10x"): + _make_layer(model_dir, 8, decode="flashinfer", num_spec=2) + + +@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"], indirect=True) +@pytest.mark.parametrize( + "state_dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"] +) +@pytest.mark.parametrize("num_heads", [16, 64]) +@torch.inference_mode() +def test_decode_step_matches_triton( + model_dir, num_heads, state_dtype, conv_state_layout +): + flashinfer = _make_layer(model_dir, num_heads, "auto", state_dtype=state_dtype) + _require_flashinfer_decode(flashinfer[0]) + triton = _make_layer(model_dir, num_heads, state_dtype=state_dtype) + before = [cache.clone() for cache in triton[0].kv_cache] + # Null rows (slot 0, also CUDA-graph padding) must not touch any cache. + slots = [3, 0, 7, 11, 0, 19, 1, 22] + meta = _non_spec_metadata([1] * len(slots), len(slots), slots, [True] * 8) + hidden = _hidden(len(slots)) + out_t = _forward(*triton, meta, hidden) + out_f = _forward(*flashinfer, meta, hidden) + + (conv_t, state_t), (conv_f, state_f) = triton[0].kv_cache, flashinfer[0].kv_cache + torch.testing.assert_close(conv_f, conv_t, atol=0, rtol=0) + tol = 2e-2 if state_dtype == torch.bfloat16 else 2e-3 + torch.testing.assert_close(state_f, state_t, atol=tol, rtol=tol) + live = torch.tensor(slots) > 0 + torch.testing.assert_close(out_f[live], out_t[live], atol=3e-2, rtol=3e-2) + untouched = [s for s in range(NUM_SLOTS) if s not in slots or s == 0] + for cache, original in zip(flashinfer[0].kv_cache, before): + torch.testing.assert_close( + cache[untouched], original[untouched], atol=0, rtol=0 + ) + + +@pytest.mark.parametrize("prefill", ["triton", "flashkda"]) +@torch.inference_mode() +def test_mixed_step_matches_triton(model_dir, workspace_init, prefill): + """The decode rows of a mixed step go through the fused decode; the + prefill rows (fresh and continued) and their states do not change.""" + num_heads = 16 + # FlashKDA is the ``auto`` prefill backend where it is available. + requested = "auto" if prefill == "flashkda" else prefill + flashinfer = _make_layer(model_dir, num_heads, "auto", requested) + _require_flashinfer_decode(flashinfer[0]) + if flashinfer[0].kda_prefill_backend != prefill: + pytest.skip("FlashKDA is not available") + triton = _make_layer(model_dir, num_heads, "triton", requested) + decode_slots = [5, 2, 20, 9] + n = len(decode_slots) + query_lens = [1] * n + [37, 1, 130] + slots = decode_slots + [12, 14, 17] + meta = _non_spec_metadata(query_lens, n, slots, [True] * n + [False, True, True]) + hidden = _hidden(sum(query_lens)) + out_t = _forward(*triton, meta, hidden) + out_f = _forward(*flashinfer, meta, hidden) + + (conv_t, state_t), (conv_f, state_f) = triton[0].kv_cache, flashinfer[0].kv_cache + torch.testing.assert_close(conv_f, conv_t, atol=0, rtol=0) + torch.testing.assert_close(out_f[n:], out_t[n:], atol=0, rtol=0) + prefill_slots = slots[n:] + torch.testing.assert_close( + state_f[prefill_slots], state_t[prefill_slots], atol=0, rtol=0 + ) + # The Triton path runs a mixed step's decode rows through the chunked + # kernel, so they differ by chunked-vs-recurrent rounding. + torch.testing.assert_close(out_f[:n], out_t[:n], atol=3e-2, rtol=3e-2) + torch.testing.assert_close( + state_f[decode_slots], state_t[decode_slots], atol=2e-2, rtol=2e-2 + ) + + +@pytest.mark.parametrize("with_prefill", [False, True], ids=["verify", "mixed"]) +@pytest.mark.parametrize( + "state_dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"] +) +@torch.inference_mode() +def test_spec_verify_rows_match_triton(model_dir, state_dtype, with_prefill): + num_heads, num_spec = 64, 2 + flashinfer = _make_layer( + model_dir, num_heads, "auto", "triton", num_spec, state_dtype + ) + _require_flashinfer_decode(flashinfer[0]) + triton = _make_layer( + model_dir, num_heads, "triton", "triton", num_spec, state_dtype + ) + # Three verify rows (the third a null row), then a fresh and a continued + # prefill. + slots = [[1, 2, 3], [4, 5, 6], [0, 0, 0], [9, 9, 9], [10, 10, 10]] + meta = _spec_metadata(num_spec, [1, 3, 1], [21, 6] if with_prefill else [], slots) + hidden = _hidden(meta.num_actual_tokens) + out_t = _forward(*triton, meta, hidden) + out_f = _forward(*flashinfer, meta, hidden) + + (conv_t, state_t), (conv_f, state_f) = triton[0].kv_cache, flashinfer[0].kv_cache + torch.testing.assert_close(conv_f, conv_t, atol=0, rtol=0) + tol = 2e-2 if state_dtype == torch.bfloat16 else 2e-3 + torch.testing.assert_close(state_f, state_t, atol=tol, rtol=tol) + live = torch.ones(meta.num_actual_tokens, dtype=torch.bool) + live[6:9] = False + torch.testing.assert_close(out_f[live], out_t[live], atol=3e-2, rtol=3e-2) + + +def _refit(layer: Glm5NextLinearAttention, seed: int) -> None: + """New conv and norm weights through their weight loaders, then the + post-load hook, as a weight reload runs them.""" + g = torch.Generator(device="cuda").manual_seed(seed) + shard_shape = (layer.local_projection_size, 1, W) + for shard_id in range(3): + weight = 0.5 * torch.randn(shard_shape, device="cuda", generator=g) + layer.conv1d.weight.weight_loader(layer.conv1d.weight, weight, shard_id) + norm = 1 + 0.1 * torch.randn(D, device="cuda", generator=g) + default_weight_loader(layer.o_norm.weight, norm) + layer.process_weights_after_loading() + + +@torch.inference_mode() +def test_decode_cuda_graph_replays_new_inputs_and_refit(model_dir): + """A captured fused decode step replays with new inputs, slots and + refitted weights: FlashInfer's conv and norm weights are refilled in place, + at the addresses the graph captured.""" + num_heads, n = 32, 6 + captured = _make_layer(model_dir, num_heads, "auto") + _require_flashinfer_decode(captured[0]) + eager = _make_layer(model_dir, num_heads, "auto") + meta = _non_spec_metadata([1] * n, n, [1, 2, 3, 4, 5, 6], [True] * n) + hidden = _hidden(n, seed=1) + _forward(*captured, meta, hidden) + _forward(*eager, meta, hidden) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _forward(*captured, meta, hidden) + + new_hidden = _hidden(n, seed=2) + hidden.copy_(new_hidden) + slots = [7, 0, 9, 10, 11, 12] + assert meta.non_spec_state_indices_tensor is not None + meta.non_spec_state_indices_tensor.copy_(torch.tensor(slots)) + for layer, _ in (captured, eager): + _refit(layer, seed=3) + graph.replay() + expected = _forward(*eager, meta, new_hidden) + live = torch.tensor(slots) > 0 + torch.testing.assert_close(out[live], expected[live], atol=0, rtol=0) + for replayed, reference in zip(captured[0].kv_cache, eager[0].kv_cache): + torch.testing.assert_close(replayed, reference, atol=0, rtol=0) diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index e904ce8b6ac1..4b10a9f0d65a 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/tests/models/glm5next/test_kda_recurrent.py @@ -302,6 +302,8 @@ def test_prefill_checkpoint_resumes_suffix(monkeypatch, dim_first, num_spec, off layer.kda_safe_gate = True layer.kda_lower_bound = -5.0 layer.kda_prefill_backend = "flashkda" + layer.kda_decode_backend = layer.kda_spec_decode_backend = "triton" + layer._norm_in_forward = False layer._conv_state_dim_first = dim_first layer.conv1d = SimpleNamespace( weight=torch.randn(width, 1, 4, device=device) * 0.1, bias=None @@ -395,7 +397,7 @@ def run(qkv, g, b, slots, has_initial, checkpoint=None): out = torch.empty( 1, qkv.shape[0], heads, D, device=device, dtype=torch.bfloat16 ) - layer._forward(qkv, g, b, out) + layer._forward(qkv, g, b, out, g[0]) return out[:, spec_tokens:] checkpoint = MambaPrefillCheckpointMetadata( diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 702d799433ce..494847c6c69d 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -12,6 +12,7 @@ from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import divide from vllm.forward_context import get_forward_context +from vllm.logger import init_logger from vllm.model_executor.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, @@ -44,14 +45,24 @@ from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader from vllm.model_executor.utils import ( maybe_disable_graph_partition, + refill_derived, set_weight_attrs, ) from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from vllm.utils.flashinfer import ( + flashinfer_fused_kda_decode, + flashinfer_packed_fused_kda_decode, + has_flashinfer_fused_kda_decode, + has_flashinfer_packed_fused_kda_decode, +) from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.worker.workspace import current_workspace_manager +logger = init_logger(__name__) + if current_platform.is_rocm(): from vllm.models.glm5next.amd.ops.third_party.kda import ( chunk_kda_with_fused_gate, @@ -193,6 +204,90 @@ def _fused_qkv_conv1d( return conv1d +def _is_flashinfer_kda_decode_supported( + num_heads: int, + head_dim: int, + conv_size: int, + query_len: int, + dtype: torch.dtype, + conv_state_dtype: torch.dtype, + state_dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + """Whether FlashInfer's fused KDA decode (``query_len`` 1) or its packed + draft-verify decode (``query_len`` 1 + num_spec) runs this layer.""" + heads: tuple[int, ...] + if query_len == 1: + has_kernel = has_flashinfer_fused_kda_decode() + heads = (8, 12, 16, 24, 32, 48, 64, 96) + else: + # The packed kernel supports only the SD conv-state layout. + has_kernel = ( + has_flashinfer_packed_fused_kda_decode() and not is_conv_state_dim_first() + ) + heads = (12, 16, 24, 32, 48, 64, 96) + return ( + has_kernel + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and num_heads in heads + and head_dim == 128 + and conv_size == 4 + and 1 <= query_len <= 8 + and dtype == torch.bfloat16 + and conv_state_dtype == torch.bfloat16 + and state_dtype in (torch.float32, torch.bfloat16) + and lower_bound is not None + and lower_bound < 0 + ) + + +def _resolve_kda_decode_backend( + backend: str, + num_heads: int, + head_dim: int, + conv_size: int, + query_len: int, + dtype: torch.dtype, + conv_state_dtype: torch.dtype, + state_dtype: torch.dtype, + lower_bound: float | None, +) -> str: + """Pick the decode kernels: FlashInfer's fused KDA decode (conv + SiLU, the + recurrence and the gated RMSNorm in one launch per layer) or the Triton + conv, recurrence and norm kernels. ``backend`` comes from + ``additional_config`` (auto / flashinfer / triton); ``auto`` picks + FlashInfer when it supports the layer. ``query_len`` is 1 for plain decode + rows and 1 + num_speculative_tokens for draft-verify rows.""" + if backend not in ("auto", "flashinfer", "triton"): + raise ValueError(f"Unsupported GLM-5.3 KDA decode backend: {backend}") + if backend == "triton": + return "triton" + if _is_flashinfer_kda_decode_supported( + num_heads, + head_dim, + conv_size, + query_len, + dtype, + conv_state_dtype, + state_dtype, + lower_bound, + ): + logger.info_once( + "Using FlashInfer fused KDA decode (%d tokens per row).", query_len + ) + return "flashinfer" + if backend == "flashinfer": + raise RuntimeError( + "FlashInfer fused KDA decode requires CUDA SM10x, 8/12/16/24/32/48/" + "64/96 local heads (12-96 and the SD conv-state layout for " + "draft-verify rows, at most 7 speculative tokens), head_dim 128, " + "conv width 4, bfloat16 activations and conv state, a float32 or " + "bfloat16 recurrent state and a bounded gate." + ) + return "triton" + + class Glm5NextLinearAttention(GatedDeltaNetAttention): head_dim: int num_heads: int @@ -358,6 +453,38 @@ def _a_log_weight_loader(param, loaded_weight): vllm_config.model_config.dtype, self.kda_lower_bound, ) + + # With speculative decoding every decode row is a draft-verify row of + # 1 + num_spec tokens (kda_spec_decode_backend); otherwise rows have + # one token (kda_decode_backend). + conv_state_dtype, recurrent_state_dtype = self.get_state_dtype() + decode_backend_key = ( + "kda_spec_decode_backend" if self.num_spec > 0 else "kda_decode_backend" + ) + decode_backend = _resolve_kda_decode_backend( + additional_config.get(decode_backend_key, "auto") + if isinstance(additional_config, dict) + else "auto", + self.local_num_heads, + self.head_dim, + self.conv_size, + 1 + self.num_spec, + vllm_config.model_config.dtype, + conv_state_dtype, + recurrent_state_dtype, + self.kda_lower_bound, + ) + self.kda_decode_backend = decode_backend if self.num_spec == 0 else "triton" + self.kda_spec_decode_backend = decode_backend if self.num_spec > 0 else "triton" + # FlashInfer's fused decode applies o_norm itself, so with it _forward + # returns normed output for every step type: a breakable CUDA graph + # replays the code after _forward as captured, whatever the live step. + self._norm_in_forward = decode_backend == "flashinfer" + # FlashInfer's [3, width, P] fp32 conv weight and fp32 norm weight, + # built in process_weights_after_loading. + self._fi_conv_weight: torch.Tensor | None = None + self._fi_norm_weight: torch.Tensor | None = None + self._flashkda_buffer_specs: ( tuple[tuple[tuple[int, ...], torch.dtype], ...] | None ) = None @@ -503,11 +630,72 @@ def forward( g1=g1, beta=beta, core_attn_out=core_attn_out, + g2=g2, ) - core_attn_out = self.o_norm(core_attn_out, g2) + if not self._norm_in_forward: + core_attn_out = self.o_norm(core_attn_out, g2) core_attn_out = core_attn_out.reshape(core_attn_out.size(1), -1) return self.o_proj(core_attn_out)[0] + def process_weights_after_loading(self) -> None: + """Build FlashInfer's conv and norm weights from the loaded weights; + reloads refill them in place.""" + if not self._norm_in_forward: + return + conv_weight = self.conv1d.weight.view( + 3, self.local_projection_size, self.conv_size + ) + refill_derived( + self, "_fi_conv_weight", conv_weight.transpose(1, 2).contiguous() + ) + refill_derived(self, "_fi_norm_weight", self.o_norm.weight.float().contiguous()) + + def _flashinfer_decode( + self, + x: torch.Tensor, + conv_state: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + state_indices: torch.Tensor, + state: torch.Tensor, + output_gate: torch.Tensor, + output: torch.Tensor, + query_start_loc: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + ) -> None: + """FlashInfer fused KDA decode: conv + SiLU with the conv-state update, + the recurrence and the gated RMSNorm, written into ``output``. Rows + with a non-positive slot are null rows. With ``query_start_loc`` the + rows are draft-verify rows (``state_indices`` [N, 1 + num_spec]).""" + if self._fi_conv_weight is None: # only when post-load processing never ran + self.process_weights_after_loading() + assert self._fi_conv_weight is not None + assert self._fi_norm_weight is not None + kwargs = dict( + x=x, + weight=self._fi_conv_weight, + conv_state=conv_state, + raw_gate=raw_gate, + raw_beta=raw_beta, + A_log=self.A_log.view(-1), + dt_bias=self.dt_bias, + state=state, + output_gate=output_gate, + norm_weight=self._fi_norm_weight, + lower_bound=self.kda_lower_bound, + norm_eps=self.o_norm.eps, + output=output, + ) + if query_start_loc is None: + flashinfer_fused_kda_decode(state_indices=state_indices, **kwargs) + else: + flashinfer_packed_fused_kda_decode( + state_indices=state_indices.contiguous(), + query_start_loc=query_start_loc, + num_accepted_tokens=num_accepted_tokens, + **kwargs, + ) + @eager_break_during_capture def _forward( self, @@ -515,6 +703,7 @@ def _forward( g1: torch.Tensor, beta: torch.Tensor, core_attn_out: torch.Tensor, + g2: torch.Tensor, ) -> None: forward_context = get_forward_context() attn_metadata_raw = forward_context.attn_metadata @@ -568,6 +757,47 @@ def _forward( ) conv_bias = self.conv1d.bias + num_decodes = attn_metadata_narrowed.num_decodes + num_prefills = attn_metadata_narrowed.num_prefills + if ( + self.kda_decode_backend == "flashinfer" + and num_prefills == 0 + and num_decodes > 0 + ): + assert non_spec_state_indices_tensor is not None + self._flashinfer_decode( + x=qkv_proj_states[:num_decodes], + conv_state=conv_state, + raw_gate=g1[:, :num_decodes], + raw_beta=beta[:, :num_decodes], + state_indices=non_spec_state_indices_tensor[:num_decodes], + state=recurrent_state, + output_gate=g2[:num_decodes], + output=core_attn_out[:, :num_decodes], + ) + return + if ( + self.kda_spec_decode_backend == "flashinfer" + and use_spec + and num_prefills == 0 + and num_decodes == 0 + ): + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + self._flashinfer_decode( + x=qkv_proj_states, + conv_state=conv_state, + raw_gate=g1, + raw_beta=beta, + state_indices=spec_state_indices_tensor, + state=recurrent_state, + output_gate=g2[:num_actual_tokens], + output=core_attn_out[:, :num_actual_tokens], + query_start_loc=spec_query_start_loc, + num_accepted_tokens=num_accepted_tokens, + ) + return + # Split projections / gating into spec (draft-verify) and non-spec token # groups when speculative decoding is active. Spec tokens carry # num_spec+1 recurrent-state columns each and are advanced with @@ -621,6 +851,32 @@ def _forward( # --- causal conv1d: non-spec path (prefill or plain decode) --- raw_qkv_ns = qkv_ns q_ns = k_ns = v_ns = None + # Mixed step: the leading decode rows (vLLM orders decodes first) go + # through the fused decode kernel; the conv skips them as null slots and + # the prefill kernel only sees the prefill rows. + num_fused_decodes = 0 + conv_cache_indices = non_spec_state_indices_tensor + if ( + self.kda_decode_backend == "flashinfer" + and num_prefills > 0 + and num_decodes > 0 + and attn_metadata_narrowed.checkpoint is None + ): + assert qkv_ns is not None and g1_ns is not None and beta_ns is not None + assert non_spec_state_indices_tensor is not None + num_fused_decodes = attn_metadata_narrowed.num_decode_tokens + self._flashinfer_decode( + x=qkv_ns[:num_fused_decodes], + conv_state=conv_state, + raw_gate=g1_ns[:, :num_fused_decodes], + raw_beta=beta_ns[:, :num_fused_decodes], + state_indices=non_spec_state_indices_tensor[:num_fused_decodes], + state=recurrent_state, + output_gate=g2[:num_fused_decodes], + output=core_attn_out[:, :num_fused_decodes], + ) + conv_cache_indices = non_spec_state_indices_tensor.clone() + conv_cache_indices[:num_decodes] = NULL_BLOCK_ID if attn_metadata_narrowed.num_prefills > 0: assert qkv_ns is not None qkv_ns = causal_conv1d_fn( @@ -630,7 +886,7 @@ def _forward( activation="silu", conv_states=conv_state, has_initial_state=has_initial_state, - cache_indices=non_spec_state_indices_tensor, + cache_indices=conv_cache_indices, query_start_loc=non_spec_query_start_loc, metadata=attn_metadata_narrowed, ).transpose(0, 1) @@ -697,17 +953,40 @@ def _rearr(x): # stays None there and the merge copy below runs as before. ns_out = None if attn_metadata_narrowed.num_prefills > 0: - assert q_ns is not None + assert q_ns is not None and k_ns is not None and v_ns is not None + assert g1_ns is not None and beta_ns is not None assert non_spec_state_indices_tensor is not None assert has_initial_state is not None + prefill_query_start_loc = non_spec_query_start_loc + prefill_state_indices = non_spec_state_indices_tensor + prefill_has_initial_state = has_initial_state + if num_fused_decodes: + # The builder precomputes the prefill tail of a mixed step. + prefill_query_start_loc = attn_metadata_narrowed.prefill_query_start_loc + prefill_state_indices = attn_metadata_narrowed.prefill_state_indices + prefill_has_initial_state = ( + attn_metadata_narrowed.prefill_has_initial_state + ) + assert prefill_query_start_loc is not None + assert prefill_state_indices is not None + assert prefill_has_initial_state is not None + q_ns = q_ns[num_fused_decodes:] + k_ns = k_ns[num_fused_decodes:] + v_ns = v_ns[num_fused_decodes:] + g1_ns = g1_ns[:, num_fused_decodes:] + beta_ns = beta_ns[:, num_fused_decodes:] initial_state = gather_initial_states( - recurrent_state, non_spec_state_indices_tensor, has_initial_state + recurrent_state, prefill_state_indices, prefill_has_initial_state ) if self.kda_prefill_backend == "flashkda": # Non-spec step: write straight into the layer output buffer # (dense token order, no merge copy). Step with spec-decode # tokens: write to the workspace buffer and scatter below. - ns_out = None if use_spec else core_attn_out[:, :num_actual_tokens] + ns_out = ( + None + if use_spec + else core_attn_out[:, num_fused_decodes:num_actual_tokens] + ) core_attn_out_non_spec, last_recurrent_state = self._flashkda_prefill( q=_rearr(q_ns), k=_rearr(k_ns), @@ -715,7 +994,7 @@ def _rearr(x): g=g1_ns, beta=beta_ns, initial_state=initial_state, - cu_seqlens=non_spec_query_start_loc, + cu_seqlens=prefill_query_start_loc, out=ns_out, checkpoint=attn_metadata_narrowed.checkpoint, raw_qkv=raw_qkv_ns, @@ -739,7 +1018,7 @@ def _rearr(x): initial_state=initial_state, output_final_state=True, use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc, + cu_seqlens=prefill_query_start_loc, safe_gate=safe_gate, lower_bound=lower_bound, ) @@ -747,7 +1026,7 @@ def _rearr(x): scatter_states( recurrent_state, last_recurrent_state, - non_spec_state_indices_tensor, + prefill_state_indices, ) elif attn_metadata_narrowed.num_decodes > 0: assert non_spec_query_start_loc is not None @@ -790,6 +1069,12 @@ def _rearr(x): else: assert core_attn_out_non_spec is not None if ns_out is None: - core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ - 0, :num_actual_tokens - ] + core_attn_out[0, num_fused_decodes:num_actual_tokens] = ( + core_attn_out_non_spec[0, : num_actual_tokens - num_fused_decodes] + ) + if self._norm_in_forward: + # Rows from the fused decode kernel are normed already. + live_out = core_attn_out[:, num_fused_decodes:num_actual_tokens] + normed = self.o_norm(live_out, g2[num_fused_decodes:num_actual_tokens]) + if normed.data_ptr() != live_out.data_ptr(): + live_out.copy_(normed) From 6a0ad6dd527004eb84dcfdfdb008f7275da36821 Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Thu, 22 Oct 2026 21:38:19 -0400 Subject: [PATCH 2/2] [Perf][GLM-5.3-Flash] KDA prefill through FlashInfer's CuTe DSL kernel kda_prefill_backend=flashinfer (--kda-prefill-backend flashinfer, already a CLI choice; GLM-5.3 rejected it) runs the chunked KDA prefill with FlashInfer's recurrent_kda(backend="cute-dsl"): chunk preparation and the recurrence in one kernel, from the raw gate and beta logits, reading and writing the recurrent state pool in place through ssm_state_indices. The gather / scatter of the FlashKDA path and its final-state buffer go away. - FlashInfer has no has-initial-state flag, so fresh sequences' slots are zeroed first (one index_fill; other rows zero the unused null slot). - A step whose prefill rows total one token uses the recurrent kernel. - The slot indices are passed contiguous: with speculative decoding they are column 0 of the [N, 1 + num_spec] block table, which FlashInfer rejects. - Steps with draft-verify rows write the prefill rows to their own buffer and scatter them like FlashKDA's; no FlashKDA workspace is involved. - Opt-in: auto stays FlashKDA. Prefix-cache checkpoints are not produced with this backend (as for Kimi K3). Supported on SM100 / SM103 / SM107 (SM107 from flashinfer-ai/flashinfer#6201, in FlashInfer 0.7.3). Error vs an fp64 reference (H64, 8 x 8192 prefill tokens): 4.2e-3, vs 5.1e-3 for FlashKDA. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Ao Tang --- tests/models/glm5next/test_kda_flashinfer.py | 111 ++++++++++- vllm/models/glm5next/common/kda.py | 189 +++++++++++++++++-- 2 files changed, 279 insertions(+), 21 deletions(-) diff --git a/tests/models/glm5next/test_kda_flashinfer.py b/tests/models/glm5next/test_kda_flashinfer.py index edfabcaa2192..945572cd82f8 100644 --- a/tests/models/glm5next/test_kda_flashinfer.py +++ b/tests/models/glm5next/test_kda_flashinfer.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""GLM-5.3-Flash KDA decode through FlashInfer's fused kernels (conv + SiLU, -the recurrence and the gated RMSNorm in one launch). +"""GLM-5.3-Flash KDA layers on FlashInfer's kernels: the fused decode (conv + +SiLU, the recurrence and the gated RMSNorm in one launch) and the CuTe DSL +chunked prefill. Each test runs the same layer, weights, caches and step once on the Triton path FlashInfer replaces and once on FlashInfer: conv states must match @@ -21,7 +22,10 @@ from vllm.forward_context import set_forward_context from vllm.model_executor.layers.mamba.mamba_utils import get_conv_state_layout from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.models.glm5next.common.kda import Glm5NextLinearAttention +from vllm.models.glm5next.common.kda import ( + Glm5NextLinearAttention, + _is_flashinfer_kda_prefill_supported, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_default_torch_dtype from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata @@ -120,6 +124,11 @@ def _require_flashinfer_decode(layer: Glm5NextLinearAttention) -> None: pytest.skip("FlashInfer fused KDA decode does not support this config") +def _require_flashinfer_prefill(state_dtype: torch.dtype) -> None: + if not _is_flashinfer_kda_prefill_supported(D, torch.bfloat16, state_dtype, -5.0): + pytest.skip("FlashInfer KDA prefill does not support this platform") + + def _hidden(num_tokens: int, seed: int = 1) -> torch.Tensor: g = torch.Generator(device="cuda").manual_seed(seed) return torch.randn(num_tokens, HIDDEN, device="cuda", generator=g).bfloat16() @@ -255,6 +264,8 @@ def test_backend_selection(model_dir): if not current_platform.is_device_capability_family(100): with pytest.raises(RuntimeError, match="SM10x"): _make_layer(model_dir, 64, decode="flashinfer") + with pytest.raises(RuntimeError, match="SM100/SM103/SM107"): + _make_layer(model_dir, 64, prefill="flashinfer") return for num_heads in (16, 64): layer, _ = _make_layer(model_dir, num_heads, decode="auto") @@ -265,6 +276,8 @@ def test_backend_selection(model_dir): assert layer.kda_spec_decode_backend == "triton" with pytest.raises(RuntimeError, match="SM10x"): _make_layer(model_dir, 8, decode="flashinfer", num_spec=2) + layer, _ = _make_layer(model_dir, 64, prefill="flashinfer") + assert layer.kda_prefill_backend == "flashinfer" @pytest.mark.parametrize("conv_state_layout", ["SD", "DS"], indirect=True) @@ -300,12 +313,14 @@ def test_decode_step_matches_triton( ) -@pytest.mark.parametrize("prefill", ["triton", "flashkda"]) +@pytest.mark.parametrize("prefill", ["triton", "flashkda", "flashinfer"]) @torch.inference_mode() def test_mixed_step_matches_triton(model_dir, workspace_init, prefill): """The decode rows of a mixed step go through the fused decode; the prefill rows (fresh and continued) and their states do not change.""" num_heads = 16 + if prefill == "flashinfer": + _require_flashinfer_prefill(torch.float32) # FlashKDA is the ``auto`` prefill backend where it is available. requested = "auto" if prefill == "flashkda" else prefill flashinfer = _make_layer(model_dir, num_heads, "auto", requested) @@ -324,10 +339,13 @@ def test_mixed_step_matches_triton(model_dir, workspace_init, prefill): (conv_t, state_t), (conv_f, state_f) = triton[0].kv_cache, flashinfer[0].kv_cache torch.testing.assert_close(conv_f, conv_t, atol=0, rtol=0) - torch.testing.assert_close(out_f[n:], out_t[n:], atol=0, rtol=0) + # Without the fused decode, FlashInfer's prefill also runs the decode rows + # (as one-token sequences), and it is not batch invariant. + tol = 1e-2 if prefill == "flashinfer" else 0 + torch.testing.assert_close(out_f[n:], out_t[n:], atol=tol, rtol=tol) prefill_slots = slots[n:] torch.testing.assert_close( - state_f[prefill_slots], state_t[prefill_slots], atol=0, rtol=0 + state_f[prefill_slots], state_t[prefill_slots], atol=tol, rtol=tol ) # The Triton path runs a mixed step's decode rows through the chunked # kernel, so they differ by chunked-vs-recurrent rounding. @@ -368,6 +386,87 @@ def test_spec_verify_rows_match_triton(model_dir, state_dtype, with_prefill): torch.testing.assert_close(out_f[live], out_t[live], atol=3e-2, rtol=3e-2) +@pytest.mark.parametrize( + "query_lens", [[40, 1, 17, 3, 96], [1]], ids=["chunked", "one-token"] +) +@pytest.mark.parametrize( + "state_dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"] +) +@torch.inference_mode() +def test_prefill_step_matches_triton(model_dir, state_dtype, query_lens): + """FlashInfer's prefill updates the state pool in place at the slot + indices, so it zeroes fresh sequences' slots itself and leaves other slots + alone. A step whose only prefill row is one token takes the recurrent + kernel (FlashInfer's prefill needs two tokens).""" + _require_flashinfer_prefill(state_dtype) + num_heads = 64 + # Fresh and continued sequences across the 16-token chunk boundary. + slots = [3, 5, 8, 13, 21][: len(query_lens)] + has_state = [False, True, True, False, True][: len(query_lens)] + meta = _non_spec_metadata(query_lens, 0, slots, has_state) + hidden = _hidden(sum(query_lens)) + flashinfer = _make_layer( + model_dir, num_heads, prefill="flashinfer", state_dtype=state_dtype + ) + before = flashinfer[0].kv_cache[1].clone() + triton = _make_layer(model_dir, num_heads, state_dtype=state_dtype) + out_t = _forward(*triton, meta, hidden).float() + out_f = _forward(*flashinfer, meta, hidden).float() + + (conv_t, state_t), (conv_f, state_f) = triton[0].kv_cache, flashinfer[0].kv_cache + torch.testing.assert_close(conv_f, conv_t, atol=0, rtol=0) + tol = 3e-2 if state_dtype == torch.bfloat16 else 1e-2 + torch.testing.assert_close(state_f[slots], state_t[slots], atol=tol, rtol=tol) + # Different chunked kernels: compare in norm. + assert (out_f - out_t).norm() < 1e-2 * out_t.norm() + # Only the sequences' slots (and the null slot 0) may change. + untouched = [s for s in range(1, NUM_SLOTS) if s not in slots] + torch.testing.assert_close(state_f[untouched], before[untouched], atol=0, rtol=0) + + +@torch.inference_mode() +def test_prefill_with_spec_verify_rows(model_dir): + """With draft-verify rows in the step, FlashInfer's prefill writes the + prefill rows to their own buffer, scattered by ``non_spec_token_indx``.""" + _require_flashinfer_prefill(torch.float32) + num_heads, num_spec = 32, 2 + slots = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [11, 11, 11], [14, 14, 14]] + meta = _spec_metadata(num_spec, [1, 2, 3], [25, 9], slots) + hidden = _hidden(meta.num_actual_tokens) + flashinfer = _make_layer(model_dir, num_heads, "triton", "flashinfer", num_spec) + triton = _make_layer(model_dir, num_heads, "triton", "triton", num_spec) + out_t = _forward(*triton, meta, hidden).float() + out_f = _forward(*flashinfer, meta, hidden).float() + assert (out_f - out_t).norm() < 1e-2 * out_t.norm() + torch.testing.assert_close( + flashinfer[0].kv_cache[1][1:], triton[0].kv_cache[1][1:], atol=1e-2, rtol=1e-2 + ) + + +@pytest.mark.parametrize("num_decodes", [0, 3], ids=["prefill", "mixed"]) +@torch.inference_mode() +def test_block_table_column_state_indices(model_dir, num_decodes): + """A step without draft rows in an MTP deployment: the builder's slot + indices are column 0 of the [N, 1 + num_spec] block table, a strided view. + It matches the same step with contiguous indices.""" + _require_flashinfer_prefill(torch.float32) + num_heads, num_spec = 32, 2 + query_lens = [1] * num_decodes + [40, 17, 96] + slots = [2, 4, 6][:num_decodes] + [9, 12, 15] + has_state = [True] * num_decodes + [False, True, True] + hidden = _hidden(sum(query_lens)) + out = {} + for width in (1, 1 + num_spec): + meta = _non_spec_metadata(query_lens, num_decodes, slots, has_state, width) + layer = _make_layer(model_dir, num_heads, "triton", "flashinfer", num_spec) + out[width] = (_forward(*layer, meta, hidden), layer[0].kv_cache) + assert meta.non_spec_state_indices_tensor is not None + assert not meta.non_spec_state_indices_tensor.is_contiguous() + torch.testing.assert_close(out[3][0], out[1][0], atol=0, rtol=0) + for strided, contiguous in zip(out[3][1], out[1][1], strict=True): + torch.testing.assert_close(strided, contiguous, atol=0, rtol=0) + + def _refit(layer: Glm5NextLinearAttention, seed: int) -> None: """New conv and norm weights through their weight loaders, then the post-load hook, as a weight reload runs them.""" diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 494847c6c69d..ab24610c0ac9 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -53,8 +53,10 @@ from vllm.utils.flashinfer import ( flashinfer_fused_kda_decode, flashinfer_packed_fused_kda_decode, + flashinfer_recurrent_kda, has_flashinfer_fused_kda_decode, has_flashinfer_packed_fused_kda_decode, + has_flashinfer_recurrent_kda, ) from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -149,14 +151,31 @@ def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor: def _resolve_kda_prefill_backend( - backend: str, head_dim: int, dtype: torch.dtype, lower_bound: float | None + backend: str, + head_dim: int, + dtype: torch.dtype, + lower_bound: float | None, + state_dtype: torch.dtype = torch.float32, ) -> str: """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on - SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton + SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate), FlashInfer's + CuTe DSL kernel (opt-in; updates the state pool in place) or the Triton ``chunk_kda_with_fused_gate`` path. ``backend`` comes from - ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" - if backend not in ("auto", "triton", "flashkda"): + ``additional_config.kda_prefill_backend`` (auto / triton / flashkda / + flashinfer).""" + if backend not in ("auto", "triton", "flashkda", "flashinfer"): raise ValueError(f"Unsupported KDA prefill backend: {backend}") + if backend == "flashinfer": + if not _is_flashinfer_kda_prefill_supported( + head_dim, dtype, state_dtype, lower_bound + ): + raise RuntimeError( + "FlashInfer KDA prefill requires CUDA SM100/SM103/SM107, " + "bfloat16 activations, a float32 or bfloat16 recurrent state, " + "head_dim=128 and a bounded KDA gate." + ) + logger.info_once("Using FlashInfer KDA prefill.") + return "flashinfer" capability = current_platform.get_device_capability() supported = ( current_platform.is_cuda() @@ -204,6 +223,68 @@ def _fused_qkv_conv1d( return conv1d +def _is_flashinfer_kda_prefill_supported( + head_dim: int, + dtype: torch.dtype, + state_dtype: torch.dtype, + lower_bound: float | None, +) -> bool: + """Whether FlashInfer's CuTe DSL chunked KDA prefill runs this layer.""" + if not current_platform.is_cuda() or not has_flashinfer_recurrent_kda(): + return False + capability = current_platform.get_device_capability() + if capability is None: + return False + return ( + (capability.major, capability.minor) in ((10, 0), (10, 3), (10, 7)) + and head_dim == 128 + and dtype == torch.bfloat16 + and state_dtype in (torch.float32, torch.bfloat16) + and lower_bound is not None + and lower_bound < 0 + ) + + +def _flashinfer_kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float, + state: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + out: torch.Tensor, +) -> None: + """FlashInfer's CuTe DSL chunked KDA prefill (raw gate and beta logits), + reading and writing the recurrent states in place at ``state_indices`` + (made contiguous: with speculative decoding they are column 0 of the + [N, 1 + num_spec] block table, which FlashInfer rejects).""" + flashinfer_recurrent_kda( + q=q.contiguous(), + k=k.contiguous(), + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + scale=q.shape[-1] ** -0.5, + initial_state=state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices.contiguous(), + output=out, + beta_is_logit=True, + backend="cute-dsl", + ) + + def _is_flashinfer_kda_decode_supported( num_heads: int, head_dim: int, @@ -445,6 +526,7 @@ def _a_log_weight_loader(param, loaded_weight): self._conv_state_dim_first = is_conv_state_dim_first() additional_config = vllm_config.additional_config + conv_state_dtype, recurrent_state_dtype = self.get_state_dtype() self.kda_prefill_backend = _resolve_kda_prefill_backend( additional_config.get("kda_prefill_backend", "auto") if isinstance(additional_config, dict) @@ -452,12 +534,12 @@ def _a_log_weight_loader(param, loaded_weight): self.head_dim, vllm_config.model_config.dtype, self.kda_lower_bound, + recurrent_state_dtype, ) # With speculative decoding every decode row is a draft-verify row of # 1 + num_spec tokens (kda_spec_decode_backend); otherwise rows have # one token (kda_decode_backend). - conv_state_dtype, recurrent_state_dtype = self.get_state_dtype() decode_backend_key = ( "kda_spec_decode_backend" if self.num_spec > 0 else "kda_decode_backend" ) @@ -696,6 +778,66 @@ def _flashinfer_decode( **kwargs, ) + def _flashinfer_prefill( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + out: torch.Tensor | None, + ) -> torch.Tensor: + """FlashInfer's chunked prefill (raw gate and beta logits), in place on + the state pool at ``state_indices``: no gather / scatter. Writes + ``out`` (a new buffer when ``None``) and returns it.""" + # FlashInfer has no has-initial-state flag: zero the slots of fresh + # sequences first (the others zero the unused null slot). + state.index_fill_( + 0, torch.where(has_initial_state, NULL_BLOCK_ID, state_indices).long(), 0 + ) + if out is None: + out = torch.empty_like(q) + if q.size(1) > 1: + assert self.kda_lower_bound is not None + _flashinfer_kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log.view(-1), + dt_bias=self.dt_bias, + lower_bound=self.kda_lower_bound, + state=state, + state_indices=state_indices, + cu_seqlens=cu_seqlens, + out=out, + ) + else: + # FlashInfer's prefill needs two tokens; one is a decode step. + fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=state, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + out=out, + sigmoid_beta=True, + a_log=self.A_log, + g_bias=self.dt_bias, + compute_gate=True, + lower_bound=self.kda_lower_bound, + ) + return out + @eager_break_during_capture def _forward( self, @@ -975,9 +1117,25 @@ def _rearr(x): v_ns = v_ns[num_fused_decodes:] g1_ns = g1_ns[:, num_fused_decodes:] beta_ns = beta_ns[:, num_fused_decodes:] - initial_state = gather_initial_states( - recurrent_state, prefill_state_indices, prefill_has_initial_state - ) + if self.kda_prefill_backend == "flashinfer": + core_attn_out_non_spec = ns_out = self._flashinfer_prefill( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + g=g1_ns, + beta=beta_ns, + state=recurrent_state, + state_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + cu_seqlens=prefill_query_start_loc, + out=None + if use_spec + else core_attn_out[:, num_fused_decodes:num_actual_tokens], + ) + else: + initial_state = gather_initial_states( + recurrent_state, prefill_state_indices, prefill_has_initial_state + ) if self.kda_prefill_backend == "flashkda": # Non-spec step: write straight into the layer output buffer # (dense token order, no merge copy). Step with spec-decode @@ -1001,7 +1159,7 @@ def _rearr(x): conv_state=conv_state, recurrent_state=recurrent_state, ) - else: + elif self.kda_prefill_backend == "triton": ( core_attn_out_non_spec, last_recurrent_state, @@ -1022,12 +1180,13 @@ def _rearr(x): safe_gate=safe_gate, lower_bound=lower_bound, ) - # Init cache - scatter_states( - recurrent_state, - last_recurrent_state, - prefill_state_indices, - ) + if self.kda_prefill_backend != "flashinfer": + # Init cache + scatter_states( + recurrent_state, + last_recurrent_state, + prefill_state_indices, + ) elif attn_metadata_narrowed.num_decodes > 0: assert non_spec_query_start_loc is not None assert non_spec_state_indices_tensor is not None