[ROCm][Perf] Add AITER FlashKDA prefill backend for GLM KDA

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

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
 
@@ -431,3 +434,115 @@ 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.conv1d = SimpleNamespace(
+        weight=0.5 * torch.randn(width, 1, 4, device=device), 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}"

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,
@@ -2672,6 +2673,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
@@ -5257,6 +5286,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()

vllm/engine/arg_utils.py

@@ -1961,8 +1961,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",
@@ -2871,12 +2872,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(