[vLLM][unified_attention] Add TILE_SIZE autotuning#7405
Draft
sarah12121212 wants to merge 55 commits into
Draft
Conversation
This reverts commit 5449fe8. Revert commit
…ntinel
The 3D (decode) path runs kernel_unified_attention (producer) followed by
reduce_segments (reducer). The reducer rebuilt its segment-validity mask by
recomputing the segment layout from TILE_SIZE
(act_num_segments = cdiv(seq_len, tiles_per_segment * TILE_SIZE)), so producer
and reducer had to agree on a single TILE_SIZE. That forced the autotuner to
pin the 3D path to a fixed tile (TILE_SIZE=16 for bf16), preventing it from
tuning while the 2D path tuned freely.
Remove the coupling by construction:
- Producer: segment programs that cover no keys now write a -inf sentinel into
their segment-max slot (and 0 into expsum) before returning, instead of
early-returning silently. Offset/mask computations are hoisted above the 3D
scoping check so the dead-segment store can address its slot.
- Reducer: derive segment validity from the data (segm_max != -inf) instead of
recomputing it from TILE_SIZE. Drop the TILE_SIZE parameter entirely. The
segm_output load stays masked by this sentinel-derived mask, so uninitialized
output slots never reach a nan * 0.
- Prune: drop the is_3d special case so 3D autotunes TILE_SIZE in {16, 32};
fp8 stays pinned to 32 via the existing rule.
This also makes the reducer correct when the segment buffers are reused across
decode steps (the serving path), where an unwritten slot otherwise holds stale
finite values, and leaves the reducer correct for any future joint tuning of
TILE_SIZE and NUM_SEGMENTS_PER_SEQ.
Patch regenerated against the pristine vLLM base (blob 257481768); verified to
apply and reverse cleanly.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Standalone XPU test backing the in-kernel-sentinel change. Forces the producer
autotuner to each legal TILE_SIZE and compares the 3D decode output against a
CPU paged-attention reference:
- sweeps KV lengths around the segment-layout boundaries
cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE): 1/15/16/17 ... 255/256/257
... 511/512/513 ... 4095/4096/4097
- bf16 covers TILE_SIZE in {16, 32}; fp8 covers 32 (16 is pruned)
- head_size in {128, 256}, block_size 64
- --reused exercises persistent segm_* buffers reused and pre-poisoned with
finite garbage across calls (the serving-path case), which is correct only
because the producer now overwrites every segment slot (live data or -inf)
each call
Measured on Intel Data Center GPU Max 1100: sweep 150/150 pass (max bf16 abs
err 0.008, fp8 0.125); reused-buffer scenario passes with no NaNs. The
TILE_SIZE=32 3D cases are the ones the previously pinned reducer could not run.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Previously the sentinel branch's @triton.autotune key was [BLOCK_SIZE, HEAD_SIZE, IS_FP8_INPUT, IS_3D, BLOCK_Q] -- coarser than the ua-3d-pipeline-autotune branch, which keys 3D tile selection on (block_size, head_size, is_fp8_input, BLOCK_Q, num_seqs, max_seqlen_k, effective_max_seqlen_k) with the dynamic shapes bucketed to powers of 2. Add the same bucketed dimensions here so steady-state shapes drawn from nearby num_seqs / max_seqlen_k / sliding-window-bounded contexts share the cache, and pick distinct configs only across power-of-2 boundaries: - introduce _power_of_2_bucket() helper - thread NUM_SEQS_BUCKET, MAX_SEQLEN_K_BUCKET, EFFECTIVE_MAX_SEQLEN_K_BUCKET into kernel_unified_attention's signature and the autotune key - compute the three bucket values (effective_max_seqlen_k clamps to sliding_window_val when one is set) in unified_attention() and forward them in the launch The sentinel branch's autotune wrapper covers both 2D and 3D, so 2D launches now also partition the cache by these buckets -- a benign over-partitioning relative to upstream Triton's stock 5-key set. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…entinel # Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention_benchmark.py
…entinel # Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch
…to ua-3d-tile-autotune-sentinel
# Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch
…entinel # Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch
# Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch
# Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch
# Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch
# Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch
sarah12121212
force-pushed
the
ua-combined-autotune
branch
5 times, most recently
from
July 10, 2026 18:31
25bdd4c to
608e30f
Compare
# Conflicts: # benchmarks/triton_kernels_benchmark/vllm/unified_attention/unified_attention.patch
sarah12121212
force-pushed
the
ua-combined-autotune
branch
from
July 10, 2026 18:41
608e30f to
1859c74
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
draft autotune pr with TILE_SIZE and BLOCK_M