From 20339d62ac596c3affbd36d870f2c0f0e1488fed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Sat, 24 Oct 2026 17:38:04 +0000 Subject: [PATCH 1/6] [ROCm][Perf] Add AITER FlashKDA prefill backend for GLM KDA MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Let --kda-prefill-backend=flashkda select AITER's FlashKDA chunked prefill for the GLM-5.3-Flash KDA layers on ROCm, where it was rejected before. The layer stores it as aiter_flashkda so the CUDA FlashKDA code, which reserves prefix-cache checkpoint blocks, is not entered. The flag's platform check and help text are updated to match; other ROCm KDA models still reject the value when their layer is built. AITER runs the prefill of steps without speculative-decode tokens. It reads and updates the paged recurrent state in place and writes straight into the layer output, so the layer skips the state gather and scatter and the merge copy. Steps that carry speculative-decode tokens keep the Triton chunk path. AITER reads cu_seqlens on the host, so the call runs under gpu_sync_allowed(), as the FLA path does for its chunk indices. The backend is validated once at layer construction and an unsupported request raises there instead of falling back. It needs VLLM_ROCM_USE_AITER=1, gfx950, an AITER release with the paged-state API (v0.1.23 or later), head_dim 128, bfloat16, a bounded gate and an fp32 recurrent state. 'auto' still resolves to Triton on ROCm, and the outputs differ from the Triton path at bf16 level (chunk size 32 vs 64). Tests: CPU tests check that flashkda resolves to aiter_flashkda on ROCm and that a step carrying speculative-decode tokens does not reach AITER. A GPU test runs the layer's prefill through AITER and through the Triton chunk path and compares the output and the recurrent state; it is skipped unless AITER FlashKDA is available. Co-authored-by: Claude Signed-off-by: Rony Leppänen --- .../glm5next/test_kda_prefill_backend.py | 51 ++++++++ tests/models/glm5next/test_kda_recurrent.py | 116 +++++++++++++++++- vllm/_aiter_ops.py | 85 +++++++++++++ vllm/engine/arg_utils.py | 12 +- vllm/models/glm5next/common/kda.py | 75 ++++++++++- 5 files changed, 329 insertions(+), 10 deletions(-) create mode 100644 tests/models/glm5next/test_kda_prefill_backend.py diff --git a/tests/models/glm5next/test_kda_prefill_backend.py b/tests/models/glm5next/test_kda_prefill_backend.py new file mode 100644 index 000000000000..15901341156d --- /dev/null +++ b/tests/models/glm5next/test_kda_prefill_backend.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-5.3-Flash KDA prefill through AITER FlashKDA on ROCm. Nothing here +needs a GPU: the platform and the AITER probe are faked.""" + +import pytest +import torch + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.models.glm5next.common import kda +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata + + +def test_rocm_flashkda_resolves_to_aiter(monkeypatch): + """``--kda-prefill-backend flashkda`` selects AITER FlashKDA on ROCm + instead of being rejected.""" + monkeypatch.setattr(kda.current_platform, "is_rocm", lambda: True) + monkeypatch.setattr( + rocm_aiter_ops, "is_kda_flashkda_prefill_available", lambda: True + ) + backend = kda._resolve_kda_prefill_backend( + "flashkda", + head_dim=128, + dtype=torch.bfloat16, + lower_bound=-5.0, + state_dtype=torch.float32, + ) + assert backend == "aiter_flashkda" + + +@pytest.mark.parametrize( + ("num_spec_decodes", "expected"), [(0, True), (1, False)], ids=["no-spec", "spec"] +) +def test_aiter_flashkda_skips_steps_with_spec_tokens(num_spec_decodes, expected): + """Steps that carry speculative-decode tokens keep the Triton chunk path.""" + spec_tokens = 4 * num_spec_decodes + metadata = GDNAttentionMetadata( + num_prefills=2, + num_prefill_tokens=74, + num_decodes=0, + num_decode_tokens=0, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=spec_tokens, + num_actual_tokens=74 + spec_tokens, + spec_sequence_masks=( + torch.tensor([True] * num_spec_decodes + [False] * 2) + if num_spec_decodes + else None + ), + ) + assert kda._aiter_flashkda_runs_prefill("aiter_flashkda", metadata) is expected diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index 06b2eafac6a2..e9387a9f1a2c 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/tests/models/glm5next/test_kda_recurrent.py @@ -284,7 +284,10 @@ def test_prefill_checkpoint_resumes_suffix(monkeypatch, dim_first, num_spec, off from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata from vllm.v1.worker.workspace import WorkspaceManager - if kda._resolve_kda_prefill_backend("auto", D, torch.bfloat16, -5) != "flashkda": + if ( + kda._resolve_kda_prefill_backend("auto", D, torch.bfloat16, -5, torch.float32) + != "flashkda" + ): pytest.skip("requires FlashKDA") import vllm._flashkda_C # noqa: F401 @@ -430,3 +433,114 @@ def run(qkv, g, b, slots, has_initial, checkpoint=None): actual[:, : length - offset], expected[:, suffix], atol=0.003, rtol=0.03 ) torch.testing.assert_close(recurrent[3], expected_state, atol=0.003, rtol=0.03) + + +@torch.inference_mode() +def test_aiter_flashkda_prefill_matches_triton(monkeypatch): + """With ``flashkda`` on ROCm the layer's prefill runs AITER FlashKDA on + the paged state and output in place, and matches the Triton chunk path + to bf16 precision (chunk size 32 vs 64).""" + from types import SimpleNamespace + + from vllm._aiter_ops import rocm_aiter_ops + from vllm.models.glm5next.common import kda + from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata + from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata + + if not rocm_aiter_ops.is_kda_flashkda_prefill_available(): + pytest.skip("requires AITER FlashKDA (VLLM_ROCM_USE_AITER=1, gfx950)") + + torch.manual_seed(7) + device = torch.device("cuda") + # 32 local heads: the default 64 KDA heads at tensor parallel size 2. + heads, lengths = 32, [37, 100] + num_tokens, num_slots = sum(lengths), 4 + width, state_len, slot_size = 3 * heads * D, 3, heads * D * D + layer = object.__new__(kda.Glm5NextLinearAttention) + torch.nn.Module.__init__(layer) + layer.prefix = "layer.0" + layer.local_num_heads = heads + layer.head_dim = D + layer.local_projection_size = heads * D + layer.kda_safe_gate = True + layer.kda_lower_bound = LOWER_BOUND + layer._conv_state_dim_first = False + layer._merged_conv_weight = 0.5 * torch.randn(width, 4, device=device) + layer.q_conv1d = SimpleNamespace(bias=None) + layer.A_log = 0.5 * torch.randn(1, 1, heads, 1, device=device) + layer.dt_bias = 0.1 * torch.randn(heads * D, device=device) + + # Padded slot strides, as in the hybrid KV cache. + conv_storage = torch.randn( + num_slots, width * state_len + 8, device=device, dtype=torch.bfloat16 + ) + recurrent_storage = 0.5 * torch.randn(num_slots, slot_size + 8, device=device) + projected = torch.randn( + num_tokens, width + heads, device=device, dtype=torch.bfloat16 + ) + raw_qkv, beta = projected[:, :width], projected[:, width:].unsqueeze(0) + gate = torch.randn(1, num_tokens, heads, D, device=device, dtype=torch.bfloat16) + + cu_cpu = torch.tensor([0, lengths[0], num_tokens], dtype=torch.int32) + nums, batch_ptr, chunk_offsets = compute_causal_conv1d_metadata( + cu_cpu, device=device + ) + metadata = GDNAttentionMetadata( + num_prefills=len(lengths), + num_prefill_tokens=num_tokens, + num_decodes=0, + num_decode_tokens=0, + num_spec_decodes=0, + num_spec_decode_tokens=0, + num_actual_tokens=num_tokens, + has_initial_state=torch.tensor([True, False], device=device), + non_spec_query_start_loc=cu_cpu.to(device), + non_spec_state_indices_tensor=torch.tensor( + [1, 2], dtype=torch.int32, device=device + ), + nums_dict=nums, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=chunk_offsets, + ) + monkeypatch.setattr( + kda, + "get_forward_context", + lambda: SimpleNamespace(attn_metadata={layer.prefix: metadata}), + ) + + prefill_calls = [] + kda_flashkda_prefill = rocm_aiter_ops.kda_flashkda_prefill + + def spy(**kwargs): + prefill_calls.append(kwargs) + return kda_flashkda_prefill(**kwargs) + + monkeypatch.setattr(rocm_aiter_ops, "kda_flashkda_prefill", staticmethod(spy)) + + def run(backend: str): + conv_pool, recurrent_pool = conv_storage.clone(), recurrent_storage.clone() + conv = conv_pool[:, : width * state_len].view(num_slots, state_len, width) + recurrent = recurrent_pool[:, :slot_size].view(num_slots, heads, D, D) + layer.kv_cache = (conv, recurrent) + layer.kda_prefill_backend = backend + out = torch.full( + (1, num_tokens, heads, D), float("nan"), device=device, dtype=torch.bfloat16 + ) + layer._forward(raw_qkv, gate, beta, out) + return out, recurrent + + triton_out, triton_recurrent = run("triton") + assert not prefill_calls + out, recurrent = run("aiter_flashkda") + assert len(prefill_calls) == 1 + + torch.testing.assert_close(out, triton_out, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(recurrent, triton_recurrent, rtol=2e-2, atol=2e-2) + written = slice(1, 3) + for actual, expected in ( + (out, triton_out), + (recurrent[written], triton_recurrent[written]), + ): + actual, expected = actual.float(), expected.float() + error = (actual - expected).norm() / expected.norm() + assert error < 2e-2, f"relative L2 error {error:.2e}" diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 6af965b49286..dc5327e79396 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -14,6 +14,7 @@ from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_aiter_sparse_attn_indexer, @@ -2624,6 +2625,34 @@ def is_kda_gluon_prefill_available(cls) -> bool: and cls._kda_gluon_prefill_importable() ) + @staticmethod + @functools.cache + def _kda_flashkda_paged_prefill_importable() -> bool: + import inspect + + try: + from aiter.ops.triton.kimi_delta_attn import chunk_kimi_delta_attn + except ImportError: + return False + # AITER releases before v0.1.23 export it without the paged + # state_cache / out arguments. + return "state_cache" in inspect.signature(chunk_kimi_delta_attn).parameters + + @classmethod + @if_aiter_supported + def is_kda_flashkda_prefill_available(cls) -> bool: + """Whether AITER's FlashKDA prefill with paged state can run here. + + Limited to gfx950, the only arch AITER's FlashKDA Gluon kernels target. + """ + from vllm.platforms.rocm import on_gfx950 + + return ( + cls._AITER_ENABLED + and on_gfx950() + and cls._kda_flashkda_paged_prefill_importable() + ) + @classmethod def get_aiter_allreduce(cls): """Return the TP device communicator's AITER custom-allreduce if it has @@ -5200,6 +5229,62 @@ def kda_chunk_prefill( has_initial_state=has_initial_state, ) + @staticmethod + def kda_flashkda_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, + cu_seqlens: torch.Tensor, + state_cache: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + out: torch.Tensor, + ) -> torch.Tensor: + """Run AITER FlashKDA chunked prefill over a paged recurrent state. + + ``q``, ``k`` and ``v`` are ``[1, T, H, 128]`` bf16, ``g`` is the raw + gate of the same shape and ``beta`` the raw ``[1, T, H]`` logits; the + gate activation and the sigmoid run in the kernel. ``A_log`` is + ``[H]`` and ``dt_bias`` ``[H * 128]``, both fp32. ``state_cache`` is the + fp32 V-first ``[slots, H, V, K]`` pool: sequence ``n`` reads and writes + row ``state_indices[n]`` in place. ``out`` must be contiguous with + ``v``'s shape. Returns ``out``. + """ + from aiter.ops.triton.kimi_delta_attn import chunk_kimi_delta_attn + + logger.info_once("Dispatching AITER FlashKDA KDA prefill.") + # AITER reads cu_seqlens on the host to plan its chunks and segments. + with gpu_sync_allowed(): + chunk_kimi_delta_attn( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=True, + lower_bound=lower_bound, + state_v_first=True, + # 32 selects the FlashKDA kernels, the only ones that take + # state_cache and out. + chunk_size=32, + cu_seqlens=cu_seqlens, + out=out, + state_cache=state_cache, + state_indices=state_indices, + has_initial_state=has_initial_state, + ) + return out + rocm_aiter_ops.register_ops_once() _sync_aiter_situv2_moe_env() diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 8fa86a95415d..971aa7ea6289 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1956,8 +1956,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: dest="kda_prefill_backend", choices=["auto", "triton", "flashkda", "flashinfer", "fused"], default=None, - help="Select KDA prefill backend. 'flashkda' is CUDA-only and " - "'fused' is ROCm-only; 'auto' picks a supported backend.", + help="Select KDA prefill backend. 'flashkda' runs FlashKDA on CUDA, " + "or AITER's FlashKDA on ROCm for GLM-5.3-Flash only; 'fused' is " + "ROCm-only; 'auto' picks a supported backend.", ) parser.add_argument( "--kda-decode-backend", @@ -2865,12 +2866,11 @@ def create_engine_config( if self.gdn_prefill_backend is not None: self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend if self.kda_prefill_backend is not None: - if ( - self.kda_prefill_backend == "flashkda" - and not current_platform.is_cuda() + if self.kda_prefill_backend == "flashkda" and not ( + current_platform.is_cuda() or current_platform.is_rocm() ): raise ValueError( - "--kda-prefill-backend=flashkda is only available on CUDA." + "--kda-prefill-backend=flashkda is only available on CUDA or ROCm." ) if self.kda_prefill_backend == "fused" and not current_platform.is_rocm(): raise ValueError( diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 68d405b40cfc..09bd08a6b7a9 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -8,6 +8,7 @@ from torch import nn from transformers import Glm5NextTextConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import divide @@ -135,14 +136,45 @@ 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, ) -> 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 ``chunk_kda_with_fused_gate`` path. ``backend`` comes from - ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" + ``additional_config.kda_prefill_backend`` (auto / triton / flashkda). + + On ROCm, ``flashkda`` runs AITER's FlashKDA kernels and resolves to + ``aiter_flashkda``, which reserves no prefix-cache checkpoint blocks; + ``auto`` keeps Triton there.""" if backend not in ("auto", "triton", "flashkda"): raise ValueError(f"Unsupported KDA prefill backend: {backend}") + if backend == "flashkda" and current_platform.is_rocm(): + reasons = [] + if not rocm_aiter_ops.is_kda_flashkda_prefill_available(): + reasons.append( + "AITER FlashKDA with paged state is unavailable (needs " + "VLLM_ROCM_USE_AITER=1, gfx950 and AITER >= v0.1.23)" + ) + if head_dim != 128: + reasons.append(f"head_dim is {head_dim}, not 128") + if dtype != torch.bfloat16: + reasons.append(f"model dtype is {dtype}, not bfloat16") + if lower_bound is None or not -5 <= lower_bound < 0: + reasons.append(f"KDA lower_bound {lower_bound} is outside [-5, 0)") + if state_dtype != torch.float32: + reasons.append( + f"recurrent state dtype is {state_dtype}, not float32 " + "(see --mamba-ssm-cache-dtype)" + ) + if reasons: + raise RuntimeError( + "KDA prefill backend 'flashkda' on ROCm: " + "; ".join(reasons) + ) + return "aiter_flashkda" capability = current_platform.get_device_capability() supported = ( current_platform.is_cuda() @@ -160,6 +192,16 @@ def _resolve_kda_prefill_backend( return "flashkda" if supported and backend != "triton" else "triton" +def _aiter_flashkda_runs_prefill(backend: str, metadata: GDNAttentionMetadata) -> bool: + """Whether AITER FlashKDA runs this step's chunked prefill. Steps that + carry speculative-decode tokens keep the Triton chunk path, as do all + steps of the other backends.""" + has_spec = ( + metadata.spec_sequence_masks is not None and metadata.num_spec_decodes > 0 + ) + return backend == "aiter_flashkda" and metadata.num_prefills > 0 and not has_spec + + class Glm5NextLinearAttention(GatedDeltaNetAttention): head_dim: int num_heads: int @@ -348,6 +390,7 @@ def _a_log_weight_loader(param, loaded_weight): self.head_dim, vllm_config.model_config.dtype, self.kda_lower_bound, + recurrent_state_dtype, ) self._flashkda_buffer_specs: ( tuple[tuple[tuple[int, ...], torch.dtype], ...] | None @@ -695,7 +738,33 @@ def _rearr(x): # layer output buffer; the chunked prefill kernel cannot, so this # stays None there and the merge copy below runs as before. ns_out = None - if attn_metadata_narrowed.num_prefills > 0: + if _aiter_flashkda_runs_prefill( + self.kda_prefill_backend, attn_metadata_narrowed + ): + assert q_ns is not None + assert non_spec_state_indices_tensor is not None + assert has_initial_state is not None + # AITER reads and updates the paged recurrent state in place, so + # there is no gather or scatter, and writes straight into the layer + # output: without spec tokens the prefill tokens are its leading + # slice. + ns_out = core_attn_out[:, :num_actual_tokens] + core_attn_out_non_spec = rocm_aiter_ops.kda_flashkda_prefill( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + g=g1_ns, + beta=beta_ns, + A_log=self.A_log.view(-1), + dt_bias=self.dt_bias, + lower_bound=lower_bound, + cu_seqlens=non_spec_query_start_loc, + state_cache=recurrent_state, + state_indices=non_spec_state_indices_tensor, + has_initial_state=has_initial_state, + out=ns_out, + ) + elif attn_metadata_narrowed.num_prefills > 0: assert q_ns is not None assert non_spec_state_indices_tensor is not None assert has_initial_state is not None From f7f9512a9edcdb7ad263d6386668cc796d274807 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Tue, 27 Oct 2026 15:18:16 +0000 Subject: [PATCH 2/6] [ROCm] KDA: read the prefill backend only on steps with prefills MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The AITER FlashKDA branch evaluated self.kda_prefill_backend on every non-spec step, decode included. A layer built without __init__, as in tests/models/glm5next/test_kda_conv_weight_refit.py, has no such attribute, so its plain decode raised AttributeError. Check num_prefills first, so decode reads the same attributes as the Triton path. Signed-off-by: Rony Leppänen Co-authored-by: Claude --- vllm/models/glm5next/common/kda.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 58173db29249..6cc436e85b95 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -738,7 +738,7 @@ def _rearr(x): # layer output buffer; the chunked prefill kernel cannot, so this # stays None there and the merge copy below runs as before. ns_out = None - if _aiter_flashkda_runs_prefill( + if attn_metadata_narrowed.num_prefills > 0 and _aiter_flashkda_runs_prefill( self.kda_prefill_backend, attn_metadata_narrowed ): assert q_ns is not None From cc2676cf707d861af99981efc379e01a7f72ab05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 30 Oct 2026 15:26:21 -0500 Subject: [PATCH 3/6] [ROCm] KDA: run the FlashKDA parity test in CI, and on a segmented sequence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MI355 GLM5Next CI step does not set VLLM_ROCM_USE_AITER, which rocm_aiter_ops reads at import, so the AITER FlashKDA parity test always skipped there. Enable AITER for the test with monkeypatch, so it is gated only on gfx950 and the FlashKDA paged-state API. Add a second case: 16 local heads (tensor parallel size 4) and a 4096-token sequence that starts from a cached state. It is long enough, with few enough heads, that AITER splits its recurrence into segments, which the 100-token case never reaches. Signed-off-by: Rony Leppänen Co-authored-by: Claude --- tests/models/glm5next/test_kda_recurrent.py | 23 ++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index e5f5d008314f..06877554e5dc 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/tests/models/glm5next/test_kda_recurrent.py @@ -436,8 +436,21 @@ def run(qkv, g, b, slots, has_initial, checkpoint=None): torch.testing.assert_close(recurrent[3], expected_state, atol=0.003, rtol=0.03) +@pytest.mark.parametrize( + "heads,lengths,has_initial", + [ + # 32 local heads: the default 64 KDA heads at tensor parallel size 2. + (32, [37, 100], [True, False]), + # Tensor parallel size 4, and a 4096-token sequence that starts from a + # cached state: long enough, with few enough heads, that AITER splits + # its recurrence into segments. + (16, [37, 4096], [False, True]), + ], +) @torch.inference_mode() -def test_aiter_flashkda_prefill_matches_triton(monkeypatch): +def test_aiter_flashkda_prefill_matches_triton( + monkeypatch, heads, lengths, has_initial +): """With ``flashkda`` on ROCm the layer's prefill runs AITER FlashKDA on the paged state and output in place, and matches the Triton chunk path to bf16 precision (chunk size 32 vs 64).""" @@ -448,13 +461,13 @@ def test_aiter_flashkda_prefill_matches_triton(monkeypatch): from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata + # VLLM_ROCM_USE_AITER is read at import, and CI does not set it. + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True) if not rocm_aiter_ops.is_kda_flashkda_prefill_available(): - pytest.skip("requires AITER FlashKDA (VLLM_ROCM_USE_AITER=1, gfx950)") + pytest.skip("requires gfx950 and AITER FlashKDA with paged state") torch.manual_seed(7) device = torch.device("cuda") - # 32 local heads: the default 64 KDA heads at tensor parallel size 2. - heads, lengths = 32, [37, 100] num_tokens, num_slots = sum(lengths), 4 width, state_len, slot_size = 3 * heads * D, 3, heads * D * D layer = object.__new__(kda.Glm5NextLinearAttention) @@ -495,7 +508,7 @@ def test_aiter_flashkda_prefill_matches_triton(monkeypatch): num_spec_decodes=0, num_spec_decode_tokens=0, num_actual_tokens=num_tokens, - has_initial_state=torch.tensor([True, False], device=device), + has_initial_state=torch.tensor(has_initial, device=device), non_spec_query_start_loc=cu_cpu.to(device), non_spec_state_indices_tensor=torch.tensor( [1, 2], dtype=torch.int32, device=device From 972c5a7be32d0008dba24c9a6507c7d0cabd6864 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Tue, 10 Nov 2026 10:49:17 -0600 Subject: [PATCH 4/6] Use single backticks in the docstrings this PR adds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit main's check-docstring-style hook rewrites double backticks in docstrings to single ones. The docstrings the PR adds (the AITER FlashKDA wrapper in _aiter_ops.py and the two KDA test docstrings) still had them, so the hook failed on the PR's files. Nine lines in three files; no code change. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Rony Leppänen --- tests/models/glm5next/test_kda_prefill_backend.py | 2 +- tests/models/glm5next/test_kda_recurrent.py | 2 +- vllm/_aiter_ops.py | 14 +++++++------- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/models/glm5next/test_kda_prefill_backend.py b/tests/models/glm5next/test_kda_prefill_backend.py index 15901341156d..7584d7a53743 100644 --- a/tests/models/glm5next/test_kda_prefill_backend.py +++ b/tests/models/glm5next/test_kda_prefill_backend.py @@ -12,7 +12,7 @@ def test_rocm_flashkda_resolves_to_aiter(monkeypatch): - """``--kda-prefill-backend flashkda`` selects AITER FlashKDA on ROCm + """`--kda-prefill-backend flashkda` selects AITER FlashKDA on ROCm instead of being rejected.""" monkeypatch.setattr(kda.current_platform, "is_rocm", lambda: True) monkeypatch.setattr( diff --git a/tests/models/glm5next/test_kda_recurrent.py b/tests/models/glm5next/test_kda_recurrent.py index 9c76f1a6656b..3b51b6ebdc54 100644 --- a/tests/models/glm5next/test_kda_recurrent.py +++ b/tests/models/glm5next/test_kda_recurrent.py @@ -451,7 +451,7 @@ def run(qkv, g, b, slots, has_initial, checkpoint=None): def test_aiter_flashkda_prefill_matches_triton( monkeypatch, heads, lengths, has_initial ): - """With ``flashkda`` on ROCm the layer's prefill runs AITER FlashKDA on + """With `flashkda` on ROCm the layer's prefill runs AITER FlashKDA on the paged state and output in place, and matches the Triton chunk path to bf16 precision (chunk size 32 vs 64).""" from types import SimpleNamespace diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index ad7a6adc9db7..fc2dd217820d 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -5401,13 +5401,13 @@ def kda_flashkda_prefill( ) -> torch.Tensor: """Run AITER FlashKDA chunked prefill over a paged recurrent state. - ``q``, ``k`` and ``v`` are ``[1, T, H, 128]`` bf16, ``g`` is the raw - gate of the same shape and ``beta`` the raw ``[1, T, H]`` logits; the - gate activation and the sigmoid run in the kernel. ``A_log`` is - ``[H]`` and ``dt_bias`` ``[H * 128]``, both fp32. ``state_cache`` is the - fp32 V-first ``[slots, H, V, K]`` pool: sequence ``n`` reads and writes - row ``state_indices[n]`` in place. ``out`` must be contiguous with - ``v``'s shape. Returns ``out``. + `q`, `k` and `v` are `[1, T, H, 128]` bf16, `g` is the raw + gate of the same shape and `beta` the raw `[1, T, H]` logits; the + gate activation and the sigmoid run in the kernel. `A_log` is + `[H]` and `dt_bias` `[H * 128]`, both fp32. `state_cache` is the + fp32 V-first `[slots, H, V, K]` pool: sequence `n` reads and writes + row `state_indices[n]` in place. `out` must be contiguous with + `v`'s shape. Returns `out`. """ from aiter.ops.triton.kimi_delta_attn import chunk_kimi_delta_attn From 03a76608aaa2cecbdef998d3fb00214fcc8cf947 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Tue, 10 Nov 2026 11:25:54 -0600 Subject: [PATCH 5/6] [ROCm] Raise ValueError for an unsupported KDA prefill backend in the shared Kimi GDN layer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shared Kimi GDN layer rejected every backend but Triton with an assert, so under `python -O` the check was dropped and `--kda-prefill-backend flashkda` (which now passes argument parsing on ROCm as well as CUDA) ran the Triton kernel without saying so. Raise ValueError with the same message. `auto` and `triton` are accepted as before; any other value was already rejected, and still is. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Rony Leppänen --- .../layers/mamba/gdn/kimi_gdn_linear_attn.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index 27860a31126a..f44e9a8357c0 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -283,10 +283,11 @@ def __init__( else "auto" ) backend = "triton" if backend == "auto" else backend - assert backend == "triton", ( - "The shared Kimi GDN layer only supports the Triton KDA " - f"prefill backend, got {backend!r}." - ) + if backend != "triton": + raise ValueError( + "The shared Kimi GDN layer only supports the Triton KDA " + f"prefill backend, got {backend!r}." + ) if not self.use_full_rank_gate: self.g_a_proj = ReplicatedLinear( self.hidden_size, From c0d91fd2144a2018f9b55e196b7a87672d424e72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Thu, 12 Nov 2026 10:14:17 -0600 Subject: [PATCH 6/6] [ROCm] Correct the comment above ns_out in the GLM KDA layer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The comment said only the plain-decode recurrent kernel writes into the layer output buffer, so `ns_out` stays None for prefill. That stopped being true with this PR: AITER FlashKDA prefill, and CUDA FlashKDA on steps without spec tokens, write into it too (`ns_out` is set in those branches below). Only the Triton chunk path leaves it None. Co-Authored-By: Claude Sonnet 5.5 Signed-off-by: Rony Leppänen --- vllm/models/glm5next/common/kda.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index dc03a22df26f..ef840899f0f7 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -741,9 +741,10 @@ def _rearr(x): # --- core attention: non-spec path (prefill or plain decode) --- core_attn_out_non_spec = None - # Only the plain-decode recurrent kernel can write straight into the - # layer output buffer; the chunked prefill kernel cannot, so this - # stays None there and the merge copy below runs as before. + # The plain-decode recurrent kernel and the FlashKDA prefill kernels (steps + # without spec tokens) write straight into the layer output buffer; the + # Triton chunk path cannot, so this stays None there and the merge copy + # below runs as before. ns_out = None if attn_metadata_narrowed.num_prefills > 0 and _aiter_flashkda_runs_prefill( self.kda_prefill_backend, attn_metadata_narrowed