[Hybrid] calling get_mamba_groups() once at MambaCopyBuffers.create() (#37318)

Signed-off-by: Francesco Fusco <ffu@zurich.ibm.com>
This commit is contained in:
Francesco Fusco
2026-03-21 10:29:43 +01:00
committed by GitHub
parent 3982bc2cd0
commit 298e510848
2 changed files with 11 additions and 5 deletions

View File

@@ -36,6 +36,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():
spec = MagicMock(block_size=64, num_speculative_blocks=0)
cache_config = MagicMock(enable_prefix_caching=True)
input_batch = MagicMock(req_ids=[])
copy_bufs = MagicMock(mamba_group_ids=[0], mamba_spec=spec)
mamba_state_idx = {
"finished": 1,
@@ -62,7 +63,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():
{},
{},
(),
MagicMock(),
copy_bufs,
)
assert mamba_state_idx == {"keep": 99}