dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[BASE] qwen3_6_scripts/paged_attn.py: CCCL shifted_output defensive init

Browse Source 
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:
muh-pipeline
2026-08-06 02:53:07 +00:00
parent 8056641f08
commit f59d30dcb2
1 changed files with 11 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
qwen3_6_scripts/paged_attn.py
View File

@@ -412,17 +412,25 @@ class PagedAttention:
else:
# Run PagedAttention V2.
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),
dtype=output.dtype,
device=output.device,
)
exp_sums = torch.empty(
exp_sums = torch.zeros(
size=(num_seqs, num_heads, max_num_partitions),
dtype=torch.float32,
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(
output,
exp_sums,