From 83d633798f7f3a2ff24316ff98088502e5259d74 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 00:13:49 +0000 Subject: [PATCH] =?UTF-8?q?fix(overflow):=20chunk=5Fsize=2064=E2=86=9216?= =?UTF-8?q?=20=E2=80=94=20CCCL=20counter=20overflow=20prevention?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit agent_radix_sort_upsweep.cuh (517 lines) key insight: UNROLL_COUNT = min(64, 255/KEYS_PER_THREAD) — limits accumulation steps to prevent unsigned char counter overflow Same principle applied to GatedDeltaNet cumsum: chunk=64 + pre_clamp_max=2.0 → worst cumsum = 128 → exp(128) = inf chunk=16 + pre_clamp_max=2.0 → worst cumsum = 32 → clamp(-20,20) safe This was the remaining NaN source: clamp at [-5,2] before cumsum was necessary but not sufficient when chunk_size=64. --- qwen3_6_scripts/qwen3_5.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 6f13eead..93fab699 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -155,7 +155,12 @@ def _torch_chunk_gated_delta_rule( value: torch.Tensor, # (batch, seq, num_heads, head_v_dim) g: torch.Tensor, # (batch, seq, num_heads) beta: torch.Tensor, # (batch, seq, num_heads) - chunk_size: int = 64, + # CCCL agent_radix_sort_upsweep overflow pattern: UNROLL_COUNT = min(64, 255/KEYS_PER_THREAD) + # prevents counter overflow by limiting accumulation steps. + # Same principle: chunk_size limits cumsum steps. With pre-clamp [-5,2]: + # chunk=64: worst cumsum = 64*2 = 128 → exp(128) = inf + # chunk=16: worst cumsum = 16*2 = 32 → clamp(-20,20) catches it + chunk_size: int = 16, initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False,