[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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user