[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,
     )

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,
+    )