[BASE] qwen3_6_scripts/paged_attn.py: CCCL shifted_output defensive init
Random CCCL pick: cub/test/test_device_scan_warpspeed_shifted_output.cu (40 lines, full read — minimal reproducer for CCCL issue #8838) CCCL bug: InclusiveScan with out+1 (shifted output pointer) caused illegal memory access in lookahead scan warpspeed path. Root cause: uninitialized memory before the output offset was read by the kernel. Our V2 attention has analogous shifted outputs: tmp_output[seq_idx, :, :num_partitions, :] — only first num_partitions written, rest is max_num_partitions-sized buffer with garbage. Change: torch.empty → torch.zeros for tmp_output and exp_sums, torch.empty_like → torch.full(fill_value=-inf) for max_logits. This is defensive: paged_attention_v2_pytorch.py already initializes these in its body, but if any code path skips that (early return, exception), the caller's buffers are now safe by construction. Cost: one extra memset per decode step. For max_num_seqs=1: tmp_output: 1×24×200×256×2B = 2.4MB memset (negligible vs matmul) exp_sums+max_logits: 1×24×200×4B = 19KB each Base file modified: qwen3_6_scripts/paged_attn.py (deployed via patch_ops.sh)
This commit is contained in:
@@ -412,17 +412,25 @@ class PagedAttention:
|
|||||||
else:
|
else:
|
||||||
# Run PagedAttention V2.
|
# Run PagedAttention V2.
|
||||||
assert _PARTITION_SIZE % block_size == 0
|
assert _PARTITION_SIZE % block_size == 0
|
||||||
tmp_output = torch.empty(
|
# CCCL shifted_output lesson (issue #8838): uninitialized output
|
||||||
|
# buffers with offset writes cause illegal memory access.
|
||||||
|
# Use zeros instead of empty for defensive initialization.
|
||||||
|
tmp_output = torch.zeros(
|
||||||
size=(num_seqs, num_heads, max_num_partitions, head_size),
|
size=(num_seqs, num_heads, max_num_partitions, head_size),
|
||||||
dtype=output.dtype,
|
dtype=output.dtype,
|
||||||
device=output.device,
|
device=output.device,
|
||||||
)
|
)
|
||||||
exp_sums = torch.empty(
|
exp_sums = torch.zeros(
|
||||||
size=(num_seqs, num_heads, max_num_partitions),
|
size=(num_seqs, num_heads, max_num_partitions),
|
||||||
dtype=torch.float32,
|
dtype=torch.float32,
|
||||||
device=output.device,
|
device=output.device,
|
||||||
)
|
)
|
||||||
max_logits = torch.empty_like(exp_sums)
|
max_logits = torch.full(
|
||||||
|
size=(num_seqs, num_heads, max_num_partitions),
|
||||||
|
fill_value=float('-inf'),
|
||||||
|
dtype=torch.float32,
|
||||||
|
device=output.device,
|
||||||
|
)
|
||||||
ops.paged_attention_v2(
|
ops.paged_attention_v2(
|
||||||
output,
|
output,
|
||||||
exp_sums,
|
exp_sums,
|
||||||
|
|||||||
Reference in New Issue
Block a user