[ROCm][Perf] Use AITER single-launch KDA decode in GLM-5.3-Flash

tests/models/glm5next/test_kda_recurrent.py

@@ -265,6 +265,208 @@ def run(sl: slice, cu_: torch.Tensor, st: torch.Tensor, out: torch.Tensor):
     assert torch.equal(full_state, ref_state)
 
 
+def _aiter_kda_decode_runnable() -> bool:
+    if not current_platform.is_rocm():
+        return False
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.platforms.rocm import get_cdna_version
+
+    return get_cdna_version() >= 4 and rocm_aiter_ops._kda_gluon_decode_importable()
+
+
+@pytest.mark.skipif(
+    not _aiter_kda_decode_runnable(),
+    reason="Needs AITER's Gluon KDA decode kernel (CDNA4+)",
+)
+@torch.inference_mode()
+def test_aiter_decode_route_follows_replayed_step(dist_init, monkeypatch):
+    """AITER's single-launch decode kernel runs on plain decode steps only.
+
+    Breakable PIECEWISE graphs replay ``_forward`` against each step's metadata,
+    so the route is chosen there. Capture on an all-decode batch, then replay a
+    mixed step (stock chain) and a decode step (one fused launch); a route fixed
+    in ``forward`` fails one of the two. The reference is a layer built without
+    AITER, with the same weights and caches.
+    """
+    from transformers import Glm5NextTextConfig
+
+    from tests.v1.attention.utils import create_vllm_config
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+    from vllm.config import CUDAGraphMode, set_current_vllm_config
+    from vllm.forward_context import ForwardContext, override_forward_context
+    from vllm.models.glm5next.common.kda import Glm5NextLinearAttention
+    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
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    # 32 heads: one rank's share of GLM-5.3-Flash's 64 at tensor parallel 2. With
+    # hidden_size = heads * D, o_proj can be the identity, so the layer output is
+    # the normed core output the tolerances below are set for.
+    heads, max_tokens, num_slots = 32, 64, 80
+    hidden_size = heads * D
+    config = Glm5NextTextConfig(
+        hidden_size=hidden_size,
+        linear_num_heads=heads,
+        linear_head_dim=D,
+        linear_lower_bound=LOWER_BOUND,
+    )
+    vllm_config = create_vllm_config(
+        model_name="Qwen/Qwen3.5-0.8B", dtype=torch.bfloat16
+    )
+    layers: dict[str, Glm5NextLinearAttention] = {}
+    try:
+        with (
+            set_current_vllm_config(vllm_config),
+            set_default_torch_dtype(torch.bfloat16),
+            torch.device(device),
+        ):
+            for index, (name, use_aiter) in enumerate((("stock", "0"), ("fused", "1"))):
+                with monkeypatch.context() as env:
+                    env.setenv("VLLM_ROCM_USE_AITER", use_aiter)
+                    rocm_aiter_ops.refresh_env_variables()
+                    layers[name] = Glm5NextLinearAttention(
+                        config, vllm_config, prefix=f"model.layers.{index}.linear_attn"
+                    )
+    finally:
+        rocm_aiter_ops.refresh_env_variables()
+
+    stock = layers["stock"]
+    for name, param in stock.named_parameters():
+        if name.endswith("conv1d.weight"):
+            param.normal_(0, 0.5)
+        elif name == "A_log":
+            param.copy_(torch.log(15 * torch.rand_like(param) + 1))
+        elif name == "dt_bias":
+            param.normal_(0, 0.1)
+        elif name == "o_norm.weight":
+            param.normal_(1, 0.1)
+        elif name == "o_proj.weight":
+            param.copy_(torch.eye(hidden_size))
+        else:  # Input projections.
+            param.normal_(0, param.shape[1] ** -0.5)
+    layers["fused"].load_state_dict(stock.state_dict())
+    conv_shape, state_shape = stock.get_state_shape()
+    saved_caches = (
+        (0.5 * torch.randn(num_slots, *conv_shape, device=device)).bfloat16(),
+        0.1 * torch.randn(num_slots, *state_shape, device=device),
+    )
+    for layer in layers.values():
+        layer.kv_cache = tuple(cache.clone() for cache in saved_caches)
+
+    def metadata(query_lens: list[int]) -> GDNAttentionMetadata:
+        """One request per entry on slots 1, 2, ..., one-token requests first
+        (the builder counts them as decodes). The state indices are a strided
+        column of the block table, as outside FULL graphs."""
+        num_requests, num_decodes = len(query_lens), query_lens.count(1)
+        block_table = torch.zeros(num_requests, 8, dtype=torch.int32, device=device)
+        block_table[:, 0] = torch.arange(1, num_requests + 1, device=device)
+        cu_seqlens = torch.tensor([0, *query_lens]).cumsum(0, dtype=torch.int32)
+        nums_dict, batch_ptr, token_chunk_offset_ptr = compute_causal_conv1d_metadata(
+            cu_seqlens, device=device
+        )
+        return GDNAttentionMetadata(
+            num_prefills=num_requests - num_decodes,
+            num_prefill_tokens=sum(query_lens) - num_decodes,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decodes,
+            num_spec_decodes=0,
+            num_spec_decode_tokens=0,
+            num_actual_tokens=sum(query_lens),
+            has_initial_state=torch.arange(num_requests, device=device) < num_decodes,
+            non_spec_query_start_loc=cu_seqlens.to(device),
+            non_spec_state_indices_tensor=block_table[:, 0],
+            nums_dict=nums_dict,
+            batch_ptr=batch_ptr,
+            token_chunk_offset_ptr=token_chunk_offset_ptr,
+        )
+
+    def forward_context(step_metadata: GDNAttentionMetadata):
+        return override_forward_context(
+            ForwardContext(
+                {},
+                {layer.prefix: step_metadata for layer in layers.values()},
+                {},
+                cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+            )
+        )
+
+    launches: list[int] = []
+    launch = rocm_aiter_ops.kda_fused_recurrent_packed_decode
+
+    def record_launch(**kwargs):
+        launches.append(kwargs["mixed_qkv"].shape[0])
+        return launch(**kwargs)
+
+    monkeypatch.setattr(
+        rocm_aiter_ops, "kda_fused_recurrent_packed_decode", record_launch
+    )
+
+    stream = torch.cuda.Stream()
+    positions = torch.zeros(max_tokens, dtype=torch.int64, device=device)
+    graphs = {}
+    for name, layer in layers.items():
+        static_hidden = torch.zeros(
+            max_tokens, hidden_size, dtype=torch.bfloat16, device=device
+        )
+        torch.accelerator.synchronize()
+        with torch.cuda.stream(stream), forward_context(metadata([1] * max_tokens)):
+            layer(static_hidden, positions)  # Warm-up, as before every capture.
+            with BreakableCUDAGraphCapture() as capture:
+                output = layer(static_hidden, positions)
+        stream.synchronize()
+        assert capture.num_eager_breaks == 1  # ``_forward`` alone.
+        graphs[name] = (capture, static_hidden, output)
+
+    def replay(name: str, step_metadata: GDNAttentionMetadata, hidden: torch.Tensor):
+        """The outputs, caches and fused launches of one replay of graph ``name``
+        from the saved caches."""
+        capture, static_hidden, output = graphs[name]
+        launches.clear()
+        with torch.cuda.stream(stream), forward_context(step_metadata):
+            for cache, saved in zip(layers[name].kv_cache, saved_caches):
+                cache.copy_(saved)
+            static_hidden.copy_(hidden)
+            capture.replay()
+        stream.synchronize()
+        num_tokens = step_metadata.num_actual_tokens
+        conv, recurrent = (cache.clone() for cache in layers[name].kv_cache)
+        return output[:num_tokens].clone(), conv, recurrent, launches[:]
+
+    # The fused kernel keeps outputs in fp32 where the stock chain rounds to bf16:
+    # 2 bf16 ulp of the largest value. State rtol is Kimi-K3's spec-decode one, as
+    # one of these 41.9M elements reaches 3.8e-3, past its packed-decode 2e-3.
+    fused_rtol = 2 * torch.finfo(torch.bfloat16).eps
+    fused_state_tol = {"atol": 2e-3, "rtol": 3e-2}
+    # (query lengths, expected fused launches (rows each), output rtol, state tol)
+    steps: list[tuple[list[int], list[int], float, dict[str, float]]] = [
+        ([1, 1, 50], [], 0, {"atol": 0, "rtol": 0}),  # Mixed: stock, bit for bit.
+        ([1, 1, 1], [3], fused_rtol, fused_state_tol),  # Decode: one fused launch.
+    ]
+    for query_lens, expected_launches, rtol, state_tol in steps:
+        step_metadata = metadata(query_lens)
+        hidden = torch.randn(max_tokens, hidden_size, device=device).bfloat16()
+        torch.accelerator.synchronize()
+        stock_out, stock_conv, stock_state, stock_launches = replay(
+            "stock", step_metadata, hidden
+        )
+        fused_out, fused_conv, fused_state, fused_launches = replay(
+            "fused", step_metadata, hidden
+        )
+
+        assert stock_launches == []
+        assert fused_launches == expected_launches, (
+            f"{query_lens}: fused launches (rows each) {fused_launches}"
+        )
+        torch.testing.assert_close(
+            fused_out, stock_out, rtol=rtol, atol=rtol * stock_out.abs().max().item()
+        )
+        torch.testing.assert_close(fused_state, stock_state, **state_tol)
+        torch.testing.assert_close(fused_conv, stock_conv, rtol=0, atol=0)
+
+
 @pytest.mark.parametrize("offset", [16, 144, 256])
 @pytest.mark.parametrize("dim_first", [False, True])
 @pytest.mark.parametrize("num_spec", [0, 3])

vllm/models/glm5next/common/kda.py

@@ -8,10 +8,12 @@
 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
 from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
@@ -63,6 +65,8 @@
         fused_recurrent_kda,
     )
 
+logger = init_logger(__name__)
+
 
 class _Glm5NextMergedColumnParallelLinear(MergedColumnParallelLinear):
     """Merged projection with multiple replicated output shards.
@@ -306,6 +310,31 @@ def __init__(
             self.tp_rank,
             prefix=f"{prefix}.conv1d",
         )
+        # [3, conv_size, dim] copy of the q|k|v conv weight in the layout AITER's
+        # decode kernel reads (built on first decode forward, after weights are
+        # loaded). See _forward.
+        self._decode_conv_weight: torch.Tensor | None = None
+        additional_config = vllm_config.additional_config
+        decode_backend = (
+            additional_config.get("kda_decode_backend", "auto")
+            if isinstance(additional_config, dict)
+            else "auto"
+        )
+        # AITER's single-launch decode needs an fp32 state, head_dim 128 and a
+        # width-4 conv. Without a lower bound the stock kernel falls back to -5.0
+        # while AITER switches to the unbounded gate, so that config keeps the
+        # stock chain, as do spec decode and any explicit --kda-decode-backend.
+        self._use_aiter_decode = (
+            decode_backend == "auto"
+            and recurrent_state_dtype == torch.float32
+            and self.head_dim == 128
+            and self.conv_size == 4
+            and config.linear_lower_bound is not None
+            and self.num_spec == 0
+            and bool(rocm_aiter_ops.is_kda_gluon_decode_available())
+        )
+        if self._use_aiter_decode:
+            logger.info_once("GLM-5.3-Flash KDA decode: AITER single-launch kernel.")
 
         self.A_log = nn.Parameter(
             torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
@@ -349,7 +378,6 @@ def _a_log_weight_loader(param, loaded_weight):
         # every _forward call (it reads an env-derived flag each time).
         self._conv_state_dim_first = is_conv_state_dim_first()
 
-        additional_config = vllm_config.additional_config
         self.kda_prefill_backend = _resolve_kda_prefill_backend(
             additional_config.get("kda_prefill_backend", "auto")
             if isinstance(additional_config, dict)
@@ -495,15 +523,21 @@ def forward(
             dtype=hidden_states.dtype,
             device=hidden_states.device,
         )
+        # AITER's single-launch decode applies the gated RMSNorm in-kernel, so an
+        # eligible layer hands the norm to _forward. The route is chosen there
+        # per step: breakable PIECEWISE graphs replay _forward, not this function.
+        out_gate = g2 if self._use_aiter_decode else None
         # Call the decorated eager break directly so host-side prefill branches
         # are not captured by PIECEWISE CUDA graphs.
         self._forward(
             qkv_proj_states=qkv,
             g1=g1,
             beta=beta,
             core_attn_out=core_attn_out,
+            out_gate=out_gate,
         )
-        core_attn_out = self.o_norm(core_attn_out, g2)
+        if out_gate is None:
+            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]
 
@@ -514,17 +548,22 @@ def _forward(
         g1: torch.Tensor,
         beta: torch.Tensor,
         core_attn_out: torch.Tensor,
+        out_gate: torch.Tensor | None = None,
     ) -> None:
+        """With ``out_gate``, also apply the gated RMSNorm that ``forward`` skips."""
         forward_context = get_forward_context()
         attn_metadata_raw = forward_context.attn_metadata
 
-        if attn_metadata_raw is None:
-            return
-
-        assert isinstance(attn_metadata_raw, dict)
-        attn_metadata_narrowed = attn_metadata_raw.get(self.prefix)
+        attn_metadata_narrowed = None
+        if attn_metadata_raw is not None:
+            assert isinstance(attn_metadata_raw, dict)
+            attn_metadata_narrowed = attn_metadata_raw.get(self.prefix)
         if attn_metadata_narrowed is None:
-            # Profile/warmup dummy runs may omit mamba-family metadata.
+            # Profile/warmup dummy runs may omit mamba-family metadata. They
+            # still norm, as forward does for other layers, so the profiled
+            # memory matches a real step.
+            if out_gate is not None:
+                core_attn_out.copy_(self.o_norm(core_attn_out, out_gate))
             return
         assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
         has_initial_state = attn_metadata_narrowed.has_initial_state
@@ -566,6 +605,42 @@ def _forward(
             self.conv1d.weight.size(0), self.conv1d.weight.size(-1)
         )
         conv_bias = self.conv1d.bias
+        if out_gate is not None and self._decode_conv_weight is None:
+            self._decode_conv_weight = (
+                conv_weights.view(3, self.local_projection_size, self.conv_size)
+                .transpose(1, 2)
+                .contiguous()
+            )
+
+        num_decodes = attn_metadata_narrowed.num_decodes
+        if (
+            out_gate is not None
+            and attn_metadata_narrowed.num_prefills == 0
+            and num_decodes > 0
+        ):
+            # One launch for the conv update, recurrence and gated RMSNorm
+            # (out_gate implies no spec decode). The [N, 1] indices: the
+            # block-table column can be strided, which the flat form rejects.
+            assert non_spec_state_indices_tensor is not None
+            assert self._decode_conv_weight is not None
+            state_indices = non_spec_state_indices_tensor[:num_decodes].unsqueeze(1)
+            rocm_aiter_ops.kda_fused_recurrent_packed_decode(
+                mixed_qkv=qkv_proj_states[:num_decodes],
+                g=g1[0, :num_decodes],
+                beta=beta[0, :num_decodes],
+                A_log=self.A_log,
+                dt_bias=self.dt_bias,
+                lower_bound=lower_bound,
+                initial_state=recurrent_state,
+                ssm_state_indices=state_indices,
+                conv_state=conv_state,
+                conv_weight=self._decode_conv_weight,
+                out_gate=out_gate[:num_decodes],
+                norm_weight=self.o_norm.weight,
+                norm_eps=self.o_norm.eps,
+                out=core_attn_out[0, :num_decodes],
+            )
+            return
 
         # Split projections / gating into spec (draft-verify) and non-spec token
         # groups when speculative decoding is active. Spec tokens carry
@@ -792,3 +867,6 @@ def _rearr(x):
                 core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
                     0, :num_actual_tokens
                 ]
+
+        if out_gate is not None:
+            core_attn_out.copy_(self.o_norm(core_attn_out, out_gate))