[Qwen3-next] remove heuristics and add radix cache kl test (#14520)

python/sglang/srt/layers/attention/fla/chunk_delta_h.py

@@ -18,16 +18,6 @@ from sglang.srt.layers.attention.fla.utils import is_nvidia_hopper
 NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16]
 
 
-@triton.heuristics(
-    {
-        "USE_G": lambda args: args["g"] is not None,
-        "USE_GK": lambda args: args["gk"] is not None,
-        "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
-        "STORE_FINAL_STATE": lambda args: args["ht"] is not None,
-        "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None,
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-    }
-)
 # @triton.autotune(
 #     configs=[
 #         triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages)
@@ -337,6 +327,12 @@ def chunk_gated_delta_rule_fwd_h(
         V=V,
         BT=BT,
         BV=32,
+        USE_G=g is not None,
+        USE_GK=gk is not None,
+        USE_INITIAL_STATE=initial_state is not None,
+        STORE_FINAL_STATE=final_state is not None,
+        SAVE_NEW_VALUE=v_new is not None,
+        IS_VARLEN=cu_seqlens is not None,
         num_warps=4,
         num_stages=2,
     )

python/sglang/srt/layers/attention/fla/chunk_o.py

@@ -16,12 +16,6 @@ BKV_LIST = [64, 128] if check_shared_mem() else [32, 64]
 NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8]
 
 
-@triton.heuristics(
-    {
-        "USE_G": lambda args: args["g"] is not None,
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-    }
-)
 # @triton.autotune(
 #     configs=[
 #         triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
@@ -172,6 +166,8 @@ def chunk_fwd_o(
         BT=BT,
         BK=128,
         BV=64,
+        USE_G=g is not None,
+        IS_VARLEN=cu_seqlens is not None,
         num_warps=4,
         num_stages=2,
     )

python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py

@@ -12,12 +12,6 @@ from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
 from sglang.srt.layers.attention.fla.op import safe_exp
 
 
-@triton.heuristics(
-    {
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-        "USE_G": lambda args: args["g_cumsum"] is not None,
-    }
-)
 # @triton.autotune(
 #     configs=[
 #         triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages)
@@ -145,6 +139,8 @@ def chunk_scaled_dot_kkt_fwd(
         K=K,
         BT=BT,
         BK=64,
+        IS_VARLEN=cu_seqlens is not None,
+        USE_G=g_cumsum is not None,
         num_warps=8,
         num_stages=3,
     )

python/sglang/srt/layers/attention/fla/cumsum.py

@@ -14,12 +14,6 @@ from sglang.srt.layers.attention.fla.utils import check_shared_mem, input_guard
 BS_LIST = [32, 64] if check_shared_mem() else [16, 32]
 
 
-@triton.heuristics(
-    {
-        "HAS_SCALE": lambda args: args["scale"] is not None,
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-    }
-)
 # @triton.autotune(
 #     configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]],
 #     key=["B", "H", "BT", "IS_VARLEN", "REVERSE"],
@@ -74,19 +68,13 @@ def chunk_local_cumsum_scalar_kernel(
     tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
 
 
-@triton.heuristics(
-    {
-        "HAS_SCALE": lambda args: args["scale"] is not None,
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-    }
-)
 @triton.autotune(
     configs=[
         triton.Config({"BS": BS}, num_warps=num_warps)
         for BS in BS_LIST
         for num_warps in [2, 4, 8]
     ],
-    key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE"],
+    key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE", "HAS_SCALE"],
 )
 @triton.jit(do_not_specialize=["T"])
 def chunk_local_cumsum_vector_kernel(
@@ -202,6 +190,8 @@ def chunk_local_cumsum_scalar(
         BT=BT,
         HEAD_FIRST=head_first,
         REVERSE=reverse,
+        HAS_SCALE=scale is not None,
+        IS_VARLEN=cu_seqlens is not None,
         num_warps=8,
         num_stages=3,
     )
@@ -253,6 +243,8 @@ def chunk_local_cumsum_vector(
         BT=BT,
         HEAD_FIRST=head_first,
         REVERSE=reverse,
+        HAS_SCALE=scale is not None,
+        IS_VARLEN=cu_seqlens is not None,
     )
     return g
 

python/sglang/srt/layers/attention/fla/fused_recurrent.py

@@ -12,13 +12,6 @@ from sglang.srt.layers.attention.fla.op import exp
 from sglang.srt.layers.attention.fla.utils import input_guard
 
 
-@triton.heuristics(
-    {
-        "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
-        "STORE_FINAL_STATE": lambda args: args["ht"] is not None,
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-    }
-)
 @triton.jit(do_not_specialize=["T"])
 def fused_recurrent_gated_delta_rule_fwd_kernel(
     q,
@@ -175,8 +168,11 @@ def fused_recurrent_gated_delta_rule_fwd(
         V=V,
         BK=BK,
         BV=BV,
+        USE_INITIAL_STATE=initial_state is not None,
+        STORE_FINAL_STATE=final_state is not None,
         IS_BETA_HEADWISE=beta.ndim == v.ndim,
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        IS_VARLEN=cu_seqlens is not None,
         IS_KDA=False,
         num_warps=num_warps,
         num_stages=num_stages,
@@ -344,18 +340,6 @@ def fused_recurrent_gated_delta_rule(
 # 3
 # When calculating token 3's attention, it should attend to token 1 (parent) and token 0 (grand-parent)
 # When calculating token 2's attention, it should attend to token 0 (parent)
-@triton.heuristics(
-    {
-        "USE_INITIAL_STATE": lambda args: args["h0_source"] is not None,
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-        "CACHE_INTERMEDIATE_STATES": lambda args: args["intermediate_states_buffer"]
-        is not None,
-        "HAS_EAGLE_TREE_CUSTOM_ATTN_MASK": lambda args: args[
-            "retrieve_parent_token_ptr"
-        ]
-        is not None,
-    }
-)
 @triton.jit(do_not_specialize=["T"])
 def fused_recurrent_gated_delta_rule_update_fwd_kernel(
     q,
@@ -603,10 +587,14 @@ def fused_recurrent_gated_delta_rule_update_fwd(
         V=V,
         BK=BK,
         BV=BV,
+        USE_INITIAL_STATE=initial_state_source is not None,
         IS_BETA_HEADWISE=beta.ndim == v.ndim,
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        IS_VARLEN=cu_seqlens is not None,
         DISABLE_STATE_UPDATE=disable_state_update,
         DISABLE_OUTPUT_CALCULATION=disable_output_calculation,
+        CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None,
+        HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None,
         num_warps=num_warps,
         num_stages=num_stages,
     )

python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py

@@ -7,12 +7,6 @@ import triton.language as tl
 from sglang.srt.layers.attention.fla.utils import input_guard
 
 
-@triton.heuristics(
-    {
-        "USE_INITIAL_STATE": lambda args: args["h0_source"] is not None,
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-    }
-)
 @triton.jit(do_not_specialize=["T"])
 def fused_sigmoid_gating_delta_rule_update_kernel(
     A_log,
@@ -224,7 +218,9 @@ def fused_sigmoid_gating_delta_rule_update(
         V=V,
         BK=BK,
         BV=BV,
+        USE_INITIAL_STATE=initial_state_source is not None,
         USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        IS_VARLEN=cu_seqlens is not None,
         num_warps=num_warps,
         num_stages=num_stages,
     )

python/sglang/srt/layers/attention/fla/kda.py

@@ -152,14 +152,6 @@ def fused_recurrent_kda(
     return o, final_state
 
 
-@triton.heuristics(
-    {
-        "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,
-        "HAS_RESIDUAL": lambda args: args["residual"] is not None,
-        "HAS_WEIGHT": lambda args: args["w"] is not None,
-        "HAS_BIAS": lambda args: args["b"] is not None,
-    }
-)
 @triton.jit
 def layer_norm_gated_fwd_kernel(
     x,  # pointer to the input
@@ -240,14 +232,6 @@ def layer_norm_gated_fwd_kernel(
     tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
 
 
-@triton.heuristics(
-    {
-        "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,
-        "HAS_RESIDUAL": lambda args: args["residual"] is not None,
-        "HAS_WEIGHT": lambda args: args["w"] is not None,
-        "HAS_BIAS": lambda args: args["b"] is not None,
-    }
-)
 @triton.jit
 def layer_norm_gated_fwd_kernel1(
     x,  # pointer to the input
@@ -377,6 +361,10 @@ def layer_norm_gated_fwd(
             BT=BT,
             ACTIVATION=activation,
             IS_RMS_NORM=is_rms_norm,
+            STORE_RESIDUAL_OUT=residual_out is not None,
+            HAS_RESIDUAL=residual is not None,
+            HAS_WEIGHT=weight is not None,
+            HAS_BIAS=bias is not None,
             num_warps=4,
         )
     else:
@@ -395,6 +383,10 @@ def layer_norm_gated_fwd(
             BD=BD,
             ACTIVATION=activation,
             IS_RMS_NORM=is_rms_norm,
+            STORE_RESIDUAL_OUT=residual_out is not None,
+            HAS_RESIDUAL=residual is not None,
+            HAS_WEIGHT=weight is not None,
+            HAS_BIAS=bias is not None,
             num_warps=4,
         )
     # residual_out is None if residual is None and residual_dtype == input_dtype
@@ -487,7 +479,6 @@ class FusedRMSNormGated(nn.Module):
         )
 
 
-@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 @triton.autotune(
     configs=[
         triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages)
@@ -495,7 +486,7 @@ class FusedRMSNormGated(nn.Module):
         for num_warps in [1, 2, 4, 8]
         for num_stages in [2, 3, 4]
     ],
-    key=["BC"],
+    key=["BC", "IS_VARLEN"],
 )
 @triton.jit(do_not_specialize=["T"])
 def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter(
@@ -598,10 +589,9 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter(
     tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))
 
 
-@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 @triton.autotune(
     configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]],
-    key=["BK", "BT"],
+    key=["BK", "BT", "IS_VARLEN"],
 )
 @triton.jit(do_not_specialize=["T"])
 def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra(
@@ -755,6 +745,7 @@ def chunk_kda_scaled_dot_kkt_fwd(
         BT=BT,
         BC=BC,
         NC=NC,
+        IS_VARLEN=cu_seqlens is not None,
     )
 
     grid = (NT, NC, B * H)
@@ -774,17 +765,11 @@ def chunk_kda_scaled_dot_kkt_fwd(
         BT=BT,
         BC=BC,
         BK=BK,
+        IS_VARLEN=cu_seqlens is not None,
     )
     return A, Aqk
 
 
-@triton.heuristics(
-    {
-        "STORE_QG": lambda args: args["qg"] is not None,
-        "STORE_KG": lambda args: args["kg"] is not None,
-        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
-    }
-)
 @triton.autotune(
     configs=[
         triton.Config({}, num_warps=num_warps, num_stages=num_stages)
@@ -979,12 +964,14 @@ def recompute_w_u_fwd(
         BT=BT,
         BK=BK,
         BV=BV,
+        STORE_QG=False,
+        STORE_KG=kg is not None,
+        IS_VARLEN=cu_seqlens is not None,
         DOT_PRECISION="ieee",
     )
     return w, u, None, kg
 
 
-@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 @triton.autotune(
     configs=[
         triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
@@ -993,7 +980,7 @@ def recompute_w_u_fwd(
         for num_warps in [2, 4, 8]
         for num_stages in [2, 3, 4]
     ],
-    key=["BT"],
+    key=["BT", "IS_VARLEN"],
 )
 @triton.jit(do_not_specialize=["T"])
 def chunk_gla_fwd_kernel_o(
@@ -1143,6 +1130,7 @@ def chunk_gla_fwd_o_gk(
         K=K,
         V=V,
         BT=BT,
+        IS_VARLEN=cu_seqlens is not None,
     )
     return o
 

python/sglang/srt/layers/attention/fla/layernorm_gated.py

@@ -50,8 +50,6 @@ def rms_norm_ref(
     return out.to(dtype)
 
 
-@triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None})
-@triton.heuristics({"HAS_Z": lambda args: args["Z"] is not None})
 @triton.jit
 def _layer_norm_fwd_1pass_kernel(
     X,  # pointer to the input
@@ -177,6 +175,8 @@ def _layer_norm_fwd(
             group_size,
             eps,
             BLOCK_N=BLOCK_N,
+            HAS_BIAS=bias is not None,
+            HAS_Z=z is not None,
             NORM_BEFORE_GATE=norm_before_gate,
             IS_RMS_NORM=is_rms_norm,
             num_warps=num_warps,

python/sglang/srt/layers/attention/fla/solve_tril.py

@@ -12,7 +12,6 @@ from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
 from sglang.srt.layers.attention.fla.utils import input_guard
 
 
-@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 # @triton.autotune(
 #     configs=[
 #         triton.Config({}, num_warps=num_warps, num_stages=num_stages)
@@ -70,7 +69,6 @@ def solve_tril_16x16_kernel(
     )
 
 
-@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 # @triton.autotune(
 #     configs=[
 #         triton.Config({}, num_warps=num_warps, num_stages=num_stages)
@@ -150,7 +148,6 @@ def merge_16x16_to_32x32_inverse_kernel(
     )
 
 
-@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 # @triton.autotune(
 #     configs=[
 #         triton.Config({}, num_warps=num_warps, num_stages=num_stages)
@@ -434,6 +431,7 @@ def solve_tril(
         T=T,
         H=H,
         BT=BT,
+        IS_VARLEN=cu_seqlens is not None,
         num_warps=1,
         num_stages=4,
     )
@@ -459,6 +457,7 @@ def solve_tril(
         T=T,
         H=H,
         BT=BT,
+        IS_VARLEN=cu_seqlens is not None,
         num_warps=4,
         num_stages=3,
     )

python/sglang/srt/layers/attention/fla/wy_fast.py

@@ -11,7 +11,6 @@ import triton.language as tl
 from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
 
 
-@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 # @triton.autotune(
 #     configs=[
 #         triton.Config({}, num_warps=num_warps, num_stages=num_stages)
@@ -147,6 +146,7 @@ def recompute_w_u_fwd(
         BT=BT,
         BK=BK,
         BV=BV,
+        IS_VARLEN=cu_seqlens is not None,
         num_warps=4,
         num_stages=3,
     )