perf(glm5next): streamline B12X sparse decode metadata

tests/models/test_glm5next_pooled_indexer.py

@@ -13,13 +13,18 @@
 from vllm.models.glm5next.nvidia.ops.glm_kpool import (
     expand_c4_block_table,
     expand_pool_ids,
+    expand_pool_ids_physical,
     gather_c4_block_table_rows,
     pool_seq_lens,
+    prepare_c4_decode_metadata,
     update_decode_pools,
 )
 from vllm.models.glm5next.nvidia.pooled_indexer import Glm5NextPooledIndexer
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.mla.b12x_mla_sparse import B12xMLASparseMetadata
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    triton_convert_req_index_to_global_index,
+)
 from vllm.v1.kv_cache_interface import MLAAttentionSpec
 
 
@@ -197,6 +202,277 @@ def test_glm53_packed_c4_metadata_uses_parent_stride() -> None:
         )
 
 
+@pytest.mark.parametrize(("rows", "requests"), [(1, 1), (7, 4), (32, 32)])
+@pytest.mark.parametrize(
+    ("dcp_size", "dcp_rank", "pool_interleave"),
+    [(1, 0, 1), (4, 2, 1), (4, 3, 2)],
+)
+def test_glm53_c4_decode_metadata_matches_reference(
+    rows: int,
+    requests: int,
+    dcp_size: int,
+    dcp_rank: int,
+    pool_interleave: int,
+) -> None:
+    device = _require_glm_gpu()
+    source_width = 5
+    subpages_per_parent = 9
+    parent_stride_pages = 37
+    source = torch.arange(
+        requests * source_width, dtype=torch.int32, device=device
+    ).reshape(requests, source_width)
+    source[0, -1] = -1
+    source[-1, 0] = 58_000_000
+    request_ids = torch.arange(rows, dtype=torch.int32, device=device) % requests
+    positions = torch.arange(rows, dtype=torch.int64, device=device) * 257 + 3
+
+    expanded = torch.empty(
+        (requests, source_width * subpages_per_parent),
+        dtype=torch.int32,
+        device=device,
+    )
+    expected_table = torch.empty(
+        (rows, source_width * subpages_per_parent),
+        dtype=torch.int32,
+        device=device,
+    )
+    expected_seq_lens = torch.empty(rows, dtype=torch.int32, device=device)
+    actual_table = torch.empty_like(expected_table)
+    actual_seq_lens = torch.empty_like(expected_seq_lens)
+
+    expand_c4_block_table(
+        source,
+        expanded,
+        rows=requests,
+        subpages_per_parent=subpages_per_parent,
+        parent_stride_pages=parent_stride_pages,
+    )
+    gather_c4_block_table_rows(expanded, request_ids, expected_table)
+    pool_seq_lens(
+        positions,
+        expected_seq_lens,
+        dcp_size=dcp_size,
+        dcp_rank=dcp_rank,
+        pool_interleave=pool_interleave,
+    )
+    prepare_c4_decode_metadata(
+        source,
+        request_ids,
+        positions,
+        actual_table,
+        actual_seq_lens,
+        subpages_per_parent=subpages_per_parent,
+        parent_stride_pages=parent_stride_pages,
+        dcp_size=dcp_size,
+        dcp_rank=dcp_rank,
+        pool_interleave=pool_interleave,
+    )
+
+    torch.testing.assert_close(actual_table, expected_table, rtol=0, atol=0)
+    torch.testing.assert_close(actual_seq_lens, expected_seq_lens, rtol=0, atol=0)
+
+
+def test_glm53_c4_decode_metadata_graph_replays_live_inputs() -> None:
+    device = _require_glm_gpu()
+    rows = 7
+    requests = 4
+    source_width = 5
+    subpages_per_parent = 9
+    parent_stride_pages = 37
+    source = torch.arange(
+        requests * source_width, dtype=torch.int32, device=device
+    ).reshape(requests, source_width)
+    request_ids = torch.arange(rows, dtype=torch.int32, device=device) % requests
+    positions = torch.arange(rows, dtype=torch.int64, device=device) * 4 + 3
+    output_table = torch.empty(
+        (rows, source_width * subpages_per_parent),
+        dtype=torch.int32,
+        device=device,
+    )
+    output_seq_lens = torch.empty(rows, dtype=torch.int32, device=device)
+
+    def prepare() -> None:
+        prepare_c4_decode_metadata(
+            source,
+            request_ids,
+            positions,
+            output_table,
+            output_seq_lens,
+            subpages_per_parent=subpages_per_parent,
+            parent_stride_pages=parent_stride_pages,
+            dcp_size=4,
+            dcp_rank=2,
+            pool_interleave=2,
+        )
+
+    prepare()
+    device_module = torch.get_device_module(device)
+    graph = device_module.CUDAGraph()
+    with device_module.graph(graph):
+        prepare()
+
+    source.add_(100)
+    source[1, -1] = -1
+    request_ids.copy_(
+        torch.tensor([3, 1, 2, 0, 3, 2, 1], dtype=torch.int32, device=device)
+    )
+    positions.add_(4096)
+    output_table.fill_(37)
+    output_seq_lens.fill_(37)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    expanded = torch.empty(
+        (requests, source_width * subpages_per_parent),
+        dtype=torch.int32,
+        device=device,
+    )
+    expected_table = torch.empty_like(output_table)
+    expected_seq_lens = torch.empty_like(output_seq_lens)
+    expand_c4_block_table(
+        source,
+        expanded,
+        rows=requests,
+        subpages_per_parent=subpages_per_parent,
+        parent_stride_pages=parent_stride_pages,
+    )
+    gather_c4_block_table_rows(expanded, request_ids, expected_table)
+    pool_seq_lens(
+        positions,
+        expected_seq_lens,
+        dcp_size=4,
+        dcp_rank=2,
+        pool_interleave=2,
+    )
+    torch.testing.assert_close(output_table, expected_table, rtol=0, atol=0)
+    torch.testing.assert_close(output_seq_lens, expected_seq_lens, rtol=0, atol=0)
+
+    allocated = torch.accelerator.memory_allocated()
+    graph.replay()
+    graph.replay()
+    torch.accelerator.synchronize()
+    assert torch.accelerator.memory_allocated() == allocated
+
+
+@pytest.mark.parametrize("rows", [1, 7, 32])
+def test_glm53_physical_pool_expansion_matches_reference(rows: int) -> None:
+    device = _require_glm_gpu()
+    block_size = 256
+    max_blocks = 80
+    requests = min(rows, 8)
+    test_positions = [0, 1, 3, 4, 255, 256, 2047, 2048, 4095, 16383]
+    positions = torch.tensor(
+        [test_positions[row % len(test_positions)] for row in range(rows)],
+        dtype=torch.int64,
+        device=device,
+    )
+    request_ids = torch.arange(rows, dtype=torch.int32, device=device) % requests
+    block_table = torch.arange(
+        requests * max_blocks, dtype=torch.int32, device=device
+    ).reshape(requests, max_blocks)
+    block_table.mul_(101).add_(7_000_000)
+    block_table[:, -1] = -1
+    pool_ids = torch.full((rows, 512), -1, dtype=torch.int32, device=device)
+    for row, position in enumerate(positions.cpu().tolist()):
+        selected = min((position + 1) // 4, 512)
+        if selected:
+            pool_ids[row, :selected] = torch.arange(
+                selected - 1, -1, -1, dtype=torch.int32, device=device
+            )
+
+    logical = torch.empty((rows, 2051), dtype=torch.int32, device=device)
+    expand_pool_ids(pool_ids, positions, logical)
+    expected, expected_counts = triton_convert_req_index_to_global_index(
+        request_ids,
+        block_table,
+        logical,
+        BLOCK_SIZE=block_size,
+        BLOCK_STRIDE_ROWS=block_size,
+        NUM_TOPK_TOKENS=2051,
+        return_valid_counts=True,
+    )
+    actual = torch.empty_like(expected)
+    actual_counts = torch.empty_like(expected_counts)
+    expand_pool_ids_physical(
+        pool_ids,
+        positions,
+        request_ids,
+        block_table,
+        actual,
+        actual_counts,
+        block_size=block_size,
+        block_stride_rows=block_size,
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(actual_counts, expected_counts, rtol=0, atol=0)
+
+
+def test_glm53_physical_pool_expansion_graph_replays_live_inputs() -> None:
+    device = _require_glm_gpu()
+    rows = 7
+    requests = 4
+    block_size = 256
+    max_blocks = 16
+    pool_ids = torch.arange(512, dtype=torch.int32, device=device).repeat(rows, 1)
+    positions = torch.full((rows,), 2047, dtype=torch.int64, device=device)
+    request_ids = torch.arange(rows, dtype=torch.int32, device=device) % requests
+    block_table = torch.arange(
+        requests * max_blocks, dtype=torch.int32, device=device
+    ).reshape(requests, max_blocks)
+    output = torch.empty((rows, 2051), dtype=torch.int32, device=device)
+    active_counts = torch.empty(rows, dtype=torch.int32, device=device)
+
+    def expand() -> None:
+        expand_pool_ids_physical(
+            pool_ids,
+            positions,
+            request_ids,
+            block_table,
+            output,
+            active_counts,
+            block_size=block_size,
+            block_stride_rows=block_size,
+        )
+
+    expand()
+    device_module = torch.get_device_module(device)
+    graph = device_module.CUDAGraph()
+    with device_module.graph(graph):
+        expand()
+
+    pool_ids.copy_(pool_ids.flip(dims=(1,)))
+    positions.add_(1)
+    request_ids.copy_(
+        torch.tensor([3, 1, 2, 0, 3, 2, 1], dtype=torch.int32, device=device)
+    )
+    block_table.add_(7_000_000)
+    output.fill_(37)
+    active_counts.fill_(37)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    logical = torch.empty_like(output)
+    expand_pool_ids(pool_ids, positions, logical)
+    expected, expected_counts = triton_convert_req_index_to_global_index(
+        request_ids,
+        block_table,
+        logical,
+        BLOCK_SIZE=block_size,
+        BLOCK_STRIDE_ROWS=block_size,
+        NUM_TOPK_TOKENS=2051,
+        return_valid_counts=True,
+    )
+    torch.testing.assert_close(output, expected, rtol=0, atol=0)
+    torch.testing.assert_close(active_counts, expected_counts, rtol=0, atol=0)
+
+    allocated = torch.accelerator.memory_allocated()
+    graph.replay()
+    graph.replay()
+    torch.accelerator.synchronize()
+    assert torch.accelerator.memory_allocated() == allocated
+
+
 def _packed_main_cache(
     *, device: torch.device, blocks: int, layers: int, block_size: int, layer: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -711,7 +987,7 @@ def test_glm53_pool_expansion_appends_only_the_incomplete_tail() -> None:
     pool_ids[1, :2] = torch.tensor([1, 0], dtype=torch.int32, device=device)
     pool_ids[2] = torch.arange(512, dtype=torch.int32, device=device)
     positions = torch.tensor([2, 7, 2052], dtype=torch.int64, device=device)
-    output = torch.empty((3, 2051), dtype=torch.int32, device=device)
+    output = torch.full((3, 2051), 37, dtype=torch.int32, device=device)
 
     expand_pool_ids(pool_ids, positions, output)
 

tests/v1/attention/test_b12x_sparse_mla_api.py

@@ -1246,6 +1246,92 @@ def index_topk_fp8(**kwargs):
     assert calls["caps"]["max_page_table_width"] == 1024
 
 
+@pytest.mark.parametrize(
+    ("route", "expected_initial_value"),
+    [("paged_fused", 17), ("paged_tiled", -1)],
+)
+def test_b12x_paged_topk_initializes_only_non_fused_routes(
+    monkeypatch, route: str, expected_initial_value: int
+) -> None:
+    observed: list[torch.Tensor] = []
+
+    plan = SimpleNamespace(
+        layout=SimpleNamespace(route=route),
+        shapes_and_dtypes=lambda: (),
+        bind=lambda **kwargs: SimpleNamespace(route=route),
+    )
+
+    def index_topk_fp8(**kwargs):
+        output = kwargs["out_indices"]
+        observed.append(output.clone())
+        output.fill_(5)
+
+    module = SimpleNamespace(
+        PAGED_INDEX_PAGE_SIZE=64,
+        index_topk_fp8=index_topk_fp8,
+    )
+    monkeypatch.setattr(b12x_indexer, "current_workspace_manager", lambda: _Workspace())
+    output = torch.full((2, 4), 17, dtype=torch.int32)
+
+    b12x_indexer._run_paged_topk(
+        module=module,
+        plan=plan,
+        q=torch.empty((2, 16, 128), dtype=torch.float8_e4m3fn),
+        weights=torch.empty((2, 16, 1), dtype=torch.float32),
+        kv_cache=torch.empty((4, 64, 132), dtype=torch.uint8),
+        seq_lens=torch.full((2,), 128, dtype=torch.int32),
+        block_table=torch.zeros((2, 2), dtype=torch.int32),
+        schedule_metadata=None,
+        active_width=None,
+        output=output,
+        scores=None,
+        topk=4,
+        shared_page_table=False,
+    )
+
+    assert torch.count_nonzero(observed[0] != expected_initial_value) == 0
+    assert torch.count_nonzero(output != 5) == 0
+
+
+def test_b12x_decode_metadata_uses_plan_capacity_for_active_width(monkeypatch) -> None:
+    decode = b12x_indexer.DeepSeekV32IndexerDecodeMetadata(
+        block_table=torch.zeros((2, 4), dtype=torch.int32),
+        seq_lens=torch.full((2,), 128, dtype=torch.int32),
+        decode_lens=torch.ones((2,), dtype=torch.int32),
+        requires_padding=False,
+        schedule_metadata=torch.empty(0, dtype=torch.int32),
+    )
+    metadata = b12x_indexer.DeepseekV32IndexerMetadata(
+        seq_lens=decode.seq_lens,
+        max_seq_len=128,
+        slot_mapping=torch.arange(2, dtype=torch.int64),
+        num_decodes=2,
+        num_decode_tokens=2,
+        num_prefills=0,
+        num_prefill_tokens=0,
+        decode=decode,
+    )
+    monkeypatch.setattr(
+        b12x_indexer.DeepseekV32IndexerMetadataBuilder,
+        "build",
+        lambda self, *args, **kwargs: metadata,
+    )
+    monkeypatch.setattr(
+        b12x_indexer,
+        "_require_b12x_indexer",
+        lambda: SimpleNamespace(uses_paged_schedule=lambda **kwargs: False),
+    )
+    builder = object.__new__(b12x_indexer.DeepseekV4B12xIndexerMetadataBuilder)
+    builder.scheduler_metadata_buffer = torch.empty(0, dtype=torch.int32)
+    builder.num_sms = 1
+
+    result = builder.build()
+
+    assert isinstance(result.decode, b12x_indexer.DeepseekV4B12xIndexerDecodeMetadata)
+    assert result.decode.active_width is None
+    assert not hasattr(builder, "active_width_buffer")
+
+
 def test_b12x_dsa_indexer_reuses_plans_and_rebinds_shared_workspace(
     monkeypatch,
 ) -> None: