From 70881864516990314eb6dc8240a917adb1f70e09 Mon Sep 17 00:00:00 2001 From: Koushik Dutta Date: Fri, 17 Jul 2026 08:28:49 +0000 Subject: [PATCH 1/7] perf(glm5next): gather full C4 cache for DCP prefill Pure-prefill DCP batches pack each rank-local CKV shard, gather the full cache in rank-major order, remap global C4 selections, and execute B12X extend attention with local query heads and global causal lengths. Decode, speculative verification, CUDA-graph capture, and ineligible batches retain the ordinary DCP attention path. The implementation excludes indexer query splitting and cross-layer asynchronous prefetch. Those mechanisms have independent topology and ownership contracts and are not required for synchronous full-CKV prefill. (cherry picked from commit d98b9b245c1430e02090e06edc4a836018f1e727) Co-authored-by: Jack Zampolin Co-authored-by: opencode (cherry picked from commit ddfa3d558079fbba8d5b874e1c1c269ec0506ce2) Signed-off-by: Martin Vit --- .../v1/attention/test_b12x_sparse_mla_api.py | 51 ++ vllm/envs.py | 15 + .../layers/attention/mla_attention.py | 7 +- vllm/v1/attention/backend.py | 4 + .../attention/backends/mla/b12x_mla_sparse.py | 581 +++++++++++++++++- 5 files changed, 635 insertions(+), 23 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 70d263f40d5a..09242f8db0cf 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -35,9 +35,12 @@ B12xMLASparseImpl, B12xMLASparseMetadata, B12xMLASparseMetadataBuilder, + _global_causal_lens_for_ckv_gather, + _is_glm_next_ckv_source_layout, _is_speculative_decode_batch, _max_speculative_decode_query_len, _selected_index_block_stride_rows, + _use_b12x_full_ckv_gather, ) from vllm.v1.attention.backends.mla.sparse_utils import _remap_tiling from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -258,6 +261,52 @@ def test_b12x_glm5_next_accepts_dcp_with_speculation(monkeypatch) -> None: assert invalid_reasons == [] +@pytest.mark.parametrize( + ("max_query_len", "is_spec_decode", "num_tokens", "expected"), + [ + (1, False, 32, False), + (6, True, 192, False), + (6, False, 192, True), + (128, False, 8192, True), + (128, False, 600000, False), + ], +) +def test_b12x_full_ckv_gather_excludes_decode_and_mtp_batches( + max_query_len: int, + is_spec_decode: bool, + num_tokens: int, + expected: bool, +) -> None: + assert ( + _use_b12x_full_ckv_gather( + enabled=True, + is_glm_next=True, + dcp_world_size=4, + max_query_len=max_query_len, + num_tokens=num_tokens, + is_spec_decode=is_spec_decode, + min_tokens=16, + max_tokens=524288, + ) + is expected + ) + + +def test_b12x_full_ckv_gather_uses_global_causal_lengths() -> None: + global_seq_lens = torch.tensor([5, 12], dtype=torch.int32) + query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32) + req_id_per_token = torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32) + + actual = _global_causal_lens_for_ckv_gather( + global_seq_lens, + query_start_loc, + req_id_per_token, + num_actual_tokens=5, + ) + + assert actual.tolist() == [4, 5, 10, 11, 12] + + def test_b12x_glm5_next_accepts_dcp_with_prefix_caching(monkeypatch) -> None: monkeypatch.setattr(b12x_mla_sparse, "get_b12x_sparse_mla", lambda: object()) with set_current_vllm_config( @@ -341,6 +390,8 @@ def test_b12x_glm5_next_selected_indices_use_physical_slots() -> None: ) == 64 ) + assert _is_glm_next_ckv_source_layout(cache, page_size=64) + assert not _is_glm_next_ckv_source_layout(cache[:, :, ::2], page_size=64) def test_sparse_index_remap_tiling_covers_glm5_next_width() -> None: diff --git a/vllm/envs.py b/vllm/envs.py index 657f484407bd..2874dbf1aa7b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -190,6 +190,9 @@ VLLM_HUMMING_USE_F16_ACCUM: bool = False VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None VLLM_B12X_MOE_FP4_FORCE_A16: bool = False + VLLM_B12X_MLA_CKV_GATHER: bool = False + VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS: int = 16 + VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS: int = 524288 VLLM_PLE_CPU_OFFLOAD: bool = False VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False @@ -1627,6 +1630,18 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_B12X_MOE_FP4_FORCE_A16": lambda: bool( int(os.getenv("VLLM_B12X_MOE_FP4_FORCE_A16", "0")) ), + # Gather DCP-sharded C4 records before B12X sparse-MLA prefill. This avoids + # query replication plus the per-rank LSE combine and is opt-in while the + # path is being qualified on GLM5Next. + "VLLM_B12X_MLA_CKV_GATHER": lambda: ( + os.getenv("VLLM_B12X_MLA_CKV_GATHER", "0").lower() in ("1", "true", "yes", "on") + ), + "VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS": lambda: int( + os.getenv("VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS", "16") + ), + "VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS": lambda: int( + os.getenv("VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS", "524288") + ), # Qwen3.8-Flash-Next only. Store PLE table payloads in CUDA-mapped host # memory unless additional_config.ple_table_memory is explicitly set. "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "0"))), diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 0a9a332becda..f9fed2d32548 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -909,6 +909,9 @@ def forward_impl( else: mqa_q = q[:num_mqa_tokens] qrep_decode = False + full_ckv_dcp = self.impl.uses_full_ckv_dcp( # type: ignore[attr-defined] + attn_metadata, num_mqa_tokens + ) mqa_output_slice = output[:num_mqa_tokens] mqa_q_nope, mqa_q_pe = mqa_q.split( @@ -998,7 +1001,7 @@ def forward_impl( if isinstance(mqa_q, tuple): # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) mqa_q = torch.cat(mqa_q, dim=-1) - if not qrep_decode: + if not qrep_decode and not full_ckv_dcp: assert self.dcp_manager.query_gather is not None mqa_q = self.dcp_manager.query_gather(mqa_q) @@ -1008,7 +1011,7 @@ def forward_impl( attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined] # correct dcp attn_out with lse. - if self.impl.dcp_world_size > 1: + if self.impl.dcp_world_size > 1 and not full_ckv_dcp: assert lse is not None assert self.dcp_manager is not None decode_metadata = getattr(attn_metadata, "decode", None) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 4fdfafaaea9a..4041dcca312c 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -985,6 +985,10 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): supports_pcp: bool = True + def uses_full_ckv_dcp(self, attn_metadata: T, num_tokens: int) -> bool: + """Whether this call attends a transient globally gathered DCP cache.""" + return False + @abstractmethod def __init__( self, diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 5748f7330e1b..af7f310df006 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -8,17 +8,21 @@ import numpy as np import torch +import torch.distributed as dist from vllm import _custom_ops as ops +from vllm import envs from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.config.cache import CacheDType from vllm.distributed import get_dcp_group +from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, SparseMLACommonMetadataBuilder, ) from vllm.platforms.interface import DeviceCapability +from vllm.triton_utils import tl, triton from vllm.utils.b12x import get_b12x_sparse_mla from vllm.v1.attention.backend import ( AttentionBackend, @@ -50,6 +54,8 @@ _GLM_NEXT_CACHE_RECORD_BYTES = 528 _GLM_NEXT_INDEX_TAIL_BYTES_PER_TOKEN = 132 // 4 +logger = init_logger(__name__) + def _is_glm_next_config(hf_config: object | None) -> bool: return getattr(hf_config, "model_type", None) in _GLM_NEXT_MODEL_TYPES @@ -139,6 +145,254 @@ def _is_speculative_decode_batch( ) +def _is_glm_next_ckv_source_layout( + kv_cache: torch.Tensor, + *, + page_size: int, +) -> bool: + return ( + kv_cache.dtype == torch.uint8 + and kv_cache.ndim == 3 + and tuple(kv_cache.shape[1:]) == (page_size, _GLM_NEXT_CACHE_RECORD_BYTES) + and kv_cache.stride(1) == _GLM_NEXT_CACHE_RECORD_BYTES + and kv_cache.stride(2) == 1 + ) + + +def _use_b12x_full_ckv_gather( + *, + enabled: bool, + is_glm_next: bool, + dcp_world_size: int, + max_query_len: int, + num_tokens: int, + is_spec_decode: bool, + min_tokens: int, + max_tokens: int, +) -> bool: + return ( + enabled + and is_glm_next + and dcp_world_size > 1 + and max_query_len > 1 + and not is_spec_decode + and num_tokens > min_tokens + and num_tokens <= max_tokens + ) + + +def _dcp_all_gather_current_stream( + group, + input_tensor: torch.Tensor, + output_tensor: torch.Tensor, +) -> None: + if not input_tensor.is_contiguous() or not output_tensor.is_contiguous(): + raise ValueError("CKV all-gather tensors must be contiguous") + if output_tensor.numel() != input_tensor.numel() * group.world_size: + raise ValueError("CKV all-gather tensors have incompatible sizes") + + communicator = getattr(group, "device_communicator", None) + pynccl_comm = getattr(communicator, "pynccl_comm", None) + if pynccl_comm is not None and not getattr(pynccl_comm, "disabled", False): + pynccl_comm.all_gather(output_tensor, input_tensor) + return + + device_group = getattr(group, "device_group", None) + if device_group is None: + device_group = getattr(communicator, "device_group", None) + if device_group is not None: + dist.all_gather_into_tensor( + output_tensor, + input_tensor, + group=device_group, + async_op=False, + ) + return + + output_tensor.copy_(group.all_gather(input_tensor, dim=0)) + + +@triton.jit +def _mask_page_table_after_nsa_len_kernel( + page_table_ptr, + nsa_len_ptr, + page_stride0, + page_stride1, + width: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + offs = tile * BLOCK_N + tl.arange(0, BLOCK_N) + valid = offs < width + nsa_len = tl.load(nsa_len_ptr + row) + tl.store( + page_table_ptr + row * page_stride0 + offs * page_stride1, + -1, + mask=valid & (offs >= nsa_len), + ) + + +def _mask_page_table_after_nsa_len( + page_table: torch.Tensor, + nsa_cache_seqlens: torch.Tensor, +) -> None: + width = page_table.shape[1] + if width == 0 or page_table.shape[0] == 0: + return + block_n = 128 + _mask_page_table_after_nsa_len_kernel[ + (page_table.shape[0], triton.cdiv(width, block_n)) + ]( + page_table, + nsa_cache_seqlens, + page_table.stride(0), + page_table.stride(1), + width, + BLOCK_N=block_n, + ) + + +def _global_causal_lens_for_ckv_gather( + global_seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + req_id_per_token: torch.Tensor, + num_actual_tokens: int, +) -> torch.Tensor: + """Return each query token's causal length in the gathered global cache.""" + num_reqs = global_seq_lens.shape[0] + qsl = query_start_loc[: num_reqs + 1].to(torch.int32) + req_ids = req_id_per_token[:num_actual_tokens].to(torch.int64) + chunk_start = qsl[:-1][req_ids] + chunk_len = (qsl[1:] - qsl[:-1])[req_ids] + full_seq = global_seq_lens[req_ids].to(torch.int32) + token_idx = torch.arange( + num_actual_tokens, + device=global_seq_lens.device, + dtype=torch.int32, + ) + return full_seq - chunk_len + (token_idx - chunk_start) + 1 + + +@triton.jit +def _map_global_topk_to_gathered_ckv_kernel( + req_id_ptr, + token_indices_ptr, + rank_req_starts_ptr, + rank_req_lens_ptr, + out_ptr, + valid_count_ptr, + starts_stride0, + starts_stride1, + lens_stride0, + lens_stride1, + ti_stride0, + ti_stride1, + out_stride0, + out_stride1, + padded_rank_tokens, + DCP_SIZE: tl.constexpr, + DCP_INTERLEAVE: tl.constexpr, + NUM_TOPK_TOKENS: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + tile = tl.program_id(1) + cols = tile * BLOCK_N + tl.arange(0, BLOCK_N) + col_mask = cols < NUM_TOPK_TOKENS + req = tl.load(req_id_ptr + row) + tok = tl.load( + token_indices_ptr + row * ti_stride0 + cols * ti_stride1, + mask=col_mask, + other=-1, + ) + owner = (tok // DCP_INTERLEAVE) % DCP_SIZE + local_idx = ( + tok // (DCP_SIZE * DCP_INTERLEAVE) + ) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE + valid_tok = col_mask & (tok >= 0) + req_start = tl.load( + rank_req_starts_ptr + owner * starts_stride0 + req * starts_stride1, + mask=valid_tok, + other=0, + ) + req_len = tl.load( + rank_req_lens_ptr + owner * lens_stride0 + req * lens_stride1, + mask=valid_tok, + other=0, + ) + valid = valid_tok & (local_idx >= 0) & (local_idx < req_len) + gathered_slot = owner * padded_rank_tokens + req_start + local_idx + valid_i32 = valid.to(tl.int32) + local_offset = tl.cumsum(valid_i32) - valid_i32 + tile_valid_count = tl.sum(valid_i32) + output_base = tl.atomic_add(valid_count_ptr + row, tile_valid_count) + tl.store( + out_ptr + row * out_stride0 + (output_base + local_offset) * out_stride1, + gathered_slot, + mask=valid, + ) + + +def _map_global_topk_to_gathered_ckv( + req_ids: torch.Tensor, + token_indices: torch.Tensor, + rank_req_starts: torch.Tensor, + rank_req_lens: torch.Tensor, + out: torch.Tensor, + valid_counts: torch.Tensor, + *, + dcp_size: int, + cp_kv_cache_interleave_size: int, + padded_rank_tokens: int, +) -> None: + if token_indices.shape != out.shape: + raise ValueError("CKV gather index output shape does not match top-k input") + if rank_req_starts.shape != rank_req_lens.shape: + raise ValueError("CKV gather request starts/lens shapes do not match") + if rank_req_starts.shape[0] != dcp_size: + raise ValueError("CKV gather request metadata does not match DCP size") + if any( + tensor.dtype != torch.int32 + for tensor in ( + req_ids, + token_indices, + rank_req_starts, + rank_req_lens, + out, + valid_counts, + ) + ): + raise TypeError("CKV gather index metadata must be int32") + + block_n = 128 + out.fill_(-1) + valid_counts.zero_() + _map_global_topk_to_gathered_ckv_kernel[ + (token_indices.shape[0], triton.cdiv(token_indices.shape[1], block_n)) + ]( + req_ids, + token_indices, + rank_req_starts, + rank_req_lens, + out, + valid_counts, + rank_req_starts.stride(0), + rank_req_starts.stride(1), + rank_req_lens.stride(0), + rank_req_lens.stride(1), + token_indices.stride(0), + token_indices.stride(1), + out.stride(0), + out.stride(1), + padded_rank_tokens, + DCP_SIZE=dcp_size, + DCP_INTERLEAVE=cp_kv_cache_interleave_size, + NUM_TOPK_TOKENS=token_indices.shape[1], + BLOCK_N=block_n, + ) + + class B12xMLASparseBackend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ @@ -301,6 +555,15 @@ class B12xMLASparseMetadata(AttentionMetadata): selector_num_accepted_tokens: torch.Tensor | None = None selector_is_prefilling: torch.Tensor | None = None is_spec_decode: bool = False + ckv_selected_indices: torch.Tensor | None = None + ckv_active_counts: torch.Tensor | None = None + dcp_rank_req_starts: torch.Tensor | None = None + dcp_rank_req_lens: torch.Tensor | None = None + dcp_local_cu_seq_lens: torch.Tensor | None = None + global_cache_seq_lens_per_req: torch.Tensor | None = None + dcp_local_total_tokens: int = 0 + dcp_padded_total_tokens: int = 0 + dcp_ckv_gather_eligible: bool = False class B12xMLASparseMetadataBuilder( @@ -330,11 +593,12 @@ def __init__( self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 scheduler_config = vllm_config.scheduler_config max_tokens = scheduler_config.max_num_batched_tokens + max_reqs = int(scheduler_config.max_num_seqs) + self._ckv_max_reqs = max_reqs self.cache_seq_lens_per_token_buffer = torch.empty( (max_tokens,), dtype=torch.int32, device=device ) if self.requires_glm_next_selector_metadata: - max_reqs = int(scheduler_config.max_num_seqs) self._capture_default_state_slot_ids = torch.arange( max_reqs, dtype=torch.int32, device=device ) @@ -350,6 +614,35 @@ def __init__( self._capture_is_prefilling = torch.zeros( max_reqs, dtype=torch.bool, device=device ) + self._ckv_gather_requested = ( + self.requires_glm_next_selector_metadata + and self.dcp_world_size > 1 + and envs.VLLM_B12X_MLA_CKV_GATHER + ) + if self._ckv_gather_requested: + hf_config = vllm_config.model_config.hf_text_config + ckv_topk_tokens = int(hf_config.index_topk) + int(hf_config.index_kpool) - 1 + self.ckv_selected_indices_buffer = torch.empty( + (max_tokens, ckv_topk_tokens), dtype=torch.int32, device=device + ) + self.ckv_active_counts_buffer = torch.empty( + (max_tokens,), dtype=torch.int32, device=device + ) + self.dcp_rank_req_lens_buffer = torch.empty( + (self.dcp_world_size, max_reqs), dtype=torch.int32, device=device + ) + self.dcp_rank_req_starts_buffer = torch.empty( + (self.dcp_world_size, max_reqs), dtype=torch.int32, device=device + ) + self.dcp_local_cu_seq_lens_buffer = torch.empty( + (max_reqs + 1,), dtype=torch.int32, device=device + ) + else: + self.ckv_selected_indices_buffer = None + self.ckv_active_counts_buffer = None + self.dcp_rank_req_lens_buffer = None + self.dcp_rank_req_starts_buffer = None + self.dcp_local_cu_seq_lens_buffer = None num_q_heads = vllm_config.model_config.get_num_attention_heads( vllm_config.parallel_config ) @@ -533,6 +826,79 @@ def _build( metadata.prefill_seq_lens_cpu = seq_lens_cpu_source[ prefill_start : prefill_start + metadata.num_prefills ].clone() + if ( + _use_b12x_full_ckv_gather( + enabled=self._ckv_gather_requested, + is_glm_next=self.requires_glm_next_selector_metadata, + dcp_world_size=self.dcp_world_size, + max_query_len=common.max_query_len, + num_tokens=num_tokens, + is_spec_decode=metadata.is_spec_decode, + min_tokens=envs.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS, + max_tokens=envs.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS, + ) + and metadata.num_decode_tokens == 0 + ): + assert self.ckv_selected_indices_buffer is not None + assert self.ckv_active_counts_buffer is not None + assert self.dcp_rank_req_lens_buffer is not None + assert self.dcp_rank_req_starts_buffer is not None + assert self.dcp_local_cu_seq_lens_buffer is not None + global_seq_lens = common.seq_lens[: common.num_reqs] + all_rank_lens = get_dcp_local_seq_lens( + global_seq_lens, + self.dcp_world_size, + dcp_rank=None, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + ).transpose(0, 1) + rank_req_lens = self.dcp_rank_req_lens_buffer[ + : self.dcp_world_size, : common.num_reqs + ] + rank_req_lens.copy_(all_rank_lens) + rank_req_starts = self.dcp_rank_req_starts_buffer[ + : self.dcp_world_size, : common.num_reqs + ] + rank_req_starts[:, 0].zero_() + if common.num_reqs > 1: + torch.cumsum(rank_req_lens[:, :-1], dim=1, out=rank_req_starts[:, 1:]) + local_cu_seq_lens = self.dcp_local_cu_seq_lens_buffer[: common.num_reqs + 1] + local_cu_seq_lens[0].zero_() + torch.cumsum( + rank_req_lens[self.dcp_rank], + dim=0, + out=local_cu_seq_lens[1:], + ) + rank_totals = rank_req_lens.sum(dim=1).tolist() + local_total_tokens = int(rank_totals[self.dcp_rank]) + page_size = int(self.kv_cache_spec.block_size) + padded_total_tokens = ( + (max(int(total) for total in rank_totals) + page_size - 1) + // page_size + * page_size + ) + max_local_capacity = ( + ( + (envs.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS + self.dcp_world_size - 1) + // self.dcp_world_size + + self._ckv_max_reqs * self.cp_kv_cache_interleave_size + + page_size + - 1 + ) + // page_size + * page_size + ) + if 0 < padded_total_tokens <= max_local_capacity: + metadata.ckv_selected_indices = self.ckv_selected_indices_buffer[ + :num_tokens + ] + metadata.ckv_active_counts = self.ckv_active_counts_buffer[:num_tokens] + metadata.dcp_rank_req_lens = rank_req_lens + metadata.dcp_rank_req_starts = rank_req_starts + metadata.dcp_local_cu_seq_lens = local_cu_seq_lens + metadata.global_cache_seq_lens_per_req = global_seq_lens + metadata.dcp_local_total_tokens = local_total_tokens + metadata.dcp_padded_total_tokens = padded_total_tokens + metadata.dcp_ckv_gather_eligible = True ( metadata.selector_state_slot_ids, metadata.selector_state_is_fresh, @@ -713,6 +1079,19 @@ def __init__( kernel_page_size = ( int(vllm_config.cache_config.block_size) if self._is_glm_next else 64 ) + self._ckv_gather_enabled = ( + self._is_glm_next + and self.dcp_world_size > 1 + and envs.VLLM_B12X_MLA_CKV_GATHER + ) + max_ckv_tokens = envs.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS + cp_kv_cache_interleave_size = int( + vllm_config.parallel_config.cp_kv_cache_interleave_size + ) + self._ckv_capacity_tokens = ( + max_ckv_tokens + self.dcp_world_size - 1 + ) // self.dcp_world_size + max_seqs * cp_kv_cache_interleave_size + self._ckv_local_capacity = 0 self._module = module self._kernel_page_size = 0 @@ -728,11 +1107,11 @@ def _set_kernel_page_size(self, kernel_page_size: int) -> None: if kernel_page_size == self._kernel_page_size: return - def make_plan(mode: str): + def make_plan(mode: str, num_q_heads: int = self._input_num_heads): max_rows = self._decode_max_rows if mode == "decode" else self._max_tokens caps_kwargs = dict( device=torch.device("cuda", torch.accelerator.current_device_index()), - num_q_heads=self._input_num_heads, + num_q_heads=num_q_heads, max_q_rows=max_rows, max_width=self._topk_tokens, dtype=torch.bfloat16, @@ -752,6 +1131,14 @@ def make_plan(mode: str): extend_plan = make_plan("extend") self._decode_plan = decode_plan self._extend_plan = extend_plan + self._ckv_extend_plan = ( + make_plan("extend", self.num_heads) if self._ckv_gather_enabled else None + ) + self._ckv_local_capacity = ( + (self._ckv_capacity_tokens + kernel_page_size - 1) + // kernel_page_size + * kernel_page_size + ) self._kernel_page_size = kernel_page_size self._reserve_planned_workspaces() @@ -840,6 +1227,85 @@ def do_kv_cache_update( slot_mapping.flatten(), ) + def uses_full_ckv_dcp( + self, + attn_metadata: B12xMLASparseMetadata, + num_tokens: int, + ) -> bool: + if torch.cuda.is_current_stream_capturing(): + return False + return ( + self._ckv_gather_enabled + and attn_metadata.dcp_ckv_gather_eligible + and attn_metadata.num_decode_tokens == 0 + and num_tokens == attn_metadata.num_actual_tokens + and 0 < attn_metadata.dcp_padded_total_tokens <= self._ckv_local_capacity + and attn_metadata.dcp_local_total_tokens + <= attn_metadata.dcp_padded_total_tokens + and all( + value is not None + for value in ( + attn_metadata.ckv_selected_indices, + attn_metadata.ckv_active_counts, + attn_metadata.dcp_rank_req_starts, + attn_metadata.dcp_rank_req_lens, + attn_metadata.dcp_local_cu_seq_lens, + attn_metadata.global_cache_seq_lens_per_req, + ) + ) + ) + + def _gather_full_ckv( + self, + kv_cache: torch.Tensor, + attn_metadata: B12xMLASparseMetadata, + local_buffer: torch.Tensor, + gathered_buffer: torch.Tensor, + ) -> torch.Tensor: + if not self.uses_full_ckv_dcp(attn_metadata, attn_metadata.num_actual_tokens): + raise RuntimeError("full CKV gather called for an ineligible batch") + if not _is_glm_next_ckv_source_layout( + kv_cache, page_size=self._kernel_page_size + ): + raise ValueError( + "GLM5Next CKV gather requires native 528-byte records; " + f"got shape={tuple(kv_cache.shape)}, stride={kv_cache.stride()}" + ) + expected_local_shape = ( + self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + expected_gathered_shape = ( + self.dcp_world_size * self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ) + if tuple(local_buffer.shape) != expected_local_shape: + raise RuntimeError("CKV local workspace has an invalid shape") + if tuple(gathered_buffer.shape) != expected_gathered_shape: + raise RuntimeError("CKV gathered workspace has an invalid shape") + + assert attn_metadata.dcp_local_cu_seq_lens is not None + local_tokens = attn_metadata.dcp_local_total_tokens + padded_tokens = attn_metadata.dcp_padded_total_tokens + if local_tokens: + ops.cp_gather_cache( + src_cache=kv_cache, + dst=local_buffer[:local_tokens], + block_table=attn_metadata.block_table, + cu_seq_lens=attn_metadata.dcp_local_cu_seq_lens, + batch_size=attn_metadata.num_reqs, + ) + if local_tokens < padded_tokens: + local_buffer[local_tokens:padded_tokens].zero_() + _dcp_all_gather_current_stream( + get_dcp_group(), + local_buffer[:padded_tokens].view(-1), + gathered_buffer[: self.dcp_world_size * padded_tokens].view(-1), + ) + return gathered_buffer.view( + -1, self._kernel_page_size, _GLM_NEXT_CACHE_RECORD_BYTES + ) + def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -860,16 +1326,46 @@ def forward_mqa( f"plan={self._kernel_page_size}" ) num_tokens = int(q[0].shape[0] if isinstance(q, tuple) else q.shape[0]) - plan = ( - self._decode_plan - if self._use_decode_plan(attn_metadata, num_tokens) - else self._extend_plan + use_ckv_gather = self.uses_full_ckv_dcp(attn_metadata, num_tokens) + if use_ckv_gather: + assert self._ckv_extend_plan is not None + plan = self._ckv_extend_plan + logger.info_once("Using full-CKV gather for GLM5Next B12X DCP prefill") + else: + plan = ( + self._decode_plan + if self._use_decode_plan(attn_metadata, num_tokens) + else self._extend_plan + ) + input_num_heads = self.num_heads if use_ckv_gather else self._input_num_heads + q_spec = ( + (self._max_tokens, input_num_heads, self._q_head_dim), + torch.bfloat16, + ) + plan_specs = plan.shapes_and_dtypes() + ckv_specs = ( + ( + ( + (self._ckv_local_capacity, _GLM_NEXT_CACHE_RECORD_BYTES), + torch.uint8, + ), + ( + ( + self.dcp_world_size * self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ), + torch.uint8, + ), + ) + if use_ckv_gather + else () ) workspaces = current_workspace_manager().get_simultaneous( - *self._workspace_specs(plan) + q_spec, *plan_specs, *ckv_specs ) q_buffer = workspaces[0] - scratch = workspaces[1:] + scratch_end = 1 + len(plan_specs) + scratch = workspaces[1:scratch_end] if isinstance(q, tuple): q_nope, q_pe = q @@ -882,20 +1378,57 @@ def forward_mqa( q_all = q_buffer[:num_tokens] q_all.copy_(q) - if int(q_all.shape[1]) != self._input_num_heads: + if int(q_all.shape[1]) != input_num_heads: raise ValueError( "B12X sparse MLA query heads do not match the planned head " - f"count: {q_all.shape[1]} != {self._input_num_heads}." + f"count: {q_all.shape[1]} != {input_num_heads}." ) assert self.topk_indices_buffer is not None topk_indices = self.topk_indices_buffer[:num_tokens] - block_stride_rows = _selected_index_block_stride_rows( - kv_c_and_k_pe_cache, - block_size=attn_metadata.block_size, - is_glm_next=self._is_glm_next, - ) - if self.dcp_world_size > 1: + kv_cache_for_run = kv_c_and_k_pe_cache + if use_ckv_gather: + local_buffer, gathered_buffer = workspaces[scratch_end:] + kv_cache_for_run = self._gather_full_ckv( + kv_c_and_k_pe_cache, + attn_metadata, + local_buffer, + gathered_buffer, + ) + assert attn_metadata.ckv_selected_indices is not None + assert attn_metadata.ckv_active_counts is not None + assert attn_metadata.dcp_rank_req_starts is not None + assert attn_metadata.dcp_rank_req_lens is not None + selected_indices = attn_metadata.ckv_selected_indices[ + :num_tokens, : topk_indices.shape[1] + ] + active_counts = attn_metadata.ckv_active_counts[:num_tokens] + _map_global_topk_to_gathered_ckv( + attn_metadata.req_id_per_token[:num_tokens], + topk_indices, + attn_metadata.dcp_rank_req_starts, + attn_metadata.dcp_rank_req_lens, + selected_indices, + active_counts, + dcp_size=self.dcp_world_size, + cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size), + padded_rank_tokens=attn_metadata.dcp_padded_total_tokens, + ) + assert attn_metadata.global_cache_seq_lens_per_req is not None + cache_seq_lens = _global_causal_lens_for_ckv_gather( + attn_metadata.global_cache_seq_lens_per_req, + attn_metadata.query_start_loc, + attn_metadata.req_id_per_token, + num_tokens, + ).contiguous() + torch.minimum(active_counts, cache_seq_lens, out=active_counts) + _mask_page_table_after_nsa_len(selected_indices, active_counts) + elif self.dcp_world_size > 1: + block_stride_rows = _selected_index_block_stride_rows( + kv_c_and_k_pe_cache, + block_size=attn_metadata.block_size, + is_glm_next=self._is_glm_next, + ) selected_indices, active_counts = triton_filter_and_convert_dcp_index( attn_metadata.req_id_per_token[:num_tokens], attn_metadata.block_table, @@ -909,6 +1442,11 @@ def forward_mqa( return_valid_counts=True, ) else: + block_stride_rows = _selected_index_block_stride_rows( + kv_c_and_k_pe_cache, + block_size=attn_metadata.block_size, + is_glm_next=self._is_glm_next, + ) selected_indices, active_counts = triton_convert_req_index_to_global_index( attn_metadata.req_id_per_token[:num_tokens], attn_metadata.block_table, @@ -919,9 +1457,10 @@ def forward_mqa( return_valid_counts=True, ) - cache_seq_lens = attn_metadata.cache_seq_lens_per_token - assert cache_seq_lens is not None - cache_seq_lens = cache_seq_lens[:num_tokens].contiguous() + if not use_ckv_gather: + cache_seq_lens = attn_metadata.cache_seq_lens_per_token + assert cache_seq_lens is not None + cache_seq_lens = cache_seq_lens[:num_tokens].contiguous() binding = plan.bind( scratch=scratch, q=q_all, @@ -932,7 +1471,7 @@ def forward_mqa( run = self._run_decode if plan is self._decode_plan else self._run_extend run_kwargs = dict( binding=binding, - kv_cache=kv_c_and_k_pe_cache, + kv_cache=kv_cache_for_run, sm_scale=self.scale, v_head_dim=self.kv_lora_rank, return_lse=self.need_to_return_lse_for_decode, From 56e3b2a3038e9091f823bbfcd912e05ea73b7766 Mon Sep 17 00:00:00 2001 From: Jack Zampolin Date: Thu, 27 Aug 2026 14:48:14 -0700 Subject: [PATCH 2/7] fix(glm5next): size CKV scratch from parallel config Co-authored-by: OpenAI Codex (cherry picked from commit d71b3cae8cdaaa2f392d16965c733d2f8e1ddaff) (cherry picked from commit 208f574034cfc38076e9b69007b78fad64a7927f) (cherry picked from commit 6bfb5fc7c9fb78c12bd88dcd361c3b1e881dd45f) Signed-off-by: Martin Vit From e8cde51315d4e212b4de1195bd162a0327b98832 Mon Sep 17 00:00:00 2001 From: Jack Zampolin Date: Thu, 27 Aug 2026 19:25:05 -0700 Subject: [PATCH 3/7] Fix full-cache DCP query selection Co-Authored-By: Claude opus (cherry picked from commit ea338d7342eb4d1c31c242b52d87be7c21c6734d) (cherry picked from commit 0cb13562a9763fc349c2d3761bea315aa0d00e4d) (cherry picked from commit c17abaeb08b4b7893a396585c19ea6845afbbd94) Signed-off-by: Martin Vit --- tests/v1/attention/test_mla_backends.py | 27 +++++++++++++++++ .../layers/attention/mla_attention.py | 29 +++++++++++++++---- 2 files changed, 50 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index cbccbfc9a2ac..c1bbded4c35a 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -27,6 +27,7 @@ MLAAttention, QueryLenSupport, _DecodeConcatQuantFP8, + _select_mqa_query, build_mla_chunked_context_metadata, ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape @@ -68,6 +69,32 @@ DEVICE_TYPE = current_platform.device_type +@pytest.mark.cpu_test +def test_full_ckv_dcp_prefers_local_query_geometry() -> None: + q = torch.zeros((2, 2, 6)) + q_dcp_replicated = torch.ones((2, 8, 6)) + + selected, replicated = _select_mqa_query( + q, + q_dcp_replicated, + num_mqa_tokens=1, + full_ckv_dcp=True, + ) + assert selected.shape == (1, 2, 6) + assert not replicated + assert torch.equal(selected, q[:1]) + + selected, replicated = _select_mqa_query( + q, + q_dcp_replicated, + num_mqa_tokens=1, + full_ckv_dcp=False, + ) + assert selected.shape == (1, 8, 6) + assert replicated + assert torch.equal(selected, q_dcp_replicated[:1]) + + @pytest.mark.parametrize( ("cache_dtype", "expected_quant_mode"), [ diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index f9fed2d32548..f806103f2042 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -346,6 +346,19 @@ def _detect_output_quant_key( return kFp8StaticTensorSym +def _select_mqa_query( + q: torch.Tensor, + q_dcp_replicated: torch.Tensor | None, + *, + num_mqa_tokens: int, + full_ckv_dcp: bool, +) -> tuple[torch.Tensor, bool]: + """Select local or replicated query geometry for MLA decode/prefill.""" + if q_dcp_replicated is not None and not full_ckv_dcp: + return q_dcp_replicated[:num_mqa_tokens], True + return q[:num_mqa_tokens], False + + def _canonicalize_sparse_mla_kv_cache_dtype( attn_backend: type[AttentionBackend], kv_cache_dtype: CacheDType, @@ -903,15 +916,19 @@ def forward_impl( ) if num_mqa_tokens > 0: - if q_dcp_replicated is not None: - mqa_q = q_dcp_replicated[:num_mqa_tokens] - qrep_decode = True - else: - mqa_q = q[:num_mqa_tokens] - qrep_decode = False full_ckv_dcp = self.impl.uses_full_ckv_dcp( # type: ignore[attr-defined] attn_metadata, num_mqa_tokens ) + # Full-CKV prefill already makes every rank's cache visible to + # its local query heads. Prefer the local projection even when + # dcp_q_replicate retained a global query for ordinary DCP decode; + # the replicated query does not fit the local-head CKV plan. + mqa_q, qrep_decode = _select_mqa_query( + q, + q_dcp_replicated, + num_mqa_tokens=num_mqa_tokens, + full_ckv_dcp=full_ckv_dcp, + ) mqa_output_slice = output[:num_mqa_tokens] mqa_q_nope, mqa_q_pe = mqa_q.split( From 5b5776145ed6bf11f0222e5dca97c02daf066c68 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 29 Aug 2026 08:51:55 +0000 Subject: [PATCH 4/7] fix(glm5next): finalize CKV geometry before profiling Hybrid cache alignment applies the resolved page size to loaded attention implementations before device-memory profiling. B12X rebuilds its plans from that immutable geometry and reserves every workspace execution slot, so the KV-cache budget includes full-CKV gather storage. KV-cache binding validates the finalized page size and fails closed when full-CKV gathering reaches binding without pre-profile finalization. The 64-to-2304-token page transition and multi-slot workspace reservation have regression coverage. Validation: uvx ruff check and format on changed files; 11 targeted pytest cases passed in the source-compatible runtime image. (cherry picked from commit 65fc70fe002c5318b4ea9ccf998353368fe2cb4a) (cherry picked from commit 87b9d1195a133a9aa328f989c63549fbb34b1031) Signed-off-by: Martin Vit --- .../v1/attention/test_b12x_sparse_mla_api.py | 58 +++++++- tests/v1/worker/test_workspace.py | 28 ++++ .../layers/attention_layer_base.py | 17 +++ vllm/platforms/interface.py | 18 ++- .../attention/backends/mla/b12x_mla_sparse.py | 133 ++++++++++++++---- vllm/v1/worker/workspace.py | 47 +++++++ 6 files changed, 266 insertions(+), 35 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 09242f8db0cf..bd3059713f05 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -420,19 +420,29 @@ def test_b12x_glm5_next_cache_writer_ignores_empty_rope() -> None: assert calls == [(kv_c, kv_cache, slots)] -def test_b12x_glm5_next_cache_bind_replans_aligned_manager_page(monkeypatch) -> None: +def test_b12x_glm5_next_cache_geometry_is_finalized_before_bind(monkeypatch) -> None: planned: list[SimpleNamespace] = [] + reservations: list[tuple[tuple[tuple[int, ...], torch.dtype], ...]] = [] monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) monkeypatch.setattr( b12x_mla_sparse, "is_workspace_manager_initialized", lambda: False, ) + monkeypatch.setattr( + b12x_mla_sparse, + "current_workspace_manager", + lambda: SimpleNamespace(reserve_all=lambda *specs: reservations.append(specs)), + ) class FakeCaps(SimpleNamespace): def __init__(self, **kwargs): super().__init__(**kwargs) + @staticmethod + def shapes_and_dtypes(): + return () + class FakeModule: Caps = FakeCaps @@ -445,7 +455,10 @@ def plan(caps): impl._is_glm_next = True impl._module = FakeModule impl._kernel_page_size = 64 + impl._kernel_page_size_finalized = False impl._input_num_heads = 64 + impl.num_heads = 16 + impl.dcp_world_size = 4 impl._max_tokens = 4096 impl._max_seqs = 4 impl._max_speculative_decode_query_len = 6 @@ -455,23 +468,60 @@ def plan(caps): impl._q_head_dim = 512 impl.kv_lora_rank = 512 impl._model_type = 1 + impl._ckv_gather_enabled = True + impl._ckv_capacity_tokens = 131200 + impl._ckv_local_capacity = 131200 impl._decode_plan = SimpleNamespace() impl._extend_plan = SimpleNamespace() + impl._ckv_extend_plan = SimpleNamespace() owner = SimpleNamespace(impl=impl, indexer=None) cache = torch.empty((2, 1, 2304, 528), dtype=torch.uint8) + MLAAttention.finalize_kv_cache_geometry( + owner, + SimpleNamespace(cache_config=SimpleNamespace(block_size=2304)), + ) MLAAttention.bind_kv_cache(owner, cache) assert owner.kv_cache.shape == (2, 2304, 528) assert impl._kernel_page_size == 2304 + assert impl._kernel_page_size_finalized assert [(caps.mode, caps.page_size) for caps in planned] == [ ("decode", 2304), ("extend", 2304), + ("extend", 2304), ] - assert [(caps.max_q_rows, caps.max_batch) for caps in planned] == [ - (24, 24), - (4096, 4096), + plan_geometry = [ + (caps.num_q_heads, caps.max_q_rows, caps.max_batch) for caps in planned ] + assert plan_geometry == [ + (64, 24, 24), + (64, 4096, 4096), + (16, 4096, 4096), + ] + assert len(reservations) == 3 + assert reservations[0] == (((4096, 64, 512), torch.bfloat16),) + assert reservations[1] == (((4096, 64, 512), torch.bfloat16),) + assert reservations[2] == ( + ((4096, 16, 512), torch.bfloat16), + ((131328, 528), torch.uint8), + ((525312, 528), torch.uint8), + ) + + with pytest.raises(RuntimeError, match="immutable after finalization"): + impl.finalize_kv_cache_geometry(64) + with pytest.raises(RuntimeError, match="does not match the finalized"): + impl.bind_kv_cache(torch.empty((2, 64, 528), dtype=torch.uint8)) + + +def test_b12x_glm5_next_full_ckv_bind_requires_geometry_finalization() -> None: + impl = object.__new__(B12xMLASparseImpl) + impl._is_glm_next = True + impl._ckv_gather_enabled = True + impl._kernel_page_size_finalized = False + + with pytest.raises(RuntimeError, match="before KV-cache memory profiling"): + impl.bind_kv_cache(torch.empty((2, 2304, 528), dtype=torch.uint8)) @pytest.mark.parametrize( diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py index 8bbe8e9cc0d1..813b8f0de021 100644 --- a/tests/v1/worker/test_workspace.py +++ b/tests/v1/worker/test_workspace.py @@ -82,6 +82,34 @@ def test_workspace_lanes_compose_with_ubatches(monkeypatch) -> None: assert len(pointers) == 4 +def test_workspace_reservation_covers_every_execution_slot(monkeypatch) -> None: + active_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0]) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + + manager.reserve_all(((257,), torch.uint8), ((1,), torch.float32)) + + assert [ + buffer.numel() if buffer is not None else 0 + for buffer in manager._current_workspaces + ] == [768, 768, 768, 768] + for ubatch_id in range(2): + active_ubatch[0] = ubatch_id + for lane in range(2): + workspace_id = ubatch_id * 2 + lane + with workspace.use_workspace_lane(lane): + (view,) = manager.get_simultaneous(((8,), torch.uint8)) + reserved = manager._current_workspaces[workspace_id] + assert reserved is not None + assert view.data_ptr() == reserved.data_ptr() + + manager.lock() + with pytest.raises(AssertionError, match="reserve_all"): + manager.reserve_all(((1024,), torch.uint8)) + + def test_workspace_lane_validation(monkeypatch) -> None: monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1) diff --git a/vllm/model_executor/layers/attention_layer_base.py b/vllm/model_executor/layers/attention_layer_base.py index a7ab98673077..337e024237f6 100644 --- a/vllm/model_executor/layers/attention_layer_base.py +++ b/vllm/model_executor/layers/attention_layer_base.py @@ -47,6 +47,23 @@ def unbind_kv_cache(self) -> None: if hasattr(impl, "_v_scale_cache"): impl._v_scale_cache = None + def finalize_kv_cache_geometry(self, vllm_config: VllmConfig) -> None: + """Apply the resolved cache geometry to a loaded backend. + + Hybrid cache alignment runs after model construction because it needs + the selected attention backend. Implementations that derive kernel + plans or persistent workspace sizes from the page geometry receive the + resolved value here, before device-memory profiling assigns the + remaining capacity to KV cache. + + Args: + vllm_config: Configuration containing the resolved cache geometry. + """ + impl = getattr(self, "impl", None) + finalize = getattr(impl, "finalize_kv_cache_geometry", None) + if callable(finalize): + finalize(vllm_config.cache_config.block_size) + @abstractmethod def get_attn_backend(self) -> type[AttentionBackend]: """Get the attention backend class for this layer.""" diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 5c041f1632b7..6fc0f663b928 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -612,7 +612,10 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: For hybrid models, also aligns block_size with mamba page sizes. """ from vllm.config.cache import CacheConfig - from vllm.config.vllm import set_current_vllm_config + from vllm.config.vllm import ( + get_layers_from_vllm_config, + set_current_vllm_config, + ) cache_config = vllm_config.cache_config model_config = vllm_config.model_config @@ -650,6 +653,19 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: if cache_config.kv_cache_dtype_skip_layers: cls._align_heterogeneous_kv_block_size(vllm_config, backend_cls) + # Model construction precedes hybrid block-size alignment. Notify + # loaded layers while memory profiling can still account for plans and + # persistent workspaces derived from the resolved cache geometry. + from vllm.model_executor.layers.attention_layer_base import ( + AttentionLayerBase, + ) + + for layer in get_layers_from_vllm_config( + vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + ).values(): + layer.finalize_kv_cache_geometry(vllm_config) + @classmethod def _align_heterogeneous_kv_block_size( cls, diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index af7f310df006..2e214d2dfc9b 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, replace from math import prod -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch @@ -1095,6 +1095,7 @@ def __init__( self._module = module self._kernel_page_size = 0 + self._kernel_page_size_finalized = not self._is_glm_next self._set_kernel_page_size(kernel_page_size) self.supports_quant_query_input = False @@ -1142,7 +1143,9 @@ def make_plan(mode: str, num_q_heads: int = self._input_num_heads): self._kernel_page_size = kernel_page_size self._reserve_planned_workspaces() - def _workspace_specs(self, plan) -> tuple[tuple[tuple[int, ...], torch.dtype], ...]: + def _base_workspace_specs( + self, plan + ) -> tuple[tuple[tuple[int, ...], torch.dtype], ...]: q_spec = ( (self._max_tokens, self._input_num_heads, self._q_head_dim), torch.bfloat16, @@ -1162,8 +1165,8 @@ def _reserve_planned_workspaces(self) -> None: if not is_workspace_manager_initialized(): return plan_specs = ( - self._workspace_specs(self._decode_plan), - self._workspace_specs(self._extend_plan), + self._base_workspace_specs(self._decode_plan), + self._base_workspace_specs(self._extend_plan), ) largest_specs = max(plan_specs, key=self._workspace_nbytes) current_workspace_manager().get_simultaneous(*largest_specs) @@ -1183,6 +1186,81 @@ def _use_decode_plan( and num_tokens <= self._decode_max_rows ) + def _workspace_specs( + self, + plan: Any, + *, + input_num_heads: int, + include_ckv: bool, + ) -> tuple[tuple[tuple[int, ...], torch.dtype], ...]: + q_spec = ( + (self._max_tokens, input_num_heads, self._q_head_dim), + torch.bfloat16, + ) + ckv_specs = ( + ( + (self._ckv_local_capacity, _GLM_NEXT_CACHE_RECORD_BYTES), + torch.uint8, + ), + ( + ( + self.dcp_world_size * self._ckv_local_capacity, + _GLM_NEXT_CACHE_RECORD_BYTES, + ), + torch.uint8, + ), + ) + return ( + q_spec, + *plan.shapes_and_dtypes(), + *(ckv_specs if include_ckv else ()), + ) + + def _reserve_attention_workspaces(self) -> None: + if not self._ckv_gather_enabled: + return + assert self._ckv_extend_plan is not None + manager = current_workspace_manager() + for plan, input_num_heads, include_ckv in ( + (self._decode_plan, self._input_num_heads, False), + (self._extend_plan, self._input_num_heads, False), + (self._ckv_extend_plan, self.num_heads, True), + ): + manager.reserve_all( + *self._workspace_specs( + plan, + input_num_heads=input_num_heads, + include_ckv=include_ckv, + ) + ) + + def finalize_kv_cache_geometry(self, kernel_page_size: int) -> None: + """Finalize kernel plans and workspace memory before KV profiling. + + GLM5Next hybrid cache alignment resolves the physical page size after + model construction. Full-CKV gather workspace depends on that value, + so every execution slot must be sized while the memory profiler can + still subtract the allocation from the KV-cache budget. + + Args: + kernel_page_size: Resolved physical KV-cache page size in tokens. + + Raises: + RuntimeError: If an established page size is changed. + """ + if not self._is_glm_next: + return + if self._kernel_page_size_finalized: + if kernel_page_size != self._kernel_page_size: + raise RuntimeError( + "B12X GLM5Next KV-cache page size is immutable after " + f"finalization: {self._kernel_page_size} != {kernel_page_size}." + ) + return + self._set_kernel_page_size(kernel_page_size) + self._reserve_attention_workspaces() + self._kernel_page_size_finalized = True + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: if self._is_glm_next: if kv_cache.ndim != 3 or int(kv_cache.shape[-1]) != 528: @@ -1192,7 +1270,20 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: f"shape={tuple(kv_cache.shape)}, stride={kv_cache.stride()}, " f"dtype={kv_cache.dtype}" ) - self._set_kernel_page_size(int(kv_cache.shape[1])) + cache_page_size = int(kv_cache.shape[1]) + if getattr(self, "_kernel_page_size_finalized", False): + if cache_page_size != self._kernel_page_size: + raise RuntimeError( + "B12X GLM5Next bound cache does not match the finalized " + f"page size: {cache_page_size} != {self._kernel_page_size}." + ) + return + if self._ckv_gather_enabled: + raise RuntimeError( + "B12X GLM5Next full-CKV gather requires page geometry " + "finalization before KV-cache memory profiling." + ) + self._set_kernel_page_size(cache_page_size) def do_kv_cache_update( self, @@ -1236,6 +1327,7 @@ def uses_full_ckv_dcp( return False return ( self._ckv_gather_enabled + and getattr(self, "_kernel_page_size_finalized", False) and attn_metadata.dcp_ckv_gather_eligible and attn_metadata.num_decode_tokens == 0 and num_tokens == attn_metadata.num_actual_tokens @@ -1338,33 +1430,14 @@ def forward_mqa( else self._extend_plan ) input_num_heads = self.num_heads if use_ckv_gather else self._input_num_heads - q_spec = ( - (self._max_tokens, input_num_heads, self._q_head_dim), - torch.bfloat16, - ) - plan_specs = plan.shapes_and_dtypes() - ckv_specs = ( - ( - ( - (self._ckv_local_capacity, _GLM_NEXT_CACHE_RECORD_BYTES), - torch.uint8, - ), - ( - ( - self.dcp_world_size * self._ckv_local_capacity, - _GLM_NEXT_CACHE_RECORD_BYTES, - ), - torch.uint8, - ), - ) - if use_ckv_gather - else () - ) - workspaces = current_workspace_manager().get_simultaneous( - q_spec, *plan_specs, *ckv_specs + workspace_specs = self._workspace_specs( + plan, + input_num_heads=input_num_heads, + include_ckv=use_ckv_gather, ) + workspaces = current_workspace_manager().get_simultaneous(*workspace_specs) q_buffer = workspaces[0] - scratch_end = 1 + len(plan_specs) + scratch_end = len(workspace_specs) - (2 if use_ckv_gather else 0) scratch = workspaces[1:scratch_end] if isinstance(q, tuple): diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index e98a05c99e4e..115ac1461686 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -178,6 +178,53 @@ def get_simultaneous( for i in range(len(shapes_and_dtypes)) ] + def reserve_all( + self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype] + ) -> None: + """Reserve one equal-size workspace for every execution slot. + + Startup code uses this method when a runtime path can execute in any + microbatch or model lane. Reserving every slot before memory profiling + prevents the first request on an otherwise unused slot from growing + device memory after the KV-cache budget has been assigned. + + Args: + *shapes_and_dtypes: Simultaneously live tensor shapes and dtypes. + + Raises: + AssertionError: If the manager is locked and any slot is too small. + """ + required_bytes = sum( + round_up(_compute_bytes(shape, dtype), 256) + for shape, dtype in shapes_and_dtypes + ) + undersized = [ + workspace_id + for workspace_id, workspace in enumerate(self._current_workspaces) + if self._workspace_size_bytes(workspace) < required_bytes + ] + if self._locked and undersized: + raise AssertionError( + "Workspace is locked but reserve_all requires " + f"{required_bytes / _MB:.2f} MB in slot(s) {undersized}." + ) + + for workspace_id in undersized: + current_workspace = self._current_workspaces[workspace_id] + self._current_workspaces[workspace_id] = None + del current_workspace + torch.accelerator.empty_cache() + self._current_workspaces[workspace_id] = torch.empty( + (required_bytes,), dtype=torch.uint8, device=self._device + ) + + if envs.VLLM_DEBUG_WORKSPACE and undersized: + logger.info( + "[WORKSPACE DEBUG] Reserved %.2f MB in execution slots %s", + required_bytes / _MB, + undersized, + ) + def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor: """Ensure workspace is allocated and large enough, return current workspace. From 20c3bdd6bb251ec57d49ef67db5b04498ae6669a Mon Sep 17 00:00:00 2001 From: Jack Zampolin Date: Fri, 28 Aug 2026 01:04:28 -0700 Subject: [PATCH 5/7] Refresh B12X CKV test fixtures Co-authored-by: OpenAI Codex (cherry picked from commit 68eb9def9610ba8c4fe2ef99d5e28788135f79d8) (cherry picked from commit 710fe52ba4373f2e7b52efdfc821b7726b302b00) --- tests/v1/attention/test_b12x_sparse_mla_api.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index bd3059713f05..5fadaee934ca 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -175,10 +175,12 @@ def test_b12x_glm5_next_cache_spec_and_layout(monkeypatch) -> None: assert invalid_reasons == [] assert unidentified == probe assert packed_by_glm_backend.state_content_bytes == 528 - assert packed_by_glm_backend.page_size_padded == 64 * 528 + 16 * 128 * 2 + assert packed_by_glm_backend.page_size_padded is None + assert packed_by_glm_backend.page_size_bytes == 64 * (528 + 33) assert packed_by_glm_backend.model_version == "glm5_next" assert packed.state_content_bytes == 528 - assert packed.page_size_padded == 64 * 528 + 16 * 128 * 2 + assert packed.page_size_padded is None + assert packed.page_size_bytes == 64 * (528 + 33) assert packed.model_version == "glm5_next" assert packed_without_config_context == packed assert layouts == (KVCacheLayout.BLHNC,) From 9e537c586dbe27ed35caa37b0cf27aacf67431fb Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 29 Aug 2026 08:58:43 +0000 Subject: [PATCH 6/7] test(glm5next): model the selector metadata contract Selector-state fixture stubs expose prefill, decode-token, and speculative-decode fields required by the sparse-MLA metadata builder. The bare builder disables full-CKV eligibility so selector-state staging tests remain isolated from DCP cache-gather behavior. (cherry picked from commit 702004ec703c2f36237b160fa2cb9c2cb01754aa) --- tests/v1/attention/test_b12x_sparse_mla_api.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 5fadaee934ca..095b47cbb94a 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -636,6 +636,7 @@ def _bare_glm_selector_metadata_builder() -> B12xMLASparseMetadataBuilder: builder = B12xMLASparseMetadataBuilder.__new__(B12xMLASparseMetadataBuilder) builder.requires_glm_next_selector_metadata = True builder.supports_draft_decode_metadata_update = True + builder._ckv_gather_requested = False builder.dcp_world_size = 1 builder._max_speculative_decode_query_len = 6 builder._capture_default_state_slot_ids = torch.arange(4, dtype=torch.int32) @@ -734,7 +735,11 @@ def test_glm_selector_metadata_builder_stages_padded_rows_and_capture( monkeypatch.setattr( SparseMLACommonMetadataBuilder, "build", - lambda *args, **kwargs: SimpleNamespace(), + lambda *args, **kwargs: SimpleNamespace( + num_prefills=0, + num_decode_tokens=0, + is_spec_decode=False, + ), ) builder = _bare_glm_selector_metadata_builder() common = SimpleNamespace( @@ -810,7 +815,11 @@ def test_glm_selector_metadata_builder_requires_complete_runtime_state( monkeypatch.setattr( SparseMLACommonMetadataBuilder, "build", - lambda *args, **kwargs: SimpleNamespace(), + lambda *args, **kwargs: SimpleNamespace( + num_prefills=0, + num_decode_tokens=0, + is_spec_decode=False, + ), ) builder = _bare_glm_selector_metadata_builder() common = SimpleNamespace( From 92d807af333bdabf2f7110de7eac62595cc4fbe3 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 29 Aug 2026 09:09:22 +0000 Subject: [PATCH 7/7] fix(glm5next): classify full-CKV batches from decode tokens Enable the synchronous full-cache gather only when B12X metadata reports zero decode tokens. This excludes ordinary and speculative decode without depending on a metadata field that the sparse MLA builder does not define, including during CUDA graph memory profiling. (cherry picked from commit eeae70219b9f86889108afa734f6de1bbd42d99e) Signed-off-by: Martin Vit --- .../v1/attention/test_b12x_sparse_mla_api.py | 18 ++++++------- .../attention/backends/mla/b12x_mla_sparse.py | 25 ++++++++----------- 2 files changed, 19 insertions(+), 24 deletions(-) diff --git a/tests/v1/attention/test_b12x_sparse_mla_api.py b/tests/v1/attention/test_b12x_sparse_mla_api.py index 095b47cbb94a..67463ebb4648 100644 --- a/tests/v1/attention/test_b12x_sparse_mla_api.py +++ b/tests/v1/attention/test_b12x_sparse_mla_api.py @@ -264,18 +264,18 @@ def test_b12x_glm5_next_accepts_dcp_with_speculation(monkeypatch) -> None: @pytest.mark.parametrize( - ("max_query_len", "is_spec_decode", "num_tokens", "expected"), + ("max_query_len", "num_decode_tokens", "num_tokens", "expected"), [ - (1, False, 32, False), - (6, True, 192, False), - (6, False, 192, True), - (128, False, 8192, True), - (128, False, 600000, False), + (1, 0, 32, False), + (6, 192, 192, False), + (6, 0, 192, True), + (128, 0, 8192, True), + (128, 0, 600000, False), ], ) def test_b12x_full_ckv_gather_excludes_decode_and_mtp_batches( max_query_len: int, - is_spec_decode: bool, + num_decode_tokens: int, num_tokens: int, expected: bool, ) -> None: @@ -286,7 +286,7 @@ def test_b12x_full_ckv_gather_excludes_decode_and_mtp_batches( dcp_world_size=4, max_query_len=max_query_len, num_tokens=num_tokens, - is_spec_decode=is_spec_decode, + num_decode_tokens=num_decode_tokens, min_tokens=16, max_tokens=524288, ) @@ -738,7 +738,6 @@ def test_glm_selector_metadata_builder_stages_padded_rows_and_capture( lambda *args, **kwargs: SimpleNamespace( num_prefills=0, num_decode_tokens=0, - is_spec_decode=False, ), ) builder = _bare_glm_selector_metadata_builder() @@ -818,7 +817,6 @@ def test_glm_selector_metadata_builder_requires_complete_runtime_state( lambda *args, **kwargs: SimpleNamespace( num_prefills=0, num_decode_tokens=0, - is_spec_decode=False, ), ) builder = _bare_glm_selector_metadata_builder() diff --git a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py index 2e214d2dfc9b..5d131145d872 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/b12x_mla_sparse.py @@ -166,7 +166,7 @@ def _use_b12x_full_ckv_gather( dcp_world_size: int, max_query_len: int, num_tokens: int, - is_spec_decode: bool, + num_decode_tokens: int, min_tokens: int, max_tokens: int, ) -> bool: @@ -175,7 +175,7 @@ def _use_b12x_full_ckv_gather( and is_glm_next and dcp_world_size > 1 and max_query_len > 1 - and not is_spec_decode + and num_decode_tokens == 0 and num_tokens > min_tokens and num_tokens <= max_tokens ) @@ -826,18 +826,15 @@ def _build( metadata.prefill_seq_lens_cpu = seq_lens_cpu_source[ prefill_start : prefill_start + metadata.num_prefills ].clone() - if ( - _use_b12x_full_ckv_gather( - enabled=self._ckv_gather_requested, - is_glm_next=self.requires_glm_next_selector_metadata, - dcp_world_size=self.dcp_world_size, - max_query_len=common.max_query_len, - num_tokens=num_tokens, - is_spec_decode=metadata.is_spec_decode, - min_tokens=envs.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS, - max_tokens=envs.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS, - ) - and metadata.num_decode_tokens == 0 + if _use_b12x_full_ckv_gather( + enabled=self._ckv_gather_requested, + is_glm_next=self.requires_glm_next_selector_metadata, + dcp_world_size=self.dcp_world_size, + max_query_len=common.max_query_len, + num_tokens=num_tokens, + num_decode_tokens=metadata.num_decode_tokens, + min_tokens=envs.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS, + max_tokens=envs.VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS, ): assert self.ckv_selected_indices_buffer is not None assert self.ckv_active_counts_buffer is not None