[Scheduler] Fix the deterministic-inference hang when a prefill chunk is shorter than its alignment

docs/docs/references/environment_variables.mdx

@@ -1725,11 +1725,6 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>FlashInfer decode split-tile size for deterministic attention.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>2048</code></td>
     </tr>
-    <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE</code></td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Triton prefill truncation alignment size for deterministic attention.</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>4096</code></td>
-    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_TRITON_DECODE_SPLIT_TILE_SIZE</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Triton decode split-tile size for deterministic attention.</td>

python/sglang/kernels/ops/attention/extend_attention.py

@@ -270,6 +270,61 @@ def _copy_unified_indices_kernel(
         tl.store(unified_kv_indices + dst_idx, vals, mask=mask)
 
 
+@triton.jit
+def _align_window_kv_kernel(
+    window_kv_indptr,
+    window_kv_indices,
+    aligned_kv_indptr,
+    pad,
+    aligned_kv_indices,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    src_start = tl.load(window_kv_indptr + pid)
+    num_keys = tl.load(window_kv_indptr + pid + 1) - src_start
+    dst_start = tl.load(aligned_kv_indptr + pid)
+    cur_pad = tl.load(pad + pid)
+    for off in range(0, num_keys + cur_pad, BLOCK):
+        j = off + tl.arange(0, BLOCK)
+        mask = j < num_keys + cur_pad
+        val = tl.load(
+            window_kv_indices + src_start + tl.maximum(j - cur_pad, 0), mask=mask
+        )
+        tl.store(aligned_kv_indices + dst_start + j, val, mask=mask)
+
+
+def align_window_kv_to_tiles(
+    window_kv_indptr: torch.Tensor,
+    window_kv_indices: torch.Tensor,
+    window_start_pos: torch.Tensor,
+    bs: int,
+    tile: int = 128,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pad each request's window KV list with copies of its first key so the list
+    starts on an absolute multiple of ``tile`` (a multiple of every BLOCK_N). The
+    unified kernel then sums a chunked prefill's keys in the same tiles as an
+    unchunked one; the copies sit before the window and are masked out."""
+    pad = window_start_pos[:bs] % tile
+    lens = window_kv_indptr[1 : bs + 1] - window_kv_indptr[:bs] + pad
+    aligned_kv_indptr = torch.zeros_like(window_kv_indptr[: bs + 1])
+    aligned_kv_indptr[1:] = torch.cumsum(lens, dim=0)
+    # Upper bound on the padded length, so no device-to-host sync is needed.
+    aligned_kv_indices = torch.empty(
+        window_kv_indices.numel() + bs * (tile - 1),
+        dtype=window_kv_indices.dtype,
+        device=window_kv_indices.device,
+    )
+    _align_window_kv_kernel[(bs,)](
+        window_kv_indptr,
+        window_kv_indices,
+        aligned_kv_indptr,
+        pad,
+        aligned_kv_indices,
+        BLOCK=128,
+    )
+    return aligned_kv_indptr, aligned_kv_indices
+
+
 def build_unified_kv_indices(
     prefix_kv_indptr: torch.Tensor,
     prefix_kv_indices: torch.Tensor,

python/sglang/srt/environ.py

@@ -1183,7 +1183,6 @@ class Envs:
     # the batch fits the chunk budget.
     SGLANG_TRITON_DENSE_PREFILL_ATTN = EnvBool(True)
     SGLANG_ENABLE_TORCH_COMPILE = EnvBool(False)
-    SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096)
     SGLANG_TRITON_DECODE_SPLIT_TILE_SIZE = EnvInt(256)
 
     # ===================================================================
@@ -1993,6 +1992,9 @@ def apply(self, old_name: str):
         note="Buffer-mode anchor pinning is always on; set "
         "SGLANG_HICACHE_BUFFER_ANCHOR_LOCK_CAP=0 to disable it."
     ),
+    "SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE": _DeprecatedEnv(
+        note="Deterministic inference on triton no longer aligns prefill chunks."
+    ),
     # Replaced by CLI flags.
     "SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE": _DeprecatedEnv(
         note="Please use '--enable-prefill-delayer' instead."

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

@@ -132,6 +132,9 @@ class ForwardMetadata:
     swa_attn_logits: Optional[torch.Tensor] = None
     # full->SWA translated out_cache_loc (SWA KV-store write target)
     swa_out_cache_loc: Optional[torch.Tensor] = None
+    # Deterministic extend: window KV lists padded to absolute tile boundaries
+    aligned_window_kv_indptr: Optional[torch.Tensor] = None
+    aligned_window_kv_indices: Optional[torch.Tensor] = None
     # Lean decode (persistent-grid partial-result buffers)
     lean_Mp: Optional[torch.Tensor] = None
     lean_Lp: Optional[torch.Tensor] = None
@@ -165,6 +168,7 @@ def __init__(
             lean_decode_seqlen_gate,
         )
         from sglang.kernels.ops.attention.extend_attention import (
+            align_window_kv_to_tiles,
             build_unified_kv_indices,
             can_use_dense_prefill_fp8,
             dense_prefill_attention_fwd,
@@ -187,6 +191,7 @@ def __init__(
             extend_attention_fwd_unified
         )
         self.build_unified_kv_indices = torch.compiler.disable(build_unified_kv_indices)
+        self.align_window_kv_to_tiles = torch.compiler.disable(align_window_kv_to_tiles)
         # Dense (non-absorbed) MLA prefill over a materialized prefix; see
         # handle_attention_triton for when the dispatcher selects it.
         self.dense_prefill_attention_fwd = torch.compiler.disable(
@@ -829,6 +834,8 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
         window_kv_indices = None
         window_num_kv_splits = None
         window_kv_offsets = None
+        aligned_window_kv_indptr = None
+        aligned_window_kv_indices = None
         swa_attn_logits = None
         spec_info = forward_batch.spec_info
         # Lean decode buffers are only allocated on the decode path below; default
@@ -1036,6 +1043,15 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                     self.device,
                     self.token_to_kv_pool,
                 )
+                if self.enable_deterministic:
+                    aligned_window_kv_indptr, aligned_window_kv_indices = (
+                        self.align_window_kv_to_tiles(
+                            window_kv_indptr,
+                            window_kv_indices,
+                            forward_batch.extend_prefix_lens[:bs] - window_kv_lens,
+                            bs,
+                        )
+                    )
 
             qo_indptr = self.qo_indptr
             qo_indptr[1 : bs + 1] = torch.cumsum(forward_batch.extend_seq_lens, dim=0)
@@ -1073,6 +1089,8 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
             window_kv_offsets,
             swa_attn_logits=swa_attn_logits,
             swa_out_cache_loc=swa_out_cache_loc,
+            aligned_window_kv_indptr=aligned_window_kv_indptr,
+            aligned_window_kv_indices=aligned_window_kv_indices,
             lean_Mp=lean_Mp,
             lean_Lp=lean_Lp,
             lean_Op=lean_Op,
@@ -2084,6 +2102,17 @@ def _forward_extend_unified(
                 window_start_pos = forward_batch.seq_lens[:bs] - window_kv_lens
             else:
                 window_start_pos = None
+            # The padding is masked only by the sliding-window mask, so a layer
+            # that attends bidirectionally keeps the unpadded list.
+            if (
+                sliding_window_size > 0
+                and self.forward_metadata.aligned_window_kv_indptr is not None
+            ):
+                prefix_kv_indptr = self.forward_metadata.aligned_window_kv_indptr
+                prefix_kv_indices = self.forward_metadata.aligned_window_kv_indices
+                window_start_pos = forward_batch.extend_prefix_lens[:bs] - (
+                    prefix_kv_indptr[1 : bs + 1] - prefix_kv_indptr[:bs]
+                )
         else:
             sliding_window_size = -1
             prefix_kv_indptr = self.forward_metadata.kv_indptr

python/sglang/srt/managers/scheduler.py

@@ -671,6 +671,7 @@ def __init__(
         # Init prefill kv split size when deterministic inference is enabled with various attention backends
         self.init_deterministic_inference_config()
         self.init_dsa_kpool_truncation_align()
+        self.check_truncation_align_fits_chunk()
 
         self.init_weight_updater()
 
@@ -1671,12 +1672,11 @@ def init_deterministic_inference_config(self):
             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),
-            "triton": ("SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE", 4096),
         }
-        # Both entries are prefill knobs (SPLIT_TILE / PREFILL_TRUNCATION):
-        # the prefill half decides.
         prefill_backend, _ = attention_backends()
         env_var, default_size = backend_sizes.get(prefill_backend, (None, None))
         self.truncation_align_size = (
@@ -1705,6 +1705,31 @@ def init_dsa_kpool_truncation_align(self):
                 self.truncation_align_size, dsa_index_kpool
             )
 
+    def check_truncation_align_fits_chunk(self):
+        """A chunk shorter than the alignment truncates to zero tokens, so a prompt
+        longer than the chunk would wait in the queue forever."""
+        align = self.truncation_align_size
+        chunk = self.chunked_prefill_size
+        if align is None or chunk is None or chunk >= align:
+            return
+        if self.disaggregation_mode == DisaggregationMode.DECODE:
+            # A decode server never prefills a prompt.
+            return
+        attn_dp_size = get_parallel().attn_dp_size
+        dp_note = (
+            f" (--chunked-prefill-size is split across {attn_dp_size} DP attention ranks)"
+            if attn_dp_size > 1
+            else ""
+        )
+        raise ValueError(
+            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."
+        )
+
     def init_request_dispatcher(self):
         self._request_dispatcher = TypeBasedDispatcher(
             [

test/registered/attention/test_triton_attention_kernels.py

@@ -11,6 +11,7 @@
 )
 from sglang.kernels.ops.attention.extend_attention import (
     _compact_extend_q_tiles_per_head,
+    align_window_kv_to_tiles,
     build_unified_kv_indices,
     extend_attention_fwd,
     extend_attention_fwd_unified,
@@ -1130,6 +1131,51 @@ def test_build_unified_kv_indices(self):
             unified_seq = unified_kv_indices[start_idx:end_idx]
             self.assertEqual(len(unified_seq), prefix_len + extend_len)
 
+    def test_unified_sliding_window_chunked_matches_unchunked(self):
+        """Deterministic inference: a chunk whose window starts off a tile boundary
+        must match the unchunked prefill bit for bit."""
+        torch.manual_seed(0)
+        device = get_device()
+        seq_len, window, prefix = 3000, 128, 2192  # window starts at 2064
+        q = torch.randn(seq_len, 8, 64, dtype=torch.bfloat16, device=device)
+        k = torch.randn(seq_len, 8, 64, dtype=torch.bfloat16, device=device)
+        v = torch.randn_like(k)
+        i32 = lambda *x: torch.tensor(x, dtype=torch.int32, device=device)
+        arange = lambda lo, hi: torch.arange(lo, hi, device=device)
+
+        def run(prefix_len):
+            win = min(prefix_len, window)
+            indptr, indices = align_window_kv_to_tiles(
+                i32(0, win),
+                arange(prefix_len - win, prefix_len),
+                i32(prefix_len - win),
+                1,
+            )
+            extend_len = seq_len - prefix_len
+            kv_indptr, kv_indices, kv_prefix_lens = build_unified_kv_indices(
+                indptr, indices, i32(0), i32(extend_len), arange(prefix_len, seq_len), 1
+            )
+            o = torch.empty_like(q[prefix_len:])
+            extend_attention_fwd_unified(
+                q[prefix_len:],
+                o,
+                k,
+                v,
+                1.0,
+                1.0,
+                i32(0, extend_len),
+                kv_indptr,
+                kv_indices,
+                kv_prefix_lens.to(torch.int32),
+                extend_len,
+                sm_scale=0.125,
+                sliding_window_size=window,
+                window_start_pos=i32(prefix_len) - (indptr[1:] - indptr[:-1]),
+            )
+            return o
+
+        self.assertTrue(torch.equal(run(prefix), run(0)[prefix:]))
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/managers/test_deterministic_chunk_align.py

@@ -0,0 +1,39 @@
+"""A per-rank chunk shorter than the deterministic prefill alignment must fail at
+startup: it truncates every chunk to zero tokens and the prompt waits forever."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.scheduler import Scheduler
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+_MOD = "sglang.srt.managers.scheduler"
+
+
+def _check(align, chunk, attn_dp_size=1, mode=DisaggregationMode.NULL):
+    scheduler = Scheduler.__new__(Scheduler)
+    scheduler.disaggregation_mode = mode
+    scheduler.truncation_align_size = align
+    scheduler.chunked_prefill_size = chunk
+    parallel = SimpleNamespace(attn_dp_size=attn_dp_size)
+    with patch(f"{_MOD}.get_parallel", return_value=parallel):
+        scheduler.check_truncation_align_fits_chunk()
+
+
+class TestDeterministicChunkAlign(CustomTestCase):
+    def test_chunk_shorter_than_alignment_fails_at_startup(self):
+        with self.assertRaisesRegex(ValueError, "at least 8192"):
+            _check(align=4096, chunk=2048, attn_dp_size=2)
+
+    def test_chunk_covering_alignment_starts(self):
+        _check(align=4096, chunk=4096)
+        _check(align=4096, chunk=2048, mode=DisaggregationMode.DECODE)
+
+
+if __name__ == "__main__":
+    unittest.main()