[Deterministic] Support Qwen3-Next model deterministic inference (#13100)

This commit is contained in:
Minglei Zhu
2025-11-14 18:01:18 -08:00
committed by GitHub
parent 8f4e18a294
commit a5be6ef98e
3 changed files with 99 additions and 1 deletions

View File

@@ -297,6 +297,15 @@ class MambaPool:
self.mamba_cache.temporal[:, free_index] = 0
def clear(self):
# Zero the entire mamba cache before resetting free_slots
# This ensures that when slots are reallocated, they start with clean state
if self.is_kda_cache:
for i in range(len(self.mamba_cache.conv)):
self.mamba_cache.conv[i].zero_()
else:
self.mamba_cache.conv.zero_()
self.mamba_cache.temporal.zero_()
self.free_slots = torch.arange(self.size, dtype=torch.int64, device=self.device)
def copy_from(self, src_index: torch.Tensor, dst_index: torch.Tensor):