From e592ee6545109abd1d1593651bf1829c0e5ae285 Mon Sep 17 00:00:00 2001 From: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> Date: Sat, 6 Dec 2025 12:11:40 -0800 Subject: [PATCH] [Qwen3-next] remove heuristics and add radix cache kl test (#14520) --- .gitignore | 3 + .../srt/layers/attention/fla/chunk_delta_h.py | 16 +- .../srt/layers/attention/fla/chunk_o.py | 8 +- .../attention/fla/chunk_scaled_dot_kkt.py | 8 +- .../sglang/srt/layers/attention/fla/cumsum.py | 18 +- .../layers/attention/fla/fused_recurrent.py | 26 +- .../fla/fused_sigmoid_gating_recurrent.py | 8 +- python/sglang/srt/layers/attention/fla/kda.py | 46 +-- .../layers/attention/fla/layernorm_gated.py | 4 +- .../srt/layers/attention/fla/solve_tril.py | 5 +- .../srt/layers/attention/fla/wy_fast.py | 2 +- python/sglang/test/kl_test_utils.py | 324 ++++++++++++++++++ test/srt/models/test_qwen3_next_models.py | 52 ++- test/srt/run_suite.py | 2 +- 14 files changed, 418 insertions(+), 104 deletions(-) create mode 100644 python/sglang/test/kl_test_utils.py diff --git a/.gitignore b/.gitignore index 532af84b6..3326fca2c 100644 --- a/.gitignore +++ b/.gitignore @@ -249,3 +249,6 @@ lmms-eval **/.serena/ ctags/ outputs/ + +# Eval Cache +.longbench_cache/ diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py index eaaba64fd..dbae27c16 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/fla/chunk_o.py b/python/sglang/srt/layers/attention/fla/chunk_o.py index 21d338d5d..bb89421eb 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_o.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py b/python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py index 7a25e68c4..e1059ac29 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/fla/cumsum.py b/python/sglang/srt/layers/attention/fla/cumsum.py index b8e3cdde1..39d2f4722 100644 --- a/python/sglang/srt/layers/attention/fla/cumsum.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/fla/fused_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_recurrent.py index ce30ec744..92c236edd 100644 --- a/python/sglang/srt/layers/attention/fla/fused_recurrent.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py index feeb7c31c..1ed66680d 100644 --- a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/fla/kda.py b/python/sglang/srt/layers/attention/fla/kda.py index de5200c86..4235a3fe9 100644 --- a/python/sglang/srt/layers/attention/fla/kda.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py index e91e34d83..ed64cc438 100644 --- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/fla/solve_tril.py b/python/sglang/srt/layers/attention/fla/solve_tril.py index 5c519507d..4aedad342 100644 --- a/python/sglang/srt/layers/attention/fla/solve_tril.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/fla/wy_fast.py b/python/sglang/srt/layers/attention/fla/wy_fast.py index fa39312df..757e56210 100644 --- a/python/sglang/srt/layers/attention/fla/wy_fast.py +++ b/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, ) diff --git a/python/sglang/test/kl_test_utils.py b/python/sglang/test/kl_test_utils.py new file mode 100644 index 000000000..0161d0273 --- /dev/null +++ b/python/sglang/test/kl_test_utils.py @@ -0,0 +1,324 @@ +import inspect +import json +import os + +import numpy as np +import requests + +from sglang.srt.utils.hf_transformers_utils import get_tokenizer + +# LongBench V2 dataset configuration +# Reference: https://github.com/THUDM/LongBench +LONGBENCH_V2_DATASET = "THUDM/LongBench-v2" +LONGBENCH_V2_SPLIT = "train" +DEFAULT_NUM_SAMPLES = 48 # Number of samples to use +DEFAULT_PROMPT_TOKENS = 3000 # Maximum number of tokens to use +CACHE_DIR = os.path.join(os.path.dirname(__file__), ".longbench_cache") + +# In-memory cache for the current session +_cached_input_ids = {} + + +def format_longbench_v2_example(example): + """Format a LongBench V2 example into a single text string (context + question only).""" + context = example.get("context", "") + question = example.get("question", "") + return f"{context} {question}" + + +def get_input_ids( + tokenizer_path, max_prompt_tokens=DEFAULT_PROMPT_TOKENS, num_samples=None +): + """Get input_ids from LongBench V2 dataset with local caching.""" + # Create cache key based on parameters + if num_samples is None: + num_samples = DEFAULT_NUM_SAMPLES + cache_key = f"{tokenizer_path}_{max_prompt_tokens}_{num_samples}" + + # Check in-memory cache first (fastest) + if cache_key in _cached_input_ids: + print( + f"Using in-memory cached data ({len(_cached_input_ids[cache_key])} prompts)" + ) + return _cached_input_ids[cache_key] + + # Check local file cache + os.makedirs(CACHE_DIR, exist_ok=True) + # Use a safe filename + safe_name = tokenizer_path.replace("/", "_").replace("\\", "_") + cache_file = os.path.join( + CACHE_DIR, f"input_ids_{safe_name}_{max_prompt_tokens}_{num_samples}.json" + ) + + if os.path.exists(cache_file): + print(f"Loading from local cache: {cache_file}") + with open(cache_file, "r") as f: + input_ids = json.load(f) + _cached_input_ids[cache_key] = input_ids + print(f"Loaded {len(input_ids)} prompts from cache") + return input_ids + + # Download from HuggingFace using streaming + try: + from datasets import load_dataset + except ImportError as exc: + raise ImportError( + "Please install the 'datasets' package: pip install datasets" + ) from exc + + tokenizer = get_tokenizer(tokenizer_path) + + print(f"Downloading {num_samples} samples from LongBench V2 (streaming)...") + dataset = load_dataset( + LONGBENCH_V2_DATASET, split=LONGBENCH_V2_SPLIT, streaming=True + ) + + input_ids = [] + for i, example in enumerate(dataset): + if len(input_ids) >= num_samples: + break + text = format_longbench_v2_example(example) + tokens = tokenizer.encode(text) + # Truncate to max_tokens + input_ids.append(tokens[:max_prompt_tokens]) + + # Save to local cache + with open(cache_file, "w") as f: + json.dump(input_ids, f) + print(f"Saved {len(input_ids)} prompts to cache: {cache_file}") + + # Also cache in memory + _cached_input_ids[cache_key] = input_ids + + return input_ids + + +def compare_kl_divergence( + input_logprobs, output_logprobs, ACC_THRESHOLDS, model_name, test_name +): + """Compare the KL divergence between input and output log probabilities.""" + kl_divs = [] + for input_logprob, output_logprob in zip(input_logprobs, output_logprobs): + input_logprob = np.array(input_logprob) + output_logprob = np.array(output_logprob) + logr = input_logprob - output_logprob + kl_approx = (np.exp(logr) - 1) - logr + kl_divs.append(np.mean(kl_approx)) + + print(f"kl_divs={kl_divs}") + avg_kl_div = sum(kl_divs) / len(kl_divs) + print(f"avg_kl_div={avg_kl_div}") + print(f"ACC_THRESHOLDS={ACC_THRESHOLDS[model_name]}") + assert avg_kl_div < ACC_THRESHOLDS[model_name]["kl_div"], ( + f"avg_kl_div={avg_kl_div} > threshold={ACC_THRESHOLDS[model_name]['kl_div']} " + f"for {model_name} {test_name}" + ) + + +# Common request helpers +def _flush_cache(base_url): + requests.post(base_url + "/flush_cache") + + +def _generate( + base_url, input_ids, max_new_tokens, return_logprob=False, logprob_start_len=-1 +): + """Send generate request and return results.""" + json_data = { + "input_ids": input_ids, + "sampling_params": { + "temperature": 1, + "max_new_tokens": max_new_tokens, + "ignore_eos": True, + }, + } + if return_logprob: + json_data.update( + { + "return_logprob": True, + "return_text_in_logprobs": False, + "logprob_start_len": logprob_start_len, + } + ) + response = requests.post(base_url + "/generate", json=json_data) + return response.json() + + +def _get_input_logprobs(base_url, new_input_ids, output_logprobs): + """Run prefill to get input logprobs matching output logprobs.""" + _flush_cache(base_url) + results = _generate( + base_url, + new_input_ids, + max_new_tokens=0, + return_logprob=True, + logprob_start_len=0, + ) + assert len(results) == len(new_input_ids) + + input_logprobs = [] + for i, result in enumerate(results): + logprob = result["meta_info"]["input_token_logprobs"] + logprob = [x[0] for x in logprob][-len(output_logprobs[i]) :] + input_logprobs.append(logprob) + return input_logprobs + + +def _extract_output_logprobs(result): + """Extract output logprobs from a result.""" + return [x[0] for x in result["meta_info"]["output_token_logprobs"]] + + +def test_input_output_logprobs_match_helper( + base_url, ACC_THRESHOLDS, model_name, max_samples=None, max_new_tokens=16000 +): + num_samples = DEFAULT_NUM_SAMPLES + if max_samples is not None and max_samples > num_samples: + num_samples = max_samples + input_ids = get_input_ids(tokenizer_path=model_name, num_samples=num_samples) + if max_samples is not None: + input_ids = input_ids[:max_samples] + print(f"Running test_input_output_logprobs_match with {len(input_ids)} prompts") + + print("Flush Cache and Running generation to get output logprobs ...") + _flush_cache(base_url) + results = _generate(base_url, input_ids, max_new_tokens, return_logprob=True) + assert len(results) == len(input_ids) + + new_input_ids = [] + output_logprobs = [] + for i, result in enumerate(results): + new_input_ids.append(input_ids[i] + result["output_ids"]) + output_logprobs.append(_extract_output_logprobs(result)) + + print("Running prefill to get input logprobs ...") + input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs) + + compare_kl_divergence( + input_logprobs, + output_logprobs, + ACC_THRESHOLDS, + model_name, + inspect.currentframe().f_code.co_name, + ) + + +def test_input_output_logprobs_match_prefill_cache_hit_helper( + base_url, ACC_THRESHOLDS, model_name, max_samples=None, max_new_tokens=8192 +): + server_info = requests.get(base_url + "/get_server_info").json() + if server_info["disable_radix_cache"]: + print("Radix cache is disabled, skipping test") + return + + num_samples = DEFAULT_NUM_SAMPLES + if max_samples is not None and max_samples > num_samples: + num_samples = max_samples + input_ids = get_input_ids(tokenizer_path=model_name, num_samples=num_samples) + if max_samples is not None: + input_ids = input_ids[:max_samples] + print( + f"Running test_input_output_logprobs_match_prefill_cache_hit with {len(input_ids)} prompts" + ) + + # Prefill to cache the input + print("Flush Cache and Prefill to cache the input ...") + _flush_cache(base_url) + _generate(base_url, input_ids, max_new_tokens=0) + + # Generate with cache hit + print("Running generation to get output logprobs ...") + results = _generate(base_url, input_ids, max_new_tokens, return_logprob=True) + assert len(results) == len(input_ids) + + new_input_ids = [] + output_logprobs = [] + for i, result in enumerate(results): + if result["meta_info"]["cached_tokens"] == 0: + print(f"Prefill cache miss for prompt {i}, skipping") + continue + new_input_ids.append(input_ids[i] + result["output_ids"]) + output_logprobs.append(_extract_output_logprobs(result)) + + assert len(new_input_ids) > 0.5 * len( + input_ids + ), f"Too few prefill cache hits: {len(new_input_ids)}/{len(input_ids)}" + + print("Flush Cache and run prefill to get input logprobs ...") + input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs) + + compare_kl_divergence( + input_logprobs, + output_logprobs, + ACC_THRESHOLDS, + model_name, + inspect.currentframe().f_code.co_name, + ) + + +def test_input_output_logprobs_match_decode_cache_hit_helper( + base_url, ACC_THRESHOLDS, model_name, max_samples=None, max_new_tokens=8192 +): + server_info = requests.get(base_url + "/get_server_info").json() + if server_info["disable_radix_cache"]: + print("Radix cache is disabled, skipping test") + return + + num_samples = DEFAULT_NUM_SAMPLES + if max_samples is not None and max_samples > num_samples: + num_samples = max_samples + first_turn_input_ids = get_input_ids( + tokenizer_path=model_name, num_samples=num_samples + ) + if max_samples is not None: + first_turn_input_ids = first_turn_input_ids[:max_samples] + print( + f"Running test_input_output_logprobs_match_decode_cache_hit with {len(first_turn_input_ids)} prompts" + ) + + # First turn: Prefill + Decode to cache + print("Flush Cache and First turn: Prefill + Decode to cache decode ...") + _flush_cache(base_url) + results = _generate( + base_url, first_turn_input_ids, max_new_tokens, return_logprob=True + ) + assert len(results) == len(first_turn_input_ids) + + tokenizer = get_tokenizer(tokenizer_name=model_name) + comma_token_id = tokenizer.encode(",") + + second_turn_input_ids = [ + first_turn_input_ids[i] + result["output_ids"] + comma_token_id + for i, result in enumerate(results) + ] + + # Second turn: should hit decode cache + print("Running generation to get output logprobs ...") + results = _generate( + base_url, second_turn_input_ids, max_new_tokens, return_logprob=True + ) + assert len(results) == len(second_turn_input_ids) + + new_input_ids = [] + output_logprobs = [] + for i, result in enumerate(results): + if result["meta_info"]["cached_tokens"] <= len(first_turn_input_ids[i]) + 1: + print(f"Decode cache miss for prompt {i}, skipping") + continue + new_input_ids.append(second_turn_input_ids[i] + result["output_ids"]) + output_logprobs.append(_extract_output_logprobs(result)) + + assert len(new_input_ids) > 0.5 * len( + second_turn_input_ids + ), f"Too few decode cache hits: {len(new_input_ids)}/{len(second_turn_input_ids)}" + + print("Flush Cache and run prefill to get input logprobs ...") + input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs) + + compare_kl_divergence( + input_logprobs, + output_logprobs, + ACC_THRESHOLDS, + model_name, + inspect.currentframe().f_code.co_name, + ) diff --git a/test/srt/models/test_qwen3_next_models.py b/test/srt/models/test_qwen3_next_models.py index c03869746..436585caf 100644 --- a/test/srt/models/test_qwen3_next_models.py +++ b/test/srt/models/test_qwen3_next_models.py @@ -3,6 +3,10 @@ from types import SimpleNamespace from sglang.srt.utils import kill_process_tree from sglang.test.few_shot_gsm8k import run_eval +from sglang.test.kl_test_utils import ( + test_input_output_logprobs_match_decode_cache_hit_helper, + test_input_output_logprobs_match_prefill_cache_hit_helper, +) from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, @@ -10,11 +14,15 @@ from sglang.test.test_utils import ( popen_launch_server, ) +QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" + +ACC_THRESHOLDS = {QWEN3_NEXT_MODEL: {"kl_div": 0.01, "gsm8k": 0.93}} + class TestQwen3Next(CustomTestCase): @classmethod def setUpClass(cls): - cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct" + cls.model = QWEN3_NEXT_MODEL cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, @@ -23,6 +31,8 @@ class TestQwen3Next(CustomTestCase): other_args=[ "--tp-size", "4", + "--chunked-prefill-size", + "2048", ], ) @@ -42,13 +52,33 @@ class TestQwen3Next(CustomTestCase): ) metrics = run_eval(args) print(f"{metrics=}") - self.assertGreater(metrics["accuracy"], 0.93) + self.assertGreaterEqual( + metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + + def test_input_output_logprobs_match_prefill_cache_hit(self): + test_input_output_logprobs_match_prefill_cache_hit_helper( + self.base_url, + ACC_THRESHOLDS, + self.model, + max_samples=16, + max_new_tokens=256, + ) + + def test_input_output_logprobs_match_decode_cache_hit(self): + test_input_output_logprobs_match_decode_cache_hit_helper( + self.base_url, + ACC_THRESHOLDS, + self.model, + max_samples=16, + max_new_tokens=256, + ) class TestQwen3NextMTP(CustomTestCase): @classmethod def setUpClass(cls): - cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct" + cls.model = QWEN3_NEXT_MODEL cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, @@ -87,13 +117,15 @@ class TestQwen3NextMTP(CustomTestCase): ) metrics = run_eval(args) print(f"{metrics=}") - self.assertGreater(metrics["accuracy"], 0.93) + self.assertGreaterEqual( + metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) class TestQwen3NextMTPTopk(CustomTestCase): @classmethod def setUpClass(cls): - cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct" + cls.model = QWEN3_NEXT_MODEL cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, @@ -132,14 +164,16 @@ class TestQwen3NextMTPTopk(CustomTestCase): ) metrics = run_eval(args) print(f"{metrics=}") - self.assertGreater(metrics["accuracy"], 0.93) + self.assertGreaterEqual( + metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) class TestQwen3NextPiecewiseCudaGraph(CustomTestCase): @classmethod def setUpClass(cls): - cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct" + cls.model = QWEN3_NEXT_MODEL cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, @@ -170,7 +204,9 @@ class TestQwen3NextPiecewiseCudaGraph(CustomTestCase): ) metrics = run_eval(args) print(f"{metrics=}") - self.assertGreater(metrics["accuracy"], 0.93) + self.assertGreaterEqual( + metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) if __name__ == "__main__": diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index e17c25fcd..cd6a95add 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -151,7 +151,7 @@ suites = { TestFile("test_eagle_dp_attention.py", 200), ], "per-commit-4-gpu": [ - TestFile("models/test_qwen3_next_models.py", 291), + TestFile("models/test_qwen3_next_models.py", 472), TestFile("test_gpt_oss_4gpu.py", 300), TestFile("test_local_attn.py", 411), TestFile("test_multi_instance_release_memory_occupation.py", 64),