From fe45af1e6f04ac46d3a1a433cbb2049b12b023b9 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:07:49 -0700 Subject: [PATCH] perf(gdn): select ReplaySSM verify loop unrolling by shape (#36970) --- .../ops/attention/cutedsl_gdn_mtp_ring.py | 40 ++++++++++++++++++- .../gdn/test_gdn_cutedsl_ring_verify.py | 14 ++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py b/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py index d4700699e..cf50bb4cb 100644 --- a/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py +++ b/python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py @@ -918,6 +918,7 @@ def gdn_wide_vec_kernel( disable_output: cutlass.Constexpr[bool], recovery_steps: cutlass.Constexpr[int], per_request_accepted_steps: cutlass.Constexpr[bool], + phase_b_unroll: cutlass.Constexpr[int], per_token_pool_scatter: cutlass.Constexpr[bool], per_token_pool_scatter_flat: cutlass.Constexpr[bool], replayssm_rawv: cute.Tensor, @@ -1680,7 +1681,7 @@ def gdn_wide_vec_kernel( _loop_limit = T_decode_const # constexpr int — propagates as constexpr for i_t_offset in cutlass.range( _loop_limit, - unroll=1, + unroll=phase_b_unroll, unroll_full=(T_decode_const <= 1) and not per_request_accepted_steps, ): if cutlass.const_expr(per_request_fused): @@ -2615,6 +2616,7 @@ def _run_wide_vec( disable_output: cutlass.Constexpr[bool], recovery_steps: cutlass.Constexpr[int], per_request_accepted_steps: cutlass.Constexpr[bool], + phase_b_unroll: cutlass.Constexpr[int], per_token_pool_scatter: cutlass.Constexpr[bool], per_token_pool_scatter_flat: cutlass.Constexpr[bool], replayssm_rawv: cute.Tensor, @@ -2668,6 +2670,7 @@ def _run_wide_vec( disable_output, recovery_steps, per_request_accepted_steps, + phase_b_unroll, per_token_pool_scatter, per_token_pool_scatter_flat, replayssm_rawv, @@ -2957,6 +2960,15 @@ def _get_bf16_mtp_config( # picking is done by `_select_wide_vec_tile_v` below. _WIDE_VEC_WORK_UNITS_THRESHOLD = 128 +# The ReplaySSM verify path uses a short Phase-B loop while also writing its +# raw ring. Two-way unrolling exposes useful ILP in the measured T=3..8 range +# without the register-pressure regressions measured at larger factors. Keep +# shapes outside that range and non-replay modes on the generic schedule. +_WIDE_VEC_PHASE_B_DEFAULT_UNROLL = 1 +_WIDE_VEC_PHASE_B_REPLAY_UNROLL = 2 +_WIDE_VEC_PHASE_B_REPLAY_MIN_SEQ_LEN = 3 +_WIDE_VEC_PHASE_B_REPLAY_MAX_SEQ_LEN = 8 + def _select_wide_vec_tile_v(B: int, HV: int) -> Optional[int]: """Pick a wide_vec tile_v by `work_units = B * HV`, or return None to @@ -2986,6 +2998,25 @@ def _select_wide_vec_tile_v(B: int, HV: int) -> Optional[int]: return None +def _select_wide_vec_phase_b_unroll( + seq_len: int, + cache_ring: bool, +) -> int: + """Select the compile-time Phase-B schedule for the wide-vector kernel. + + Use the tuned schedule only for the measured ReplaySSM ring-write domain; + retain the generic one-way schedule as an explicit fallback. + """ + if ( + cache_ring + and _WIDE_VEC_PHASE_B_REPLAY_MIN_SEQ_LEN + <= seq_len + <= _WIDE_VEC_PHASE_B_REPLAY_MAX_SEQ_LEN + ): + return _WIDE_VEC_PHASE_B_REPLAY_UNROLL + return _WIDE_VEC_PHASE_B_DEFAULT_UNROLL + + # ============================================================================== # PYTHON ENTRY (wide_vec) — called from gated_delta_rule and gated_delta_rule_mtp # ============================================================================== @@ -3227,6 +3258,11 @@ def gated_delta_rule_mtp_wide_vec( ), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}" assert ssm_state_indices.device == q.device + phase_b_unroll = _select_wide_vec_phase_b_unroll( + seq_len=T_val, + cache_ring=cache_ring, + ) + # Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size # agnostic); padded/strided pool keeps real pool_size/stride in the key. contiguous_pool = initial_state_source.is_contiguous() @@ -3257,6 +3293,7 @@ def gated_delta_rule_mtp_wide_vec( disable_output, recovery_steps, per_request_accepted_steps, + phase_b_unroll, per_token_pool_scatter, per_token_pool_scatter_flat, cache_ring, @@ -3355,6 +3392,7 @@ def gated_delta_rule_mtp_wide_vec( disable_output, recovery_steps, per_request_accepted_steps, + phase_b_unroll, per_token_pool_scatter, per_token_pool_scatter_flat, rawv_, diff --git a/test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py b/test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py index 9c827275a..01bbd9160 100644 --- a/test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py +++ b/test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py @@ -17,7 +17,10 @@ import unittest import torch -from sglang.kernels.ops.attention.cutedsl_gdn_mtp_ring import gated_delta_rule_mtp +from sglang.kernels.ops.attention.cutedsl_gdn_mtp_ring import ( + _select_wide_vec_phase_b_unroll, + gated_delta_rule_mtp, +) from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating from sglang.kernels.ops.attention.fla.gdn_replayssm_spec_fold import ( commit_gdn_replayssm_fold_all_layers, @@ -86,6 +89,15 @@ def _verify(gating, inputs, state, slots, rings=None, disable_state_update=True) class TestGdnCuteDSLRingVerify(CustomTestCase): + def test_phase_b_unroll_selector(self): + select = _select_wide_vec_phase_b_unroll + for seq_len in range(3, 9): + with self.subTest(seq_len=seq_len, cache_ring=True): + self.assertEqual(select(seq_len, cache_ring=True), 2) + for seq_len, cache_ring in ((2, True), (9, True), (6, False)): + with self.subTest(seq_len=seq_len, cache_ring=cache_ring): + self.assertEqual(select(seq_len, cache_ring=cache_ring), 1) + def _run(self, B): gating, inputs, state0, slots, rings = _case(B)