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 14, 05:30 UTC
Optimize triton unified attention performance for sliding window attention (#24390)
Signed-off-by: zixi-qi <qizixi@meta.com>
qizixi committed
Sep 19, 2025 at 12:07 UTC
a2a5f79e09075f50943750a2376b3b5d95a47ac5
2 files changed
+25
-3
tests/kernels/attention/test_triton_unified_attention.py
+1
-1
@@ -83,7 +83,7 @@ def ref_paged_attn(
83
@pytest.mark.parametrize("num_heads", NUM_HEADS)
84
@pytest.mark.parametrize("head_size", HEAD_SIZES)
85
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
86
-@pytest.mark.parametrize("sliding_window", [None, 256])
86
+@pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
87
@pytest.mark.parametrize("dtype", DTYPES)
88
@pytest.mark.parametrize("soft_cap", [None, 50.0])
89
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
vllm/attention/ops/triton_unified_attention.py
+24
-2
@@ -184,8 +184,30 @@ def kernel_unified_attention_2d(
184
# this prefix can be skipped)
185
num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)
186
187
- # iterate through tiles
188
- for j in range(0, num_tiles):
187
+ # ---- Sliding-window tile pruning --------------------
188
+ # Default: keep previous global behavior
189
+ tile_start = 0
190
+ tile_end = num_tiles
191
+ if SLIDING_WINDOW > 0:
192
+ # Query rows covered by this Q-block
193
+ qpos_lo = q_block_local_idx * BLOCK_Q
194
+ qpos_hi = tl.minimum(
195
+ qpos_lo + (BLOCK_M - 1) // num_queries_per_kv,
196
+ cur_batch_query_len - 1,
197
+ )
198
+ # For sliding window, each query position q can only attend to
199
+ # keys in the range [q_abs - SLIDING_WINDOW + 1, q_abs]
200
+ # where q_abs = context_len + q
201
+ # The union of allowed key positions for this Q-block is:
202
+ # [context_len + qpos_lo - SLIDING_WINDOW + 1, context_len + qpos_hi]
203
+ first_allowed_key = context_len + qpos_lo - SLIDING_WINDOW + 1
204
+ last_allowed_key = context_len + qpos_hi
205
+ # Convert to tile indices and clamp
206
+ tile_start = tl.maximum(0, first_allowed_key // TILE_SIZE)
207
+ tile_end = tl.minimum((last_allowed_key // TILE_SIZE) + 1, num_tiles)
208
+
209
+ # iterate through tiles (now limited to the sliding window range)
210
+ for j in range(tile_start, tile_end):
211
seq_offset = j * TILE_SIZE + offs_t
212
tile_mask = seq_offset < max_seq_prefix_len
213