[Perf][GLM-5.3-Flash][SM90] Shape-aware launch configuration for recurrent KDA

tests/models/glm5next/test_kda_recurrent.py

@@ -16,7 +16,10 @@
 if current_platform.is_rocm():
     from vllm.models.glm5next.amd.ops.third_party.kda import fused_recurrent_kda
 else:
-    from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda
+    from vllm.models.glm5next.nvidia.ops.third_party.kda import (
+        fused_recurrent_kda,
+        kernels,
+    )
 
 pytestmark = pytest.mark.skipif(
     not current_platform.is_cuda_alike(), reason="Requires CUDA or ROCm"
@@ -164,6 +167,32 @@ def test_fused_recurrent_kda_strided_inputs_bit_identical_to_contiguous(
     torch.testing.assert_close(state, state_ref, rtol=0, atol=0)
 
 
+@pytest.mark.skipif(
+    current_platform.is_rocm(),
+    reason="The SM90 launch dispatch only exists in the NVIDIA implementation",
+)
+@pytest.mark.parametrize(
+    ("num_seqs", "query_len"), [(1, 1), (7, 1), (3, 3)], ids=["1x1", "7x1", "3x3"]
+)
+@torch.inference_mode()
+def test_fused_recurrent_kda_sm90_launch_config_bit_identical(
+    monkeypatch, num_seqs: int, query_len: int
+):
+    """The SM90 dispatch (BV=16, small-grid num_stages=2) must be bit-identical
+    to the default configuration: num_warps stays 1, so the per-element
+    reduction order is unchanged."""
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    inputs, state = make_inputs(num_seqs, query_len, device)
+    state_ref = state.clone()
+    monkeypatch.setattr(kernels, "_IS_SM90", False)
+    out_ref = run_kernel(inputs, state_ref)
+    monkeypatch.setattr(kernels, "_IS_SM90", True)
+    out = run_kernel(inputs, state)
+    torch.testing.assert_close(out, out_ref, rtol=0, atol=0)
+    torch.testing.assert_close(state, state_ref, rtol=0, atol=0)
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda_alike(),
     reason="Both implementations require zero-copy token-strided layouts",

vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py

@@ -32,6 +32,16 @@
 BT_LIST_AUTOTUNE = [32, 64, 128]
 NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
 
+_IS_SM90: bool | None = None
+
+
+def _is_sm90() -> bool:
+    global _IS_SM90
+    if _IS_SM90 is None:
+        _IS_SM90 = torch.cuda.is_available()
+        _IS_SM90 = _IS_SM90 and torch.cuda.get_device_capability() == (9, 0)
+    return _IS_SM90
+
 
 def fused_recurrent_kda_fwd(
     q: torch.Tensor,
@@ -61,6 +71,18 @@ def fused_recurrent_kda_fwd(
     assert NK == 1, "NK > 1 is not supported yet"
     num_stages = 3
     num_warps = 1
+    # Shape-aware launch configuration, tuned on H20 (SM90) with K=V=128:
+    # a wider V tile (BV=16) amortizes the per-program q/k/gate loads, and
+    # stages=2 only pays off when the grid is small; large grids keep the
+    # default 3 so non-H20 SM90 parts see no untuned stage change. num_warps
+    # stays 1, so the reduction order is unchanged and results are
+    # bit-identical to the default configuration. Other platforms keep the
+    # default above.
+    if _is_sm90():
+        BV = min(next_power_of_2(V), 16)
+        NV = cdiv(V, BV)
+        if N * HV <= 32:
+            num_stages = 2
 
     if compute_gate:
         assert a_log is not None and g_bias is not None, (