diff --git a/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py index 8f7c34200..3d6758db1 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/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 diff --git a/test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py b/test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py index fd5791f85..f66a89320 100644 --- a/test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py +++ b/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):