[ROCm][Perf][Kimi-K3] Use AITER Gluon chunk_kda for KDA prefill (#59552)

Signed-off-by: Omar Muhammad <omar.muhammad@amd.com>
Signed-off-by: Nicholas Susanto <nicholas.susanto@amd.com>
Co-authored-by: Nicholas Susanto <nicholas.susanto@amd.com>
Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>

tests/models/kimi_k3/test_amd_kda_checkpoint.py

@@ -482,7 +482,7 @@ def _spec_for(monkeypatch, **attrs) -> MambaSpec:
     layer = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention)
     for name, value in dict(
         use_prefill_checkpoint=True,
-        use_fused_chunk=True,
+        kda_prefill_backend="fused",
         num_spec=0,
         use_safe_gate=True,
     ).items():
@@ -506,7 +506,8 @@ def test_opt_in_publishes_chunk_walk_alignment(monkeypatch) -> None:
     "attr,value",
     [
         ("use_prefill_checkpoint", False),
-        ("use_fused_chunk", False),
+        ("kda_prefill_backend", "triton"),
+        ("kda_prefill_backend", "gluon"),
         ("use_safe_gate", False),
     ],
 )

tests/models/kimi_k3/test_amd_kda_chunk.py

@@ -77,7 +77,7 @@ def _run(inp: dict, use_fused: bool, out: torch.Tensor | None = None, **kw):
         output_final_state=True,
         use_qk_l2norm_in_kernel=True,
         cu_seqlens=inp["cu"],
-        use_fused_chunk=use_fused,
+        backend=kw.pop("backend", "fused" if use_fused else "triton"),
         out=out,
         **kw,
     )
@@ -112,6 +112,27 @@ def test_fused_chunk_matches_triton(seqlens: list[int]) -> None:
     torch.testing.assert_close(state_fused, state_ref, rtol=1e-3, atol=1e-3)
 
 
+@pytest.mark.parametrize(
+    "seqlens",
+    [[1], [64], [1024], [4096], [513, 64, 1, 1200], [512] * 6],
+)
+def test_gluon_chunk_matches_triton(seqlens: list[int]) -> None:
+    """AITER's Gluon chunk_kda must match the Triton chunk path too."""
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    if not rocm_aiter_ops.is_kda_gluon_prefill_available():
+        pytest.skip("AITER Gluon chunk_kda is not available here")
+
+    inp = _inputs(seqlens, seed=len(seqlens))
+    o_ref, state_ref = _run(inp, use_fused=False)
+    o_gluon, state_gluon = _run(inp, use_fused=False, backend="gluon")
+
+    torch.testing.assert_close(o_gluon.float(), o_ref.float(), rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(state_gluon, state_ref, rtol=5e-3, atol=5e-3)
+    state_rel = (state_gluon - state_ref).norm() / state_ref.norm()
+    assert state_rel < 5e-3, f"final state relative L2 error {state_rel:.2e}"
+
+
 @pytest.mark.parametrize("use_fused", [False, True])
 def test_out_argument_receives_the_result(use_fused: bool) -> None:
     """``out`` is a promise about where the result lands.
@@ -648,7 +669,7 @@ def _run_paged(inp: dict, cache, idx, warm, use_fused: bool, **kw):
         has_initial_state=warm,
         use_qk_l2norm_in_kernel=True,
         cu_seqlens=inp["cu"],
-        use_fused_chunk=use_fused,
+        backend="fused" if use_fused else "triton",
         **kw,
     )
 

vllm/_aiter_ops.py

@@ -2606,6 +2606,27 @@ class rocm_aiter_ops:
 
         return cls.is_linear_enabled() and on_gfx950()
 
+    @staticmethod
+    @functools.cache
+    def _kda_gluon_prefill_importable() -> bool:
+        try:
+            from aiter.ops.triton.attention.chunk_kda import chunk_kda  # noqa: F401
+
+            return True
+        except ImportError:
+            return False
+
+    @classmethod
+    @if_aiter_supported
+    def is_kda_gluon_prefill_available(cls) -> bool:
+        from vllm.platforms.rocm import get_cdna_version
+
+        return (
+            cls._AITER_ENABLED
+            and get_cdna_version() >= 4
+            and cls._kda_gluon_prefill_importable()
+        )
+
     @classmethod
     def get_aiter_allreduce(cls):
         """Return the TP device communicator's AITER custom-allreduce if it has
@@ -5111,6 +5132,50 @@ class rocm_aiter_ops:
             layer_input.view(*outer_shape, hidden_size),
         )
 
+    @staticmethod
+    def kda_chunk_prefill(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor | None,
+        lower_bound: float | None,
+        cu_seqlens: torch.Tensor | None,
+        chunk_indices: torch.Tensor | None,
+        chunk_offsets: torch.Tensor | None,
+        scale: float,
+        out: torch.Tensor | None,
+        initial_state: torch.Tensor | None,
+        output_final_state: bool,
+        state_cache: torch.Tensor | None,
+        state_indices: torch.Tensor | None,
+        has_initial_state: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from aiter.ops.triton.attention.chunk_kda import chunk_kda
+
+        return chunk_kda(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            lower_bound=lower_bound,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            chunk_offsets=chunk_offsets,
+            scale=scale,
+            out=out,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            state_cache=state_cache,
+            state_indices=state_indices,
+            has_initial_state=has_initial_state,
+        )
+
 
 rocm_aiter_ops.register_ops_once()
 _sync_aiter_situv2_moe_env()

vllm/model_executor/warmup/kimi_k3_triton_warmup.py

@@ -19,9 +19,11 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
-def _get_kda_layer(worker: Worker) -> KimiK3DeltaAttention | None:
+def _get_kda_layer(
+    worker: Worker, module: str = "vllm.models.kimi_k3.nvidia.kda"
+) -> KimiK3DeltaAttention | None:
     # Kimi model construction already imports kda. Avoid importing it here.
-    kda = sys.modules.get("vllm.models.kimi_k3.nvidia.kda")
+    kda = sys.modules.get(module)
     if kda is None:
         return None
 
@@ -172,9 +174,59 @@ def _warm_recurrent_kda(
         )
 
 
+def _warm_gluon_chunk_kda(layer, input_dtype: torch.dtype) -> None:
+    from vllm.models.kimi_k3.amd.ops.kda_prefill import (
+        chunk_kda_prefill,
+        get_gluon_chunk_kda_warmup_profiles,
+    )
+
+    if layer.kda_prefill_backend != "gluon":
+        return
+    kv_cache = layer.kv_cache
+    if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) < 2:
+        return
+    state = kv_cache[1]
+    if not isinstance(state, torch.Tensor) or not state.numel():
+        return
+
+    logger.info("Warming up Kimi-K3 Gluon chunk KDA prefill kernels.")
+    h, d = int(layer.local_num_heads), int(layer.head_dim)
+    for num_seqs, seq_len in get_gluon_chunk_kda_warmup_profiles():
+        num_tokens = num_seqs * seq_len
+        shape = (1, num_tokens, h, d)
+        chunk_kda_prefill(
+            q=torch.zeros(shape, dtype=input_dtype, device=state.device),
+            k=torch.zeros(shape, dtype=input_dtype, device=state.device),
+            v=torch.zeros(shape, dtype=input_dtype, device=state.device),
+            raw_g=torch.zeros(shape, dtype=input_dtype, device=state.device),
+            raw_beta=torch.zeros(
+                (1, num_tokens, h), dtype=input_dtype, device=state.device
+            ),
+            A_log=layer.A_log,
+            g_bias=layer.dt_bias,
+            lower_bound=layer.gate_lower_bound,
+            state_cache=state,
+            state_indices=torch.zeros(num_seqs, dtype=torch.int32, device=state.device),
+            has_initial_state=torch.zeros(
+                num_seqs, dtype=torch.bool, device=state.device
+            ),
+            use_qk_l2norm_in_kernel=True,
+            cu_seqlens=torch.arange(
+                0, num_tokens + 1, seq_len, dtype=torch.int32, device=state.device
+            ),
+            backend="gluon",
+            out=torch.empty(shape, dtype=input_dtype, device=state.device),
+        )
+
+
 @torch.inference_mode()
 def kimi_k3_triton_warmup(worker: Worker) -> None:
     """Warm Kimi-K3 Triton kernels reachable by this server."""
+    if current_platform.is_rocm():
+        layer = _get_kda_layer(worker, "vllm.models.kimi_k3.amd.kda")
+        if layer is not None:
+            _warm_gluon_chunk_kda(layer, worker.model_config.dtype)
+        return
     if not current_platform.is_cuda():
         return
 

vllm/models/kimi_k3/amd/kda.py

@@ -95,7 +95,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
         assert isinstance(spec, MambaSpec)
         # Only the fused path exports the checkpoint snapshots
         can_checkpoint = (
-            self.use_prefill_checkpoint and self.use_fused_chunk and self.use_safe_gate
+            self.use_prefill_checkpoint
+            and self.kda_prefill_backend == "fused"
+            and self.use_safe_gate
         )
         if can_checkpoint and vllm_config.cache_config.mamba_cache_mode == "align":
             logger.info_once("Kimi-K3 KDA prefill checkpoint enabled.")
@@ -242,13 +244,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                 "The fused KDA chunk kernel requires gfx950 and a build that "
                 "includes it."
             )
-        self.use_fused_chunk = backend == "fused" or (
-            backend == "auto" and is_fused_kda_chunk_supported()
-        )
-        logger.info_once(
-            "Kimi-K3 KDA prefill backend: %s",
-            "fused" if self.use_fused_chunk else "triton",
-        )
+        self.kda_prefill_backend = backend
+        if backend == "auto":
+            if self.use_safe_gate and rocm_aiter_ops.is_kda_gluon_prefill_available():
+                self.kda_prefill_backend = "gluon"
+            elif is_fused_kda_chunk_supported():
+                self.kda_prefill_backend = "fused"
+            else:
+                self.kda_prefill_backend = "triton"
+        logger.info_once("Kimi-K3 KDA prefill backend: %s", self.kda_prefill_backend)
         # TODO: Remove the extra config arg once tested
         self.use_prefill_checkpoint = (
             additional_config.get("kda_prefill_checkpoint", True)
@@ -694,7 +698,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                     cu_seqlens=prefill_query_start_loc,
                     chunk_indices=m.chunk_indices,
                     chunk_offsets=m.chunk_offsets,
-                    use_fused_chunk=self.use_fused_chunk,
+                    backend=self.kda_prefill_backend,
                     out=core_attn_out[:, nd_tok:num_actual_tokens] if direct else None,
                     checkpoint_state=checkpoint_state,
                     checkpoint_offsets=checkpoint_offsets,

vllm/models/kimi_k3/amd/ops/kda_prefill.py

@@ -9,6 +9,7 @@ path.
 
 import torch
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.logger import init_logger
 from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
     gather_initial_states,
@@ -24,6 +25,11 @@ from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
 logger = init_logger(__name__)
 
 
+# Consumed by kimi_k3_triton_warmup.py during kernel_warmup().
+def get_gluon_chunk_kda_warmup_profiles() -> tuple[tuple[int, int], ...]:
+    return ((1, 64), (2, 64), (4, 64), (8, 64), (16, 64), (1, 4096), (2, 2048))
+
+
 def chunk_kda_prefill(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -40,7 +46,7 @@ def chunk_kda_prefill(
     cu_seqlens: torch.Tensor | None = None,
     chunk_indices: torch.Tensor | None = None,
     chunk_offsets: torch.Tensor | None = None,
-    use_fused_chunk: bool = False,
+    backend: str = "triton",
     out: torch.Tensor | None = None,
     checkpoint_state: torch.Tensor | None = None,
     checkpoint_offsets: torch.Tensor | None = None,
@@ -68,9 +74,10 @@ def chunk_kda_prefill(
         cu_seqlens: int32 cumulative sequence lengths.
         chunk_indices: precomputed chunk indices for ``cu_seqlens``.
         chunk_offsets: int32 per-sequence first chunk index.
-        use_fused_chunk: request the two-kernel ROCm path. It is used only when
-            every one of its preconditions holds; otherwise the Triton path
-            runs unchanged.
+        backend: ``"triton"`` (the vendored chunk path), ``"fused"`` (the
+            two-kernel ROCm path; used only when every one of its
+            preconditions holds, otherwise the Triton path runs unchanged) or
+            ``"gluon"`` (AITER's ``chunk_kda``).
         out: buffer the result must land in. Honoured by both backends, so the
             caller can hand in a slice of its own output and skip a copy.
         checkpoint_state: destination for mid-prefill recurrent state
@@ -89,6 +96,11 @@ def chunk_kda_prefill(
         The output and, when requested, the final recurrent state.
 
     """
+    if backend not in ("triton", "fused", "gluon"):
+        raise ValueError(
+            "chunk_kda_prefill backend must be 'triton', 'fused' or 'gluon', "
+            f"got {backend!r}."
+        )
     if scale is None:
         scale = k.shape[-1] ** -0.5
 
@@ -96,7 +108,7 @@ def chunk_kda_prefill(
     # needs the raw projections and a bounded gate rather than the pre-normalized
     # tensors the Triton path takes.
     fused = (
-        use_fused_chunk
+        backend == "fused"
         and use_qk_l2norm_in_kernel
         and cu_seqlens is not None
         and lower_bound is not None
@@ -116,6 +128,28 @@ def chunk_kda_prefill(
         if state_indices is None or has_initial_state is None:
             raise ValueError("state_cache needs state_indices and has_initial_state")
 
+    if backend == "gluon":
+        return rocm_aiter_ops.kda_chunk_prefill(
+            q=q,
+            k=k,
+            v=v,
+            g=raw_g,
+            beta=raw_beta,
+            A_log=A_log,
+            dt_bias=g_bias,
+            lower_bound=lower_bound,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            chunk_offsets=chunk_offsets,
+            scale=scale,
+            out=out,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            state_cache=state_cache,
+            state_indices=state_indices,
+            has_initial_state=has_initial_state,
+        )
+
     # Only the fused walk addresses the paged rows directly. The Triton path
     # gathers the rows it needs and scatters the results back here, so callers
     # pass the cache the same way for either backend.