[AMD] Tune Qwen3.5 TP4 GDN recurrent launch on gfx950 (#39987)

python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py

@@ -5,6 +5,33 @@ import triton
 import triton.language as tl
 
 from sglang.kernels.jit.utils import is_arch_support_pdl
+from sglang.srt.utils import is_gfx95_supported, is_hip
+
+_is_hip = is_hip()
+_is_gfx95 = is_gfx95_supported()
+
+
+def _select_recurrent_launch_config(
+    n: int,
+    h: int,
+    hv: int,
+    k: int,
+    v: int,
+    is_kda: bool,
+) -> tuple[int, int]:
+    """Select the value tile and warp count for recurrent GDN."""
+    if (
+        _is_hip
+        and _is_gfx95
+        and not is_kda
+        and 0 < n <= 32
+        and h == 4
+        and hv == 16
+        and k == 128
+        and v == 128
+    ):
+        return (8, 4) if n == 1 else (16, 2)
+    return min(triton.next_power_of_2(v), 32), 1
 
 
 @triton.jit(do_not_specialize=["T"])
@@ -401,11 +428,11 @@ def fused_sigmoid_gating_delta_rule_update(
     stride_a = a.stride()[1] if a.ndim == 4 else a.stride()[-2]
     HV = v.shape[2]
     N = B if cu_seqlens is None else len(cu_seqlens) - 1
-    BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)
+    BV, num_warps = _select_recurrent_launch_config(N, H, HV, K, V, is_kda)
+    BK = triton.next_power_of_2(K)
     NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
     assert NK == 1, "NK > 1 is not supported yet"
     num_stages = 3
-    num_warps = 1
 
     if scale is None:
         scale = k.shape[-1] ** -0.5

test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py

@@ -11,6 +11,7 @@ import sys
 import pytest
 import torch
 
+from sglang.srt.utils import is_gfx95_supported
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 try:
@@ -19,6 +20,7 @@ try:
         fused_recurrent_gated_delta_rule_update,
     )
     from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
+        _select_recurrent_launch_config,
         fused_sigmoid_gating_delta_rule_update,
     )
 
@@ -180,6 +182,57 @@ def test_fused_gdn_mtp_precision(N: int, T: int):
     torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
 
 
+@pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernel not available")
+@pytest.mark.parametrize("N", [1, 3, 16])
+def test_qwen35_tp4_fused_gdn_mtp_precision(N: int):
+    """Exercise the gfx950 TP4 launch shape against the reference path."""
+    T, H, HV, K, V = 4, 4, 16, 128, 128
+    A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
+        N, T, H, HV, K, V
+    )
+
+    out_ref = run_reference(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state.clone(),
+        indices,
+        cu_seqlens,
+        disable_state_update=True,
+    )
+    out_fused = run_fused_mtp(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state.clone(),
+        indices,
+        cu_seqlens,
+        disable_state_update=True,
+    )
+
+    torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
+
+
+@pytest.mark.skipif(
+    not (torch.version.hip and is_gfx95_supported()), reason="requires AMD gfx95"
+)
+def test_qwen35_tp4_launch_config_is_narrow():
+    assert _select_recurrent_launch_config(1, 4, 16, 128, 128, False) == (8, 4)
+    assert _select_recurrent_launch_config(3, 4, 16, 128, 128, False) == (16, 2)
+    assert _select_recurrent_launch_config(32, 4, 16, 128, 128, False) == (16, 2)
+    assert _select_recurrent_launch_config(33, 4, 16, 128, 128, False) == (32, 1)
+    assert _select_recurrent_launch_config(3, 8, 32, 128, 128, False) == (32, 1)
+    assert _select_recurrent_launch_config(3, 4, 16, 128, 128, True) == (32, 1)
+
+
 @pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernels not available")
 @pytest.mark.parametrize("N", [1, 16, 128])
 def test_mtp_single_step_decode(N: int):