[Scheduler] Align deterministic flashinfer chunks to the q-tile grid, not the split-KV tile

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

@@ -68,6 +68,10 @@ class AttentionBackend(ABC):
     data_type: Optional[torch.dtype] = None
     kv_cache_dtype: Optional[torch.dtype] = None
 
+    # Grid a prefill chunk must start on to reproduce the unchunked prefill bit
+    # for bit; None for backends that tile keys at absolute positions.
+    prefill_chunk_align: Optional[int] = None
+
     # Wrapper backends (e.g. HybridLinearAttnBackend) set this to their child
     # backends; leaves keep None. Lets generic code (metadata glue graph)
     # enumerate every backend whose python-side forward_metadata must be

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

@@ -406,6 +406,7 @@ def __init__(
         self.token_to_kv_pool = swa_backend.token_to_kv_pool
         self.req_to_token_pool = swa_backend.req_to_token_pool
         self.kv_index_translator = swa_backend.kv_index_translator
+        self.prefill_chunk_align = swa_backend.prefill_chunk_align
         # SWA latent expansion uses host sequence-length mirrors.
         self.needs_cpu_seq_lens = True
         self._dp_rebuilt_batch_id: int | None = None

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

@@ -291,6 +291,10 @@ def fast_prefill_plan(
     self._plan_info = self._cached_module.plan(*args)
 
 
+# fa2 q-tile width; chunks must start on this grid to be bit-exact.
+PREFILL_Q_TILE_TOKENS = 128
+
+
 class FlashInferAttnBackend(AttentionBackend):
     """Flashinfer attention kernels."""
 
@@ -423,6 +427,7 @@ def __init__(
                 "SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE", 2048
             )
             self.disable_cuda_graph_kv_split = True
+            self.prefill_chunk_align = PREFILL_Q_TILE_TOKENS
             envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.set(2048 * 1024 * 1024)
 
         self.use_paged = envs.SGLANG_FLASHINFER_USE_PAGED.get()

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

@@ -38,6 +38,7 @@ def __init__(
         self.token_to_kv_pool = model_runner.token_to_kv_pool
         self.req_to_token_pool = model_runner.req_to_token_pool
         self.kv_index_translator = model_runner.kv_index_translator
+        self.prefill_chunk_align = prefill_backend.prefill_chunk_align
         self.spec_attn_is_decode = get_spec().speculative_attention_mode == "decode"
         self.spec_attn_is_prefill = get_spec().speculative_attention_mode == "prefill"
         # Gates the FutureMap's per-step seq_lens D2H (decide_needs_cpu_seq_lens

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

@@ -1194,6 +1194,7 @@ def __init__(
         self.token_to_kv_pool = full_attn_backend.token_to_kv_pool
         self.req_to_token_pool = full_attn_backend.req_to_token_pool
         self.kv_index_translator = full_attn_backend.kv_index_translator
+        self.prefill_chunk_align = full_attn_backend.prefill_chunk_align
         self.max_context_len = getattr(full_attn_backend, "max_context_len", None)
         self.needs_cpu_seq_lens = (
             full_attn_backend.needs_cpu_seq_lens

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

@@ -1868,6 +1868,7 @@ def __init__(
         self.dense = dense_backend
         self.sparse = sparse_backend
         self.kv_index_translator = dense_backend.kv_index_translator
+        self.prefill_chunk_align = dense_backend.prefill_chunk_align
         self.sparse_layer_ids = sparse_layer_ids
         # Let the sparse decode reuse the dense paged backend (page table + workspace).
         self.sparse.dense_backend = dense_backend

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

@@ -24,6 +24,7 @@ def __init__(self, primary: AttentionBackend, children: List[AttentionBackend]):
         self.token_to_kv_pool = primary.token_to_kv_pool
         self.req_to_token_pool = primary.req_to_token_pool
         self.kv_index_translator = primary.kv_index_translator
+        self.prefill_chunk_align = primary.prefill_chunk_align
         self.extend_dummy_seqs_capped_by_req_pool = getattr(
             primary, "extend_dummy_seqs_capped_by_req_pool", False
         )

python/sglang/srt/managers/scheduler.py

@@ -30,7 +30,6 @@
 
 from sglang.srt.runtime_context import (
     SpawnRanks,
-    attention_backends,
     get_context,
     get_device,
     get_disagg,
@@ -343,7 +342,6 @@
     freeze_gc,
     get_available_gpu_memory,
     get_bool_env_var,
-    get_int_env_var,
     is_cuda,
     is_hip,
     is_mps,
@@ -668,8 +666,8 @@ def __init__(
         # Init Ngram Embedding
         self.maybe_init_ngram_embedding()
 
-        # Init prefill kv split size when deterministic inference is enabled with various attention backends
-        self.init_deterministic_inference_config()
+        # Init prefill chunk alignment for attention backends that need it (deterministic flashinfer)
+        self.init_prefill_chunk_align()
         self.init_dsa_kpool_truncation_align()
         self.check_truncation_align_fits_chunk()
 
@@ -1666,21 +1664,10 @@ def maybe_init_ngram_embedding(self):
         )
         self.use_ngram_embedding = self.tp_worker.model_config.use_ngram_embedding
 
-    def init_deterministic_inference_config(self):
-        """Initialize deterministic inference configuration for different attention backends."""
-        if not get_exec().deterministic.enable_deterministic_inference:
-            self.truncation_align_size = None
-            return
-
-        # Only flashinfer needs chunks on its prefill split tiles; triton tiles
-        # keys at absolute positions (align_window_kv_to_tiles).
-        backend_sizes = {
-            "flashinfer": ("SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE", 4096),
-        }
-        prefill_backend, _ = attention_backends()
-        env_var, default_size = backend_sizes.get(prefill_backend, (None, None))
+    def init_prefill_chunk_align(self):
+        """Take the prefill chunk grid the attention backend asks for."""
         self.truncation_align_size = (
-            get_int_env_var(env_var, default_size) if env_var else None
+            self.tp_worker.model_runner.attn_backend.prefill_chunk_align
         )
 
     def init_dsa_kpool_truncation_align(self):
@@ -1725,9 +1712,7 @@ def check_truncation_align_fits_chunk(self):
             f"Chunked prefill is aligned to {align} tokens, but each rank's "
             f"chunked prefill size is {chunk}{dp_note}; a prompt longer than "
             f"{chunk} tokens could never be scheduled. Raise --chunked-prefill-size "
-            f"to at least {align * attn_dp_size}, or, with the flashinfer backend "
-            f"under deterministic inference, lower "
-            f"SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE."
+            f"to at least {align * attn_dp_size}."
         )
 
     def init_request_dispatcher(self):

test/registered/attention/test_flashinfer_deterministic_chunk_invariance.py

@@ -0,0 +1,98 @@
+"""Deterministic fa2 prefill chunked on the 128-token grid must match the
+unchunked prefill bit for bit; a chunk off the grid must not."""
+
+import unittest
+
+import torch
+from flashinfer import BatchPrefillWithPagedKVCacheWrapper
+
+from sglang.srt.layers.attention.flashinfer_backend import PREFILL_Q_TILE_TOKENS
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
+
+HEAD_DIM = 128
+SPLIT_TILE = 4096
+SEQ_LEN = 12 * PREFILL_Q_TILE_TOKENS
+DTYPE = torch.bfloat16
+
+
+def _prefill(wrapper, q, k, v, q_start, q_len, kv_len, num_qo_heads, num_kv_heads):
+    device = q.device
+    wrapper.plan(
+        torch.tensor([0, q_len], dtype=torch.int32, device=device),
+        torch.tensor([0, kv_len], dtype=torch.int32, device=device),
+        torch.arange(kv_len, dtype=torch.int32, device=device),
+        torch.ones(1, dtype=torch.int32, device=device),
+        num_qo_heads,
+        num_kv_heads,
+        HEAD_DIM,
+        1,
+        causal=True,
+        q_data_type=DTYPE,
+        kv_data_type=DTYPE,
+        fixed_split_size=SPLIT_TILE,
+    )
+    return wrapper.run(q[q_start : q_start + q_len], (k, v))
+
+
+class TestFlashinferDeterministicChunkInvariance(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        workspace = torch.empty(512 * 1024 * 1024, dtype=torch.uint8, device="cuda")
+        cls.wrapper = BatchPrefillWithPagedKVCacheWrapper(
+            workspace, "NHD", backend="fa2"
+        )
+
+    def _chunked_equals_full(self, *, num_qo_heads, num_kv_heads, bounds):
+        gen = torch.Generator(device="cuda").manual_seed(0)
+        q = torch.randn(
+            SEQ_LEN, num_qo_heads, HEAD_DIM, dtype=DTYPE, device="cuda", generator=gen
+        )
+        k = torch.randn(
+            SEQ_LEN, num_kv_heads, HEAD_DIM, dtype=DTYPE, device="cuda", generator=gen
+        )
+        v = torch.randn(
+            SEQ_LEN, num_kv_heads, HEAD_DIM, dtype=DTYPE, device="cuda", generator=gen
+        )
+        args = (self.wrapper, q, k, v)
+        full = _prefill(*args, 0, SEQ_LEN, SEQ_LEN, num_qo_heads, num_kv_heads)
+        outs = []
+        start = 0
+        for end in list(bounds) + [SEQ_LEN]:
+            outs.append(
+                _prefill(*args, start, end - start, end, num_qo_heads, num_kv_heads)
+            )
+            start = end
+        return torch.equal(torch.cat(outs), full)
+
+    def test_chunks_on_tile_grid_match_unchunked(self):
+        for num_qo_heads, num_kv_heads in ((8, 2), (8, 8)):
+            for step in (1, 2, 3):
+                bounds = range(
+                    step * PREFILL_Q_TILE_TOKENS, SEQ_LEN, step * PREFILL_Q_TILE_TOKENS
+                )
+                with self.subTest(gqa=num_qo_heads // num_kv_heads, step=step):
+                    self.assertTrue(
+                        self._chunked_equals_full(
+                            num_qo_heads=num_qo_heads,
+                            num_kv_heads=num_kv_heads,
+                            bounds=bounds,
+                        )
+                    )
+
+    def test_chunk_off_tile_grid_differs(self):
+        for num_qo_heads, num_kv_heads, bound in ((8, 2, 100), (8, 8, 64)):
+            with self.subTest(gqa=num_qo_heads // num_kv_heads, bound=bound):
+                self.assertFalse(
+                    self._chunked_equals_full(
+                        num_qo_heads=num_qo_heads,
+                        num_kv_heads=num_kv_heads,
+                        bounds=(bound,),
+                    )
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()