From 2070927e00257d923758438c6230de69dc5e882d Mon Sep 17 00:00:00 2001 From: Alex Nails Date: Mon, 24 Aug 2026 00:58:52 -0700 Subject: [PATCH] [Triton] Bound the sliding-window extend-attention KV loop: -86.6% on SWA layers, -9.4% prefill GPU, bit-identical (#34462) Co-authored-by: Claude Opus 5 (1M context) --- .../sglang/kernels/ops/attention/extend_attention.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/python/sglang/kernels/ops/attention/extend_attention.py b/python/sglang/kernels/ops/attention/extend_attention.py index 7ffc2bde3..1e32f1eb4 100644 --- a/python/sglang/kernels/ops/attention/extend_attention.py +++ b/python/sglang/kernels/ops/attention/extend_attention.py @@ -605,7 +605,16 @@ def _fwd_kernel( else tl.minimum(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M) ) extend_end = 0 if SKIP_EXTEND else cur_block_m_end - for start_n in range(0, extend_end, BLOCK_N): + # The mask below keeps (q, kv) iff q <= kv + SLIDING_WINDOW_SIZE, so no tile + # under this floor can hold an unmasked element -- tight for any BLOCK_M/BLOCK_N. + # SKIP_TILE already made those tiles no-ops, so bounding the loop is + # bit-identical and drops their cross-wave tl.max reduction. + extend_start = 0 + if SLIDING_WINDOW_SIZE > 0: + extend_start = ( + tl.maximum(cur_block_m * BLOCK_M - SLIDING_WINDOW_SIZE, 0) // BLOCK_N + ) * BLOCK_N + for start_n in range(extend_start, extend_end, BLOCK_N): start_n = tl.multiple_of(start_n, BLOCK_N) mask_n = (start_n + offs_n) < cur_block_m_end