[fix/correctness] mamba_cache: safe swap via clone, not in-place fancy indexing

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.
This commit is contained in:
muh
2026-08-06 06:33:21 +00:00
parent 32fd4299b3
commit cf245adff9

View File

@@ -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