| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for GDNAttentionMetadataBuilder.build() — specifically the |
| 4 | reclassification of non-spec decodes as prefills when spec decodes exist. |
| 5 | Covers the fix for https://github.com/vllm-project/vllm/issues/34845. |
| 6 | """ |
| 7 | |
| 8 | from dataclasses import dataclass |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from tests.v1.attention.utils import ( |
| 14 | BatchSpec, |
| 15 | create_common_attn_metadata, |
| 16 | create_vllm_config, |
| 17 | ) |
| 18 | from vllm.config import SpeculativeConfig |
| 19 | from vllm.config.compilation import CUDAGraphMode |
| 20 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 21 | from vllm.v1.attention.backends.gdn_attn import ( |
| 22 | GDNAttentionMetadata, |
| 23 | GDNAttentionMetadataBuilder, |
| 24 | ) |
| 25 | from vllm.v1.attention.backends.utils import mamba_get_block_table_tensor |
| 26 | from vllm.v1.kv_cache_interface import MambaSpec |
| 27 | |
| 28 | BLOCK_SIZE = 16 |
| 29 | DEVICE = torch.device("cpu") |
| 30 | |
| 31 | |
| 32 | @dataclass |
| 33 | class GDNBuildTestCase: |
| 34 | """Specification for a GDN metadata builder classification test.""" |
| 35 | |
| 36 | seq_lens: list[int] |
| 37 | query_lens: list[int] |
| 38 | num_decode_draft_tokens: list[int] | None # None = no spec config |
| 39 | num_speculative_tokens: int |
| 40 | expected_num_decodes: int |
| 41 | expected_num_prefills: int |
| 42 | expected_num_prefill_tokens: int |
| 43 | expected_num_spec_decodes: int |
| 44 | |
| 45 | |
| 46 | GDN_BUILD_TEST_CASES = { |
| 47 | # The original #34845 crash: non-spec query_len=1 + spec decode |
| 48 | "mixed_decode_and_spec_decode": GDNBuildTestCase( |
| 49 | seq_lens=[65, 20], |
| 50 | query_lens=[1, 3], |
| 51 | num_decode_draft_tokens=[-1, 2], |
| 52 | num_speculative_tokens=2, |
| 53 | expected_num_decodes=0, |
| 54 | expected_num_prefills=1, |
| 55 | expected_num_prefill_tokens=1, |
| 56 | expected_num_spec_decodes=1, |
| 57 | ), |
| 58 | # All requests are spec decodes — no reclassification needed |
| 59 | "pure_spec_decode": GDNBuildTestCase( |
| 60 | seq_lens=[50, 30], |
| 61 | query_lens=[3, 3], |
| 62 | num_decode_draft_tokens=[2, 2], |
| 63 | num_speculative_tokens=2, |
| 64 | expected_num_decodes=0, |
| 65 | expected_num_prefills=0, |
| 66 | expected_num_prefill_tokens=0, |
| 67 | expected_num_spec_decodes=2, |
| 68 | ), |
| 69 | # Padded (CUDA graph) sequences trail the spec decodes |
| 70 | "pure_spec_decode_with_padding": GDNBuildTestCase( |
| 71 | seq_lens=[50, 30, 16], |
| 72 | query_lens=[3, 3, 0], |
| 73 | num_decode_draft_tokens=[2, 2, -1], |
| 74 | num_speculative_tokens=2, |
| 75 | expected_num_decodes=0, |
| 76 | expected_num_prefills=0, |
| 77 | expected_num_prefill_tokens=0, |
| 78 | expected_num_spec_decodes=2, |
| 79 | ), |
| 80 | # No speculative config at all — standard decode path |
| 81 | "pure_regular_decode": GDNBuildTestCase( |
| 82 | seq_lens=[40, 30, 20], |
| 83 | query_lens=[1, 1, 1], |
| 84 | num_decode_draft_tokens=None, |
| 85 | num_speculative_tokens=0, |
| 86 | expected_num_decodes=3, |
| 87 | expected_num_prefills=0, |
| 88 | expected_num_prefill_tokens=0, |
| 89 | expected_num_spec_decodes=0, |
| 90 | ), |
| 91 | # No speculative config, decode alongside prefill |
| 92 | "regular_decode_with_prefill": GDNBuildTestCase( |
| 93 | seq_lens=[40, 100], |
| 94 | query_lens=[1, 50], |
| 95 | num_decode_draft_tokens=None, |
| 96 | num_speculative_tokens=0, |
| 97 | expected_num_decodes=1, |
| 98 | expected_num_prefills=1, |
| 99 | expected_num_prefill_tokens=50, |
| 100 | expected_num_spec_decodes=0, |
| 101 | ), |
| 102 | # Multi-token prefill alongside spec decode — no decode to reclassify |
| 103 | "spec_decode_with_real_prefill": GDNBuildTestCase( |
| 104 | seq_lens=[100, 20], |
| 105 | query_lens=[50, 3], |
| 106 | num_decode_draft_tokens=[-1, 2], |
| 107 | num_speculative_tokens=2, |
| 108 | expected_num_decodes=0, |
| 109 | expected_num_prefills=1, |
| 110 | expected_num_prefill_tokens=50, |
| 111 | expected_num_spec_decodes=1, |
| 112 | ), |
| 113 | # All three types in one batch — decode gets reclassified |
| 114 | "prefill_decode_and_spec_decode": GDNBuildTestCase( |
| 115 | seq_lens=[100, 65, 20], |
| 116 | query_lens=[50, 1, 3], |
| 117 | num_decode_draft_tokens=[-1, -1, 2], |
| 118 | num_speculative_tokens=2, |
| 119 | expected_num_decodes=0, |
| 120 | expected_num_prefills=2, |
| 121 | expected_num_prefill_tokens=51, |
| 122 | expected_num_spec_decodes=1, |
| 123 | ), |
| 124 | # Multiple non-spec query_len=1 requests all reclassified |
| 125 | "multiple_decodes_reclassified": GDNBuildTestCase( |
| 126 | seq_lens=[40, 50, 60, 20], |
| 127 | query_lens=[1, 1, 1, 3], |
| 128 | num_decode_draft_tokens=[-1, -1, -1, 2], |
| 129 | num_speculative_tokens=2, |
| 130 | expected_num_decodes=0, |
| 131 | expected_num_prefills=3, |
| 132 | expected_num_prefill_tokens=3, |
| 133 | expected_num_spec_decodes=1, |
| 134 | ), |
| 135 | # Zero-length padded sequence excluded from counts |
| 136 | "zero_length_padding_with_spec": GDNBuildTestCase( |
| 137 | seq_lens=[16, 65, 20], |
| 138 | query_lens=[0, 1, 3], |
| 139 | num_decode_draft_tokens=[-1, -1, 2], |
| 140 | num_speculative_tokens=2, |
| 141 | expected_num_decodes=0, |
| 142 | expected_num_prefills=1, |
| 143 | expected_num_prefill_tokens=1, |
| 144 | expected_num_spec_decodes=1, |
| 145 | ), |
| 146 | } |
| 147 | |
| 148 | |
| 149 | def _create_gdn_builder( |
| 150 | num_speculative_tokens: int = 0, |
| 151 | full_cuda_graph: bool = False, |
| 152 | ) -> GDNAttentionMetadataBuilder: |
| 153 | """Create a GDNAttentionMetadataBuilder with minimal config.""" |
| 154 | vllm_config = create_vllm_config( |
| 155 | model_name="Qwen/Qwen3.5-0.8B", |
| 156 | block_size=BLOCK_SIZE, |
| 157 | ) |
| 158 | vllm_config.compilation_config.cudagraph_mode = ( |
| 159 | CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE |
| 160 | ) |
| 161 | if num_speculative_tokens > 0: |
| 162 | vllm_config.speculative_config = SpeculativeConfig( |
| 163 | method="ngram", |
| 164 | num_speculative_tokens=num_speculative_tokens, |
| 165 | ) |
| 166 | mamba_spec = MambaSpec( |
| 167 | block_size=BLOCK_SIZE, |
| 168 | shapes=((16, 64),), |
| 169 | dtypes=(torch.float16,), |
| 170 | ) |
| 171 | return GDNAttentionMetadataBuilder( |
| 172 | kv_cache_spec=mamba_spec, |
| 173 | layer_names=["layer.0"], |
| 174 | vllm_config=vllm_config, |
| 175 | device=DEVICE, |
| 176 | ) |
| 177 | |
| 178 | |
| 179 | def _build( |
| 180 | builder: GDNAttentionMetadataBuilder, |
| 181 | batch_spec: BatchSpec, |
| 182 | num_decode_draft_tokens: list[int] | None = None, |
| 183 | block_table: torch.Tensor | None = None, |
| 184 | ) -> GDNAttentionMetadata: |
| 185 | """Build GDN attention metadata, optionally with spec-decode kwargs.""" |
| 186 | common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE) |
| 187 | if block_table is not None: |
| 188 | common = common.replace(block_table_tensor=block_table) |
| 189 | kwargs: dict = {} |
| 190 | if num_decode_draft_tokens is not None: |
| 191 | kwargs["num_decode_draft_tokens_cpu"] = torch.tensor( |
| 192 | num_decode_draft_tokens, dtype=torch.int32 |
| 193 | ) |
| 194 | kwargs["num_accepted_tokens"] = torch.ones( |
| 195 | batch_spec.batch_size, dtype=torch.int32, device=DEVICE |
| 196 | ) |
| 197 | return builder.build(common_prefix_len=0, common_attn_metadata=common, **kwargs) |
| 198 | |
| 199 | |
| 200 | @pytest.mark.parametrize( |
| 201 | "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys() |
| 202 | ) |
| 203 | def test_gdn_build_classification(test_case: GDNBuildTestCase): |
| 204 | """Test that GDN metadata builder classifies requests correctly.""" |
| 205 | builder = _create_gdn_builder(test_case.num_speculative_tokens) |
| 206 | batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens) |
| 207 | meta = _build(builder, batch, test_case.num_decode_draft_tokens) |
| 208 | |
| 209 | assert meta.num_decodes == test_case.expected_num_decodes |
| 210 | assert meta.num_prefills == test_case.expected_num_prefills |
| 211 | assert meta.num_prefill_tokens == test_case.expected_num_prefill_tokens |
| 212 | assert meta.num_spec_decodes == test_case.expected_num_spec_decodes |
| 213 | if meta.spec_state_indices_tensor is not None: |
| 214 | assert len(meta.spec_state_indices_tensor) == meta.num_spec_decodes |
| 215 | |
| 216 | |
| 217 | @pytest.mark.parametrize("mamba_cache_mode", ["none", "align"]) |
| 218 | @pytest.mark.parametrize("full_cuda_graph", [False, True]) |
| 219 | @pytest.mark.parametrize( |
| 220 | "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys() |
| 221 | ) |
| 222 | def test_update_block_table_matches_build( |
| 223 | test_case: GDNBuildTestCase, full_cuda_graph: bool, mamba_cache_mode: str |
| 224 | ): |
| 225 | """update_block_table() on another group's metadata matches build().""" |
| 226 | batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens) |
| 227 | src, dst, ref = ( |
| 228 | _create_gdn_builder(test_case.num_speculative_tokens, full_cuda_graph) |
| 229 | for _ in range(3) |
| 230 | ) |
| 231 | for builder in (src, dst, ref): |
| 232 | builder.vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode |
| 233 | common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE) |
| 234 | if mamba_cache_mode == "align": |
| 235 | # MRV2 precomputes these for every Mamba group each step. |
| 236 | dst.mamba_aligned_state_indices = ref.mamba_aligned_state_indices = ( |
| 237 | mamba_get_block_table_tensor( |
| 238 | common.block_table_tensor, common.seq_lens, ref.kv_cache_spec, "align" |
| 239 | ) |
| 240 | ) |
| 241 | draft_tokens = test_case.num_decode_draft_tokens |
| 242 | expected = _build(ref, batch, draft_tokens, common.block_table_tensor) |
| 243 | source = _build(src, batch, draft_tokens) |
| 244 | fields = ( |
| 245 | "spec_state_indices_tensor", |
| 246 | "non_spec_state_indices_tensor", |
| 247 | "prefill_state_indices", |
| 248 | ) |
| 249 | source_indices = [getattr(source, f) for f in fields] |
| 250 | source_indices = [t if t is None else t.clone() for t in source_indices] |
| 251 | meta = dst.update_block_table( |
| 252 | source, common.block_table_tensor, common.slot_mapping |
| 253 | ) |
| 254 | |
| 255 | for field, source_index in zip(fields, source_indices): |
| 256 | actual = getattr(meta, field) |
| 257 | torch.testing.assert_close(actual, getattr(expected, field)) |
| 258 | # The source group's indices are untouched. |
| 259 | torch.testing.assert_close(getattr(source, field), source_index) |
| 260 | # FULL graph state indices land in this group's own buffers. |
| 261 | if full_cuda_graph and meta.num_prefills == 0 and actual is not None: |
| 262 | assert actual.data_ptr() == getattr(dst, field).data_ptr() |
| 263 | |
| 264 | |
| 265 | def test_has_initial_state_after_reclassification(): |
| 266 | """After reclassification, num_prefills > 0 so the prefill kernel path |
| 267 | should compute has_initial_state. For the reclassified request with |
| 268 | context_lens > 0, the corresponding entry must be True.""" |
| 269 | builder = _create_gdn_builder(num_speculative_tokens=2) |
| 270 | batch = BatchSpec(seq_lens=[65, 20], query_lens=[1, 3]) |
| 271 | meta = _build(builder, batch, num_decode_draft_tokens=[-1, 2]) |
| 272 | |
| 273 | assert meta.num_prefills > 0, "reclassification should produce prefills" |
| 274 | assert meta.has_initial_state is not None |
| 275 | # req0 has context_lens = 65 - 1 = 64 > 0, so has_initial_state[0] = True |
| 276 | assert meta.has_initial_state[0].item() is True |
| 277 | |
| 278 | |
| 279 | def test_full_cudagraph_spec_metadata_uses_request_count(): |
| 280 | """FULL cudagraph token padding must not pad request-indexed metadata.""" |
| 281 | num_speculative_tokens = 3 |
| 282 | builder = _create_gdn_builder( |
| 283 | num_speculative_tokens=num_speculative_tokens, |
| 284 | full_cuda_graph=True, |
| 285 | ) |
| 286 | batch = BatchSpec(seq_lens=[80, 96], query_lens=[4, 4]) |
| 287 | meta = _build(builder, batch, num_decode_draft_tokens=[3, 3]) |
| 288 | |
| 289 | assert meta.num_spec_decodes == batch.batch_size |
| 290 | assert meta.num_spec_decode_tokens == batch.compute_num_tokens() |
| 291 | assert meta.spec_state_indices_tensor is not None |
| 292 | assert meta.spec_state_indices_tensor.shape == ( |
| 293 | batch.batch_size, |
| 294 | num_speculative_tokens + 1, |
| 295 | ) |
| 296 | assert meta.spec_sequence_masks is not None |
| 297 | assert meta.spec_sequence_masks.shape == (batch.batch_size,) |
| 298 | assert meta.spec_query_start_loc is not None |
| 299 | assert meta.spec_query_start_loc.shape == (batch.batch_size + 1,) |
| 300 | assert meta.num_accepted_tokens is not None |
| 301 | assert meta.num_accepted_tokens.shape == (batch.batch_size,) |
| 302 | |
| 303 | |
| 304 | def _build_non_spec( |
| 305 | batch: BatchSpec, |
| 306 | is_prefilling: list[bool] | None, |
| 307 | full_cuda_graph: bool = False, |
| 308 | ): |
| 309 | common_attn_metadata = create_common_attn_metadata( |
| 310 | batch, BLOCK_SIZE, DEVICE |
| 311 | ).replace( |
| 312 | is_prefilling=None |
| 313 | if is_prefilling is None |
| 314 | else torch.tensor(is_prefilling, dtype=torch.bool) |
| 315 | ) |
| 316 | builder = _create_gdn_builder(full_cuda_graph=full_cuda_graph) |
| 317 | return builder, common_attn_metadata, builder.build(0, common_attn_metadata) |
| 318 | |
| 319 | |
| 320 | @pytest.mark.parametrize( |
| 321 | ("seq_len", "query_len", "is_prefilling", "num_prefills"), |
| 322 | [ |
| 323 | pytest.param(1, 1, True, 1, id="first-chunk"), |
| 324 | pytest.param(65, 1, True, 0, id="resumed-chunk"), |
| 325 | pytest.param(0, 0, True, 0, id="padding"), |
| 326 | pytest.param(1, 1, None, 0, id="missing-prefill-flag"), |
| 327 | ], |
| 328 | ) |
| 329 | def test_one_token_chunk_classification( |
| 330 | seq_len: int, |
| 331 | query_len: int, |
| 332 | is_prefilling: bool | None, |
| 333 | num_prefills: int, |
| 334 | ): |
| 335 | """Only a real first chunk with a prefill flag needs state initialization.""" |
| 336 | _, _, meta = _build_non_spec( |
| 337 | BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]), |
| 338 | is_prefilling=None if is_prefilling is None else [False, is_prefilling], |
| 339 | ) |
| 340 | |
| 341 | assert meta.num_prefills == num_prefills |
| 342 | assert meta.num_decodes == 2 - num_prefills |
| 343 | assert meta.num_prefill_tokens == num_prefills |
| 344 | assert meta.num_decode_tokens == 1 + query_len - num_prefills |
| 345 | if num_prefills: |
| 346 | assert meta.has_initial_state is not None |
| 347 | assert meta.has_initial_state.tolist() == [True, False] |
| 348 | else: |
| 349 | assert meta.has_initial_state is None |
| 350 | |
| 351 | |
| 352 | def test_one_token_first_chunk_excludes_padding(): |
| 353 | """Neither padding requests nor padding tokens count as prefill work.""" |
| 354 | common = create_common_attn_metadata( |
| 355 | BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]), |
| 356 | BLOCK_SIZE, |
| 357 | DEVICE, |
| 358 | ).replace( |
| 359 | is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool), |
| 360 | num_actual_tokens=4, |
| 361 | ) |
| 362 | meta = _create_gdn_builder().build(0, common) |
| 363 | |
| 364 | assert meta.num_decodes == 1 |
| 365 | assert meta.num_prefills == 1 |
| 366 | assert meta.num_decode_tokens == 1 |
| 367 | assert meta.num_prefill_tokens == 1 |
| 368 | |
| 369 | |
| 370 | def test_cudagraph_capture_batch_stays_decode_only(): |
| 371 | """Capture rows have no history, but must still select decode kernels.""" |
| 372 | batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4) |
| 373 | builder, common_attn_metadata, _ = _build_non_spec( |
| 374 | batch, [False] * 4, full_cuda_graph=True |
| 375 | ) |
| 376 | meta = builder.build_for_cudagraph_capture(common_attn_metadata) |
| 377 | |
| 378 | assert meta.num_prefills == 0 |
| 379 | assert meta.num_decodes == 4 |
| 380 | assert meta.has_initial_state is None |
| 381 | staged = meta.non_spec_state_indices_tensor |
| 382 | assert staged is not None |
| 383 | assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr() |
| 384 | torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0]) |
| 385 | |
| 386 | |
| 387 | def _create_checkpoint_builder_and_batch( |
| 388 | num_spec: int, prefix_match_unit: int = 16 |
| 389 | ) -> tuple[GDNAttentionMetadataBuilder, CommonAttentionMetadata, dict]: |
| 390 | vllm_config = create_vllm_config( |
| 391 | model_name="Qwen/Qwen3.5-0.8B", |
| 392 | block_size=BLOCK_SIZE, |
| 393 | ) |
| 394 | if num_spec > 0: |
| 395 | vllm_config.speculative_config = SpeculativeConfig( |
| 396 | method="ngram", |
| 397 | num_speculative_tokens=num_spec, |
| 398 | ) |
| 399 | vllm_config.cache_config.mamba_cache_mode = "align" |
| 400 | vllm_config.cache_config.prefix_match_unit = prefix_match_unit |
| 401 | mamba_spec = MambaSpec( |
| 402 | block_size=64, |
| 403 | shapes=((16, 64),), |
| 404 | dtypes=(torch.float16,), |
| 405 | mamba_cache_mode="align", |
| 406 | num_speculative_blocks=num_spec, |
| 407 | num_prefill_checkpoint_blocks=1, |
| 408 | prefill_checkpoint_alignment=16, |
| 409 | ) |
| 410 | builder = GDNAttentionMetadataBuilder( |
| 411 | kv_cache_spec=mamba_spec, |
| 412 | layer_names=["layer.0"], |
| 413 | vllm_config=vllm_config, |
| 414 | device=DEVICE, |
| 415 | ) |
| 416 | batch = BatchSpec( |
| 417 | seq_lens=[65, 100, 100], query_lens=[3 if num_spec else 1, 100, 99] |
| 418 | ) |
| 419 | common = create_common_attn_metadata(batch, 64, DEVICE, arange_block_indices=True) |
| 420 | table = common.block_table_tensor + 1 |
| 421 | if num_spec: |
| 422 | table = torch.cat([table, torch.zeros(3, num_spec, dtype=table.dtype)], dim=1) |
| 423 | common = common.replace(block_table_tensor=table) |
| 424 | kwargs = {} |
| 425 | if num_spec: |
| 426 | kwargs = dict( |
| 427 | num_decode_draft_tokens_cpu=torch.tensor([2, -1, -1], dtype=torch.int32), |
| 428 | num_accepted_tokens=torch.ones(3, dtype=torch.int32), |
| 429 | ) |
| 430 | return builder, common, kwargs |
| 431 | |
| 432 | |
| 433 | @pytest.mark.parametrize("num_spec", [0, 3]) |
| 434 | @pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)]) |
| 435 | def test_checkpoint_metadata_preserves_non_spec_order( |
| 436 | num_spec, prefix_match_unit, expected_offset |
| 437 | ): |
| 438 | """Only eligible non-spec rows get a checkpoint in their reserved page.""" |
| 439 | from vllm.v1.attention.backends.utils import NULL_BLOCK_ID |
| 440 | |
| 441 | builder, common, kwargs = _create_checkpoint_builder_and_batch( |
| 442 | num_spec, prefix_match_unit |
| 443 | ) |
| 444 | actual = builder.build(0, common, **kwargs) |
| 445 | assert actual.checkpoint is not None |
| 446 | offsets = [expected_offset, 0] if num_spec else [0, expected_offset, 0] |
| 447 | slots = [common.block_table_tensor[1, 0].item(), NULL_BLOCK_ID] |
| 448 | if not num_spec: |
| 449 | slots.insert(0, NULL_BLOCK_ID) |
| 450 | torch.testing.assert_close( |
| 451 | actual.checkpoint.checkpoint_offsets, torch.tensor(offsets, dtype=torch.int32) |
| 452 | ) |
| 453 | torch.testing.assert_close( |
| 454 | actual.checkpoint.state_indices, torch.tensor(slots, dtype=torch.int32) |
| 455 | ) |
| 456 | |
| 457 | |
| 458 | @pytest.mark.parametrize("num_spec", [0, 3]) |
| 459 | def test_update_block_table_regathers_checkpoint(num_spec): |
| 460 | """Checkpoint pages come from the target group's block table, not the |
| 461 | source group's (otherwise the target group's checkpoint page is never |
| 462 | written while the prefix cache treats it as valid).""" |
| 463 | src, common, kwargs = _create_checkpoint_builder_and_batch(num_spec) |
| 464 | dst, _, _ = _create_checkpoint_builder_and_batch(num_spec) |
| 465 | ref, _, _ = _create_checkpoint_builder_and_batch(num_spec) |
| 466 | other_table = common.block_table_tensor + 100 |
| 467 | other = common.replace(block_table_tensor=other_table) |
| 468 | dst.mamba_aligned_state_indices = mamba_get_block_table_tensor( |
| 469 | other_table, other.seq_lens, dst.kv_cache_spec, "align" |
| 470 | ) |
| 471 | |
| 472 | source = src.build(0, common, **kwargs) |
| 473 | expected = ref.build(0, other, **kwargs) |
| 474 | actual = dst.update_block_table(source, other_table, other.slot_mapping) |
| 475 | |
| 476 | assert source.checkpoint is not None and expected.checkpoint is not None |
| 477 | assert actual.checkpoint is not None |
| 478 | torch.testing.assert_close( |
| 479 | actual.checkpoint.state_indices, expected.checkpoint.state_indices |
| 480 | ) |
| 481 | torch.testing.assert_close( |
| 482 | actual.checkpoint.checkpoint_offsets, expected.checkpoint.checkpoint_offsets |
| 483 | ) |
| 484 | assert not torch.equal( |
| 485 | actual.checkpoint.state_indices, source.checkpoint.state_indices |
| 486 | ) |