[CI Bug] Fix kimi test `AssertionError: Aligned Mamba state indices must be precomputed` (#53766)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 25, 2026 at 15:06 UTC
0e30bd62fedfad509e045b01e4aed1ef82ff6dbf
1 file changed
+20
-4
tests/models/kimi_k3/test_kda_metadata.py
+20
-4
@@ -130,14 +130,22 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
130
common_attn_metadata = create_common_attn_metadata(
131
batch, BLOCK_SIZE, device, arange_block_indices=True
132
).replace(is_prefilling=torch.tensor([True, True]))
133
- actual = _make_builder(
133
+ builder = _make_builder(
134
KimiK3KDAMetadataBuilder,
135
num_speculative_tokens=0,
136
full_cuda_graph=False,
137
mamba_cache_mode="align",
138
num_prefill_checkpoint_blocks=1,
139
device=device,
140
- ).build(0, common_attn_metadata)
140
+ )
141
+ assert isinstance(builder, KimiK3KDAMetadataBuilder)
142
+ builder.mamba_aligned_state_indices = mamba_get_block_table_tensor(
143
+ common_attn_metadata.block_table_tensor,
144
+ common_attn_metadata.seq_lens,
145
+ builder.kv_cache_spec,
146
+ "align",
147
+ )
148
+ actual = builder.build(0, common_attn_metadata)
149
150
assert actual.checkpoint is not None
151
torch.testing.assert_close(
@@ -159,14 +167,22 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():
167
common_attn_metadata = create_common_attn_metadata(
168
batch, BLOCK_SIZE, device, arange_block_indices=True
169
).replace(is_prefilling=torch.tensor([True]))
162
- actual = _make_builder(
170
+ builder = _make_builder(
171
KimiK3KDAMetadataBuilder,
172
num_speculative_tokens=0,
173
full_cuda_graph=False,
174
mamba_cache_mode="align",
175
num_prefill_checkpoint_blocks=1,
176
device=device,
169
- ).build(0, common_attn_metadata)
177
+ )
178
+ assert isinstance(builder, KimiK3KDAMetadataBuilder)
179
+ builder.mamba_aligned_state_indices = mamba_get_block_table_tensor(
180
+ common_attn_metadata.block_table_tensor,
181
+ common_attn_metadata.seq_lens,
182
+ builder.kv_cache_spec,
183
+ "align",
184
+ )
185
+ actual = builder.build(0, common_attn_metadata)
186
187
assert actual.checkpoint is None
188