diff --git a/qwen3_6_scripts/mamba_cache.py b/qwen3_6_scripts/mamba_cache.py index 696bc493..3fcb739b 100644 --- a/qwen3_6_scripts/mamba_cache.py +++ b/qwen3_6_scripts/mamba_cache.py @@ -10,7 +10,7 @@ class MambaCacheManager: def __init__(self, dtype, num_mamba_layers, max_batch_size, conv_state_shape, temporal_state_shape): - conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) + + conv_state = torch.zeros(size=(num_mamba_layers, max_batch_size) + conv_state_shape, dtype=dtype, device="cuda") @@ -209,10 +209,6 @@ class MambaCacheManager: for req_id in finished_seq_groups_req_ids: if req_id in self.mamba_cache_indices_mapping: seq_mapping = self.mamba_cache_indices_mapping.pop(req_id) - # Zero out released cache slots so new requests start from - # clean state. Without this, a new prefill reuses the stale - # temporal_state of the previous (unrelated) request, causing - # the GDN delta-rule to diverge → NaN → GPU memory corruption. for cache_idx in seq_mapping.values(): for cache_t in self.mamba_cache: cache_t[:, cache_idx].zero_()