A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 09:00 UTC
[Kernel] Optimize Sliding Window Attention in 3D Triton Kernel (#31984)
Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com>
jvlunteren committed
Jan 10, 2026 at 19:13 UTC
b8bf5c45bbbb1249bd1656557603da71e9d4c3b9
1 file changed
+26
-3
vllm/v1/attention/ops/triton_unified_attention.py
+26
-3
@@ -545,10 +545,33 @@ def kernel_unified_attention_3d(
545
# this prefix can be skipped)
546
num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)
547
548
- # iterate through tiles within current segment
548
+ # ---- Sliding-window tile pruning --------------------
549
+ # Default: keep previous global behavior
550
+ tile_start = 0
551
+ tile_end = num_tiles
552
+ # TODO(Isotr0py): sliding window pruning with image bidirectional mask
553
+ if SLIDING_WINDOW > 0 and not USE_MM_PREFIX:
554
+ # Query rows covered by this Q-block
555
+ qpos_lo = q_block_local_idx * BLOCK_Q
556
+ qpos_hi = tl.minimum(
557
+ qpos_lo + (BLOCK_M - 1) // num_queries_per_kv,
558
+ cur_batch_query_len - 1,
559
+ )
560
+ # For sliding window, each query position q can only attend to
561
+ # keys in the range [q_abs - SLIDING_WINDOW + 1, q_abs]
562
+ # where q_abs = context_len + q
563
+ # The union of allowed key positions for this Q-block is:
564
+ # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi]
565
+ first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1
566
+ last_allowed_key = context_len + qpos_hi
567
+ # Convert to tile indices and clamp
568
+ tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)
569
+ tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles)
570
+
571
+ # iterate through tiles (now limited to the sliding window range)
572
for j in range(
550
- segm_idx * tiles_per_segment,
551
- min((segm_idx + 1) * tiles_per_segment, num_tiles),
573
+ max(segm_idx * tiles_per_segment, tile_start),
574
+ min((segm_idx + 1) * tiles_per_segment, tile_end),
575
):
576
seq_offset = j * TILE_SIZE + offs_t
577
tile_mask = seq_offset < max_seq_prefix_len