Fix padded positions in breakable CUDA Graph attention (#33253)

TestBreakableCUDAGraphBasic and all NVIDIA CI tests pass.
This commit is contained in:
YAMY
2026-08-05 15:53:28 -07:00
committed by GitHub
parent 7bc90ab394
commit 990a446773
3 changed files with 56 additions and 0 deletions
@@ -351,9 +351,12 @@ def _unified_attention_with_output_impl(
kwargs["topk_indices"] = topk_indices[:real_query_num_tokens]
original_out_cache_loc = forward_batch.out_cache_loc
original_positions = forward_batch.positions
# Keep the original ForwardBatch object and only narrow cache locations for
# this backend call so model/backend state is still written to the same batch.
forward_batch.out_cache_loc = original_out_cache_loc[:real_query_num_tokens]
if original_positions is not None:
forward_batch.positions = original_positions[:real_query_num_tokens]
# Store pre-allocated output for FA backend to write directly into.
# Must slice to real_query_num_tokens to match the narrowed query shape —
@@ -370,6 +373,7 @@ def _unified_attention_with_output_impl(
**kwargs,
)
forward_batch.out_cache_loc = original_out_cache_loc
forward_batch.positions = original_positions
lse = None
if return_lse: