refactor linear memory pool (#13004)

2025-11-17 16:24:29 +08:00
parent f3e9336dcb
commit a3e4fe4b41
5 changed files with 89 additions and 113 deletions
@@ -320,8 +320,8 @@ class TestMamba(unittest.TestCase):
        kv_indices, last_node = result.device_indices, result.last_device_node
        assert req9.mamba_pool_idx is not None
        assert torch.all(
-            mamba_pool.mamba_cache.conv[:, req9.mamba_pool_idx]
-            == mamba_pool.mamba_cache.conv[:, last_node.mamba_value]
+            mamba_pool.mamba_cache.conv[0][:, req9.mamba_pool_idx]
+            == mamba_pool.mamba_cache.conv[0][:, last_node.mamba_value]
        )
        assert torch.all(
            mamba_pool.mamba_cache.temporal[:, req9.mamba_pool_idx]