From f59d30dcb285b65670583802c36e386b3b806c41 Mon Sep 17 00:00:00 2001 From: muh-pipeline Date: Thu, 6 Aug 2026 02:53:07 +0000 Subject: [PATCH] [BASE] qwen3_6_scripts/paged_attn.py: CCCL shifted_output defensive init MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- qwen3_6_scripts/paged_attn.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/qwen3_6_scripts/paged_attn.py b/qwen3_6_scripts/paged_attn.py index 3d98fe47..1b107131 100644 --- a/qwen3_6_scripts/paged_attn.py +++ b/qwen3_6_scripts/paged_attn.py @@ -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,