From 0e30bd62fedfad509e045b01e4aed1ef82ff6dbf Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:06:14 -0400 Subject: [PATCH] [CI Bug] Fix kimi test `AssertionError: Aligned Mamba state indices must be precomputed` (#53766) Signed-off-by: yewentao256 --- tests/models/kimi_k3/test_kda_metadata.py | 24 +++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index bab3bc0094..8431705ec6 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -130,14 +130,22 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary(): common_attn_metadata = create_common_attn_metadata( batch, BLOCK_SIZE, device, arange_block_indices=True ).replace(is_prefilling=torch.tensor([True, True])) - actual = _make_builder( + builder = _make_builder( KimiK3KDAMetadataBuilder, num_speculative_tokens=0, full_cuda_graph=False, mamba_cache_mode="align", num_prefill_checkpoint_blocks=1, device=device, - ).build(0, common_attn_metadata) + ) + assert isinstance(builder, KimiK3KDAMetadataBuilder) + builder.mamba_aligned_state_indices = mamba_get_block_table_tensor( + common_attn_metadata.block_table_tensor, + common_attn_metadata.seq_lens, + builder.kv_cache_spec, + "align", + ) + actual = builder.build(0, common_attn_metadata) assert actual.checkpoint is not None torch.testing.assert_close( @@ -159,14 +167,22 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset(): common_attn_metadata = create_common_attn_metadata( batch, BLOCK_SIZE, device, arange_block_indices=True ).replace(is_prefilling=torch.tensor([True])) - actual = _make_builder( + builder = _make_builder( KimiK3KDAMetadataBuilder, num_speculative_tokens=0, full_cuda_graph=False, mamba_cache_mode="align", num_prefill_checkpoint_blocks=1, device=device, - ).build(0, common_attn_metadata) + ) + assert isinstance(builder, KimiK3KDAMetadataBuilder) + builder.mamba_aligned_state_indices = mamba_get_block_table_tensor( + common_attn_metadata.block_table_tensor, + common_attn_metadata.seq_lens, + builder.kv_cache_spec, + "align", + ) + actual = builder.build(0, common_attn_metadata) assert actual.checkpoint is None