[Fix] attention: Make FlashInfer prefill bit-exact on radix-cache hits

docs/docs/advanced_features/deterministic_inference.mdx

@@ -49,7 +49,7 @@ The following table shows feature compatibility for deterministic inference acro
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>**FlashInfer**</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅ Yes</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>✅ Yes</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌ No</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>✅ Yes</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>✅ Yes</td>
     </tr>
     <tr>

python/sglang/srt/arg_groups/attention_hook.py

@@ -583,6 +583,7 @@ def handle_multi_item_scoring(server_args: Any):
 
 def handle_deterministic_inference(server_args: Any):
     from sglang.srt.server_args import (
+        MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND,
         RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND,
     )
 
@@ -645,9 +646,9 @@ def handle_deterministic_inference(server_args: Any):
 
         attention_backend = resolved_view(server_args).attention_backend
         if is_deepseek_model:
-            if attention_backend not in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND:
+            if attention_backend not in MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND:
                 raise ValueError(
-                    f"Currently only {RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with absorbed-MLA models. But you're using {attention_backend}."
+                    f"Currently only {MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with absorbed-MLA models. But you're using {attention_backend}."
                 )
             if attention_backend == "fa4" and not get_platform().is_sm100_or_sm110:
                 raise ValueError(

python/sglang/srt/arg_groups/choices.py

@@ -246,7 +246,15 @@
     DETERMINISTIC_ATTENTION_BACKEND_CHOICES.extend
 )
 
-RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["ascend", "fa3", "fa4", "triton"]
+MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["ascend", "fa3", "fa4", "triton"]
+
+RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = [
+    "ascend",
+    "flashinfer",
+    "fa3",
+    "fa4",
+    "triton",
+]
 add_radix_supported_deterministic_attention_backend_choices = (
     RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND.extend
 )

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -35,6 +35,10 @@
     resolve_window_id_space,
     resolve_window_kv_pool,
 )
+from sglang.srt.layers.attention.flashinfer_deterministic import (
+    build_prefill_pad_indices,
+    deterministic_prefill_min_q,
+)
 from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
     KVCacheAttentionAccessKind,
 )
@@ -172,6 +176,8 @@ class PrefillMetadata:
     extend_no_prefix: bool
     multi_item_params: Optional[MultiItemScoringParams] = None
     swa_out_cache_loc: Optional[torch.Tensor] = None
+    prefill_q_pad_gather_idx: Optional[torch.Tensor] = None
+    prefill_q_pad_out_idx: Optional[torch.Tensor] = None
 
 
 # Reuse this workspace buffer across all flashinfer wrappers
@@ -1094,6 +1100,8 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                 extend_no_prefix,
                 multi_item_params,
                 swa_out_cache_loc=swa_out_cache_loc,
+                prefill_q_pad_gather_idx=self.indices_updater_prefill.prefill_q_pad_gather_idx,
+                prefill_q_pad_out_idx=self.indices_updater_prefill.prefill_q_pad_out_idx,
             )
 
     def init_cuda_graph_state(
@@ -1431,8 +1439,14 @@ def forward_extend(
                 not layer.is_cross_attention
                 and layer.attn_type != AttentionType.ENCODER_ONLY
             )
+            q_3d = q.view(-1, layer.tp_q_head_num, layer.head_dim)
+            pad_gather_idx = self.forward_metadata.prefill_q_pad_gather_idx
+            if pad_gather_idx is not None:
+                q_3d = torch.cat(
+                    [q_3d, q_3d.new_zeros((1, *q_3d.shape[1:]))], dim=0
+                ).index_select(0, pad_gather_idx)
             o = prefill_wrapper_paged.forward(
-                q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                q_3d,
                 kv_cache,
                 causal=causal,
                 sm_scale=layer.scaling,
@@ -1442,6 +1456,8 @@ def forward_extend(
                 k_scale=layer.k_scale_float,
                 v_scale=layer.v_scale_float,
             )
+            if pad_gather_idx is not None:
+                o = o.index_select(0, self.forward_metadata.prefill_q_pad_out_idx)
         else:
             # If `k`/`v` are not explicitly provided, fall back to the KV cache stored in
             # `self.token_to_kv_pool` for this layer. This enables attention over
@@ -1925,6 +1941,21 @@ def __init__(self, model_runner: ModelRunner, attn_backend: FlashInferAttnBacken
         self._swa_space = attn_backend._swa_space
         self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
 
+        self.prefill_cta_min_q = (
+            deterministic_prefill_min_q(
+                self.num_qo_heads, self.num_kv_heads, self.head_dim
+            )
+            if attn_backend.enable_deterministic
+            else 1
+        )
+        self.prefill_q_pad_enabled = (
+            attn_backend.enable_deterministic
+            and attn_backend.dispatch_reason is None
+            and not attn_backend.enable_mis
+        )
+        self.prefill_q_pad_gather_idx: Optional[torch.Tensor] = None
+        self.prefill_q_pad_out_idx: Optional[torch.Tensor] = None
+
         # Dispatch the update function
         if self.attn_backend.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
             self.update = self.update_sliding_window
@@ -2003,6 +2034,7 @@ def update_single_wrapper(
             fixed_split_size=fixed_split_size,
             multi_item_params=multi_item_params,
             seq_lens_cpu=seq_lens_cpu,
+            extend_prefix_lens_cpu=extend_prefix_lens_cpu,
             custom_kv_indices=custom_kv_indices,
             plan=plan,
         )
@@ -2211,6 +2243,48 @@ def update_cross_attention(
                 plan=plan,
             )
 
+    def _can_pad_prefill_q(
+        self,
+        wrapper_paged: BatchPrefillWithPagedKVCacheWrapper,
+        use_ragged: bool,
+        use_sliding_window_kv_pool: bool,
+        custom_kv_indices: Optional[torch.Tensor],
+        cross_attention_custom_mask: Optional[torch.Tensor],
+        multi_item_params: Optional[MultiItemScoringParams],
+    ) -> bool:
+        if not self.prefill_q_pad_enabled or self.prefill_cta_min_q <= 1:
+            return False
+        if wrapper_paged.is_cuda_graph_enabled:
+            return False
+        if use_ragged or use_sliding_window_kv_pool:
+            return False
+        if custom_kv_indices is not None or cross_attention_custom_mask is not None:
+            return False
+        if multi_item_params is not None and multi_item_params.is_enabled():
+            return False
+        return not is_in_tc_piecewise_cuda_graph()
+
+    def _prefill_q_lens_cpu(
+        self,
+        q_lens: torch.Tensor,
+        seq_lens_cpu: Optional[torch.Tensor],
+        extend_prefix_lens_cpu: Optional[List[int]],
+    ) -> List[int]:
+        bs = q_lens.shape[0]
+        if (
+            extend_prefix_lens_cpu is not None
+            and seq_lens_cpu is not None
+            and len(extend_prefix_lens_cpu) == bs
+            and len(seq_lens_cpu) == bs
+        ):
+            return [
+                int(seq_len) - int(prefix_len)
+                for seq_len, prefix_len in zip(
+                    seq_lens_cpu.tolist(), extend_prefix_lens_cpu
+                )
+            ]
+        return q_lens.detach().to("cpu", torch.int64).tolist()
+
     def call_begin_forward(
         self,
         wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper,
@@ -2230,12 +2304,15 @@ def call_begin_forward(
         multi_item_params: Optional[MultiItemScoringParams] = None,
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
         seq_lens_cpu: Optional[torch.Tensor] = None,
+        extend_prefix_lens_cpu: Optional[List[int]] = None,
         custom_kv_indices: Optional[torch.Tensor] = None,
         window: Optional[WindowRead] = None,
         *,
         plan: IterationIds,
     ):
         bs = len(seq_lens)
+        self.prefill_q_pad_gather_idx = None
+        self.prefill_q_pad_out_idx = None
         # Unified SWA wrapper-0: gather from the swa canonical directly -- its
         # entries are already swa-side physical ids, so the in-place
         # full->swa translate below must not run on top of them.
@@ -2276,7 +2353,33 @@ def call_begin_forward(
                     kv_start_idx=kv_start_idx,
                     kind=read_kind,
                 )
-            qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
+            q_lens = seq_lens - prefix_lens
+            if self._can_pad_prefill_q(
+                wrapper_paged=wrapper_paged,
+                use_ragged=use_ragged,
+                use_sliding_window_kv_pool=use_sliding_window_kv_pool,
+                custom_kv_indices=custom_kv_indices,
+                cross_attention_custom_mask=cross_attention_custom_mask,
+                multi_item_params=multi_item_params,
+            ):
+                q_lens_cpu = self._prefill_q_lens_cpu(
+                    q_lens, seq_lens_cpu, extend_prefix_lens_cpu
+                )
+                gather_idx, out_idx = build_prefill_pad_indices(
+                    q_lens_cpu, self.prefill_cta_min_q, seq_lens.device
+                )
+                if gather_idx is not None:
+                    self.prefill_q_pad_gather_idx = gather_idx
+                    self.prefill_q_pad_out_idx = out_idx
+                    q_lens = torch.tensor(
+                        [
+                            max(length, self.prefill_cta_min_q) if length > 0 else 0
+                            for length in q_lens_cpu
+                        ],
+                        dtype=q_lens.dtype,
+                        device=q_lens.device,
+                    )
+            qo_indptr[1 : bs + 1] = torch.cumsum(q_lens, dim=0)
             qo_indptr = qo_indptr[: bs + 1]
 
             custom_mask = cross_attention_custom_mask

python/sglang/srt/layers/attention/flashinfer_deterministic.py

@@ -0,0 +1,43 @@
+from typing import List, Optional, Sequence, Tuple
+
+import torch
+
+
+def deterministic_prefill_min_q(
+    num_qo_heads: int, num_kv_heads: int, head_dim: int
+) -> int:
+    gqa_group = max(num_qo_heads // num_kv_heads, 1)
+    if head_dim >= 512:
+        threshold = 32
+    elif head_dim >= 256:
+        threshold = 16
+    else:
+        threshold = 64
+    return threshold // gqa_group + 1
+
+
+def build_prefill_pad_indices(
+    real_lens: Sequence[int], min_q: int, device: torch.device
+) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+    padded_lens: List[int] = [max(real, min_q) if real > 0 else 0 for real in real_lens]
+    real_total = sum(real_lens)
+    padded_total = sum(padded_lens)
+    if padded_total == real_total:
+        return None, None
+
+    gather: List[int] = []
+    out: List[int] = []
+    real_off = 0
+    pad_off = 0
+    for real, padded in zip(real_lens, padded_lens):
+        lead = padded - real
+        gather.extend([real_total] * lead)
+        gather.extend(range(real_off, real_off + real))
+        out.extend(range(pad_off + lead, pad_off + padded))
+        real_off += real
+        pad_off += padded
+
+    return (
+        torch.tensor(gather, dtype=torch.long, device=device),
+        torch.tensor(out, dtype=torch.long, device=device),
+    )

python/sglang/srt/server_args.py

@@ -111,6 +111,7 @@ def _real_kv_hash_modes():
     GRAMMAR_BACKEND_CHOICES,
     LINEAR_ATTN_KERNEL_BACKEND_CHOICES,
     LOAD_FORMAT_CHOICES,
+    MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND,
     MOE_RUNNER_BACKEND_CHOICES,
     MXFP8_MOE_RUNNER_BACKEND_CHOICES,
     QUANTIZATION_CHOICES,

python/sglang/test/test_deterministic.py

@@ -623,8 +623,12 @@ def test_deterministic(args):
         cached_token_data = cached_logprobs[0]
         cached_logprob = cached_token_data[0]
         cached_token_id = cached_token_data[1]
+        cached_tokens = cached_response["meta_info"].get("cached_tokens", 0)
 
-        print(f"✓ Generated with cache:")
+        print(f"✓ Generated with cache (cached_tokens={cached_tokens}):")
+        if cached_tokens == 0:
+            print("TEST FAILED - Radix cache did not hit")
+            return [0]
         print(f"  Token ID: {cached_token_id}")
         print(f"  Logprob:  {cached_logprob:.10f}")
 

test/registered/attention/test_deterministic.py

@@ -15,16 +15,18 @@
     register_cuda_ci,
     register_xpu_ci,
 )
+from sglang.test.test_deterministic import BenchArgs, test_deterministic
 from sglang.test.test_deterministic_utils import (
     COMMON_SERVER_ARGS,
     TestDeterministicBase,
 )
 from sglang.test.test_utils import (
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
+    DEFAULT_URL_FOR_TEST,
     is_in_amd_ci,
 )
 
-register_cuda_ci(est_time=285, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=385, stage="base-b", runner_config="1-gpu-large")
 register_amd_ci(est_time=278, suite="stage-b-test-1-gpu-small-amd")
 register_xpu_ci(est_time=207, suite="stage-b-test-1-gpu-xpu")
 
@@ -47,6 +49,25 @@ def get_server_args(cls):
         return args
 
 
+@unittest.skipIf(_is_xpu, "CUDA runner only")
+@unittest.skipIf(is_in_amd_ci(), "Skip for AMD CI.")
+class TestFlashinferDeterministicRadixCache(TestDeterministicBase):
+    @classmethod
+    def get_server_args(cls):
+        return list(COMMON_SERVER_ARGS) + [
+            "--attention-backend",
+            "flashinfer",
+            "--disable-prefill-cuda-graph",
+        ]
+
+    def test_radix_cache(self):
+        args = BenchArgs()
+        args.host, args.port = self._extract_host_and_port(DEFAULT_URL_FOR_TEST)
+        args.test_mode = "radix_cache"
+        results = test_deterministic(args)
+        self.assertEqual(results, [1])
+
+
 @unittest.skipIf(_is_xpu, "CUDA runner only")
 @unittest.skipIf(is_in_amd_ci(), "Skip for AMD CI.")
 class TestFa3Deterministic(TestDeterministicBase):