From 4f88c20fa8c69de0e5dd702e8d7d2c51165801a0 Mon Sep 17 00:00:00 2001 From: MadeBy561 <155126856+MadeBy561@users.noreply.github.com> Date: Sun, 30 Aug 2026 22:15:01 -0400 Subject: [PATCH 1/3] perf(indexer): skip fused decode initialization Let the B12X fused paged route own output padding and bind the plan capacity for its unused active-width scalar. Keep explicit output initialization on every non-fused route. Co-authored-by: OpenAI Codex (cherry picked from commit d18ed5301fcbe7e2b277556e580fd4555c71c187) --- tests/models/test_glm5next_pooled_indexer.py | 2 +- .../v1/attention/test_b12x_sparse_mla_api.py | 86 +++++++++++++++++++ .../models/deepseek_v4/nvidia/b12x_indexer.py | 21 +++-- vllm/models/glm5next/nvidia/pooled_indexer.py | 4 +- 4 files changed, 99 insertions(+), 14 deletions(-) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index 289c5752b4a6..a1913b42990b 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -711,7 +711,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) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 67463ebb4648..1620fa534b2e 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/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: diff --git a/vllm/models/deepseek_v4/nvidia/b12x_indexer.py b/vllm/models/deepseek_v4/nvidia/b12x_indexer.py index 9d95acbecab5..0a65cff283fe 100644 --- a/vllm/models/deepseek_v4/nvidia/b12x_indexer.py +++ b/vllm/models/deepseek_v4/nvidia/b12x_indexer.py @@ -52,9 +52,6 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: max_page_table_width=block_table_width, page_size=_INDEX_PAGE_SIZE, ) - self.active_width_buffer = torch.zeros( - (1,), dtype=torch.int32, device=self.device - ) def _supports_native_decode(self, next_n: int) -> bool: return True @@ -99,15 +96,13 @@ def build(self, *args, **kwargs) -> DeepseekV32IndexerMetadata: self.num_sms, out=self.scheduler_metadata_buffer, ) - active_width = ( - int(metadata.max_seq_len) + int(self.compress_ratio) - 1 - ) // int(self.compress_ratio) - self.active_width_buffer.fill_(active_width) decode_fields = vars(decode).copy() decode_fields["schedule_metadata"] = schedule_metadata metadata.decode = DeepseekV4B12xIndexerDecodeMetadata( **decode_fields, - active_width=self.active_width_buffer, + # Fused decode does not consume active_width; other routes bind + # the immutable plan-capacity tensor when no override is given. + active_width=None, ) return metadata @@ -201,6 +196,13 @@ def _run_paged_topk( topk: int, shared_page_table: bool, ) -> None: + route = getattr(plan, "route", None) + if route is None: + route = getattr(getattr(plan, "layout", None), "route", None) + # The fused decode route owns every output slot, including -1 padding. + # Other routes retain caller initialization through this shared entry point. + if route != "paged_fused": + output.fill_(-1) if shared_page_table: _assert_prefill_route(plan) scratch = current_workspace_manager().get_simultaneous(*plan.shapes_and_dtypes()) @@ -355,7 +357,6 @@ def run_paged_topk( raise ValueError( "B12x C4 scores must be float32 with the same shape as output" ) - output.fill_(-1) _run_paged_topk( module=self._b12x_indexer, plan=self._plan_paged_topk( @@ -427,7 +428,6 @@ def forward( block_table = chunk.block_table[:1, :active_pages].expand( int(q_chunk.shape[0]), active_pages ) - output.fill_(-1) _run_paged_topk( module=self._b12x_indexer, plan=self._plan_paged_topk( @@ -465,7 +465,6 @@ def forward( ) num_tokens = metadata.num_decode_tokens output = self.topk_indices_buffer[:num_tokens, : self.topk_tokens] - output.fill_(-1) active_width = getattr(decode, "active_width", None) _run_paged_topk( module=self._b12x_indexer, diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 65d6409dda67..af85681f601a 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -457,7 +457,6 @@ def forward( pool_interleave=self.pool_interleave, ) pool_ids = self.pool_topk_indices_buffer[:rows] - pool_ids.fill_(-1) pool_scores = self._pool_scores[:rows] if self.dcp_world_size > 1 else None if decode_rows: @@ -546,7 +545,8 @@ def forward( ) output = self.topk_indices_buffer[:rows] - output.fill_(-1) + if live_rows < rows: + output[live_rows:].fill_(-1) expand_pool_ids(pool_ids[:live_rows], positions[:live_rows], output[:live_rows]) return output From a1111e41470f6e2e0a084c69caaa1bb3113b2ac1 Mon Sep 17 00:00:00 2001 From: MadeBy561 <155126856+MadeBy561@users.noreply.github.com> Date: Sun, 30 Aug 2026 22:23:06 -0400 Subject: [PATCH 2/3] perf(glm5next): fuse C4 decode metadata Co-authored-by: OpenAI Codex (cherry picked from commit 2ae5b5808a9c614ea840f49474a4b00234fad08a) --- tests/models/test_glm5next_pooled_indexer.py | 153 ++++++++++++++++++ vllm/models/glm5next/nvidia/ops/glm_kpool.py | 116 +++++++++++++ vllm/models/glm5next/nvidia/pooled_indexer.py | 58 ++++--- 3 files changed, 307 insertions(+), 20 deletions(-) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index a1913b42990b..c1c8290415b3 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -15,6 +15,7 @@ expand_pool_ids, gather_c4_block_table_rows, pool_seq_lens, + prepare_c4_decode_metadata, update_decode_pools, ) from vllm.models.glm5next.nvidia.pooled_indexer import Glm5NextPooledIndexer @@ -197,6 +198,158 @@ 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 + + def _packed_main_cache( *, device: torch.device, blocks: int, layers: int, block_size: int, layer: int ) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/vllm/models/glm5next/nvidia/ops/glm_kpool.py b/vllm/models/glm5next/nvidia/ops/glm_kpool.py index 918f4e8ee621..1f62255fdc89 100644 --- a/vllm/models/glm5next/nvidia/ops/glm_kpool.py +++ b/vllm/models/glm5next/nvidia/ops/glm_kpool.py @@ -676,6 +676,121 @@ def gather_c4_block_table_rows( ) +@triton.jit +def _prepare_c4_decode_metadata_kernel( + source, + request_ids, + positions, + output_table, + output_seq_lens, + rows, + source_width, + output_width, + source_stride, + output_stride, + parent_stride_pages, + dcp_size, + dcp_rank, + pool_interleave, + SUBPAGES_PER_PARENT: tl.constexpr, + POOL_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + linear = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + total = rows * output_width + mask = linear < total + row = linear // output_width + output_col = linear - row * output_width + request = tl.load(request_ids + row, mask=mask, other=0).to(tl.int64) + source_col = output_col // SUBPAGES_PER_PARENT + child_page = output_col - source_col * SUBPAGES_PER_PARENT + parent_page = tl.load( + source + request * source_stride + source_col, + mask=mask & (source_col < source_width), + other=-1, + ).to(tl.int64) + child_page_id = parent_page * parent_stride_pages.to(tl.int64) + child_page + child_page_id = tl.where(parent_page >= 0, child_page_id, -1) + tl.store( + output_table + row * output_stride + output_col, + child_page_id, + mask=mask, + ) + + first_column = mask & (output_col == 0) + position = tl.load(positions + row, mask=first_column, other=-1).to(tl.int64) + global_pool_len = (position + 1) // POOL_SIZE + rounds = global_pool_len // (dcp_size * pool_interleave) + remainder = global_pool_len % (dcp_size * pool_interleave) + remainder = tl.maximum(remainder - dcp_rank * pool_interleave, 0) + remainder = tl.minimum(remainder, pool_interleave) + local_pool_len = rounds * pool_interleave + remainder + tl.store(output_seq_lens + row, local_pool_len, mask=first_column) + + +def prepare_c4_decode_metadata( + source: torch.Tensor, + request_ids: torch.Tensor, + positions: torch.Tensor, + output_table: torch.Tensor, + output_seq_lens: torch.Tensor, + *, + subpages_per_parent: int, + parent_stride_pages: int, + dcp_size: int = 1, + dcp_rank: int = 0, + pool_interleave: int = 1, +) -> None: + """Gather expanded C4 page rows and write their local sequence lengths.""" + if source.dtype != torch.int32 or request_ids.dtype != torch.int32: + raise TypeError("GLM block tables and request IDs must use int32") + if positions.dtype != torch.int64: + raise TypeError("GLM positions must use int64") + if output_table.dtype != torch.int32 or output_seq_lens.dtype != torch.int32: + raise TypeError("GLM decode metadata outputs must use int32") + if source.ndim != 2 or int(source.shape[1]) < 1: + raise ValueError("GLM parent block table must be non-empty and rank two") + if request_ids.ndim != 1: + raise ValueError("GLM decode request IDs must be rank one") + rows = int(request_ids.shape[0]) + if positions.shape != (rows,): + raise ValueError("GLM selector positions must have one entry per row") + if subpages_per_parent < 1 or parent_stride_pages < 1: + raise ValueError("GLM C4 child-page geometry must be positive") + expected_width = int(source.shape[1]) * subpages_per_parent + if output_table.shape != (rows, expected_width): + raise ValueError("GLM decode block-table output has the wrong contract") + if output_seq_lens.shape != (rows,): + raise ValueError("GLM pool sequence lengths must have shape [rows]") + if dcp_size < 1 or not 0 <= dcp_rank < dcp_size: + raise ValueError("GLM pool DCP rank must be within the DCP world") + if pool_interleave < 1: + raise ValueError("GLM pool interleave must be positive") + if rows: + block = 256 + total = rows * expected_width + _prepare_c4_decode_metadata_kernel[(triton.cdiv(total, block),)]( + source, + request_ids, + positions, + output_table, + output_seq_lens, + rows, + int(source.shape[1]), + expected_width, + int(source.stride(0)), + int(output_table.stride(0)), + parent_stride_pages, + dcp_size, + dcp_rank, + pool_interleave, + SUBPAGES_PER_PARENT=subpages_per_parent, + POOL_SIZE=_POOL_SIZE, + BLOCK=block, + num_warps=4, + ) + + @triton.jit def _pool_seq_lens_kernel( positions, @@ -803,5 +918,6 @@ def expand_pool_ids( "fwht128_quant_fp8", "gather_c4_block_table_rows", "pool_seq_lens", + "prepare_c4_decode_metadata", "update_decode_pools", ] diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index af85681f601a..492cede1affd 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -34,6 +34,7 @@ fwht128_quant_fp8, gather_c4_block_table_rows, pool_seq_lens, + prepare_c4_decode_metadata, update_decode_pools, ) @@ -441,31 +442,48 @@ def forward( model_block_size=self.block_size, parent_stride_pages=self._parent_stride_pages, ) - expand_c4_block_table( - main_metadata.block_table[:num_reqs, : self._parent_table_width], - self._pool_block_table, - rows=num_reqs, - subpages_per_parent=self._subpages_per_parent, - parent_stride_pages=self._parent_stride_pages, - ) + parent_table = main_metadata.block_table[:num_reqs, : self._parent_table_width] seq_lens = self._pool_seq_lens[:live_rows] - pool_seq_lens( - positions[:live_rows], - seq_lens, - dcp_size=self.dcp_world_size, - dcp_rank=self.dcp_rank, - pool_interleave=self.pool_interleave, - ) + decode_only = decode_rows == live_rows + decode_table = self._decode_block_table[:decode_rows] + if decode_only: + prepare_c4_decode_metadata( + parent_table, + main_metadata.req_id_per_token[:decode_rows], + positions[:decode_rows], + decode_table, + seq_lens, + subpages_per_parent=self._subpages_per_parent, + parent_stride_pages=self._parent_stride_pages, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + pool_interleave=self.pool_interleave, + ) + else: + expand_c4_block_table( + parent_table, + self._pool_block_table, + rows=num_reqs, + subpages_per_parent=self._subpages_per_parent, + parent_stride_pages=self._parent_stride_pages, + ) + pool_seq_lens( + positions[:live_rows], + seq_lens, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + pool_interleave=self.pool_interleave, + ) pool_ids = self.pool_topk_indices_buffer[:rows] pool_scores = self._pool_scores[:rows] if self.dcp_world_size > 1 else None if decode_rows: - decode_table = self._decode_block_table[:decode_rows] - gather_c4_block_table_rows( - self._pool_block_table, - main_metadata.req_id_per_token[:decode_rows], - decode_table, - ) + if not decode_only: + gather_c4_block_table_rows( + self._pool_block_table, + main_metadata.req_id_per_token[:decode_rows], + decode_table, + ) self.indexer_op.run_paged_topk( q=q_fp8[:decode_rows], weights=weights[:decode_rows], From 52006e858766911cebccc8a7b2b21a260f5c97f2 Mon Sep 17 00:00:00 2001 From: MadeBy561 <155126856+MadeBy561@users.noreply.github.com> Date: Sun, 30 Aug 2026 22:30:23 -0400 Subject: [PATCH 3/3] perf(glm5next): emit physical pooled indices Co-authored-by: OpenAI Codex (cherry picked from commit adb69eac865a1a37081fa4edb9f7599a351f7aac) --- tests/models/test_glm5next_pooled_indexer.py | 123 ++++++++++++++++++ vllm/models/glm5next/nvidia/ops/glm_kpool.py | 123 ++++++++++++++++++ vllm/models/glm5next/nvidia/pooled_indexer.py | 22 +++- .../attention/backends/mla/b12x_mla_sparse.py | 12 ++ 4 files changed, 279 insertions(+), 1 deletion(-) diff --git a/tests/models/test_glm5next_pooled_indexer.py b/tests/models/test_glm5next_pooled_indexer.py index c1c8290415b3..ac1c6e6633df 100644 --- a/tests/models/test_glm5next_pooled_indexer.py +++ b/tests/models/test_glm5next_pooled_indexer.py @@ -13,6 +13,7 @@ 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, @@ -21,6 +22,9 @@ 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 @@ -350,6 +354,125 @@ def prepare() -> None: 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]: diff --git a/vllm/models/glm5next/nvidia/ops/glm_kpool.py b/vllm/models/glm5next/nvidia/ops/glm_kpool.py index 1f62255fdc89..71832885d593 100644 --- a/vllm/models/glm5next/nvidia/ops/glm_kpool.py +++ b/vllm/models/glm5next/nvidia/ops/glm_kpool.py @@ -912,9 +912,132 @@ def expand_pool_ids( ) +@triton.jit +def _expand_pool_ids_physical_kernel( + pool_ids, + positions, + request_ids, + block_table, + output, + active_counts, + pool_stride, + block_table_stride, + output_stride, + max_num_blocks, + HISTORY_TOKENS: tl.constexpr, + OUTPUT_WIDTH: tl.constexpr, + POOL_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_STRIDE_ROWS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + column = tile * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + mask = column < OUTPUT_WIDTH + sequence_length = tl.load(positions + row).to(tl.int64) + 1 + complete_pools = sequence_length // POOL_SIZE + tail_start = complete_pools * POOL_SIZE + history = column < HISTORY_TOKENS + pool_column = column // POOL_SIZE + pool_offset = column % POOL_SIZE + pool_id = tl.load( + pool_ids + row * pool_stride + pool_column, + mask=mask & history, + other=-1, + ).to(tl.int64) + history_value = tl.where(pool_id >= 0, pool_id * POOL_SIZE + pool_offset, -1) + tail_offset = column - HISTORY_TOKENS + tail_count = sequence_length - tail_start + in_tail = (tail_offset >= 0) & (tail_offset < tail_count) + logical_token = tl.where( + history, + history_value, + tl.where(in_tail, tail_start + tail_offset, -1), + ) + + request = tl.load(request_ids + row).to(tl.int64) + block_id = logical_token // BLOCK_SIZE + in_block = logical_token - block_id * BLOCK_SIZE + valid = (logical_token >= 0) & (block_id < max_num_blocks) + page = tl.load( + block_table + request * block_table_stride + block_id, + mask=mask & valid, + other=-1, + ).to(tl.int64) + physical_token = tl.where( + valid & (page >= 0), + page * BLOCK_STRIDE_ROWS + in_block, + -1, + ).to(tl.int32) + tl.store(output + row * output_stride + column, physical_token, mask=mask) + + # The selector returns one valid ID for every complete pool until top-k is + # saturated, followed by the 0--3 unpooled tail tokens. Every column tile + # computes the same scalar; only tile zero publishes it. + selected_pools = tl.minimum(complete_pools, HISTORY_TOKENS // POOL_SIZE) + active_count = selected_pools * POOL_SIZE + tail_count + tl.store(active_counts + row, active_count, mask=tile == 0) + + +def expand_pool_ids_physical( + pool_ids: torch.Tensor, + positions: torch.Tensor, + request_ids: torch.Tensor, + block_table: torch.Tensor, + output: torch.Tensor, + active_counts: torch.Tensor, + *, + block_size: int, + block_stride_rows: int, +) -> None: + """Expand pooled selections directly into physical main-cache token slots.""" + if pool_ids.dtype != torch.int32 or request_ids.dtype != torch.int32: + raise TypeError("GLM pool selections and request IDs must use int32") + if positions.dtype != torch.int64: + raise TypeError("GLM expansion positions must use int64") + if block_table.dtype != torch.int32: + raise TypeError("GLM expansion block table must use int32") + if output.dtype != torch.int32 or active_counts.dtype != torch.int32: + raise TypeError("GLM physical-selection outputs must use int32") + if pool_ids.ndim != 2 or int(pool_ids.shape[1]) != 512: + raise ValueError("GLM pool selection must have shape [rows, 512]") + rows = int(pool_ids.shape[0]) + if positions.shape != (rows,) or request_ids.shape != (rows,): + raise ValueError("GLM expansion metadata must have one entry per row") + if block_table.ndim != 2 or int(block_table.shape[1]) < 1: + raise ValueError("GLM expansion block table must be non-empty and rank two") + if output.shape != (rows, 2051) or active_counts.shape != (rows,): + raise ValueError("GLM physical-selection outputs have the wrong contract") + if block_size <= 0 or block_stride_rows < block_size: + raise ValueError("GLM physical token block geometry is invalid") + if rows: + block_cols = 128 + _expand_pool_ids_physical_kernel[(rows, triton.cdiv(2051, block_cols))]( + pool_ids, + positions, + request_ids, + block_table, + output, + active_counts, + int(pool_ids.stride(0)), + int(block_table.stride(0)), + int(output.stride(0)), + int(block_table.shape[1]), + HISTORY_TOKENS=2048, + OUTPUT_WIDTH=2051, + POOL_SIZE=_POOL_SIZE, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + BLOCK_COLS=block_cols, + num_warps=4, + ) + + __all__ = [ "expand_c4_block_table", "expand_pool_ids", + "expand_pool_ids_physical", "fwht128_quant_fp8", "gather_c4_block_table_rows", "pool_seq_lens", diff --git a/vllm/models/glm5next/nvidia/pooled_indexer.py b/vllm/models/glm5next/nvidia/pooled_indexer.py index 492cede1affd..bbc9da20f1dd 100644 --- a/vllm/models/glm5next/nvidia/pooled_indexer.py +++ b/vllm/models/glm5next/nvidia/pooled_indexer.py @@ -31,6 +31,7 @@ from .ops.glm_kpool import ( expand_c4_block_table, expand_pool_ids, + expand_pool_ids_physical, fwht128_quant_fp8, gather_c4_block_table_rows, pool_seq_lens, @@ -202,6 +203,11 @@ def __init__( ), persistent=False, ) + self.register_buffer( + "physical_active_counts_buffer", + torch.empty(self.max_tokens, dtype=torch.int32, device=device), + persistent=False, + ) self.register_buffer( "_q_scale", torch.empty( @@ -565,7 +571,21 @@ def forward( output = self.topk_indices_buffer[:rows] if live_rows < rows: output[live_rows:].fill_(-1) - expand_pool_ids(pool_ids[:live_rows], positions[:live_rows], output[:live_rows]) + if decode_only and self.dcp_world_size == 1: + expand_pool_ids_physical( + pool_ids[:live_rows], + positions[:live_rows], + main_metadata.req_id_per_token[:live_rows], + main_metadata.block_table, + output[:live_rows], + self.physical_active_counts_buffer[:live_rows], + block_size=self.block_size, + block_stride_rows=self.block_size, + ) + else: + expand_pool_ids( + pool_ids[:live_rows], positions[:live_rows], output[:live_rows] + ) return output def snapshot_speculative_interval_starts(self) -> None: diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 5d131145d872..a241403bf683 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -1012,6 +1012,9 @@ def __init__( vllm_config = get_current_vllm_config() hf_config = vllm_config.model_config.hf_text_config self._is_glm_next = _is_glm_next_config(hf_config) + self._glm_physical_active_counts = getattr( + indexer, "physical_active_counts_buffer", None + ) self.supports_mtp_with_cp_non_trivial_interleave_size = self._is_glm_next if self._is_glm_next: if recipe_error := _glm_next_recipe_error(hf_config): @@ -1493,6 +1496,15 @@ def forward_mqa( ).contiguous() torch.minimum(active_counts, cache_seq_lens, out=active_counts) _mask_page_table_after_nsa_len(selected_indices, active_counts) + elif ( + self._is_glm_next + and self.dcp_world_size == 1 + and int(attn_metadata.num_prefills) == 0 + and int(attn_metadata.num_decode_tokens) == num_tokens + and self._glm_physical_active_counts is not None + ): + selected_indices = topk_indices + active_counts = self._glm_physical_active_counts[:num_tokens] elif self.dcp_world_size > 1: block_stride_rows = _selected_index_block_stride_rows( kv_c_and_k_pe_cache,