From cf245adff956d57c166f5652a164a881b3275f1c Mon Sep 17 00:00:00 2001 From: muh Date: Thu, 6 Aug 2026 06:33:21 +0000 Subject: [PATCH] [fix/correctness] mamba_cache: safe swap via clone, not in-place fancy indexing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CCCL source: catch2_test_device_copy_batched.cu CCCL pattern: DeviceCopy::Batched always uses separate src/dst buffers with shuffled destination offsets. Never does in-place scatter. Bug: _swap_mamba_cache used cache[:, [to,from]] = cache[:, [from,to]] PyTorch advanced indexing assignment has undefined evaluation order when src and dst overlap — this can corrupt DeltaNet conv_state and temporal_state during decode, causing silent numerical errors. Fix: explicit temp = clone(from), copy(to→from), copy(tmp→to). Three CUDA memcpy calls instead of one potentially-racy fancy index. This affects every decode step of every DeltaNet layer (alternating layers in Qwen3.6). Corrupt temporal_state → wrong attention output → garbage text or NaN propagation. --- qwen3_6_scripts/mamba_cache.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/qwen3_6_scripts/mamba_cache.py b/qwen3_6_scripts/mamba_cache.py index 7537b9e9..8a3795f1 100644 --- a/qwen3_6_scripts/mamba_cache.py +++ b/qwen3_6_scripts/mamba_cache.py @@ -70,10 +70,15 @@ class MambaCacheManager: return tuple(buffer[:, :batch_size] for buffer in self.mamba_cache) def _swap_mamba_cache(self, from_index: int, to_index: int): + # CCCL DeviceCopy::Batched uses separate src/dst buffers — never + # in-place scatter. PyTorch advanced indexing assignment + # cache[:, [a,b]] = cache[:, [b,a]] has undefined evaluation order. + # Use explicit temp clone for correctness. assert len(self.mamba_cache) > 0 for cache_t in self.mamba_cache: - cache_t[:, [to_index,from_index]] = \ - cache_t[:, [from_index,to_index]] + tmp = cache_t[:, from_index].clone() + cache_t[:, from_index].copy_(cache_t[:, to_index]) + cache_t[:, to_index].copy_(tmp) def _copy_mamba_cache(self, from_index: int, to_index: int): assert len(self.mamba_cache) > 0