From 73d2a8cf86c46e21af723ffd34ce0d56668fdd14 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:42:33 +0800 Subject: [PATCH 001/139] [2/2][Model Runner V2] FULL CUDA graph capture for microbatched steps (DBO) (#51700) Signed-off-by: specture724 Co-authored-by: Codex --- tests/v1/worker/test_gpu_ubatch_slicing.py | 163 ++++++++++++++++++++- vllm/config/vllm.py | 5 +- vllm/v1/worker/gpu/cudagraph_utils.py | 146 ++++++++++++++---- vllm/v1/worker/gpu/dp_utils.py | 46 ++++-- vllm/v1/worker/gpu/model_runner.py | 7 +- vllm/v1/worker/gpu/ubatch_utils.py | 120 +++++++++++---- 6 files changed, 407 insertions(+), 80 deletions(-) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index f54e746710..e735953d8b 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -20,13 +20,22 @@ import pytest import torch from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata -from vllm.config import CUDAGraphMode, ModelConfig, ParallelConfig, VllmConfig +from vllm.config import ( + CompilationConfig, + CUDAGraphMode, + ModelConfig, + ParallelConfig, + VllmConfig, +) from vllm.forward_context import create_forward_context from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils -from vllm.v1.worker.gpu import dp_utils -from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu import cudagraph_utils, dp_utils +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, +) from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.ubatch_utils import ( @@ -272,6 +281,7 @@ def _sync_dp( num_tokens_per_rank: list[int], uniform_token_count_per_rank: list[int] | None = None, allow_ubatching: bool = True, + cudagraph_manager: Any = None, ) -> tuple[BatchExecutionDescriptor, dp_utils.DPSyncState | None]: """Run the DP handshake with the all-reduce stubbed out. @@ -295,7 +305,7 @@ def _sync_dp( patch.object(dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)), ): return dp_utils.sync_cudagraph_and_dp_padding( - cudagraph_manager=None, + cudagraph_manager=cudagraph_manager, desired_batch_desc=BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens_per_rank[0], @@ -718,3 +728,148 @@ def test_slicing_drops_stale_dcp_metadata_when_dcp_is_off(): ] assert all(ubatch.dcp_local_seq_lens is None for ubatch in ubatches) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU") +def test_capturable_run_replays_as_a_cudagraph(): + """Microbatched graph replay reads updated persistent input buffers.""" + vllm_config = VllmConfig( + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0), + parallel_config=ParallelConfig( + enable_dbo=True, all2all_backend="deepep_low_latency" + ), + ) + device = torch.device("cuda:0") + runner = _make_execution_runner(vllm_config) + + num_tokens = 16 + input_ids = torch.arange(num_tokens, device=device) + positions = torch.arange(num_tokens, device=device) + model_inputs = { + "input_ids": input_ids, + "positions": positions, + "inputs_embeds": None, + "intermediate_tensors": None, + } + ubatch_state = _make_ubatch_state( + vllm_config, + [ + UBatchSlice(slice(0, 4), slice(0, 8)), + UBatchSlice(slice(4, 8), slice(8, 16)), + ], + ) + + model = _YieldingModel([]) + graph = torch.cuda.CUDAGraph() + finish = runner.begin_capturable_run( + model, model_inputs, ubatch_state, for_capture=True + ) + with torch.cuda.graph(graph, stream=runner.capture_stream): + captured = finish() + + input_ids.fill_(3) + positions.fill_(5) + graph.replay() + torch.accelerator.synchronize() + + expected = torch.full( + (num_tokens, 1), 3 * 2 + 5, dtype=torch.float32, device=device + ) + torch.testing.assert_close(captured, expected) + + +def _request_slices(input_batch: InputBatch) -> list[slice]: + return [s.request_slice for s in create_ubatch_slices(input_batch, 2)] + + +def test_captured_split_survives_a_replay_with_enough_requests(): + """Replay preserves the captured split only if real requests reach it.""" + buffers = _make_buffers() + captured = _request_slices(InputBatch.make_dummy(16, 16, buffers)) + + def replay(num_reqs: int) -> list[slice]: + return _request_slices( + _make_input_batch( + [1] * num_reqs, + [128] * num_reqs, + buffers, + num_reqs_padded=16, + num_tokens_padded=16, + ) + ) + + assert replay(16) == captured + assert replay(9) == captured + assert replay(4) != captured + + +def _make_cudagraph_manager(capture_sizes: list[int]) -> CudaGraphManager: + vllm_config = VllmConfig( + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0), + parallel_config=ParallelConfig( + enable_dbo=True, + all2all_backend="deepep_low_latency", + dbo_decode_token_threshold=DECODE_THRESHOLD, + dbo_prefill_token_threshold=PREFILL_THRESHOLD, + ), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.FULL, + cudagraph_capture_sizes=capture_sizes, + ), + ) + with patch.object( + cudagraph_utils, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ): + manager = CudaGraphManager( + vllm_config, + torch.device("cuda:0"), + CUDAGraphMode.FULL, + decode_query_len=1, + ubatch_runner=cast(UBatchRunner, object()), + ) + manager._graphs_captured = True + return manager + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool") +def test_microbatched_graphs_are_only_offered_to_uniform_batches(): + """Mixed batches cannot reuse graphs captured with uniform query lengths.""" + manager = _make_cudagraph_manager([64, 128]) + + desc = manager.dispatch( + num_reqs=64, + num_tokens=64, + uniform_token_count=1, + num_active_loras=0, + num_ubatches=2, + ) + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.num_ubatches == 2 + + desc = manager.dispatch( + num_reqs=8, + num_tokens=64, + uniform_token_count=None, + num_active_loras=0, + num_ubatches=2, + ) + assert desc.cg_mode == CUDAGraphMode.NONE + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a graph pool") +def test_microbatched_graph_needs_every_rank_to_reach_the_split(): + """Fall back to eager on all ranks if any rank cannot reach the captured split.""" + manager = _make_cudagraph_manager([64, 128]) + uniform = [1, 1] + + desc, _ = _sync_dp([64, 64], uniform, cudagraph_manager=manager) + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.num_ubatches == 2 + + # Rank 1 pads the group up to a 128-token graph, but rank 0 only reaches 33. + desc, dp_sync = _sync_dp([33, 128], uniform, cudagraph_manager=manager) + assert desc.cg_mode == CUDAGraphMode.NONE + assert desc.num_ubatches == 2 + assert dp_sync is not None and dp_sync.eager diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index fbe0213855..2fcc144af0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3034,8 +3034,9 @@ class VllmConfig: unsupported.append("dual batch overlap with multimodal models") if model_config is not None and model_config.is_hybrid: unsupported.append("dual batch overlap with hybrid models") - if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: - unsupported.append("dual batch overlap with CUDA graphs") + if self.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE: + # DBO captures FULL graphs only. + unsupported.append("dual batch overlap with PIECEWISE CUDA graphs") if self.is_mm_encoder_only: unsupported.append("dual batch overlap with encoder only models") diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 12354f3c22..d197dbbc6d 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -4,7 +4,8 @@ import gc import itertools from collections import defaultdict from collections.abc import Callable, Iterable -from dataclasses import dataclass +from contextlib import ExitStack +from dataclasses import dataclass, replace from itertools import groupby, product from typing import TYPE_CHECKING, Any, NamedTuple, Protocol @@ -42,11 +43,13 @@ from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches if TYPE_CHECKING: from vllm.v1.worker.gpu.model_runner import GPUModelRunner from vllm.v1.worker.gpu.pcp_manager import PCPManager + from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner logger = init_logger(__name__) @@ -144,6 +147,7 @@ class CudaGraphManager: decode_query_len: int, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + ubatch_runner: "UBatchRunner | None" = None, ): self.vllm_config = vllm_config self.device = device @@ -153,6 +157,8 @@ class CudaGraphManager: self.cudagraph_mode = cudagraph_mode self.decode_query_len = decode_query_len self.varlen_decode = varlen_decode + # DBO supports FULL CUDA graphs only. + self.ubatch_runner = ubatch_runner self.dp_size = vllm_config.parallel_config.data_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size @@ -212,6 +218,33 @@ class CudaGraphManager: # Counts above the largest captured case clamp to it. return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case) + def _maybe_ubatch_twin( + self, desc: BatchExecutionDescriptor + ) -> BatchExecutionDescriptor | None: + """Return a microbatched capture candidate when eligible. + + Uniform query lengths preserve the captured request split. Use the DP + dispatch thresholds so all ranks generate the same candidates. + """ + if self.ubatch_runner is None or desc.cg_mode != CUDAGraphMode.FULL: + return None + if desc.num_reqs is None: + return None + uniform_token_count, remainder = divmod(desc.num_tokens, desc.num_reqs) + if remainder or desc.uniform_token_count not in (None, uniform_token_count): + return None + parallel_config = self.vllm_config.parallel_config + num_ubatches = get_num_ubatches(parallel_config) + if desc.num_tokens < num_ubatches: + return None + if not check_ubatch_thresholds( + parallel_config, desc.num_tokens, uniform_decode=True + ): + return None + return replace( + desc, num_ubatches=num_ubatches, uniform_token_count=uniform_token_count + ) + def _init_candidates(self) -> None: """Build priority-ordered candidate lists for each token count.""" capture_sizes = self.compilation_config.cudagraph_capture_sizes @@ -300,6 +333,12 @@ class CudaGraphManager: if desc not in descs_by_mode[decode_mode]: descs_by_mode[decode_mode].append(desc) + ubatch_desc = self._maybe_ubatch_twin(desc) + if ubatch_desc is not None and ( + ubatch_desc not in descs_by_mode[decode_mode] + ): + descs_by_mode[decode_mode].append(ubatch_desc) + # recoverSSM cannot capture a dummy query wider than its workspace. if mixed_mode and ( not self.vllm_config.cache_config.use_kda_recoverssm @@ -321,6 +360,10 @@ class CudaGraphManager: ) descs_by_mode[mixed_mode].append(desc) + ubatch_desc = self._maybe_ubatch_twin(desc) + if ubatch_desc is not None: + descs_by_mode[mixed_mode].append(ubatch_desc) + for mode, descs in descs_by_mode.items(): descs.sort(key=lambda d: d.num_tokens, reverse=True) self._capture_descs[mode] = descs @@ -344,6 +387,13 @@ class CudaGraphManager: def needs_capture(self) -> bool: return len(self._capture_descs) > 0 + def _capture_stream(self, desc: BatchExecutionDescriptor) -> torch.cuda.Stream: + """Capture on the stream used by the microbatch threads.""" + if desc.num_ubatches > 1: + assert self.ubatch_runner is not None + return self.ubatch_runner.capture_stream + return current_stream() + @torch.inference_mode() def capture( self, @@ -359,7 +409,10 @@ class CudaGraphManager: because attention backends may mutate or lazily initialize metadata during warmup. """ - with graph_capture(device=self.device): + with graph_capture(device=self.device), ExitStack() as stack: + if self.ubatch_runner is not None: + # Join parked threads on failure to avoid blocking later captures. + stack.callback(self.ubatch_runner.abort_pending_run) # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger # activations so FULL activations should fit in already allocated # buffers in the graph pool. @@ -387,7 +440,9 @@ class CudaGraphManager: # Capture logger.debug( - "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc + "CG Capture: mode=%s, batch_desc=%s", + desc.cg_mode.name, + desc, ) if ( desc.cg_mode == CUDAGraphMode.PIECEWISE @@ -415,7 +470,7 @@ class CudaGraphManager: torch.accelerator.synchronize() free_before = torch.accelerator.get_memory_info()[0] with torch.cuda.graph( - graph, self.pool, stream=current_stream() + graph, self.pool, stream=self._capture_stream(desc) ): forward_fn(CUDAGraphMode.NONE) # Join offloader's copy stream after forward to avoid @@ -429,6 +484,7 @@ class CudaGraphManager: self._capture_mem_samples.append(free_before - free_after) self.graphs[desc] = graph compilation_counter.num_cudagraph_captured += 1 + self._graphs_captured = True def captured_token_counts(self) -> list[int]: @@ -510,6 +566,7 @@ class ModelCudaGraphManager(CudaGraphManager): decode_query_len: int, lora_capture_cases: list[int] | None = None, varlen_decode: bool = False, + ubatch_runner: "UBatchRunner | None" = None, ): super().__init__( vllm_config, @@ -518,6 +575,7 @@ class ModelCudaGraphManager(CudaGraphManager): decode_query_len, lora_capture_cases=lora_capture_cases, varlen_decode=varlen_decode, + ubatch_runner=ubatch_runner, ) self.hidden_states: torch.Tensor | None = None self.aux_hidden_states: list[torch.Tensor] = [] @@ -554,6 +612,35 @@ class ModelCudaGraphManager(CudaGraphManager): "Set VLLM_USE_BREAKABLE_CUDAGRAPH=1 or cudagraph_mode=NONE/FULL." ) + def store_capture_output(num_tokens: int, model_output: Any) -> None: + """Copy outputs to persistent buffers, allocating on first use.""" + if self.is_last_pp_rank: + # Last PP rank (common case). + if self.use_aux_hidden_state_outputs: + hidden_states, aux_hidden_states = model_output + else: + hidden_states = model_output + aux_hidden_states = [] + if self.hidden_states is None: + self.hidden_states = torch.empty_like(hidden_states) + self.hidden_states[:num_tokens] = hidden_states + if self.use_aux_hidden_state_outputs and not self.aux_hidden_states: + self.aux_hidden_states = [ + torch.empty_like(x) for x in aux_hidden_states + ] + for i, aux in enumerate(aux_hidden_states): + self.aux_hidden_states[i][:num_tokens] = aux + else: + # Non-last PP rank. + assert isinstance(model_output, IntermediateTensors) + intermediate_tensors = model_output + if self.intermediate_tensors is None: + self.intermediate_tensors = IntermediateTensors.empty_like( + intermediate_tensors + ) + for k, v in intermediate_tensors.tensors.items(): + self.intermediate_tensors[k][:num_tokens] = v + def create_forward_fn( desc: BatchExecutionDescriptor, warmup: bool, @@ -583,6 +670,30 @@ class ModelCudaGraphManager(CudaGraphManager): assert intermediate_tensors is not None model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens] + if desc.num_ubatches > 1: + # Prepare and park threads before capture; finish runs inside it. + assert self.ubatch_runner is not None + ubatch_state = self.ubatch_runner.prepare( + InputBatch.make_dummy(num_reqs, num_tokens, input_buffers), + block_tables.get_dummy_block_tables(num_reqs), + block_tables.get_dummy_slot_mappings(num_tokens), + cg_mode=CUDAGraphMode.FULL, + for_capture=True, + ) + # Capture with dummy rows marked as padding. + input_buffers.is_padding.fill_(True) + finish = self.ubatch_runner.begin_capturable_run( + model, model_inputs, ubatch_state, for_capture=True + ) + + def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None: + assert cg_mode != CUDAGraphMode.PIECEWISE, ( + "DBO does not support PIECEWISE cudagraphs" + ) + store_capture_output(num_tokens, finish()) + + return ubatch_forward_fn + attn_metadata, slot_mappings = prepare_inputs_to_capture( num_reqs, num_tokens, @@ -629,32 +740,7 @@ class ModelCudaGraphManager(CudaGraphManager): # model outputs. No need to keep track of the hidden states. return None - if self.is_last_pp_rank: - # Last PP rank (common case). - if self.use_aux_hidden_state_outputs: - hidden_states, aux_hidden_states = model_output - else: - hidden_states = model_output - aux_hidden_states = [] - if self.hidden_states is None: - self.hidden_states = torch.empty_like(hidden_states) - self.hidden_states[:num_tokens] = hidden_states - if self.use_aux_hidden_state_outputs and not self.aux_hidden_states: - self.aux_hidden_states = [ - torch.empty_like(x) for x in aux_hidden_states - ] - for i, aux in enumerate(aux_hidden_states): - self.aux_hidden_states[i][:num_tokens] = aux - else: - # Non-last PP rank. - assert isinstance(model_output, IntermediateTensors) - intermediate_tensors = model_output - if self.intermediate_tensors is None: - self.intermediate_tensors = IntermediateTensors.empty_like( - intermediate_tensors - ) - for k, v in intermediate_tensors.tensors.items(): - self.intermediate_tensors[k][:num_tokens] = v + store_capture_output(num_tokens, model_output) return forward_fn diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index b7c73789de..80309bc009 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -112,23 +112,45 @@ def sync_cudagraph_and_dp_padding( int(num_tokens_across_dp.min()), uniform_decode=uniform_decode_across_dp, ): - # Microbatching is all-or-nothing: the expert all-to-all is - # collective, so every rank splits and pads to the same token count. - # A rank too small to fill a microbatch does no work in it, like a - # dummy run. Microbatched steps run eager; nothing is captured yet. + # Expert all-to-all requires every rank to split and pad equally. + # Empty microbatches run as dummy batches. ubatch_num_tokens = int(num_tokens_across_dp.max()) - return BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=ubatch_num_tokens, - num_reqs=num_reqs, - num_ubatches=get_num_ubatches(parallel_config), - ), DPSyncState( + num_ubatches = get_num_ubatches(parallel_config) + ubatch_desc = None + if cudagraph_manager is not None: + # Match a FULL microbatched CUDA graph, falling back to an eager + # descriptor if no graph matches. + ubatch_desc = cudagraph_manager.dispatch( + num_reqs, + ubatch_num_tokens, + synced_uniform_token_count, + num_active_loras=num_active_loras, + num_ubatches=num_ubatches, + ) + if 2 * int(num_tokens_across_dp.min()) < ubatch_desc.num_tokens: + # If one rank has an empty second microbatch, run without + # CUDA graphs. + ubatch_desc = None + if ubatch_desc is None: + ubatch_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=ubatch_num_tokens, + num_reqs=num_reqs, + num_ubatches=num_ubatches, + ) + # Refresh the token count to include CUDA graph padding. + ubatch_num_tokens = ubatch_desc.num_tokens + num_reqs = int(num_reqs_across_dp.max()) + if ubatch_desc.cg_mode == CUDAGraphMode.FULL: + assert ubatch_desc.num_reqs is not None + num_reqs = ubatch_desc.num_reqs + return ubatch_desc, DPSyncState( num_tokens_across_dp=torch.full_like( num_tokens_across_dp, ubatch_num_tokens ), uniform_token_count=synced_uniform_token_count, - eager=True, - num_reqs=int(num_reqs_across_dp.max()), + eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE, + num_reqs=num_reqs, ) synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item())) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9d7918126f..d66fc6c90f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -723,6 +723,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, varlen_decode=self.adaptive_verification is not None, + ubatch_runner=self.ubatch_runner, ) if self.cache_config.kv_sharing_fast_prefill and self.pcp_manager is None: self.fast_prefill = FastPrefillHelper( @@ -1775,7 +1776,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): assert self.ubatch_runner is not None assert block_tables is not None and slot_mappings is not None ubatch_state = self.ubatch_runner.prepare( - input_batch, block_tables, slot_mappings + input_batch, + block_tables, + slot_mappings, + cg_mode=batch_desc.cg_mode, + for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, ) elif not (dummy_run and skip_attn_for_dummy_run): assert slot_mappings is not None diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 8007e2dbb4..e3cade2d0b 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -3,6 +3,8 @@ """Microbatching (DBO) helpers for the V2 GPU model runner.""" import threading +from collections.abc import Callable +from contextlib import ExitStack from dataclasses import replace from typing import Any, NamedTuple @@ -317,8 +319,11 @@ class UBatchRunner: for _ in range(self.num_ubatches) ] self.comm_stream = torch.cuda.Stream(device=device) + # Threads and graph capture must use the same compute stream. + self.capture_stream = torch.cuda.Stream(device=device) # The microbatch threads plus the thread that starts them. self.ready_barrier = threading.Barrier(self.num_ubatches + 1) + self._pending_finish: Callable[[], Any] | None = None self.sm_control = create_sm_control_context(self.parallel_config) def prepare( @@ -326,11 +331,16 @@ class UBatchRunner: input_batch: InputBatch, block_tables: tuple[torch.Tensor, ...], slot_mappings: torch.Tensor, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + for_capture: bool = False, ) -> UBatchState: """Split the batch into the microbatches the step will run on. Attention metadata is built per microbatch and carried in the forward contexts the threads install, not in the caller's. + + FULL uses padded sizes; `for_capture` refreshes metadata from dummy + block tables. """ ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches) @@ -355,12 +365,13 @@ class UBatchRunner: attn_metadata.append( self.model_state.prepare_attn( ubatch, - CUDAGraphMode.NONE, + cg_mode, ubatch_block_tables, ubatch_slot_mappings, self.attn_groups, self.kv_cache_config, ubatch_idx=i, + for_capture=for_capture, ) ) slot_mappings_by_layer.append( @@ -415,14 +426,34 @@ class UBatchRunner: model_inputs: dict[str, Any], ubatch_state: UBatchState, ) -> Any: + finish = self.begin_capturable_run(model, model_inputs, ubatch_state) + return finish() + + def begin_capturable_run( + self, + model: Any, + model_inputs: dict[str, Any], + ubatch_state: UBatchState, + for_capture: bool = False, + ) -> Callable[[], Any]: + """Start microbatch threads and return a callback to run and join them. + + For capture, start threads outside the graph and call the callback inside. + `for_capture` selects the capture stream and initializes cuBLAS workspaces + before capture. + + Call the callback exactly once on this thread, or call `abort_pending_run` + on failure to release parked threads. + """ ubatch_slices = ubatch_state.slices assert len(ubatch_slices) == len(ubatch_state.forward_contexts) assert len(ubatch_slices) == self.num_ubatches + compute_stream = self.capture_stream if for_capture else current_stream() ubatch_contexts = make_ubatch_contexts( num_micro_batches=self.num_ubatches, comm_stream=self.comm_stream, - compute_stream=current_stream(), + compute_stream=compute_stream, forward_contexts=ubatch_state.forward_contexts, ready_barrier=self.ready_barrier, ) @@ -433,6 +464,13 @@ class UBatchRunner: @torch.inference_mode() def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None: try: + # New threads must select the worker's device. + torch.accelerator.set_device_index(self.device.index) + if for_capture: + # Allocate cuBLAS workspaces before capture. + for stream in (compute_stream, self.comm_stream): + with torch.cuda.stream(stream): + torch.cuda.current_blas_handle() with ubatch_context: outputs[ubatch_context.id] = model(**inputs) except BaseException as e: # noqa: BLE001 @@ -445,33 +483,56 @@ class UBatchRunner: # is deliberately out of scope here. errors[ubatch_context.id] = e - # The threads manage the forward context themselves; clear it here so - # it is restored correctly once they are done. - with override_forward_context(None), self.sm_control: - threads = [] - for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices): - thread = threading.Thread( - target=run_ubatch, - args=( - ubatch_context, - slice_model_inputs(model_inputs, ubatch_slice.token_slice), - ), - ) - threads.append(thread) - thread.start() + # Threads own the forward context; keep it cleared and SMs reserved + # until finish joins them. + stack = ExitStack() + stack.enter_context(override_forward_context(None)) + stack.enter_context(self.sm_control) - # Wait for every thread to reach its context, then start the first. - self.ready_barrier.wait() - ubatch_contexts[0].cpu_wait_event.set() - for thread in threads: - thread.join() + threads = [] + for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices): + thread = threading.Thread( + target=run_ubatch, + args=( + ubatch_context, + slice_model_inputs(model_inputs, ubatch_slice.token_slice), + ), + ) + threads.append(thread) + thread.start() - if errors: - failed = min(errors) - raise RuntimeError( - f"Microbatch {failed} of {self.num_ubatches} failed" - ) from errors[failed] - return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)]) + # Park all threads before the caller starts capture. + self.ready_barrier.wait() + + def finish() -> Any: + try: + ubatch_contexts[0].cpu_wait_event.set() + for thread in threads: + thread.join() + finally: + self._pending_finish = None + stack.close() + + if errors: + failed = min(errors) + raise RuntimeError( + f"Microbatch {failed} of {self.num_ubatches} failed" + ) from errors[failed] + return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)]) + + self._pending_finish = finish + return finish + + def abort_pending_run(self) -> None: + """Release and join parked threads after a failed capture.""" + if self._pending_finish is None: + return + try: + self._pending_finish() + except Exception: + logger.exception("Failed to join microbatch threads after an abort") + finally: + self._pending_finish = None def maybe_build_ubatch_runner( @@ -491,10 +552,7 @@ def maybe_build_ubatch_runner( if not parallel_config.use_ubatching or parallel_config.data_parallel_size <= 1: return None - logger.info_once( - "Dual batch overlap is enabled. Microbatched steps run without " - "CUDA graphs on the V2 model runner." - ) + logger.info_once("Dual batch overlap is enabled.") return UBatchRunner( vllm_config, device, model_state, attn_groups, kv_cache_config, max_num_reqs ) From 10e8d5614c50ad180acbe61878c93d98ca7098a9 Mon Sep 17 00:00:00 2001 From: Alexander Lee Date: Mon, 14 Sep 2026 01:49:44 -0400 Subject: [PATCH 002/139] [Bugfix][Beam Search] Respect skip_special_tokens during decoding (#56211) Signed-off-by: lxy-alexander Co-authored-by: OpenAI Codex --- tests/samplers/test_beam_search_online.py | 38 ++++++++++++++++++- .../generate/beam_search/offline.py | 4 +- .../generate/beam_search/online.py | 4 +- .../openai/chat_completion/protocol.py | 1 + .../entrypoints/openai/completion/protocol.py | 1 + vllm/sampling_params.py | 1 + 6 files changed, 46 insertions(+), 3 deletions(-) diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index 14481f79c6..102e66dcbe 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -11,8 +11,13 @@ from vllm.sampling_params import BeamSearchParams class _Tokenizer: eos_token_id = 0 + special_token_id = eos_token_id - def decode(self, token_ids: list[int]) -> str: + def decode(self, token_ids: list[int], skip_special_tokens: bool = False) -> str: + if skip_special_tokens: + token_ids = [ + token_id for token_id in token_ids if token_id != self.special_token_id + ] return " ".join(str(token_id) for token_id in token_ids) @@ -72,3 +77,34 @@ async def test_beam_search_handles_extra_logprob_candidates() -> None: assert outputs[0].outputs[0].finish_reason == "stop" assert outputs[0].outputs[0].token_ids == [] assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("skip_special_tokens", "expected_text"), + [ + pytest.param(True, "", id="skip"), + pytest.param(False, "0", id="keep"), + ], +) +async def test_beam_search_respects_skip_special_tokens( + skip_special_tokens: bool, expected_text: str +) -> None: + prompt = { + "type": "token", + "prompt": "prompt", + "prompt_token_ids": [1], + } + params = BeamSearchParams( + beam_width=1, + max_tokens=1, + ignore_eos=True, + skip_special_tokens=skip_special_tokens, + ) + + outputs = [ + output async for output in _Serving().beam_search(prompt, "request", params) + ] + + assert outputs[0].outputs[0].text == expected_text + assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id] diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index f412f5aebe..b919435f65 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -186,7 +186,9 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): best_beams = sorted_completed[:beam_width] for beam in best_beams: - beam.text = tokenizer.decode(beam.tokens) + beam.text = tokenizer.decode( + beam.tokens, skip_special_tokens=params.skip_special_tokens + ) outputs.append(BeamSearchOutput(sequences=best_beams)) diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index 02ff61d629..ac1a585e8f 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -197,7 +197,9 @@ class BeamSearchOnlineMixin(ABC): tokens = beam.tokens[tokenized_length:-1] else: tokens = beam.tokens[tokenized_length:] - beam.text = tokenizer.decode(tokens) + beam.text = tokenizer.decode( + tokens, skip_special_tokens=params.skip_special_tokens + ) yield RequestOutput( request_id=request_id, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index b7b292e835..badd5a0344 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -659,6 +659,7 @@ class ChatCompletionRequest(OpenAIBaseModel): temperature=temperature, length_penalty=self.length_penalty, include_stop_str_in_output=self.include_stop_str_in_output, + skip_special_tokens=self.skip_special_tokens, ) def extract_structured_outputs(self) -> StructuredOutputsParams | None: diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 6d7a49bacd..bf0a94eea7 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -303,6 +303,7 @@ class CompletionRequest(OpenAIBaseModel): temperature=temperature, length_penalty=self.length_penalty, include_stop_str_in_output=self.include_stop_str_in_output, + skip_special_tokens=self.skip_special_tokens, ) def extract_structured_outputs(self) -> StructuredOutputsParams | None: diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 52a65aefbc..c0ab45ced7 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1313,6 +1313,7 @@ class BeamSearchParams( length_penalty: float = 1.0 include_stop_str_in_output: bool = False structured_outputs: StructuredOutputsParams | None = None + skip_special_tokens: bool = True def __post_init__(self) -> None: _verify_num_sequences(self.beam_width, "beam_width") From cf1584f373d06f03c5662608fb80cf66c03fcb1d Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 14 Sep 2026 14:17:50 +0800 Subject: [PATCH 003/139] [Refactor] Normalize DeepSeek V4.1 model package naming (#56741) Signed-off-by: Bugen Zhao --- .buildkite/test_areas/distributed.yaml | 2 +- .github/mergify.yml | 5 ++--- tests/distributed/test_engram_dp_shard.py | 10 +++++----- tests/kernels/core/test_fused_q_kv_rmsnorm.py | 8 ++++---- tests/kernels/test_compressor_kv_cache.py | 10 +++++----- tests/kernels/test_engram.py | 12 ++++++------ tests/kernels/test_fused_indexer_q_rope_quant.py | 6 +++--- tests/kernels/test_mhc_kernels.py | 6 +++--- tests/models/test_deepseek_v4_mega_moe.py | 4 ++-- tests/models/test_dspark_mla.py | 2 +- tests/quantization/test_fp8.py | 12 ++++++------ .../test_indexer_deepseek_v4_slot_mapping.py | 2 +- .../model_executor/layers/quantization/__init__.py | 2 +- vllm/model_executor/models/registry.py | 4 ++-- .../{deepseek_v4_1 => deepseek_v41}/__init__.py | 0 .../amd/__init__.py | 0 .../{deepseek_v4_1 => deepseek_v41}/amd/dspark.py | 0 .../{deepseek_v4_1 => deepseek_v41}/amd/model.py | 4 ++-- .../{deepseek_v4_1 => deepseek_v41}/amd/rocm.py | 6 +++--- .../amd/vl_model.py | 0 .../{deepseek_v4_1 => deepseek_v41}/attention.py | 14 +++++++------- .../common/__init__.py | 0 .../common/engram.py | 0 .../common/mm_preprocess.py | 0 .../common/ops/__init__.py | 0 .../common/ops/cache_utils.py | 0 .../common/ops/fused_compress_quant_cache.py | 0 .../common/ops/indexer_k_store.py | 0 .../common/ops/query_quant.py | 0 .../{deepseek_v4_1 => deepseek_v41}/common/rope.py | 0 .../{deepseek_v4_1 => deepseek_v41}/compressor.py | 2 +- .../nvidia/__init__.py | 0 .../nvidia/dspark.py | 0 .../nvidia/engram.py | 6 +++--- .../nvidia/flashinfer_sparse.py | 6 +++--- .../nvidia/flashmla.py | 6 +++--- .../nvidia/model.py | 6 +++--- .../nvidia/model_state.py | 0 .../nvidia/vl_model.py | 0 .../quant_config.py | 0 .../{deepseek_v4_1 => deepseek_v41}/sparse_mla.py | 0 vllm/transformers_utils/configs/deepseek_v41.py | 2 +- vllm/v1/attention/backends/mla/sparse_swa.py | 2 +- vllm/v1/attention/backends/registry.py | 4 ++-- 44 files changed, 71 insertions(+), 72 deletions(-) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/dspark.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/model.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/rocm.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/amd/vl_model.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/attention.py (98%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/engram.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/mm_preprocess.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/cache_utils.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/fused_compress_quant_cache.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/indexer_k_store.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/ops/query_quant.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/common/rope.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/compressor.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/__init__.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/dspark.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/engram.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/flashinfer_sparse.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/flashmla.py (98%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/model.py (99%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/model_state.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/nvidia/vl_model.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/quant_config.py (100%) rename vllm/models/{deepseek_v4_1 => deepseek_v41}/sparse_mla.py (100%) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 070664ac7a..3ea26ce89c 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -208,7 +208,7 @@ steps: - tests/distributed/test_utils - tests/distributed/test_engram_dp_shard.py - vllm/config/engram.py - - vllm/models/deepseek_v4_1/ + - vllm/models/deepseek_v41/ commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 diff --git a/.github/mergify.yml b/.github/mergify.yml index ba3a69ee35..1267f0264d 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -343,9 +343,8 @@ pull_request_rules: conditions: - label != stale - or: - # The 4.1 sources spell the model three ways: deepseek_v4_1 (models), - # deepseek_v41 (most) and deepseekv41 (tool parsers). - - files~=^vllm/models/deepseek_v4_1/ + # The 4.1 sources use deepseek_v41 and deepseekv41 (tool parsers). + - files~=^vllm/models/deepseek_v41/ - files~=^vllm/(?:parser|tokenizers|reasoning|transformers_utils/configs)/deepseek_v41.*\.py - files~=^vllm/tool_parsers/deepseekv41.*\.py - files~=^rust/src/(?:chat|parser)/src/.*deepseek_v41 diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index ebb0254289..35ffc659d4 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -39,9 +39,9 @@ from vllm.distributed.parallel_state import ( initialize_model_parallel, ) from vllm.forward_context import set_forward_context -from vllm.models.deepseek_v4_1.common.engram import EngramLayout -from vllm.models.deepseek_v4_1.nvidia import engram as engram_ops -from vllm.models.deepseek_v4_1.nvidia.engram import ( +from vllm.models.deepseek_v41.common.engram import EngramLayout +from vllm.models.deepseek_v41.nvidia import engram as engram_ops +from vllm.models.deepseek_v41.nvidia.engram import ( Engram, ParallelEngramEmbedding, engram_head_shard_rank, @@ -473,8 +473,8 @@ def _check_dummy_hash_model_forward( vllm_config, engram, head_sizes, weight, scales, dp_shared_memory ): """A metadata-less replica joins DP sharding but skips a shared lookup.""" - from vllm.models.deepseek_v4_1.common.engram import NgramHashState - from vllm.models.deepseek_v4_1.nvidia import model as model_ops + from vllm.models.deepseek_v41.common.engram import NgramHashState + from vllm.models.deepseek_v41.nvidia import model as model_ops dp_rank = vllm_config.parallel_config.data_parallel_rank if dp_shared_memory and dp_rank != 1: diff --git a/tests/kernels/core/test_fused_q_kv_rmsnorm.py b/tests/kernels/core/test_fused_q_kv_rmsnorm.py index 7e5e410487..3d0e2ac11e 100644 --- a/tests/kernels/core/test_fused_q_kv_rmsnorm.py +++ b/tests/kernels/core/test_fused_q_kv_rmsnorm.py @@ -125,7 +125,7 @@ def test_fused_q_kv_rmsnorm_quant_matches_separate(num_tokens, q_size, dtype): from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( mxfp8_e4m3_quantize, ) - from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + from vllm.models.deepseek_v41.common.ops.query_quant import ( fused_q_kv_rmsnorm_quant, ) @@ -198,15 +198,15 @@ def test_shared_query_quant_preserves_projection( mxfp8_e4m3_quantize, ) from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic - from vllm.models.deepseek_v4_1.attention import ( + from vllm.models.deepseek_v41.attention import ( DeepseekV4Attention, DeepseekV4Indexer, ) - from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + from vllm.models.deepseek_v41.common.ops.query_quant import ( can_fuse_query_quant, fused_q_kv_rmsnorm_quant, ) - from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config + from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config cls: ( type[FlashInferCutedslMxfp8LinearKernel] diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 082a60b7a5..0f1976cf68 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -65,7 +65,7 @@ def test_v41_fused_save_compress_and_insert( a Torch reference with BF16 rounding before RoPE. Replays change both inputs and history so stale scratch/cache values cannot satisfy the test. """ - from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import ( + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( fused_save_compress_norm, rope_quant_insert, ) @@ -226,7 +226,7 @@ def test_v41_rope_insert_plain_row(compress_ratio: int, store_fp8: bool): Non-boundary positions, negative slots and untouched rows must keep their prior contents. """ - from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import ( + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( rope_quant_insert, ) @@ -301,7 +301,7 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): capacity`` and keep padding tokens at PAD.""" from unittest.mock import MagicMock - from vllm.models.deepseek_v4_1.compressor import CompressorMetadataBuilder + from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.kv_cache_interface import CircularBufferSpec @@ -351,11 +351,11 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph): """Both reused-event joins must publish this forward's states and cache rows.""" from vllm.forward_context import ForwardContext, override_forward_context - from vllm.models.deepseek_v4_1.attention import ( + from vllm.models.deepseek_v41.attention import ( DeepseekV4Attention, DeepseekV4Indexer, ) - from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor + from vllm.models.deepseek_v41.compressor import DeepseekCompressor torch.manual_seed(43) raw = torch.randn(19, 1024, device="cuda") diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index c7e08ced82..c2a7607cba 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -7,15 +7,15 @@ from types import SimpleNamespace import pytest import torch -from vllm.models.deepseek_v4_1.common import engram as engram_ops -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common import engram as engram_ops +from vllm.models.deepseek_v41.common.engram import ( Engram as CommonEngram, ) -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common.engram import ( NgramHashState, ) -from vllm.models.deepseek_v4_1.nvidia import engram as nvidia_engram_ops -from vllm.models.deepseek_v4_1.nvidia.engram import Engram, ParallelEngramEmbedding +from vllm.models.deepseek_v41.nvidia import engram as nvidia_engram_ops +from vllm.models.deepseek_v41.nvidia.engram import Engram, ParallelEngramEmbedding from vllm.platforms import current_platform @@ -503,7 +503,7 @@ def test_lookback_window_reproduces_single_instance(runner, capture): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") def test_v2_model_state_gathers_lookback_window(): """The window is gathered on device from the runner's token history.""" - from vllm.models.deepseek_v4_1.nvidia.model_state import DeepseekV41ModelState + from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState depth, max_num_reqs, max_model_len = 3, 4, 16 state = DeepseekV41ModelState.__new__(DeepseekV41ModelState) diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 73b147584c..8770882ddf 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -359,7 +359,7 @@ def test_indexer_k_inserts_only_valid_groups_into_padded_pages( num_tokens, compress_ratio, use_fp4, cache_dtype ): """Insert group ends and preserve skipped slots, graph rows, and page padding.""" - from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( indexer_k_norm_rope_store, ) @@ -442,7 +442,7 @@ def test_indexer_k_inserts_only_valid_groups_into_padded_pages( @torch.inference_mode() def test_indexer_k_cuda_graph_replay_reads_current_projection(use_fp4): """Replay must consume updated keys and slot mapping through the captured API.""" - from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( indexer_k_norm_rope_store, ) @@ -501,7 +501,7 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat ``block_size > 1``, so a row-major store would feed the indexer permuted key bytes and randomize its top-k. Assert the write/read pair is exact. """ - from vllm.models.deepseek_v4_1.common.ops.indexer_k_store import ( + from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( indexer_k_norm_rope_store, ) from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 7f0515d11a..46f311dfe6 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -32,7 +32,7 @@ from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4DecoderLayer, DeepseekV4Model, ) -from vllm.models.deepseek_v4_1.nvidia.model import ( +from vllm.models.deepseek_v41.nvidia.model import ( DeepseekV4DecoderLayer as DeepseekV41DecoderLayer, ) from vllm.platforms import current_platform @@ -104,7 +104,7 @@ def test_hc_collapse_custom_op_supports_compile(): @pytest.mark.parametrize("num_tokens", [1, 7]) def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): """Match reference DSparkBlock.forward_head's hc_pre before its RMSNorm.""" - from vllm.models.deepseek_v4_1.nvidia import dspark + from vllm.models.deepseek_v41.nvidia import dspark set_random_seed(0) hidden_size, hc_mult = 5120, 4 @@ -296,7 +296,7 @@ def test_deepseek_v41_decoder_mixes_match_torch( return post, res, decoder.attn_norm(collapsed), pre monkeypatch.setattr( - "vllm.models.deepseek_v4_1.nvidia.model.mhc_pre_delayed_tilelang", + "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang", reference, ) expected = decoder(x, positions, None, **kwargs) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index d3d9ccabad..0d3c2b6583 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -18,8 +18,8 @@ from vllm.models.deepseek_v4.nvidia.model import ( ) from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs -from vllm.models.deepseek_v4_1.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID -from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4MoE as DeepseekV41MoE +from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID +from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE from vllm.platforms import current_platform from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 0d1507a033..d30eb9e7c9 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -216,7 +216,7 @@ def test_v41_dspark_loads_linear_scales( """Checkpoint ``.scale`` maps to the quant method's scale parameter and loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static loader (see tests/quantization/test_modelopt.py), not in load_weights.""" - from vllm.models.deepseek_v4_1.nvidia import dspark + from vllm.models.deepseek_v41.nvidia import dspark mxfp8 = scale_dtype != torch.float32 scale_name = "weight_scale" if mxfp8 else "weight_scale_inv" diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index a91c99bf53..65e119b083 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -77,8 +77,8 @@ def test_deepseek_v41_mxfp8_scale_loading( MergedColumnParallelLinear, RowParallelLinear, ) - from vllm.models.deepseek_v4_1 import quant_config as quant_module - from vllm.models.deepseek_v4_1.nvidia import model as model_module + from vllm.models.deepseek_v41 import quant_config as quant_module + from vllm.models.deepseek_v41.nvidia import model as model_module default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16) quant_config = quant_module.DeepseekV4FP8Config( @@ -214,8 +214,8 @@ def test_deepseek_v41_vl_mapper_routes_linear_scales( """The VL wrapper must map ``.scale`` keys to the parameter the linear quant method registers, as the text model does. A hardcoded ``weight_scale_inv`` raised KeyError for native MXFP8 checkpoints.""" - from vllm.models.deepseek_v4_1.nvidia import model as model_module - from vllm.models.deepseek_v4_1.nvidia import vl_model as vl_module + from vllm.models.deepseek_v41.nvidia import model as model_module + from vllm.models.deepseek_v41.nvidia import vl_model as vl_module vllm_config = SimpleNamespace( quant_config=SimpleNamespace(weight_block_size=weight_block_size) @@ -318,8 +318,8 @@ def test_mxfp8_bmm_loads_and_projects_grouped_weights( compute_fp8_einsum_recipe, deep_gemm_fp8_o_proj, ) - from vllm.models.deepseek_v4_1.nvidia.model import DeepseekV4Model - from vllm.models.deepseek_v4_1.quant_config import DeepseekV4FP8Config + from vllm.models.deepseek_v41.nvidia.model import DeepseekV4Model + from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config from vllm.utils.deep_gemm import is_deep_gemm_supported if not is_deep_gemm_supported() or not current_platform.is_device_capability_family( diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 49462ea8fb..87df1e409c 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -8,7 +8,7 @@ import torch from tests.v1.attention.utils import create_vllm_config from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend -from vllm.models.deepseek_v4_1.sparse_mla import ( +from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend, ) from vllm.v1.attention.backend import CommonAttentionMetadata diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index fd1d047074..57dc237131 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -120,7 +120,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: # The v4.1 class is the v4 one extended to accept model_type # "deepseek_v41" and its 32x32 MXFP8 linear layout. V4.1 has # platform-specific implementations for both CUDA and ROCm. - from vllm.models.deepseek_v4_1 import ( + from vllm.models.deepseek_v41 import ( DeepseekV4FP8Config as DeepseekV41FP8Config, ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 69dd88f349..1a5f8399b6 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -377,7 +377,7 @@ _MULTIMODAL_MODELS = { "DeepseekV4ForConditionalGeneration", ), "DeepseekV41ForCausalLM": ( - "vllm.models.deepseek_v4_1", + "vllm.models.deepseek_v41", "DeepseekV41ForCausalLM", ), "Dots3NoteForCausalLM": ( @@ -645,7 +645,7 @@ _SPECULATIVE_DECODING_MODELS = { "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), "DSparkV41DraftModel": ( - "vllm.models.deepseek_v4_1", + "vllm.models.deepseek_v41", "DSparkDeepseekV4ForCausalLM", ), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), diff --git a/vllm/models/deepseek_v4_1/__init__.py b/vllm/models/deepseek_v41/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/__init__.py rename to vllm/models/deepseek_v41/__init__.py diff --git a/vllm/models/deepseek_v4_1/amd/__init__.py b/vllm/models/deepseek_v41/amd/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/amd/__init__.py rename to vllm/models/deepseek_v41/amd/__init__.py diff --git a/vllm/models/deepseek_v4_1/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py similarity index 100% rename from vllm/models/deepseek_v4_1/amd/dspark.py rename to vllm/models/deepseek_v41/amd/dspark.py diff --git a/vllm/models/deepseek_v4_1/amd/model.py b/vllm/models/deepseek_v41/amd/model.py similarity index 99% rename from vllm/models/deepseek_v4_1/amd/model.py rename to vllm/models/deepseek_v41/amd/model.py index 898d3ab7a0..eaed15214f 100644 --- a/vllm/models/deepseek_v4_1/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -54,8 +54,8 @@ from vllm.models.common.ops.sequence_parallel import ( from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) -from vllm.models.deepseek_v4_1.amd.rocm import DeepseekV41ROCMAiterMLAAttention -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention +from vllm.models.deepseek_v41.attention import DeepseekV4Attention from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum diff --git a/vllm/models/deepseek_v4_1/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py similarity index 99% rename from vllm/models/deepseek_v4_1/amd/rocm.py rename to vllm/models/deepseek_v41/amd/rocm.py index 9271e7ae5a..802a59e340 100644 --- a/vllm/models/deepseek_v4_1/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -13,9 +13,9 @@ from vllm.distributed import ( ) from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention -from vllm.models.deepseek_v4_1.common.ops import dequantize_and_gather_k_cache -from vllm.models.deepseek_v4_1.sparse_mla import ( +from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache +from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, DeepseekV4SparseMLAMetadataBuilder, diff --git a/vllm/models/deepseek_v4_1/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py similarity index 100% rename from vllm/models/deepseek_v4_1/amd/vl_model.py rename to vllm/models/deepseek_v41/amd/vl_model.py diff --git a/vllm/models/deepseek_v4_1/attention.py b/vllm/models/deepseek_v41/attention.py similarity index 98% rename from vllm/models/deepseek_v4_1/attention.py rename to vllm/models/deepseek_v41/attention.py index 9bb53526c9..131cf4d97e 100644 --- a/vllm/models/deepseek_v4_1/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -26,7 +26,7 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.models.common.ops import fused_q_kv_rmsnorm -from vllm.models.deepseek_v4_1.common.ops import ( +from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant, indexer_k_norm_rope_store, @@ -49,8 +49,8 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.utils import extract_layer_index -from vllm.models.deepseek_v4_1.common.rope import build_deepseek_v4_rope -from vllm.models.deepseek_v4_1.compressor import DeepseekCompressor +from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope +from vllm.models.deepseek_v41.compressor import DeepseekCompressor from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, @@ -156,7 +156,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ``DeepseekV4FlashInferSM120Attention`` / ``DeepseekV4FlashInferMLAAttention`` (CUDA) or ``DeepseekV41ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific - deepseek_v4_1 model module. The base is never instantiated directly. + deepseek_v41 model module. The base is never instantiated directly. """ # Provided by the platform subclass. @@ -536,7 +536,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): "FLASHMLA_SPARSE_DSV41", "ROCM_FLASHMLA_SPARSE_DSV4", ): - from vllm.models.deepseek_v4_1.common.ops.cache_utils import ( + from vllm.models.deepseek_v41.common.ops.cache_utils import ( _COMBINE_TOPK_SWA_INDICES_KERNEL, ) @@ -581,7 +581,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): @cached_property def _can_fuse_query_quant(self) -> bool: - from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + from vllm.models.deepseek_v41.common.ops.query_quant import ( can_fuse_query_quant, ) @@ -600,7 +600,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): """ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) if self.q_lora_rank % 32 == 0 and self._can_fuse_query_quant: - from vllm.models.deepseek_v4_1.common.ops.query_quant import ( + from vllm.models.deepseek_v41.common.ops.query_quant import ( fused_q_kv_rmsnorm_quant, ) diff --git a/vllm/models/deepseek_v4_1/common/__init__.py b/vllm/models/deepseek_v41/common/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/__init__.py rename to vllm/models/deepseek_v41/common/__init__.py diff --git a/vllm/models/deepseek_v4_1/common/engram.py b/vllm/models/deepseek_v41/common/engram.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/engram.py rename to vllm/models/deepseek_v41/common/engram.py diff --git a/vllm/models/deepseek_v4_1/common/mm_preprocess.py b/vllm/models/deepseek_v41/common/mm_preprocess.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/mm_preprocess.py rename to vllm/models/deepseek_v41/common/mm_preprocess.py diff --git a/vllm/models/deepseek_v4_1/common/ops/__init__.py b/vllm/models/deepseek_v41/common/ops/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/__init__.py rename to vllm/models/deepseek_v41/common/ops/__init__.py diff --git a/vllm/models/deepseek_v4_1/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/cache_utils.py rename to vllm/models/deepseek_v41/common/ops/cache_utils.py diff --git a/vllm/models/deepseek_v4_1/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/fused_compress_quant_cache.py rename to vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py diff --git a/vllm/models/deepseek_v4_1/common/ops/indexer_k_store.py b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/indexer_k_store.py rename to vllm/models/deepseek_v41/common/ops/indexer_k_store.py diff --git a/vllm/models/deepseek_v4_1/common/ops/query_quant.py b/vllm/models/deepseek_v41/common/ops/query_quant.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/ops/query_quant.py rename to vllm/models/deepseek_v41/common/ops/query_quant.py diff --git a/vllm/models/deepseek_v4_1/common/rope.py b/vllm/models/deepseek_v41/common/rope.py similarity index 100% rename from vllm/models/deepseek_v4_1/common/rope.py rename to vllm/models/deepseek_v41/common/rope.py diff --git a/vllm/models/deepseek_v4_1/compressor.py b/vllm/models/deepseek_v41/compressor.py similarity index 99% rename from vllm/models/deepseek_v4_1/compressor.py rename to vllm/models/deepseek_v41/compressor.py index d09d97d04c..44440fcb61 100644 --- a/vllm/models/deepseek_v4_1/compressor.py +++ b/vllm/models/deepseek_v41/compressor.py @@ -12,7 +12,7 @@ from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear -from vllm.models.deepseek_v4_1.common.ops.fused_compress_quant_cache import ( +from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( fused_save_compress_norm, rope_quant_insert, ) diff --git a/vllm/models/deepseek_v4_1/nvidia/__init__.py b/vllm/models/deepseek_v41/nvidia/__init__.py similarity index 100% rename from vllm/models/deepseek_v4_1/nvidia/__init__.py rename to vllm/models/deepseek_v41/nvidia/__init__.py diff --git a/vllm/models/deepseek_v4_1/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py similarity index 100% rename from vllm/models/deepseek_v4_1/nvidia/dspark.py rename to vllm/models/deepseek_v41/nvidia/dspark.py diff --git a/vllm/models/deepseek_v4_1/nvidia/engram.py b/vllm/models/deepseek_v41/nvidia/engram.py similarity index 99% rename from vllm/models/deepseek_v4_1/nvidia/engram.py rename to vllm/models/deepseek_v41/nvidia/engram.py index b141b6c349..624a36c214 100644 --- a/vllm/models/deepseek_v4_1/nvidia/engram.py +++ b/vllm/models/deepseek_v41/nvidia/engram.py @@ -23,16 +23,16 @@ from vllm.distributed.parallel_state import GroupCoordinator from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.utils import set_weight_attrs -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common.engram import ( DEAD_ID, EngramLayout, _engram_head_shard_weight_loader, _engram_select_rows, ) -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common.engram import ( Engram as BaseEngram, ) -from vllm.models.deepseek_v4_1.common.engram import ( +from vllm.models.deepseek_v41.common.engram import ( ParallelEngramEmbedding as BaseParallelEngramEmbedding, ) from vllm.utils.platform_utils import is_uva_available diff --git a/vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py similarity index 99% rename from vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py rename to vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 802ddede9c..7bed50bde7 100644 --- a/vllm/models/deepseek_v4_1/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -13,12 +13,12 @@ from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( compute_fp8_einsum_recipe, deep_gemm_fp8_o_proj, ) -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention -from vllm.models.deepseek_v4_1.common.ops import ( +from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.common.ops import ( build_flashinfer_mixed_sparse_indices, compute_global_topk_indices_and_lens, ) -from vllm.models.deepseek_v4_1.sparse_mla import ( +from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLAMetadata, DeepseekV4SparseMLABackend, DeepseekV4SparseMLAMetadataBuilder, diff --git a/vllm/models/deepseek_v4_1/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py similarity index 98% rename from vllm/models/deepseek_v4_1/nvidia/flashmla.py rename to vllm/models/deepseek_v41/nvidia/flashmla.py index 86d8fb2031..ec834aa80d 100644 --- a/vllm/models/deepseek_v4_1/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -10,13 +10,13 @@ from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( compute_fp8_einsum_recipe, deep_gemm_fp8_o_proj, ) -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention -from vllm.models.deepseek_v4_1.common.ops import ( +from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.common.ops import ( combine_topk_swa_indices, compute_global_topk_indices_and_lens, dequantize_and_gather_k_cache, ) -from vllm.models.deepseek_v4_1.sparse_mla import ( +from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4FlashMLABackend, DeepseekV4FlashMLAMetadata, DeepseekV41SparseSWAMetadataBuilder, diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py similarity index 99% rename from vllm/models/deepseek_v4_1/nvidia/model.py rename to vllm/models/deepseek_v41/nvidia/model.py index 0868089da6..df12a9d9bb 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -63,12 +63,12 @@ from vllm.models.deepseek_v4.nvidia.model import ( from vllm.models.deepseek_v4.nvidia.model import ( make_deepseek_v4_expert_params_mapping, ) -from vllm.models.deepseek_v4_1.attention import DeepseekV4Attention -from vllm.models.deepseek_v4_1.nvidia.flashinfer_sparse import ( +from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, ) -from vllm.models.deepseek_v4_1.nvidia.flashmla import DeepseekV4FlashMLAAttention +from vllm.models.deepseek_v41.nvidia.flashmla import DeepseekV4FlashMLAAttention from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv diff --git a/vllm/models/deepseek_v4_1/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py similarity index 100% rename from vllm/models/deepseek_v4_1/nvidia/model_state.py rename to vllm/models/deepseek_v41/nvidia/model_state.py diff --git a/vllm/models/deepseek_v4_1/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py similarity index 100% rename from vllm/models/deepseek_v4_1/nvidia/vl_model.py rename to vllm/models/deepseek_v41/nvidia/vl_model.py diff --git a/vllm/models/deepseek_v4_1/quant_config.py b/vllm/models/deepseek_v41/quant_config.py similarity index 100% rename from vllm/models/deepseek_v4_1/quant_config.py rename to vllm/models/deepseek_v41/quant_config.py diff --git a/vllm/models/deepseek_v4_1/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py similarity index 100% rename from vllm/models/deepseek_v4_1/sparse_mla.py rename to vllm/models/deepseek_v41/sparse_mla.py diff --git a/vllm/transformers_utils/configs/deepseek_v41.py b/vllm/transformers_utils/configs/deepseek_v41.py index 0be36d2480..8f33245cd6 100644 --- a/vllm/transformers_utils/configs/deepseek_v41.py +++ b/vllm/transformers_utils/configs/deepseek_v41.py @@ -12,7 +12,7 @@ class DeepseekV41Config(PretrainedConfig): tower under ``vision_config``. The vLLM model code (ported from ``deepseek_v4``) reads flat attributes, so both sub-configs are flattened onto the top level here: text fields are exposed as-is, vision fields - with the ``vision_*`` naming used by ``deepseek_v4_1.common.vision``. + with the ``vision_*`` naming used by ``deepseek_v41.common.vision``. """ model_type = "deepseek_v41" diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index f47554d81f..62c3921b5e 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -49,7 +49,7 @@ _LAYER_TYPE_C128A = "c128a" # v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but # their compressed page block is block_size // ratio, so each gets its own plan. # v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see -# deepseek_v4_1/sparse_mla.py. +# deepseek_v41/sparse_mla.py. _LAYER_TYPE_C1A = "c1a" _LAYER_TYPE_C2A = "c2a" diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 085f2358fa..571566db1f 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -108,10 +108,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): # layer). Separate names from DSV4 so a V4.1 model never resolves the # V4.0 backend classes through this enum. FLASHMLA_SPARSE_DSV41 = ( - "vllm.models.deepseek_v4_1.sparse_mla.DeepseekV4FlashMLABackend" + "vllm.models.deepseek_v41.sparse_mla.DeepseekV4FlashMLABackend" ) FLASHINFER_MLA_SPARSE_DSV41 = ( - "vllm.models.deepseek_v4_1.nvidia.flashinfer_sparse." + "vllm.models.deepseek_v41.nvidia.flashinfer_sparse." "DeepseekV4FlashInferMLASparseBackend" ) B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend" From dbf49dad11e30ec5d464be74b715a9a9e31c2c5a Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Mon, 14 Sep 2026 14:28:13 +0800 Subject: [PATCH 004/139] [Fast Start] Use GPU uuid as socket folder identifier (#56669) Signed-off-by: Isotr0py --- .../model_loader/test_weight_cache.py | 8 ++-- .../model_loader/weight_cache/daemon.py | 26 ++++--------- .../model_loader/weight_cache/ipc_loader.py | 27 ++++--------- .../model_loader/weight_cache/protocol.py | 38 ++++++------------- 4 files changed, 30 insertions(+), 69 deletions(-) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index fa0b6b4183..1195b8b09d 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -7,6 +7,7 @@ warm restarts (weights mapped from the daemon via CUDA IPC) must both serve identical outputs. """ +import glob import os import shutil import subprocess @@ -21,7 +22,6 @@ import pytest from vllm import SamplingParams from vllm.assets.image import ImageAsset -from vllm.model_executor.model_loader.weight_cache.protocol import get_socket_path from vllm.platforms import current_platform DAEMON_TIMEOUT_S = 600 @@ -79,9 +79,7 @@ class WeightCacheDaemon: # Poll for the socket files rather than a log line: model loading can # pull in JIT compilers that swap the process's stderr and swallow # everything logged afterwards, making log-based readiness flaky. - expected = [ - get_socket_path(gpu_id, self.socket_dir) for gpu_id in range(self.tp_size) - ] + pattern = os.path.join(self.socket_dir, "vllm_weight_cache_*.sock") deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: if self._proc.poll() is not None: @@ -89,7 +87,7 @@ class WeightCacheDaemon: f"Weight cache daemon exited with {self._proc.returncode}:\n" f"{self._logs()}" ) - if all(os.path.exists(path) for path in expected): + if len(glob.glob(pattern)) >= self.tp_size: return time.sleep(1.0) raise TimeoutError( diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 96cedbe1e9..b19fb96272 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -60,6 +60,7 @@ from vllm.distributed import ( ensure_model_parallel_initialized, init_distributed_environment, ) +from vllm.engine.arg_utils import EngineArgs from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model_loader from vllm.model_executor.model_loader.utils import process_weights_after_loading @@ -70,13 +71,15 @@ from vllm.model_executor.model_loader.weight_cache.protocol import ( check_ipc_platform_support, check_ipc_quant_support, ensure_private_socket_dir, - get_physical_device_id, + get_current_device_uuid, get_socket_path, recv_msg, send_msg, verify_peer_is_owner, ) from vllm.platforms import current_platform +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.network_utils import get_distributed_init_method, get_open_port from vllm.utils.torch_utils import set_default_torch_dtype logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon") @@ -209,7 +212,7 @@ class WeightCacheDaemon: ready_callback: Invoked once the socket is bound and listening, so the launcher can report overall readiness. """ - socket_path = self._socket_path() + socket_path = self._socket_path ensure_private_socket_dir( os.path.dirname(socket_path), strict_perms=self.socket_dir is None ) @@ -268,12 +271,9 @@ class WeightCacheDaemon: ) from e return lock_fd + @property def _socket_path(self) -> str: - device_index = torch.accelerator.current_device_index() - gpu_id = get_physical_device_id(device_index) - if gpu_id is None: - gpu_id = device_index - return get_socket_path(gpu_id, self.socket_dir) + return get_socket_path(get_current_device_uuid(), self.socket_dir) def _handle_connection(self, conn: socket.socket) -> None: request = recv_msg(conn) @@ -304,7 +304,7 @@ class WeightCacheDaemon: "status": "ok", "entries": self.entries, "aliases": self.aliases, - "gpu_uuid": self._gpu_uuid(), + "gpu_uuid": get_current_device_uuid(), }, ) logger.info_once( @@ -322,12 +322,6 @@ class WeightCacheDaemon: logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank) send_msg(conn, {"status": "ok"}) - def _gpu_uuid(self) -> str: - props = torch.cuda.get_device_properties( - torch.accelerator.current_device_index() - ) - return str(props.uuid) - def _run_daemon( tp_rank: int, @@ -359,10 +353,6 @@ def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None: def main() -> None: - from vllm.engine.arg_utils import EngineArgs - from vllm.utils.argparse_utils import FlexibleArgumentParser - from vllm.utils.network_utils import get_distributed_init_method, get_open_port - parser = FlexibleArgumentParser( description="Launch weight cache daemons (one per TP rank)." ) diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index be2e5ebae4..6c3fb742cd 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -12,6 +12,10 @@ import torch.nn as nn from vllm.config import ModelConfig, VllmConfig from vllm.config.load import LoadConfig +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) from vllm.logger import init_logger from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader @@ -27,7 +31,7 @@ from vllm.model_executor.model_loader.weight_cache.protocol import ( WeightCacheUnavailableError, check_ipc_platform_support, check_ipc_quant_support, - get_physical_device_id, + get_current_device_uuid, get_socket_path, recv_msg, send_msg, @@ -56,7 +60,7 @@ class IpcModelLoader(BaseModelLoader): Extra config keys (via --model-loader-extra-config): - socket_path: explicit daemon socket path. Defaults to a per-GPU path - derived from the physical GPU id. + derived from the physical GPU uuid. - socket_dir: directory containing the daemon sockets. - mode: "zero_copy" (default) or "copy". - fallback: fall back to disk loading when the daemon is unavailable or @@ -261,11 +265,6 @@ class IpcModelLoader(BaseModelLoader): def _fetch_entries( self, model_config: ModelConfig ) -> tuple[dict[str, TensorEntry], dict[str, str]]: - from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, - ) - cache_config = WeightCacheKey.from_model_config( model_config, tp_size=get_tensor_model_parallel_world_size(), @@ -317,22 +316,12 @@ class IpcModelLoader(BaseModelLoader): def _resolve_socket_path(self) -> str: if self.socket_path is not None: return self.socket_path - device_index = torch.accelerator.current_device_index() - gpu_id = get_physical_device_id(device_index) - if gpu_id is None: - raise WeightCacheUnavailableError( - "Cannot infer the physical GPU id from CUDA_VISIBLE_DEVICES; " - "pass socket_path via --model-loader-extra-config" - ) - return get_socket_path(gpu_id, self.socket_dir) + return get_socket_path(get_current_device_uuid(), self.socket_dir) def _check_gpu_uuid(self, daemon_uuid: str | None) -> None: if daemon_uuid is None: return - props = torch.cuda.get_device_properties( - torch.accelerator.current_device_index() - ) - local_uuid = str(props.uuid) + local_uuid = get_current_device_uuid() if daemon_uuid != local_uuid: raise CacheConfigMismatchError( f"Daemon GPU {daemon_uuid} != engine GPU {local_uuid}; " diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index fff1fcca2a..00d3dbfa8b 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -23,14 +23,19 @@ from dataclasses import dataclass, fields from typing import Any import torch +from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME import vllm.version from vllm.config import ModelConfig from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase +from vllm.model_executor.model_loader.weight_utils import ( + filter_duplicate_safetensors_files, +) from vllm.platforms import current_platform from vllm.utils.hashing import safe_hash -SOCKET_NAME_TEMPLATE = "vllm_weight_cache_gpu{gpu_id}.sock" +SOCKET_NAME_TEMPLATE = "vllm_weight_cache_{gpu_uuid}.sock" SOCKET_DIR_TEMPLATE = "vllm_weight_cache_{uid}" _LEN_STRUCT = struct.Struct("!Q") @@ -102,20 +107,9 @@ def check_ipc_quant_support(model: torch.nn.Module) -> None: ) -def get_physical_device_id(device_index: int) -> int | None: - """Map a local CUDA device index to the physical GPU id. - - Returns None if CUDA_VISIBLE_DEVICES contains non-integer entries - (e.g. GPU UUIDs), in which case an explicit socket path is required. - """ - visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if not visible: - return device_index - entries = visible.split(",") - try: - return int(entries[device_index]) - except (IndexError, ValueError): - return None +def get_current_device_uuid() -> str: + """UUID of the physical GPU backing the current accelerator device.""" + return current_platform.get_device_uuid(torch.accelerator.current_device_index()) def get_socket_dir(socket_dir: str | None = None) -> str: @@ -132,9 +126,9 @@ def get_socket_dir(socket_dir: str | None = None) -> str: ) -def get_socket_path(gpu_id: int, socket_dir: str | None = None) -> str: +def get_socket_path(gpu_uuid: str, socket_dir: str | None = None) -> str: directory = get_socket_dir(socket_dir) - return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_id=gpu_id)) + return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_uuid=gpu_uuid)) def ensure_private_socket_dir(directory: str, strict_perms: bool = True) -> None: @@ -245,12 +239,6 @@ def hash_checkpoint(model: str) -> str | None: """ if not os.path.isdir(model): return None - from transformers.utils import SAFE_WEIGHTS_INDEX_NAME - - from vllm.model_executor.model_loader.weight_utils import ( - filter_duplicate_safetensors_files, - ) - files = glob.glob(os.path.join(model, "*.safetensors")) if os.path.isfile(os.path.join(model, SAFE_WEIGHTS_INDEX_NAME)): files = filter_duplicate_safetensors_files( @@ -337,8 +325,6 @@ class TensorEntry: @classmethod def from_tensor(cls, tensor: torch.Tensor, kind: str) -> "TensorEntry": - from torch.multiprocessing.reductions import reduce_tensor - tensor = tensor.detach() if tensor.is_cuda: _, ipc_args = reduce_tensor(tensor) @@ -349,8 +335,6 @@ class TensorEntry: if self.ipc_args is None: assert self.cpu_tensor is not None return self.cpu_tensor - from torch.multiprocessing.reductions import rebuild_cuda_tensor - args = list(self.ipc_args) # Index 6 of the args from reduce_tensor is the device index. It must # be retargeted to the local index since the daemon and the engine may From 663d7f679eda78a8d48613dfede8a4b4bcff2b74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hung=20Hoang=20=28=E9=BB=84=E5=85=B4=EF=BC=89?= Date: Mon, 14 Sep 2026 13:55:21 +0700 Subject: [PATCH 005/139] [Bugfix] Fix --lora-modules name=path parsing when path contains '=' (#53353) Signed-off-by: hungh --- vllm/entrypoints/launchers/cli_args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index c0dd64164e..0da96b43e1 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -47,7 +47,7 @@ class LoRAParserAction(argparse.Action): if item in [None, ""]: # Skip if item is None or empty string continue if "=" in item and "," not in item: # Old format: name=path - name, path = item.split("=") + name, path = item.split("=", 1) lora_list.append(LoRAModulePath(name, path)) else: # Assume JSON format try: From e6b4e47d2d894a1a0815fbf03669adb1b3405083 Mon Sep 17 00:00:00 2001 From: Alec <35311602+alec-flowers@users.noreply.github.com> Date: Mon, 14 Sep 2026 00:05:33 -0700 Subject: [PATCH 006/139] [Bugfix][Rust Frontend] Preserve selected-token logprob mode (#56406) Co-authored-by: OpenAI Codex Signed-off-by: Alec Flowers --- rust/src/server/src/grpc/convert.rs | 98 ++++++++++++++++++++++++----- 1 file changed, 84 insertions(+), 14 deletions(-) diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 61eb21ef70..aea04495eb 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -193,7 +193,7 @@ fn build_sampling_params( if let Some(r) = response { if r.output_logprobs { let (count, token_ids) = candidate_logprob_spec(r.output_candidates.as_ref()); - params.logprobs = Some(count); + params.logprobs = count; params.logprob_token_ids = token_ids; } if r.prompt_logprobs { @@ -210,7 +210,7 @@ fn build_sampling_params( )); } let (count, _) = candidate_logprob_spec(r.prompt_candidates.as_ref()); - params.prompt_logprobs = Some(count); + params.prompt_logprobs = count; } } @@ -220,18 +220,22 @@ fn build_sampling_params( /// Map the proto `CandidateTokens` selector to a `(logprobs_count, /// logprob_token_ids)` pair. /// -/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability -/// - `all` → `(-1, None)` — return the full vocabulary -/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the -/// count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs -/// are carried via `logprob_token_ids` on `SamplingParams`) -/// - absent → `(1, None)` — just the sampled/scored token -fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option>) { +/// - `top_n(k)` → `(Some(k), None)` — return top-k candidates by probability +/// - `all` → `(Some(-1), None)` — return the full vocabulary +/// - nonempty `token_ids` → `(None, Some(ids))` — let the engine derive the count +/// - absent or empty `token_ids` → `(Some(0), None)` — return only the sampled/scored token +fn candidate_logprob_spec( + candidates: Option<&pb::CandidateTokens>, +) -> (Option, Option>) { match candidates.and_then(|c| c.select.as_ref()) { - Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None), - Some(pb::candidate_tokens::Select::All(true)) => (-1, None), - Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())), - _ => (1, None), + Some(pb::candidate_tokens::Select::TopN(n)) => (Some(*n as i32), None), + Some(pb::candidate_tokens::Select::All(true)) => (Some(-1), None), + Some(pb::candidate_tokens::Select::TokenIds(ids)) if !ids.ids.is_empty() => { + // Match Python HTTP: selected IDs have their own limit, independent + // of the numeric top-logprobs count and its max_logprobs cap. + (None, Some(ids.ids.clone())) + } + _ => (Some(0), None), } } @@ -518,7 +522,10 @@ impl ResponseOpts { #[cfg(test)] mod tests { use vllm_engine_core_client::protocol::output::StopReason; - use vllm_text::{FinishReason, Finished, Prompt}; + use vllm_text::{ + FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params, + }; + use vllm_tokenizer::test_utils::TestTokenizer; use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason}; use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request}; @@ -590,6 +597,69 @@ mod tests { assert_eq!(text.sampling_params.seed, Some(0)); } + #[test] + fn output_logprob_selectors_survive_request_lowering() { + use pb::candidate_tokens::Select; + + let selected_ids: Vec = (100..121).collect(); + let cases = [ + ( + Some(Select::TokenIds(pb::TokenIds { + ids: selected_ids.clone(), + })), + None, + Some(selected_ids), + ), + ( + Some(Select::TokenIds(pb::TokenIds { + ids: vec![198, 198], + })), + None, + Some(vec![198, 198]), + ), + ( + Some(Select::TokenIds(pb::TokenIds { ids: vec![] })), + Some(0), + None, + ), + (None, Some(0), None), + (Some(Select::TopN(0)), Some(0), None), + (Some(Select::TopN(2)), Some(2), None), + ]; + for (select, expected_count, expected_ids) in cases { + let req = pb::GenerateRequest { + response: Some(pb::ResponseOptions { + output_logprobs: true, + output_candidates: select.map(|select| pb::CandidateTokens { + select: Some(select), + }), + ..Default::default() + }), + ..base_request() + }; + let text = + to_text_request(req, false, &["test-model".to_string()]).expect("convert ok"); + let params = lower_sampling_params( + text.sampling_params, + SamplingHints::default(), + SamplingLimits { + max_model_len: 32, + max_logprobs: 20, + model_vocab_size: 512, + tokenizer_vocab_size: 512, + }, + 1, + &TestTokenizer::new(), + ) + .expect("logprob selector should lower to an engine request"); + + assert_eq!( + (params.logprobs, params.logprob_token_ids), + (expected_count, expected_ids) + ); + } + } + #[test] fn bypass_prefix_cache_maps_to_skip_reading_prefix_cache() { let req = pb::GenerateRequest { From 78e84261ab2ffb691a8079d698d1f3f96ecb7635 Mon Sep 17 00:00:00 2001 From: JooHo Lee <96564470+BWAAEEEK@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:42:15 +0900 Subject: [PATCH 007/139] [Performance] Add Triton kernel for Gemma3n sparse GELU (#48498) Signed-off-by: BWAAEEEK Co-authored-by: Andreas Karatzas Co-authored-by: Cyrus Leung --- .../kernels/benchmark_gelu_and_mul_sparse.py | 90 +++++++++ tests/compile/passes/ir/test_lowering.py | 34 ++++ tests/compile/test_aot_compile.py | 65 +++++++ tests/kernels/ir/test_activation.py | 183 ++++++++++++++++++ tests/test_config.py | 16 ++ vllm/config/kernel.py | 3 + vllm/ir/ops/__init__.py | 3 +- vllm/ir/ops/activation.py | 34 ++++ vllm/kernels/__init__.py | 4 +- vllm/kernels/triton/__init__.py | 4 + vllm/kernels/triton/activation.py | 110 +++++++++++ vllm/model_executor/layers/activation.py | 18 +- vllm/platforms/cuda.py | 5 +- vllm/platforms/rocm.py | 5 +- vllm/platforms/xpu.py | 2 +- 15 files changed, 555 insertions(+), 21 deletions(-) create mode 100644 benchmarks/kernels/benchmark_gelu_and_mul_sparse.py create mode 100644 tests/kernels/ir/test_activation.py create mode 100644 vllm/ir/ops/activation.py create mode 100644 vllm/kernels/triton/activation.py diff --git a/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py b/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py new file mode 100644 index 0000000000..9f8d976121 --- /dev/null +++ b/benchmarks/kernels/benchmark_gelu_and_mul_sparse.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import itertools + +import torch + +import vllm.kernels # noqa: F401 +from vllm import ir +from vllm.benchmarks.lib.utils import default_vllm_config +from vllm.triton_utils import triton +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed + +NUM_TOKENS = [1, 2, 8, 32, 128, 512, 2048] +INTERMEDIATE_SIZES = [2048, 4096, 8192, 13824, 16384, 32768] +CONFIGS = list(itertools.product(NUM_TOKENS, INTERMEDIATE_SIZES)) +STD_MULTIPLIER = 1.6448536269514722 + + +@default_vllm_config() +def benchmark_gelu_and_mul_sparse( + num_tokens: int, + intermediate_size: int, + provider: str, + dtype: torch.dtype, +): + set_random_seed(0) + torch.set_default_device("cuda") + x = torch.randn(num_tokens, 2 * intermediate_size, dtype=dtype) + native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn + triton_impl = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn + args = (x, STD_MULTIPLIER, "tanh") + + expected = native(*args) + actual = triton_impl(*args) + torch.testing.assert_close( + actual, + expected, + **ir.ops.gelu_and_mul_sparse.get_tolerance(actual.dtype), + ) + + if provider == "triton": + fn = lambda: triton_impl(*args) + elif provider == "native_compiled": + compiled = torch.compile(native, fullgraph=True) + compiled(*args) + fn = lambda: compiled(*args) + else: + fn = lambda: native(*args) + + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + fn, quantiles=[0.5, 0.2, 0.8] + ) + return ms, max_ms, min_ms + + +if __name__ == "__main__": + parser = FlexibleArgumentParser( + description="Benchmark the Gemma3n sparse GELU activation." + ) + parser.add_argument( + "--dtype", + type=str, + choices=["half", "bfloat16", "float"], + default="bfloat16", + ) + args = parser.parse_args() + dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype] + + perf_report = triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens", "intermediate_size"], + x_vals=CONFIGS, + line_arg="provider", + line_vals=["triton", "native_compiled", "native_eager"], + line_names=["Triton", "Native (torch.compile)", "Native (eager)"], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="ms", + plot_name="gelu-and-mul-sparse-performance", + args={}, + ) + ) + perf_report( + lambda num_tokens, intermediate_size, provider: ( + benchmark_gelu_and_mul_sparse( + num_tokens, intermediate_size, provider, dtype + ) + ) + ).run(print_data=True) diff --git a/tests/compile/passes/ir/test_lowering.py b/tests/compile/passes/ir/test_lowering.py index b7ca55e7d1..0ac605e796 100644 --- a/tests/compile/passes/ir/test_lowering.py +++ b/tests/compile/passes/ir/test_lowering.py @@ -11,6 +11,7 @@ from vllm.compilation.passes.ir.lowering_pass import ( ) from vllm.config import get_current_vllm_config from vllm.ir import ops +from vllm.model_executor.layers.activation import GeluAndMulSparse from vllm.platforms import current_platform from ...backend import TestBackend @@ -34,6 +35,15 @@ class Model(nn.Module): return x6 + 3.0 +class SparseActivationModel(nn.Module): + def __init__(self): + super().__init__() + self.activation = GeluAndMulSparse(0.95, "tanh") + + def forward(self, x): + return self.activation(x) + + @pytest.mark.parametrize("rms_provider", ops.rms_norm.supported_providers()) def test_lowering_rms_norm(rms_provider, default_vllm_config): torch.set_default_device(current_platform.device_type) @@ -67,3 +77,27 @@ def test_lowering_rms_norm(rms_provider, default_vllm_config): torch.testing.assert_close(output_unlowered, output) torch.testing.assert_close(output_unlowered, output2) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_lowering_gelu_and_mul_sparse(default_vllm_config): + torch.set_default_device("cuda") + lowering_pass = VllmIRLoweringPass(get_current_vllm_config()) + backend = TestBackend(lowering_pass) + model = SparseActivationModel() + x = torch.randn(8, 32768, dtype=torch.bfloat16) + + with ( + ops.gelu_and_mul_sparse.set_priority(["triton", "native"]), + ir.enable_torch_wrap(True), + ): + compiled_model = torch.compile(model, backend=backend, fullgraph=True) + output = compiled_model(x) + + selected = lowering_pass.selected_impls["gelu_and_mul_sparse"] + assert selected == {"gelu_and_mul_sparse": "triton"} + + expected = ops.gelu_and_mul_sparse.impls["native"].impl_fn( + x, 1.6448536269514722, "tanh" + ) + torch.testing.assert_close(output, expected, rtol=1.6e-2, atol=1e-3) diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py index b3a71b4c30..90c132bee1 100644 --- a/tests/compile/test_aot_compile.py +++ b/tests/compile/test_aot_compile.py @@ -13,6 +13,8 @@ import pytest import torch import vllm.envs as envs +import vllm.kernels # noqa: F401 to register kernels +from vllm import ir from vllm.compilation.backends import VllmBackend from vllm.compilation.caching import ( StandaloneCompiledArtifacts, @@ -28,6 +30,9 @@ from vllm.config import ( ) from vllm.envs import disable_envs_cache from vllm.forward_context import set_forward_context +from vllm.ir import ops +from vllm.model_executor.layers.activation import GeluAndMulSparse +from vllm.platforms import current_platform from vllm.utils.torch_utils import is_torch_equal_or_newer from ..utils import create_new_process_for_each_test @@ -77,6 +82,18 @@ class CompiledModTuple(torch.nn.Module): return reference_fn_tuple(x) +@support_torch_compile(dynamic_arg_dims={"x": 0}) +class SparseActivationAotModel(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.up_proj = torch.nn.Linear(32, 256, bias=False) + self.activation = GeluAndMulSparse(0.95, "tanh") + self.down_proj = torch.nn.Linear(128, 16, bias=False) + + def forward(self, x: torch.Tensor): + return self.down_proj(self.activation(self.up_proj(x))) + + def make_vllm_config() -> VllmConfig: return VllmConfig( compilation_config=CompilationConfig( @@ -175,6 +192,54 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch): assert torch.allclose(ret, expected) +@pytest.mark.skipif( + not current_platform.is_cuda() or not is_torch_equal_or_newer("2.10.0"), + reason="requires CUDA and torch 2.10", +) +def test_sparse_activation_standalone_aot_preserves_dependencies( + monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path +): + def native_reference(model: SparseActivationAotModel, x: torch.Tensor): + projected = torch.nn.functional.linear(x, model.up_proj.weight) + activated = ops.gelu_and_mul_sparse.impls["native"].impl_fn( + projected, model.activation.std_multiplier, "tanh" + ) + return torch.nn.functional.linear(activated, model.down_proj.weight) + + with monkeypatch.context() as m: + m.setenv("VLLM_USE_AOT_COMPILE", "1") + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1") + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1") + m.setenv("VLLM_CACHE_ROOT", str(vllm_tmp_cache / "vllm_cache")) + + vllm_config = make_vllm_config() + torch.manual_seed(0) + inputs = [ + torch.randn(4, 32, device="cuda", dtype=torch.bfloat16), + torch.randn(7, 32, device="cuda", dtype=torch.bfloat16), + ] + + with ( + use_vllm_config(vllm_config), + vllm_config.kernel_config.ir_op_priority.set_priority(), + ir.enable_torch_wrap(True), + ): + compiled = SparseActivationAotModel(vllm_config=vllm_config).to( + device="cuda", dtype=torch.bfloat16 + ) + compiled.requires_grad_(False) + expected = [native_reference(compiled, x) for x in inputs] + actual = [compiled(x) for x in inputs] + + for result, reference in zip(actual, expected): + assert torch.isfinite(result).all() + torch.testing.assert_close( + result, + reference, + **ops.gelu_and_mul_sparse.get_tolerance(result.dtype), + ) + + @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch): from torch._subclasses import FakeTensorMode diff --git a/tests/kernels/ir/test_activation.py b/tests/kernels/ir/test_activation.py new file mode 100644 index 0000000000..a36ff8ee33 --- /dev/null +++ b/tests/kernels/ir/test_activation.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm.kernels # noqa: F401 +from tests.ir.ir_test_utils import assert_close, clone_args +from vllm import ir +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON + +gelu_and_mul_sparse_native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn + + +def test_gelu_and_mul_sparse_registration(): + expected = { + "native": True, + "triton": HAS_TRITON and current_platform.is_cuda(), + } + actual = { + provider: impl.supported + for provider, impl in ir.ops.gelu_and_mul_sparse.impls.items() + } + assert actual == expected + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("num_tokens", [1, 8, 64]) +@pytest.mark.parametrize("intermediate_size", [2048, 4096, 8192, 13824, 16384, 32768]) +def test_gelu_and_mul_sparse_triton( + dtype: torch.dtype, num_tokens: int, intermediate_size: int +): + torch.manual_seed(0) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=num_tokens, + hidden_size=intermediate_size, + dtype=dtype, + ) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + assert impl.supports_args(*args) + + expected = gelu_and_mul_sparse_native(*clone_args(args)) + actual = impl.impl_fn(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + mask_mismatch = torch.count_nonzero((actual == 0) != (expected == 0)) + assert mask_mismatch / actual.numel() < 1e-3 + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_second_seed(): + torch.manual_seed(1) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=32, + hidden_size=8192, + dtype=torch.bfloat16, + ) + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args)) + expected = gelu_and_mul_sparse_native(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("intermediate_size", [1, 31, 128]) +def test_gelu_and_mul_sparse_triton_small_sizes( + dtype: torch.dtype, intermediate_size: int +): + torch.manual_seed(0) + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=5, + hidden_size=intermediate_size, + dtype=dtype, + ) + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args)) + expected = gelu_and_mul_sparse_native(*clone_args(args)) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + assert torch.equal(actual == 0, expected == 0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_nd_and_empty(): + torch.set_default_device("cuda") + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + x = torch.randn(2, 3, 16384, dtype=torch.bfloat16) + args = (x, 1.6448536269514722, "tanh") + actual = impl.impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + assert actual.shape == (2, 3, 8192) + assert_close(ir.ops.gelu_and_mul_sparse, actual, expected) + + empty = torch.empty(0, 16384, dtype=torch.bfloat16) + empty_out = impl.impl_fn(empty, 1.6448536269514722, "tanh") + assert empty_out.shape == (0, 8192) + assert empty_out.dtype == empty.dtype + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_gelu_and_mul_sparse_triton_constant_rows(dtype: torch.dtype): + torch.set_default_device("cuda") + gate = torch.full((4, 8192), 3.0, dtype=dtype) + up = torch.randn_like(gate) + x = torch.cat((gate, up), dim=-1) + args = (x, 1.6448536269514722, "tanh") + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + assert torch.count_nonzero(actual) == 0 + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")]) +def test_gelu_and_mul_sparse_triton_nonfinite_gate(nonfinite: float): + torch.set_default_device("cuda") + x = torch.randn(2, 16384, dtype=torch.bfloat16) + x[0, 0] = nonfinite + args = (x, 1.6448536269514722, "tanh") + + actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args) + expected = gelu_and_mul_sparse_native(*args) + torch.testing.assert_close(actual, expected, equal_nan=True) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_fallback_conditions(): + torch.set_default_device("cuda") + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + x = torch.randn(4, 16384, dtype=torch.bfloat16) + + assert impl.supports_args(x, 1.0, "tanh") + assert not impl.supports_args(x, 1.0, "none") + assert not impl.supports_args(x[:, ::2], 1.0, "tanh") + assert not impl.supports_args(x[:, :-1], 1.0, "tanh") + assert not impl.supports_args( + torch.randn(1, 65538, dtype=torch.bfloat16), 1.0, "tanh" + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_dispatch_and_opcheck(): + torch.set_default_device("cuda") + args = ir.ops.gelu_and_mul_sparse.generate_inputs( + num_tokens=8, + hidden_size=4096, + dtype=torch.bfloat16, + ) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + with ir.ops.gelu_and_mul_sparse.set_priority(["triton", "native"]): + dispatched = ir.ops.gelu_and_mul_sparse(*args) + torch.library.opcheck(torch.ops.vllm_ir.gelu_and_mul_sparse, args) + + direct = impl.impl_fn(*args) + torch.testing.assert_close(dispatched, direct, rtol=0.0, atol=0.0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider") +def test_gelu_and_mul_sparse_triton_cudagraph(): + torch.set_default_device("cuda") + x = torch.randn(16, 32768, dtype=torch.bfloat16) + impl = ir.ops.gelu_and_mul_sparse.impls["triton"] + + for _ in range(3): + expected = impl.impl_fn(x, 1.6448536269514722, "tanh") + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = impl.impl_fn(x, 1.6448536269514722, "tanh") + graph.replay() + torch.accelerator.synchronize() + + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) diff --git a/tests/test_config.py b/tests/test_config.py index c2629a75ab..4d39b56d26 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3322,6 +3322,22 @@ def test_ir_op_priority_default(): assert priority_config.fused_add_rms_norm == ["native"] +@pytest.mark.parametrize("mode", [CompilationMode.NONE, CompilationMode.VLLM_COMPILE]) +@pytest.mark.parametrize("backend", ["inductor", "eager"]) +def test_ir_op_platform_defaults_support_sparse_gelu(mode, backend): + """Worker initialization must not select an unregistered sparse GELU provider.""" + from vllm import ir + + config = SimpleNamespace( + compilation_config=CompilationConfig(mode=mode, backend=backend) + ) + priority = current_platform.get_default_ir_op_priority(config) + expected = ["triton", "native"] if current_platform.is_cuda() else ["native"] + + with priority.set_priority(): + assert ir.ops.gelu_and_mul_sparse.get_priority() == expected + + def test_ir_op_priority_str(): """Test that passing a comma-delimited string works""" from vllm.config.kernel import IrOpPriorityConfig diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index d44e83bf67..4505e0cfd9 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -34,6 +34,9 @@ class IrOpPriorityConfig: fused_add_rms_norm: list[str] = Field(default_factory=list) """Priority list for vllm.ir.ops.fused_add_rms_norm""" + gelu_and_mul_sparse: list[str] = Field(default_factory=list) + """Priority list for vllm.ir.ops.gelu_and_mul_sparse""" + def compute_hash(self) -> str: """ Produces a hash unique to the pass configuration. diff --git a/vllm/ir/ops/__init__.py b/vllm/ir/ops/__init__.py index d4d71afef7..13e2b7dda9 100644 --- a/vllm/ir/ops/__init__.py +++ b/vllm/ir/ops/__init__.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .activation import gelu_and_mul_sparse from .layernorm import fused_add_rms_norm, rms_norm -__all__ = ["rms_norm", "fused_add_rms_norm"] +__all__ = ["rms_norm", "fused_add_rms_norm", "gelu_and_mul_sparse"] diff --git a/vllm/ir/ops/activation.py b/vllm/ir/ops/activation.py new file mode 100644 index 0000000000..149197d735 --- /dev/null +++ b/vllm/ir/ops/activation.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn.functional as F +from torch import Tensor + +from ..op import register_op + + +@register_op +def gelu_and_mul_sparse( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> Tensor: + """Apply Gaussian sparsification, GELU, and gated multiplication.""" + d = x.shape[-1] // 2 + gate = x[..., :d] + # Statistics intentionally remain local to each tensor-parallel shard. + mean = torch.mean(gate, dim=-1, keepdim=True) + std = torch.std(gate, dim=-1, keepdim=True, unbiased=False) + sparse_gate = F.relu(gate - (mean + std * std_multiplier)) + return F.gelu(sparse_gate, approximate=approximate) * x[..., d:] + + +@gelu_and_mul_sparse.register_input_generator +def _gelu_and_mul_sparse_input_generator( + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + std_multiplier: float = 1.6448536269514722, + approximate: str = "tanh", +) -> tuple: + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype) + return x, std_multiplier, approximate diff --git a/vllm/kernels/__init__.py b/vllm/kernels/__init__.py index e83e2772f0..4821ec43ec 100644 --- a/vllm/kernels/__init__.py +++ b/vllm/kernels/__init__.py @@ -2,6 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Kernel implementations for vLLM.""" -from . import aiter_ops, oink_ops, vllm_c +from . import aiter_ops, oink_ops, triton, vllm_c -__all__ = ["vllm_c", "aiter_ops", "oink_ops"] +__all__ = ["vllm_c", "aiter_ops", "oink_ops", "triton"] diff --git a/vllm/kernels/triton/__init__.py b/vllm/kernels/triton/__init__.py index 6626213145..3a38335050 100644 --- a/vllm/kernels/triton/__init__.py +++ b/vllm/kernels/triton/__init__.py @@ -1,3 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Triton kernel implementations.""" + +from . import activation + +__all__ = ["activation"] diff --git a/vllm/kernels/triton/activation.py b/vllm/kernels/triton/activation.py new file mode 100644 index 0000000000..2bda1f51b4 --- /dev/null +++ b/vllm/kernels/triton/activation.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import Tensor +from torch.library import triton_op, wrap_triton + +from vllm import ir +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + +_MAX_INTERMEDIATE_SIZE = 32768 +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32} + + +@triton.jit +def _gelu_and_mul_sparse_kernel( + x_ptr, + output_ptr, + d: tl.constexpr, + std_multiplier, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < d + row_start = row * 2 * d + + gate = tl.load(x_ptr + row_start + offsets, mask=mask, other=0.0).to(tl.float32) + # TODO: Investigate delaying this load to reduce spills for large reductions + # without regressing smaller shapes. + up = tl.load(x_ptr + row_start + d + offsets, mask=mask, other=0.0).to(tl.float32) + + input_dtype = x_ptr.dtype.element_ty + mean = tl.sum(gate, axis=0) / d + centered = tl.where(mask, gate - mean, 0.0) + variance = tl.sum(centered * centered, axis=0) / d + # Match the input-dtype rounding at each native PyTorch op boundary. + mean = mean.to(input_dtype).to(tl.float32) + std = tl.sqrt(variance).to(input_dtype).to(tl.float32) + scaled_std = (std * std_multiplier).to(input_dtype).to(tl.float32) + cutoff = (mean + scaled_std).to(input_dtype).to(tl.float32) + sparse_gate = (gate - cutoff).to(input_dtype).to(tl.float32) + sparse_gate = tl.where(sparse_gate < 0.0, 0.0, sparse_gate) + + inner = 0.7978845608028654 * ( + sparse_gate + 0.044715 * sparse_gate * sparse_gate * sparse_gate + ) + tanh = 2.0 * tl.sigmoid(2.0 * inner) - 1.0 + activated = 0.5 * sparse_gate * (1.0 + tanh) + activated = activated.to(input_dtype).to(tl.float32) + tl.store(output_ptr + row * d + offsets, activated * up, mask=mask) + + +def _supports_gelu_and_mul_sparse( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> bool: + del std_multiplier + if ( + x.device.type != "cuda" + or x.dtype not in _SUPPORTED_DTYPES + or not x.is_contiguous() + or x.ndim == 0 + or approximate != "tanh" + ): + return False + last_dim = x.shape[-1] + return ( + last_dim > 0 and last_dim % 2 == 0 and last_dim // 2 <= _MAX_INTERMEDIATE_SIZE + ) + + +@triton_op( + "vllm::gelu_and_mul_sparse_triton", + mutates_args=(), +) +def _gelu_and_mul_sparse_triton_op(x: Tensor, std_multiplier: float) -> Tensor: + d = x.shape[-1] // 2 + output = torch.empty((*x.shape[:-1], d), dtype=x.dtype, device=x.device) + if output.numel() == 0: + return output + + n_rows = x.numel() // x.shape[-1] + + def grid(meta): + return (n_rows,) + + block_size = triton.next_power_of_2(d) + num_warps = min(max(block_size // 256, 1), 8) + wrap_triton(_gelu_and_mul_sparse_kernel)[grid]( + x, + output, + d=d, + std_multiplier=std_multiplier, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return output + + +@ir.ops.gelu_and_mul_sparse.register_impl( + "triton", + supported=HAS_TRITON and current_platform.is_cuda(), + supports_args=_supports_gelu_and_mul_sparse, +) +def gelu_and_mul_sparse_triton( + x: Tensor, std_multiplier: float, approximate: str = "none" +) -> Tensor: + assert approximate == "tanh" + return _gelu_and_mul_sparse_triton_op(x, std_multiplier) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 159c521cee..f159d2dff8 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -8,6 +8,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from vllm import ir from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.platforms import CpuArchEnum, current_platform @@ -332,24 +333,11 @@ class GeluAndMulSparse(CustomOp): raise ValueError("activation_sparsity is 0.0. Please use GeluAndMul.") target_sparsity_tensor = torch.tensor(activation_sparsity, dtype=torch.float32) normal_dist = torch.distributions.normal.Normal(0, 1) - self.std_multiplier = normal_dist.icdf(target_sparsity_tensor) - - def _gaussian_topk(self, x: torch.Tensor) -> torch.Tensor: - """Get % sparse percentile of the Gaussian distribution.""" - # NOTE(rob): for TP>1, we could all-gather to get the means/std. - # But we do not do this because in expectation they are the same - # and in practice the eval scores are good without gathering. - mean = torch.mean(x, dim=-1, keepdim=True) - std = torch.std(x, dim=-1, keepdim=True, unbiased=False) - cutoff_x = mean + std * self.std_multiplier - return nn.functional.relu(x - cutoff_x) + self.std_multiplier = normal_dist.icdf(target_sparsity_tensor).item() def forward_native(self, x: torch.Tensor) -> torch.Tensor: """PyTorch-native implementation equivalent to forward().""" - d = x.shape[-1] // 2 - out = self._gaussian_topk(x[..., :d]) - out = F.gelu(out, approximate=self.approximate) - return out * x[..., d:] + return ir.ops.gelu_and_mul_sparse(x, self.std_multiplier, self.approximate) def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 6d9444567a..98bddbde29 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -755,7 +755,10 @@ class CudaPlatformBase(Platform): rms_norm = ["oink"] + default return IrOpPriorityConfig.with_default( - default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm + default, + rms_norm=rms_norm, + fused_add_rms_norm=rms_norm, + gelu_and_mul_sparse=["triton", "native"], ) @classmethod diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index a9cf534716..b3b7fe8b08 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -1161,7 +1161,10 @@ class RocmPlatform(Platform): rms_norm = default return IrOpPriorityConfig.with_default( - default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm + default, + rms_norm=rms_norm, + fused_add_rms_norm=rms_norm, + gelu_and_mul_sparse=["native"], ) @classmethod diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 8cb29fbdbd..6920c55e71 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -491,7 +491,7 @@ class XPUPlatform(Platform): using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE default = ["native"] if using_inductor else ["vllm_c", "native"] - return IrOpPriorityConfig.with_default(default) + return IrOpPriorityConfig.with_default(default, gelu_and_mul_sparse=["native"]) @classmethod def device_count(cls) -> int: From 23cfaad49701c497def53552b23317335431f72a Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Mon, 14 Sep 2026 15:59:51 +0800 Subject: [PATCH 008/139] [CI] Update entrypoints CI (#56763) Signed-off-by: wang.yuqi --- .buildkite/test-amd.yaml | 12 ++++++------ .buildkite/test_areas/entrypoints.yaml | 14 +++++++------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index efe0a0f8bc..78130fa3c7 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1067,7 +1067,7 @@ steps: - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - pytest -v -s entrypoints/scale_out -- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD +- label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)" # TBD timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1080,9 +1080,10 @@ steps: - tests/entrypoints/openai commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness -- label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD +- label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false @@ -1096,7 +1097,6 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/chat_completion - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - label: ":amd: (MI300) Entrypoints Integration (API Server Generate)" # TBD timeout_in_minutes: 65 @@ -3217,7 +3217,7 @@ steps: - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc - pytest -v -s entrypoints/scale_out -- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 1)" # TBD +- label: ":amd: (MI355) Entrypoints Integration (OpenAI API completion)" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3230,9 +3230,10 @@ steps: - tests/entrypoints/openai commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness -- label: ":amd: (MI355) Entrypoints Integration (API Server OpenAI - Part 2)" # TBD +- label: ":amd: (MI355) Entrypoints Integration (OpenAI API chat_completion)" # TBD timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false @@ -3246,7 +3247,6 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/chat_completion - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - label: ":amd: (MI355) Entrypoints Integration (API Server Generate)" # TBD timeout_in_minutes: 40 diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 6ad579f82a..00091707d5 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -73,9 +73,9 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server OpenAI - Part 1)" +- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (OpenAI API completion)" device: h200_35gb - key: entrypoints-integration-api-server-openai-part-1 + key: entrypoints-integration-api-server-openai-completion timeout_in_minutes: 68 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -84,19 +84,20 @@ steps: - tests/entrypoints/openai commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness mirror: amd: - label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 1)" + label: ":amd: (MI300) Entrypoints Integration (OpenAI API completion)" dind: false device: mi300_1 timeout_in_minutes: 65 depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server OpenAI - Part 2)" +- label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (OpenAI API chat_completion)" device: h200_35gb - key: entrypoints-integration-api-server-openai-part-2 + key: entrypoints-integration-api-server-openai-chat_completion timeout_in_minutes: 83 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -106,10 +107,9 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s entrypoints/openai/chat_completion - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py mirror: amd: - label: ":amd: (MI300) Entrypoints Integration (API Server OpenAI - Part 2)" + label: ":amd: (MI300) Entrypoints Integration (OpenAI API chat_completion)" dind: false device: mi300_1 timeout_in_minutes: 70 From b443c1cc4e12171be0b0bbb7cc107919bda4634f Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Mon, 14 Sep 2026 16:36:19 +0800 Subject: [PATCH 009/139] [Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill (1.7-3.8x faster than the Triton chunk path) (#55737) Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/models/glm5next/nvidia/kda.py | 172 ++++++++++++++++++++++++----- 1 file changed, 144 insertions(+), 28 deletions(-) diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index 192327cf4a..436b9685cb 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -37,6 +37,7 @@ from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.worker.workspace import current_workspace_manager if current_platform.is_rocm(): from vllm.models.glm5next.amd.ops.third_party.kda import ( @@ -123,6 +124,32 @@ def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor: return x.float().sigmoid() +def _resolve_kda_prefill_backend( + backend: str, head_dim: int, dtype: torch.dtype, lower_bound: float | None +) -> str: + """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on + SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton + ``chunk_kda_with_fused_gate`` path. ``backend`` comes from + ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" + if backend not in ("auto", "triton", "flashkda"): + raise ValueError(f"Unsupported KDA prefill backend: {backend}") + capability = current_platform.get_device_capability() + supported = ( + current_platform.is_cuda() + and capability is not None + and capability.major in (9, 10, 12) + and head_dim == 128 + and dtype == torch.bfloat16 + and lower_bound is not None + ) + if backend == "flashkda" and not supported: + raise RuntimeError( + "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, head_dim=128 " + "and a bounded KDA gate." + ) + return "flashkda" if supported and backend != "triton" else "triton" + + class Glm5NextLinearAttention(GatedDeltaNetAttention): head_dim: int num_heads: int @@ -291,6 +318,85 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention): # every _forward call (it reads an env-derived flag each time). self._conv_state_dim_first = is_conv_state_dim_first() + additional_config = vllm_config.additional_config + self.kda_prefill_backend = _resolve_kda_prefill_backend( + additional_config.get("kda_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto", + self.head_dim, + vllm_config.model_config.dtype, + self.kda_lower_bound, + ) + self._flashkda_buffer_specs: ( + tuple[tuple[tuple[int, ...], torch.dtype], ...] | None + ) = None + if self.kda_prefill_backend == "flashkda": + import vllm._flashkda_C # noqa: F401 + + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + max_seqs = vllm_config.scheduler_config.max_num_seqs + workspace_size = torch.ops._flashkda_C.get_workspace_size( + max_tokens, self.local_num_heads, max_seqs + ) + self._flashkda_buffer_specs = ( + ( + (max_seqs, self.local_num_heads, self.head_dim, self.head_dim), + self.get_state_dtype()[1], + ), + ((workspace_size,), torch.uint8), + # Output buffer for steps that also carry spec-decode tokens: + # the non-spec tokens are then scattered by non_spec_token_indx. + ( + (1, max_tokens, self.local_num_heads, self.head_dim), + vllm_config.model_config.dtype, + ), + ) + + def _flashkda_prefill( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + out: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fused KDA chunked prefill (FlashKDA). Takes the raw gate logits ``g`` + and raw ``beta`` logits, l2-normalizes q/k in-kernel and applies the + bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, + matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the + attention output into ``out`` (a workspace buffer when ``None``) and + returns ``(out, final_state)``.""" + assert self._flashkda_buffer_specs is not None + final_state, workspace, workspace_out = ( + current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs) + ) + final_state = final_state[: initial_state.shape[0]] + if out is None: + out = workspace_out[:, : q.shape[1]] + # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided. + torch.ops._flashkda_C.fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta, + self.head_dim**-0.5, + out, + workspace, + self.A_log.view(-1), + self.dt_bias.view(-1, self.head_dim), + self.kda_lower_bound, + initial_state.contiguous(), + final_state, + cu_seqlens.contiguous(), + None, + None, + ) + return out, final_state + def forward( self, hidden_states: torch.Tensor, @@ -542,26 +648,42 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention): initial_state = gather_initial_states( recurrent_state, non_spec_state_indices_tensor, has_initial_state ) - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = chunk_kda_with_fused_gate( - q=_rearr(q_ns), - k=_rearr(k_ns), - v=_rearr(v_ns), - raw_g=g1_ns, - # Chunk path wants the pre-sigmoided fp32 beta (its kernels - # don't sigmoid); beta_ns is raw bf16 from forward. - beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), - A_log=self.A_log, - g_bias=self.dt_bias, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc, - safe_gate=safe_gate, - lower_bound=lower_bound, - ) + if self.kda_prefill_backend == "flashkda": + # Non-spec step: write straight into the layer output buffer + # (dense token order, no merge copy). Step with spec-decode + # tokens: write to the workspace buffer and scatter below. + ns_out = None if use_spec else core_attn_out[:, :num_actual_tokens] + core_attn_out_non_spec, last_recurrent_state = self._flashkda_prefill( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + g=g1_ns, + beta=beta_ns, + initial_state=initial_state, + cu_seqlens=non_spec_query_start_loc, + out=ns_out, + ) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + raw_g=g1_ns, + # Chunk path wants the pre-sigmoided fp32 beta (its + # kernels don't sigmoid); beta_ns is raw bf16 from forward. + beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), + A_log=self.A_log, + g_bias=self.dt_bias, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) # Init cache scatter_states( recurrent_state, @@ -600,14 +722,8 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention): # --- merge spec / non-spec outputs back into token order --- if use_spec and core_attn_out_non_spec is not None: assert core_attn_out_spec is not None - merged = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_non_spec.dtype, - device=core_attn_out_non_spec.device, - ) - merged.index_copy_(1, spec_token_indx, core_attn_out_spec) - merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) - core_attn_out[0, :num_actual_tokens] = merged.squeeze(0) + core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) + core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) elif use_spec: assert core_attn_out_spec is not None if spec_out is None: From dc89fdfb0ede2918a884faae771a49bbbe7b6cf9 Mon Sep 17 00:00:00 2001 From: Chinmay-Kulkarni-AMD Date: Mon, 14 Sep 2026 14:46:48 +0530 Subject: [PATCH 010/139] [CPU][Profiler] Group torch profiler tables by input shape when record_shapes is on (#56016) Signed-off-by: Chinmay Kulkarni Co-authored-by: Li, Jiang --- vllm/profiler/wrapper.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index 877a390d6d..fb89631986 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -19,6 +19,7 @@ import vllm.version from vllm.config import ProfilerConfig from vllm.config.profiler import _is_uri_path from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -259,9 +260,14 @@ class TorchProfilerWrapper(WorkerProfiler): sort_key: str, row_limit: int | None = None, ) -> str: + group_by_input_shape = ( + current_platform.is_cpu() + and self.profiler_config.torch_profiler_record_shapes + ) + averages = self.profiler.key_averages(group_by_input_shape=group_by_input_shape) if row_limit is None: # use profiler default row limit of 100 - return self.profiler.key_averages().table(sort_by=sort_key) - return self.profiler.key_averages().table( + return averages.table(sort_by=sort_key) + return averages.table( sort_by=sort_key, row_limit=row_limit, ) From 79f0be21ff4044d53f0a6f911422b9f52064639c Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Mon, 14 Sep 2026 17:30:28 +0800 Subject: [PATCH 011/139] [Bugfix][CPU] Fix DeepSeek-R1 (FP8 MLA + MoE) correctness on CPU backend (#56773) Signed-off-by: jiang1.li --- vllm/model_executor/kernels/linear/scaled_mm/cpu.py | 2 ++ vllm/model_executor/layers/fused_moe/experts/cpu_moe.py | 1 + 2 files changed, 3 insertions(+) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index 4ef1af0a70..92ea073a1a 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -261,6 +261,8 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if getattr(layer, "_cpu_skip_gemm_dispatch", False): + return # Skip the base class process (FP8 padding / fnuz normalization) # which is GPU-oriented. Instead, VNNI-prepack weights for AMX. params = self._get_layer_params(layer) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 44a5bda21e..f7d82c691b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -468,6 +468,7 @@ class CPUExpertsFp8(mk.FusedMoEExpertsModular): RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, RoutingMethodType.DeepseekV4, + RoutingMethodType.DeepSeekV3, ] @staticmethod From 238cb2b191630e8b02965adc772e3236d26936fe Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Mon, 14 Sep 2026 17:36:39 +0800 Subject: [PATCH 012/139] [Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat (#55738) Signed-off-by: Jared Wen Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/attention/test_mla_backends.py | 20 +++++++++++ .../v1/attention/test_mla_prefill_selector.py | 34 +++++++++++++++++++ .../v1/attention/test_sparse_mla_backends.py | 29 ++++++++++++++++ .../layers/attention/mla_attention.py | 4 +++ .../layers/attention/sparse_mla_attention.py | 2 ++ .../backends/mla/prefill/flash_attn.py | 7 ++++ 6 files changed, 96 insertions(+) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index ee5737a995..f724949f54 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -25,6 +25,7 @@ from vllm.config.vllm import set_current_vllm_config from vllm.model_executor.layers.attention import mla_attention as mla_attention_module from vllm.model_executor.layers.attention.mla_attention import ( MLAAttention, + MLACommonBaseImpl, QueryLenSupport, _DecodeConcatQuantFP8, _use_masked_mha, @@ -166,6 +167,25 @@ def test_glm5_flashinfer_masked_mha_routing( ) +@pytest.mark.parametrize("qk_rope_head_dim", [64, 0], ids=["rope", "nope"]) +def test_concat_k_nope_k_pe_matches_torch_cat(qk_rope_head_dim): + """The K concat used by the MLA prefill context loop must equal torch.cat of + k_nope with the broadcast k_pe; with no RoPE part it returns k_nope itself + instead of allocating and copying.""" + torch.manual_seed(0) + num_tokens, num_heads, qk_nope_head_dim = 5, 4, 256 + k_nope = torch.randn(num_tokens, num_heads, qk_nope_head_dim, dtype=torch.bfloat16) + k_pe = torch.randn(num_tokens, 1, qk_rope_head_dim, dtype=torch.bfloat16) + impl = SimpleNamespace(_use_flashinfer_concat_mla_k=False) + + k = MLACommonBaseImpl._concat_k_nope_k_pe(impl, k_nope, k_pe) + + expected = torch.cat([k_nope, k_pe.expand(-1, num_heads, -1)], dim=-1) + assert k.shape == (num_tokens, num_heads, qk_nope_head_dim + qk_rope_head_dim) + torch.testing.assert_close(k, expected, rtol=0, atol=0) + assert (k.data_ptr() == k_nope.data_ptr()) == (qk_rope_head_dim == 0) + + def test_masked_mha_routing_is_dimension_specific(): assert _use_masked_mha( backend_name="FLASHMLA_SPARSE", diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index d5856ab6bd..098eddfbf6 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -313,6 +313,40 @@ class TestBackendValidation: ) assert invalid_reasons == [] + def test_flash_attn_accepts_glm53_flash_nope_dimensions(self): + """(256, 0, 256) runs the same FA kernels as GLM-5's (192, 64, 256); a + RoPE-carrying 320-wide query does not.""" + try: + from vllm.v1.attention.backends.mla.prefill.flash_attn import ( + FlashAttnPrefillBackend, + ) + except ImportError: + pytest.skip("MLA prefill backend not available") + return + + capability = DeviceCapability(major=10, minor=0) + + def validate(qk_rope_head_dim: int) -> list[str]: + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=MLADimensions( + qk_nope_head_dim=256, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=256, + ), + ) + with patch.object( + FlashAttnPrefillBackend, "is_available", return_value=True + ): + return FlashAttnPrefillBackend.validate_configuration( + capability, selector_config + ) + + assert validate(qk_rope_head_dim=0) == [] + invalid_reasons = validate(qk_rope_head_dim=64) + assert len(invalid_reasons) == 1 + assert "supported MLA dimensions" in invalid_reasons[0] + @pytest.mark.skipif( not current_platform.is_cuda_alike(), diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 0766584f54..2ca839d59c 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import ( GLOBAL_TOPK_MASK_MAX_BYTES, SparseMLACommonImpl, SparseMLAPrefillMetadata, + _is_masked_mha_available, _masked_mha_workspace_fits, _topk_mask_shape, _use_dense_mha_prefill, @@ -1165,6 +1166,33 @@ def test_masked_mha_workspace_guards_long_routing_policy( ) +@pytest.mark.parametrize( + ("model_dims", "kv_cache_dtype", "fa_version", "expected"), + [ + pytest.param((128, 512, 128, 64, 128), "auto", 4, True, id="deepseek_v32"), + pytest.param((64, 512, 192, 64, 256), "auto", 4, True, id="glm5"), + pytest.param((64, 512, 256, 0, 256), "auto", 4, True, id="glm53_flash_nope"), + pytest.param((64, 512, 256, 0, 256), "fp8", 4, False, id="quantized_kv"), + pytest.param((64, 512, 256, 0, 256), "auto", 3, False, id="no_fa4"), + pytest.param((64, 512, 256, 64, 256), "auto", 4, False, id="rope_320"), + pytest.param((128, 512, 256, 0, 256), "auto", 4, False, id="wrong_heads"), + ], +) +def test_is_masked_mha_available_model_dims( + monkeypatch, model_dims, kv_cache_dtype, fa_version, expected +): + """The allow-list gates masked MHA per exact model geometry: the DeepSeek-V3.2, + GLM-5 and NoPE GLM-5.3-Flash layouts on an SM100-family GPU with FA4 and an + unquantized KV cache, nothing else.""" + import vllm.model_executor.layers.attention.sparse_mla_attention as mod + + monkeypatch.setattr( + mod.current_platform, "is_device_capability_family", lambda family: True + ) + monkeypatch.setattr(mod, "get_flash_attn_version", lambda **kwargs: fa_version) + assert _is_masked_mha_available(*model_dims, kv_cache_dtype) is expected + + def test_masked_mha_workspace_fits_accounts_for_batch_and_context(): """Request count and context chunk length are independent multipliers.""" base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048) @@ -1198,6 +1226,7 @@ PREFILL_BATCH_SPECS = { [ pytest.param(128, 128, 64, 128, id="deepseek_hd192_v128"), pytest.param(64, 192, 64, 256, id="glm5_hd256_v256"), + pytest.param(64, 256, 0, 256, id="glm53_flash_nope_hd256_v256"), ], ) def test_sparse_backend_prefill_correctness( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1c2e5b3962..874c43f8cd 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2871,6 +2871,10 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]): Returns: Tensor of shape [..., nope_dim + pe_dim] """ + if k_pe.shape[-1] == 0: + # NoPE MLA: nothing to append, so no copy either. + return k_nope + k = torch.empty( (*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]), dtype=k_nope.dtype, diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 79ca612032..aed5afdcf9 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -111,6 +111,8 @@ def _is_masked_mha_available( if model_dims not in ( (128, 512, 128, 64, 128), (64, 512, 192, 64, 256), + # GLM-5.3-Flash: NoPE, qk_head_dim 256 == the (192, 64, 256) kernel. + (64, 512, 256, 0, 256), ): return False qk_head_dim = qk_nope_head_dim + qk_rope_head_dim diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 3e0d544195..3af623b1f1 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -327,6 +327,13 @@ class FlashAttnPrefillBackend(MLAPrefillBackend): qk_rope_head_dim=64, v_head_dim=128, ), + # GLM5Next NoPE layout: qk_head_dim 256 + 0 and v_head_dim 256 run + # the same kernels as the (192, 64, 256) DeepSeek-V3.2 layout. + MLADimensions( + qk_nope_head_dim=256, + qk_rope_head_dim=0, + v_head_dim=256, + ), ] def __init__( From dc2e8f1157173d0f0fc09712f6897e9f6eee52a0 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Mon, 14 Sep 2026 11:40:25 +0200 Subject: [PATCH 013/139] [ROCm][Perf] W4A16: keep skinny GEMM zero-points packed 4-bit (#54965) Signed-off-by: Matthias Gehre --- csrc/rocm/skinny_gemms_int4.cu | 59 ++++++++++++----- csrc/rocm/torch_bindings.cpp | 3 +- .../quantization/test_rdna_hybrid_w4a16.py | 65 ++++++++++++++++--- .../mixed_precision/rdna_hybrid_w4a16.py | 35 +++++----- 4 files changed, 117 insertions(+), 45 deletions(-) diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index 8b7077d14e..a20ca41d41 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -102,6 +102,18 @@ __device__ inline unsigned int min__(uint32_t a, uint32_t b) { return min(a, b); } +// Reads row `row` of packed zero points [M/8, num_groups] uint32. +// +// The parameters must be signed and the indexing must use shift/mask: with +// unsigned parameters and row/8, row%8, the (A_CHUNK=32, UNRL=8) +// instantiation needs 256 VGPRs and 68 bytes of scratch instead of 239. +__device__ __forceinline__ uint32_t zp_nibble(const uint32_t* zp_packed, + const int row, const int group, + const int num_groups) { + const uint32_t w = zp_packed[(row >> 3) * num_groups + group]; + return (w >> (4 * (row & 7))) & 0xFu; +} + // W4A16 skinny GEMM kernel: packed int4 weights, fp16/bf16 activations // Targets the "sml" case where activations fit in LDS. // A_CHUNK: number of K-elements processed per thread per step. @@ -116,7 +128,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By, const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, - const scalar_t* zero_points, + const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, const int CuCount) { constexpr int max_lds_len = LDS_SIZE / 2; @@ -257,7 +269,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) if constexpr (!std::is_same_v) { if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) { uint32_t group_idx = k_ / GROUP_SIZE; - scalar_t zp = zero_points[(m + y) * num_groups + group_idx]; + scalar_t zp = __float2s(static_cast( + zp_nibble(zero_points, m + y, group_idx, num_groups))); #pragma unroll for (uint32_t b = 0; b < A_CHUNK; b++) { cvtB.h[b] = cvtB.h[b] - zp; @@ -280,8 +293,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } if constexpr (HAS_ZERO_POINTS) { uint32_t group_idx_zp = k_ / GROUP_SIZE; - float zp_f = __s2float( - zero_points[(m + y) * num_groups + group_idx_zp]); + float zp_f = static_cast( + zp_nibble(zero_points, m + y, group_idx_zp, num_groups)); partial -= (128.0f + zp_f) * act_sum; } else { partial -= 136.0f * act_sum; @@ -335,7 +348,7 @@ __global__ void wvSplitK_int4_hf_sml_(const int K, const int M, const int Bx, const int By, const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, - const scalar_t* zero_points, + const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, const int CuCount) { @@ -352,7 +365,7 @@ template ) { if constexpr (HAS_ZERO_POINTS && GROUP_SIZE > 0) { uint32_t group_idx = k_ / GROUP_SIZE; - scalar_t zp = zero_points[(m + y) * num_groups + group_idx]; + scalar_t zp = __float2s(static_cast( + zp_nibble(zero_points, m + y, group_idx, num_groups))); #pragma unroll for (uint32_t b = 0; b < A_CHUNK; b++) { cvtB.h[b] = cvtB.h[b] - zp; @@ -531,8 +545,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) } if constexpr (HAS_ZERO_POINTS) { uint32_t group_idx_zp = k_ / GROUP_SIZE; - float zp_f = __s2float( - zero_points[(m + y) * num_groups + group_idx_zp]); + float zp_f = static_cast( + zp_nibble(zero_points, m + y, group_idx_zp, num_groups)); partial -= (128.0f + zp_f) * act_sum; } else { partial -= 136.0f * act_sum; @@ -596,7 +610,7 @@ __global__ void wvSplitK_int4_hf_(const int K, const int M, const int Bx, const int By, const uint8_t* B_packed, const scalar_t* __restrict__ A, const scalar_t* scale, - const scalar_t* zero_points, + const uint32_t* zero_points, const scalar_t* __restrict__ BIAS, scalar_t* C, const int _WvPrGrp, const int CuCount) { @@ -621,7 +635,8 @@ static int mindiv_int4(int N, int div1, int div2) { // in_a: packed int4 weights [M, K/2] (int8) or [M, K/8] (int32) // in_b: activations [N, K] (fp16/bf16) // in_scale: group scales [M, K/group_size] (fp16/bf16) -// in_zero_points: optional raw zero points [M, K/group_size] (fp16/bf16) +// in_zero_points: optional raw zero points [M/8, K/group_size] (int32, +// 8 uint4 per word along dim 0) // If provided, kernel dequants as (nibble - zp_raw) * scale (asymmetric). // If absent, kernel dequants as (nibble - 8) * scale (symmetric uint4b8). // group_size: 32, 64, or 128 @@ -664,14 +679,22 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, "Scale must be [M, K/group_size] = [", M_in, ", ", num_groups, "] but got [", in_scale.size(0), ", ", in_scale.size(1), "]"); if (in_zero_points.has_value()) { - TORCH_CHECK(in_zero_points->dtype() == in_b.dtype(), - "Zero points dtype must match activation dtype"); + // The kernel reads the words as uint32, so either signedness is accepted. + TORCH_CHECK(in_zero_points->dtype() == at::kInt || + in_zero_points->dtype() == at::kUInt32, + "Zero points must be int32 or uint32 (packed 8x uint4 along " + "dim 0), got ", + in_zero_points->dtype()); TORCH_CHECK(in_zero_points->dim() == 2, - "Zero points must be 2D [M, K/group_size], got shape ", + "Zero points must be 2D [M/8, K/group_size], got shape ", in_zero_points->sizes()); - TORCH_CHECK(in_zero_points->size(0) == M_in && + TORCH_CHECK(M_in % 8 == 0, + "M must be divisible by 8 for packed zero points, got ", M_in); + TORCH_CHECK(in_zero_points->is_contiguous(), + "Zero points must be contiguous"); + TORCH_CHECK(in_zero_points->size(0) == M_in / 8 && in_zero_points->size(1) == num_groups, - "Zero points must be [M, K/group_size] = [", M_in, ", ", + "Zero points must be [M/8, K/group_size] = [", M_in / 8, ", ", num_groups, "] but got [", in_zero_points->size(0), ", ", in_zero_points->size(1), "]"); } @@ -769,9 +792,9 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, const fptype* aptr = reinterpret_cast(in_b.data_ptr()); const fptype* sptr = reinterpret_cast(in_scale.data_ptr()); - const fptype* zpptr = + const uint32_t* zpptr = in_zero_points.has_value() - ? reinterpret_cast(in_zero_points->data_ptr()) + ? reinterpret_cast(in_zero_points->data_ptr()) : nullptr; const fptype* biasptr = (in_bias.has_value() && in_bias->numel() > 0) diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index a8114ce196..14d1822707 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -31,7 +31,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK); // W4A16 grouped skinny GEMM: packed int4 weights, per-group scales, - // optional zero points for asymmetric quantization + // optional zero points [M/8, K/group_size] int32 for asymmetric + // quantization rocm_ops.def( "wvSplitK_int4_g(Tensor in_a, Tensor in_b, Tensor in_scale, " "Tensor? in_zero_points, Tensor? in_bias, int CuCount, " diff --git a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py index 562d1ef2ef..fe3db8dab3 100644 --- a/tests/kernels/quantization/test_rdna_hybrid_w4a16.py +++ b/tests/kernels/quantization/test_rdna_hybrid_w4a16.py @@ -37,6 +37,20 @@ MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE # --------------------------------------------------------------------------- +def _pack_zp_rows_for_kernel(zp_nkg: torch.Tensor) -> torch.Tensor: + """Pack raw uint4 zero points along N: [N, G] int32 -> [N//8, G] int32. + + Row n's nibble lands in word[n//8] at bits 4*(n%8). + """ + assert zp_nkg.dtype == torch.int32 + N, G = zp_nkg.shape + assert N % 8 == 0 + shifts = (torch.arange(8, device=zp_nkg.device, dtype=torch.int32) * 4)[:, None] + return torch.sum( + (zp_nkg.view(N // 8, 8, G) & 0xF) << shifts, dim=1, dtype=torch.int32 + ).contiguous() + + def _rdna_hybrid_w4a16_reference( x_mk: torch.Tensor, w_int4_nk: torch.Tensor, @@ -50,8 +64,8 @@ def _rdna_hybrid_w4a16_reference( x_mk: [M, K] fp16/bf16 w_int4_nk: [N, K] int32 with raw uint4 values in [0, 15] scales_nkg: [N, K//G] fp16/bf16 - zp_nkg: [N, K//G] fp16/bf16 raw zero points (already in act dtype), - or None for symmetric (uint4b8, dequant subtracts 8) + zp_nkg: [N, K//G] int32 raw zero points in [0, 15], or None for + symmetric (uint4b8, dequant subtracts 8) """ G = group_size N, K = w_int4_nk.shape @@ -110,13 +124,16 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M) 0.05 * torch.rand((N, K // group_size), device=device, dtype=torch.float32) ).to(dtype) - # Optional raw zero points [N, K//G] in act dtype. + # Optional raw zero points, [N, K//G] int32 for the reference and packed + # along N for the op. if has_zp: zp_nkg = torch.randint( 0, 16, (N, K // group_size), device=device, dtype=torch.int32 - ).to(dtype) + ) + w_zp = _pack_zp_rows_for_kernel(zp_nkg) else: zp_nkg = None + w_zp = None from vllm.utils.platform_utils import num_compute_units @@ -124,7 +141,7 @@ def test_rdna_hybrid_w4a16_apply_matches_reference(dtype, group_size, has_zp, M) x_mk, w_q, scales_nkg, - zp_nkg, + w_zp, None, # bias num_compute_units(), group_size, @@ -332,7 +349,7 @@ def test_rdna_hybrid_w4a16_process_weights_symmetric_repack(group_size, dist_ini @pytest.mark.parametrize("group_size", SUPPORTED_GROUP_SIZES) def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_init): - """uint4 (asymmetric): zero points unpacked to raw values in act dtype.""" + """uint4 (asymmetric): zero points are packed [N//8, K//G] int32.""" if not torch.cuda.is_available(): pytest.skip("CUDA/HIP device not available") @@ -374,10 +391,11 @@ def test_rdna_hybrid_w4a16_process_weights_asymmetric_repack(group_size, dist_in ) kernel.process_weights_after_loading(layer) - # Zero-points: unpacked to [N, K//G], cast to act dtype, raw values [0..15]. - assert layer.weight_zero_point.dtype == torch.float16 - assert tuple(layer.weight_zero_point.shape) == (N, K // G) - expected_zp = zeros_int4_gn.t().to(torch.float16) # [N, K//G] + # Zero-points: [N//8, K//G] int32, row n's nibble at word[n//8] + # bits 4*(n%8). + assert layer.weight_zero_point.dtype == torch.int32 + assert tuple(layer.weight_zero_point.shape) == (N // 8, K // G) + expected_zp = _pack_zp_rows_for_kernel(zeros_int4_gn.t().contiguous()) torch.testing.assert_close(layer.weight_zero_point, expected_zp) # Quantized weights match symmetric path's layout regardless of zp. @@ -518,3 +536,30 @@ def test_rdna_hybrid_w4a16_dispatch(dtype, M, K, N, G): ref = _hip_skinny_reference(a, w_int4_nk, scales, group_size=G, zp_bias=8) torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2) + + +@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only") +@pytest.mark.skipif( + not hasattr(torch.ops, "_rocm_C") + or not hasattr(torch.ops._rocm_C, "wvSplitK_int4_g"), + reason="wvSplitK_int4_g not built", +) +def test_wvsplitk_int4_g_rejects_unpacked_zero_points(): + """Act-dtype zero points raise instead of being read as packed words. + + Both layouts are 2D with compatible extents, so only the dtype separates + them; misreading one as the other returns wrong numbers silently. + """ + import vllm._custom_ops as ops + from vllm.utils.platform_utils import num_compute_units + + K, N, G, M = 256, 64, 128, 1 + a = torch.randn((M, K), device=device, dtype=torch.float16) + w = torch.randint(0, 255, (N, K // 2), device=device, dtype=torch.uint8).view( + torch.int8 + ) + scales = torch.rand((N, K // G), device=device, dtype=torch.float16) + zp_unpacked = torch.zeros((N, K // G), device=device, dtype=torch.float16) + + with pytest.raises(RuntimeError, match="Zero points must be int32 or uint32"): + ops.wvSplitK_int4_g(w, a, scales, num_compute_units(), G, zp_unpacked, None) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py index 9964c891b6..eb36308bcc 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/rdna_hybrid_w4a16.py @@ -77,7 +77,7 @@ def _triton_w4a16_skinny_fmt_kernel( a_ptr, # [M, K] fp16/bf16 activations b_ptr, # [N, K//8] int32 packed (ExLlama shuffle, K is packed dim) scales_ptr, # [N, K//G] fp16/bf16 scales (skinny layout) - zp_ptr, # [N, K//G] fp16/bf16 raw zero-points (when HAS_ZP=True) + zp_ptr, # [N//8, K//G] int32 zero-points (when HAS_ZP=True) c_ptr, # [M, N] fp16/bf16 output # Dimensions M, @@ -103,8 +103,9 @@ def _triton_w4a16_skinny_fmt_kernel( | (val[1]<<16) | (val[3]<<20) | (val[5]<<24) | (val[7]<<28) Scales are [N, K//G] (skinny layout, NOT transposed). - When HAS_ZP=True, raw zero-points zp_raw are loaded from zp_ptr [N, K//G] - and subtracted directly: (nibble - zp_raw) * scale. + When HAS_ZP=True, zp_ptr holds [N//8, K//G] int32 with row n's raw + zero-point at word[n//8] bits 4*(n%8), and dequant is + (nibble - zp_raw) * scale. When HAS_ZP=False, only the constant ZP_BIAS is subtracted (symmetric). """ pid_m = tl.program_id(0) @@ -148,9 +149,10 @@ def _triton_w4a16_skinny_fmt_kernel( scales = tl.load(scale_ptrs, mask=scale_mask, other=1.0) if HAS_ZP: - zp_ptrs = zp_ptr + offs_n * num_groups + group_idx - zp_raw = tl.load(zp_ptrs, mask=scale_mask, other=0.0) - b_fp = (b.to(scales.dtype) - zp_raw[:, None]) * scales[:, None] + zp_ptrs = zp_ptr + (offs_n // 8) * num_groups + group_idx + zp_word = tl.load(zp_ptrs, mask=scale_mask, other=0) + zp_raw = (zp_word >> (4 * (offs_n % 8))) & 0xF + b_fp = (b - zp_raw[:, None]).to(scales.dtype) * scales[:, None] else: b_fp = (b - ZP_BIAS).to(scales.dtype) * scales[:, None] @@ -192,7 +194,7 @@ def triton_w4a16_skinny_fmt_gemm( scales: torch.Tensor, # [N, K//G] fp16/bf16 group_size: int, zp_bias: int = 8, - zp: torch.Tensor | None = None, # [N, K//G] per-group zero-points + zp: torch.Tensor | None = None, # [N//8, K//G] int32 zero-points ) -> torch.Tensor: """ Fused W4A16 GEMM reading from skinny weight format [N, K//8]. @@ -203,8 +205,8 @@ def triton_w4a16_skinny_fmt_gemm( scales: Per-group scales [N, K//G], same dtype as a. group_size: Quantization group size (resolved from -1 to K by caller). zp_bias: Constant zero bias (default 8 for unsigned int4). - zp: Raw per-group zero-points [N, K//G] (asymmetric), - stored as zp_raw in activation dtype. When provided, + zp: Raw per-group zero-points [N//8, K//G] int32, row n at + word[n//8] bits 4*(n%8) (asymmetric). When provided, dequant is (nibble - zp_raw) * scale. Returns: @@ -225,8 +227,9 @@ def triton_w4a16_skinny_fmt_gemm( ) if zp is not None: assert zp.is_contiguous(), "Zero-points must be contiguous" - assert zp.shape == (N, num_groups), ( - f"zp shape mismatch: {zp.shape} vs ({N}, {num_groups})" + assert N % 8 == 0, f"N must be divisible by 8 for packed zp, got {N}" + assert zp.shape == (N // 8, num_groups), ( + f"zp shape mismatch: {zp.shape} vs ({N // 8}, {num_groups})" ) has_zp = zp is not None @@ -381,7 +384,10 @@ def _rdna_hybrid_w4a16_apply_impl( cu_count: int, group_size: int, ) -> torch.Tensor: - """Dispatch between skinny GEMM and Triton based on batch size M.""" + """Dispatch between skinny GEMM and Triton based on batch size M. + + ``w_zp`` is [N//8, K//G] int32 for asymmetric layers, None for symmetric. + """ import vllm._custom_ops as ops M = x_2d.shape[0] @@ -520,10 +526,7 @@ class RDNAHybridW4A16LinearKernel(MPLinearKernel): assert self.w_zp_name is not None w_zp_raw = getattr(layer, self.w_zp_name) permute_param_layout_(w_zp_raw, input_dim=1, output_dim=0, packed_dim=0) - zp_unpacked = unpack_quantized_values_into_int32( - w_zp_raw.data, c.weight_type, packed_dim=0 - ) - w_zp = zp_unpacked.to(c.act_type).contiguous() + w_zp = w_zp_raw.data.contiguous() self._transform_param(layer, self.w_zp_name, lambda x: w_zp) self._transform_param(layer, self.w_q_name, lambda x: w_q_skinny) From 5236bef721223a45d52b3649530e88b5e824e0d9 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 02:52:45 -0700 Subject: [PATCH 014/139] [CI] Keep OTel bytecode out of mounted checkouts (#56766) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/scripts/ci-otel/ci_otel.sh | 16 ++++++---- .../scripts/ci-otel/tests/test_ci_otel.py | 30 +++++++++++++++++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/.buildkite/scripts/ci-otel/ci_otel.sh b/.buildkite/scripts/ci-otel/ci_otel.sh index 78b18207ad..c9897e732e 100755 --- a/.buildkite/scripts/ci-otel/ci_otel.sh +++ b/.buildkite/scripts/ci-otel/ci_otel.sh @@ -26,7 +26,8 @@ if [ -z "${CI_INFRA_OTEL_RUNTIME_DIR}" ] || [ -z "${_CI_INFRA_OTEL_PYTHON}" ] || [ ! -f "${CI_INFRA_OTEL_DIR}/ci_otel.py" ] || ! mkdir -p "${CI_INFRA_OTEL_SPOOL_DIR}" || - ! PYTHONPATH="${CI_INFRA_OTEL_DIR}${PYTHONPATH:+:${PYTHONPATH}}" \ + ! PYTHONDONTWRITEBYTECODE=1 \ + PYTHONPATH="${CI_INFRA_OTEL_DIR}${PYTHONPATH:+:${PYTHONPATH}}" \ "${_CI_INFRA_OTEL_PYTHON}" -c "import ci_otel" >/dev/null 2>&1; then echo "vLLM CI OTel: tracing unavailable; test command will run normally" >&2 || : if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then @@ -58,9 +59,9 @@ fi _ci_otel_python() { if command -v timeout >/dev/null 2>&1; then - timeout 2s "${_CI_INFRA_OTEL_PYTHON}" "$@" + PYTHONDONTWRITEBYTECODE=1 timeout 2s "${_CI_INFRA_OTEL_PYTHON}" "$@" else - "${_CI_INFRA_OTEL_PYTHON}" "$@" + PYTHONDONTWRITEBYTECODE=1 "${_CI_INFRA_OTEL_PYTHON}" "$@" fi } @@ -86,7 +87,8 @@ ci_otel_start() { if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] && [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] && command -v nvidia-smi >/dev/null 2>&1; then - "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" /dev/null 2>&1; then - timeout 4s "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || : + PYTHONDONTWRITEBYTECODE=1 timeout 4s \ + "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || : else - "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || : + PYTHONDONTWRITEBYTECODE=1 \ + "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_otel.py" flush || : fi if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then rm -rf -- "${CI_INFRA_OTEL_RUNTIME_DIR}" || : diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index ba54f9f273..ae51010f9e 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -6,6 +6,7 @@ from __future__ import annotations import json import os import shlex +import shutil import subprocess import sys import threading @@ -396,6 +397,35 @@ def test_shell_wrapper_preserves_failure_status(tmp_path): assert result.returncode == 1 +def test_shell_wrapper_does_not_write_bytecode_to_helper_tree(tmp_path): + helper_dir = tmp_path / "ci-otel" + shutil.copytree( + SCRIPTS_DIR, + helper_dir, + ignore=shutil.ignore_patterns("__pycache__"), + ) + shell = ( + f'. "{helper_dir / "ci_otel.sh"}"; ' + f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0" + ) + + result = subprocess.run( + ["/bin/sh", "-c", shell], + check=False, + capture_output=True, + text=True, + env={ + **os.environ, + "CI_INFRA_OTEL_DIR": str(helper_dir), + "CI_INFRA_GPU_SAMPLING": "0", + "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"), + }, + ) + + assert result.returncode == 0, result.stderr + assert list(helper_dir.rglob("*.pyc")) == [] + + def test_ci_otel_run_records_command_and_preserves_status(tmp_path): shell = ( f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' From 2c7ee87223f0c94f614a0b8600f6f193e6e4388f Mon Sep 17 00:00:00 2001 From: BruceLoveDecimal <39156883+BruceLoveDecimal@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:55:39 +0800 Subject: [PATCH 015/139] [Bugfix][Mooncake] Report failed remote KV loads to the scheduler (#50984) Signed-off-by: liuqihao Co-authored-by: liuqihao Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../v1/mooncake/mooncake_connector.py | 80 +++++++++++++++---- 1 file changed, 66 insertions(+), 14 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 4009af7009..4451f0566d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import logging +import queue import threading import time from collections import defaultdict +from collections.abc import Collection from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from enum import IntEnum @@ -433,6 +435,9 @@ class PullReqMeta: expire_time: float = float("inf") # Designed for one D pairing to multiple P pull_tasks_count: int = 0 + # Set once any worker reports a failure, so a success from another worker + # for the same request is not counted afterwards. + failed: bool = False @dataclass @@ -582,6 +587,11 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA): assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_block_ids_with_load_errors(self) -> set[int]: + """Get block IDs whose remote KV load failed.""" + assert self.connector_worker is not None + return self.connector_worker.get_block_ids_with_load_errors() + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: assert self.connector_worker is not None assert isinstance(self._connector_metadata, MooncakeConnectorMetadata) @@ -1034,6 +1044,8 @@ class MooncakeConnectorWorker: self.finished_sending_reqs: set[ReqId] = set() self.finished_recving_reqs: set[ReqId] = set() + # Written from the receiver loop, drained from the worker thread. + self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue() self.xfer_stats = MooncakeKVConnectorStats() @@ -1907,12 +1919,11 @@ class MooncakeConnectorWorker: ret_msg = await sock.recv() response = self._xfer_resp_decoder.decode(ret_msg) if response.status == MooncakeXferResponseStatus.ERROR: - logger.error( - "Error happens during transferring kvcache for %s: %s", + self._handle_failed_recv( + pull_metas, req_ids, - response.err_msg, + response.err_msg or "transfer error", ) - self.xfer_stats.record_failed_recv() return self.process_pulling_result(response, pull_metas) if response.status == MooncakeXferResponseStatus.FINISH: @@ -1920,10 +1931,50 @@ class MooncakeConnectorWorker: except zmq.ContextTerminated: logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.") except Exception as e: - logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e) - self.xfer_stats.record_failed_recv() + self._handle_failed_recv(pull_metas, req_ids, f"transfer failed: {e}") return + def _handle_failed_recv( + self, + pull_metas: dict[ReqId, PullReqMeta], + req_ids: Collection[ReqId], + reason: str, + ) -> None: + """Report a failed remote KV load so the scheduler can fail or recompute it.""" + failed: list[ReqId] = [] + for req_id in req_ids: + pull_meta = pull_metas.get(req_id) + if pull_meta is None or pull_meta.failed: + continue + pull_meta.failed = True + failed.append(req_id) + self.xfer_stats.record_failed_recv() + + invalid = {b for group in pull_meta.local_block_ids for b in group} + if not invalid: + # A pull with no local blocks only asks P to release its blocks + # for a request that never reached the scheduler (see + # AsyncLLM.notify_kv_transfer_request_rejected, which submits an + # abort_immediately request just to run request_finished). No D + # request is waiting on a load, and reporting one here would trip + # the scheduler's `assert req_id in self.requests`. + continue + self._invalid_block_ids.put(invalid) + self.finished_recving_reqs.add(pull_meta.d_req_id) + + if failed: + logger.error("pulling kv_caches for %s failed: %s", failed, reason) + + def get_block_ids_with_load_errors(self) -> set[int]: + """Drain the blocks whose remote KV load failed since the last call.""" + result: set[int] = set() + while True: + try: + result.update(self._invalid_block_ids.get_nowait()) + except queue.Empty: + break + return result + def process_pulling_result( self, response: MooncakeXferResponse, @@ -1933,6 +1984,8 @@ class MooncakeConnectorWorker: for req_id in ok_reqs: pull_meta = pull_metas[req_id] + if pull_meta.failed: + continue # No race because we are in async loop. pull_meta.pull_tasks_count -= 1 if pull_meta.pull_tasks_count == 0: @@ -1942,10 +1995,8 @@ class MooncakeConnectorWorker: logger.debug("pulling kv_caches for %s finished", ok_reqs) if response.err_reqs: - logger.error( - "pulling kv_caches for %s failed: %s", - response.err_reqs, - response.err_msg, + self._handle_failed_recv( + pull_metas, response.err_reqs, response.err_msg or "unknown error" ) async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): @@ -2023,10 +2074,11 @@ class MooncakeConnectorWorker: await self._pending_bootstrap_queries[remote_bootstrap_addr].wait() if remote_engine_id not in self._remote_agents: - logger.error( - "Failed to find remote engine_id %s from bootstrap server %s", - remote_engine_id, - remote_bootstrap_addr, + self._handle_failed_recv( + pull_metas, + list(pull_metas), + f"remote engine_id {remote_engine_id} not found from bootstrap " + f"server {remote_bootstrap_addr}", ) return From 934b1fcbcbbac6e68da7a4246928ac797fe54690 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Mon, 14 Sep 2026 11:02:14 +0100 Subject: [PATCH 016/139] [Frontend] Replace `VLLM_ENABLE_SCALE_OUT_ENDPOINTS` with `--enable-scale-out` (#55176) Signed-off-by: Martin Hickey --- docs/serving/online_serving/README.md | 7 +-- docs/serving/online_serving/derenderer.md | 3 +- docs/serving/online_serving/renderer.md | 6 +-- docs/usage/security.md | 18 +++---- examples/scale_out/example_mm_serve.py | 6 +-- examples/scale_out/token_generation_client.py | 2 +- rust/src/cmd/src/cli.rs | 10 ++++ rust/src/cmd/src/cli/tests.rs | 6 +++ rust/src/server/src/config.rs | 2 + rust/src/server/src/routes.rs | 28 ++-------- rust/src/server/src/routes/tests.rs | 48 ++++++++++------- tests/entrypoints/cohere/test_api_router.py | 23 ++++---- tests/entrypoints/cohere/test_chat_v2.py | 3 +- .../derender/test_derender_parity.py | 7 +-- .../render/test_render_multimodal.py | 9 +--- tests/entrypoints/scale_out/test_factories.py | 53 ++++++------------- .../test_return_routed_experts.py | 7 +-- .../test_serving_multimodal_tokens.py | 7 +-- .../token_in_token_out/test_serving_tokens.py | 7 +-- tests/test_envs.py | 35 ------------ vllm/entrypoints/cohere/api_router.py | 10 ++-- vllm/entrypoints/launchers/cli_args.py | 7 +++ vllm/entrypoints/scale_out/factories.py | 24 +++------ vllm/envs.py | 17 ------ 24 files changed, 130 insertions(+), 215 deletions(-) diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 5c93ca1787..1f5d4bdb6a 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -121,12 +121,7 @@ For further details on profiling vLLM, please refer to [this page](../../contrib ## Scale-Out APIs -Scale-out APIs are disabled by default on `vllm serve`. The environment -variable accepts only `0` or `1`; set `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1` to -register the endpoints below. The -dedicated `vllm launch render` and `vllm serve --tokens-only` modes are explicit -opt-ins and enable their required endpoints when the variable is unset; an -explicit value of `0` is rejected for those modes. +Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out` to register the endpoints below. The dedicated `vllm launch render` and `vllm serve --tokens-only` modes always register their required endpoints regardless of `--enable-scale-out`. ### Tokens IN <> Tokens OUT APIs diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md index d536498f5c..298510062a 100644 --- a/docs/serving/online_serving/derenderer.md +++ b/docs/serving/online_serving/derenderer.md @@ -60,8 +60,7 @@ Launch the two servers first: ```bash vllm launch render meta-llama/Llama-3.2-1B-Instruct --port 8100 -VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve \ - meta-llama/Llama-3.2-1B-Instruct --tokens-only --port 8200 +vllm serve meta-llama/Llama-3.2-1B-Instruct --tokens-only --port 8200 ``` ```python diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index 1f5e893570..9be23142c4 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -7,16 +7,14 @@ Our renderer API is designed to disaggregate the render phase(preprocessing) and - Tokens-in / tokens-out engine: Make the engine a pure token-in / token-out service, decoupled from request preprocessing. The dedicated `vllm launch render` server always exposes the `/render` and -`/derender` endpoints when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS` is unset or set to -`1`. An explicit value of `0` conflicts with the renderer command and is -rejected at startup. +`/derender` endpoints. Scale-out endpoints, including `/render`, `/derender`, and `/inference/v1/generate`, are disabled by default on a standard inference server. To expose them with `vllm serve`, opt in explicitly: ```bash -VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve +vllm serve --enable-scale-out ``` ## API Reference diff --git a/docs/usage/security.md b/docs/usage/security.md index 559f587026..3c7f1e3f9b 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -158,16 +158,16 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/models` - List available models - `/v1/chat/completions` - Chat completions - `/v1/chat/completions/batch` - Batch chat completions -- `/v1/chat/completions/render` - Render chat completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) -- `/v1/chat/completions/derender` - Derender chat completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/chat/completions/render` - Render chat completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) +- `/v1/chat/completions/derender` - Derender chat completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/completions` - Text completions -- `/v1/completions/render` - Render completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) -- `/v1/completions/derender` - Derender completion requests (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/completions/render` - Render completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) +- `/v1/completions/derender` - Derender completion requests (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/embeddings` - Generate embeddings - `/v1/audio/transcriptions` - Audio transcription - `/v1/audio/translations` - Audio translation - `/v1/messages` - Anthropic-compatible messages API -- `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/messages/count_tokens` - Count tokens for Anthropic messages - `/v1/responses` - Create a response - `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) @@ -177,8 +177,8 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/rerank` - Reranking API - `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) - `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`) -- `/inference/v1/generate` - Generate completions (available when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or with `--tokens-only` unless explicitly disabled) -- `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires both `VLLM_ENABLE_COHERE_API=1` and `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`) +- `/inference/v1/generate` - Generate completions (available when `--enable-scale-out` is set, or with `--tokens-only`) +- `/cohere/v2/chat/render` - Render Cohere Chat v2 requests (requires `VLLM_ENABLE_COHERE_API=1` and either `--enable-scale-out` or `--tokens-only`) - `/v2/embed` - Cohere Embed API - `/v2/rerank` - Cohere Rerank API @@ -206,11 +206,11 @@ The following endpoints **do not require authentication** even when `--api-key` - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF - `/update_weights` - Update model weights (can alter model behavior) - `/get_world_size` - Get distributed world size -- `/abort_requests` - Abort in-flight requests (available with `--tokens-only` unless `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0`) +- `/abort_requests` - Abort in-flight requests (available with `--tokens-only`) **Utility endpoints:** -- `/tokenize` - Tokenize text (not disabled by `VLLM_ENABLE_SCALE_OUT_ENDPOINTS`) +- `/tokenize` - Tokenize text (not gated by `--enable-scale-out`) - `/detokenize` - Detokenize tokens - `/health` - Health check - `/ping` - SageMaker health check diff --git a/examples/scale_out/example_mm_serve.py b/examples/scale_out/example_mm_serve.py index 157a32d3be..558a4f9302 100644 --- a/examples/scale_out/example_mm_serve.py +++ b/examples/scale_out/example_mm_serve.py @@ -13,9 +13,9 @@ zero client-side transformation. Launch the server first: - VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ - vllm serve Qwen/Qwen3-VL-2B-Instruct \ - --dtype bfloat16 --max-model-len 4096 --enforce-eager + vllm serve Qwen/Qwen3-VL-2B-Instruct \ + --dtype bfloat16 --max-model-len 4096 --enforce-eager \ + --enable-scale-out Then run this script: diff --git a/examples/scale_out/token_generation_client.py b/examples/scale_out/token_generation_client.py index 922018f43d..a4792a68f9 100644 --- a/examples/scale_out/token_generation_client.py +++ b/examples/scale_out/token_generation_client.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Call a token-in/token-out server started with: -VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve Qwen/Qwen3-0.6B +vllm serve Qwen/Qwen3-0.6B --enable-scale-out """ import httpx diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index f7fcf30022..70b28018a3 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -329,6 +329,15 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub enable_request_id_headers: bool, + /// Register the scale-out `/inference/v1/generate` endpoint. + #[arg( + long, + default_missing_value = "true", + num_args = 0..=1 + )] + #[serde(default)] + pub enable_scale_out: bool, + /// If provided, the server will require one of these keys to be presented /// in the Authorization header. #[educe(Debug(ignore))] @@ -577,6 +586,7 @@ impl SharedRuntimeArgs { enable_log_requests: self.enable_log_requests, enable_prompt_tokens_details: self.enable_prompt_tokens_details, enable_request_id_headers: self.enable_request_id_headers, + enable_scale_out: self.enable_scale_out, } } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 884da135d7..411d256295 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -195,6 +195,7 @@ fn serve_args_forward_python_flags_with_separator() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( @@ -1030,6 +1031,7 @@ fn frontend_args_accept_json() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( @@ -1629,6 +1631,7 @@ fn serve_args_accept_handshake_aliases() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, disable_log_stats: false, served_model_name: [], allowed_origins: JsonStringList( @@ -1780,6 +1783,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, }, cors: CorsConfig { allow_origins: [ @@ -1868,6 +1872,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, }, cors: CorsConfig { allow_origins: [ @@ -1977,6 +1982,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present enable_log_requests: false, enable_prompt_tokens_details: false, enable_request_id_headers: false, + enable_scale_out: false, }, cors: CorsConfig { allow_origins: [ diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 4928cddf31..0f0deac86d 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -55,6 +55,8 @@ pub struct ApiServerOptions { pub enable_prompt_tokens_details: bool, /// When `true`, set `X-Request-Id` on every HTTP response. pub enable_request_id_headers: bool, + /// When `true`, register the scale-out `/inference/v1/generate` route. + pub enable_scale_out: bool, } /// CORS settings mirroring Python's `CORSMiddleware`; the default is permissive. diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index e7e91938bb..41e4987882 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -46,33 +46,14 @@ fn runtime_lora_updating_enabled() -> bool { .is_some_and(|value| matches!(value.trim().to_lowercase().as_str(), "1" | "true")) } -fn parse_scale_out_endpoints_flag(value: Option<&str>) -> Result { - let Some(value) = value else { - return Ok(false); - }; - if value.is_empty() { - return Ok(false); - } - - match value.trim() { - "0" => Ok(false), - "1" => Ok(true), - _ => Err("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1".to_string()), - } -} - -fn scale_out_endpoints_enabled() -> bool { - let value = std::env::var("VLLM_ENABLE_SCALE_OUT_ENDPOINTS").ok(); - parse_scale_out_endpoints_flag(value.as_deref()).unwrap_or_else(|message| panic!("{message}")) -} - /// Build the minimal OpenAI-compatible router for one configured model. pub fn build_router(state: Arc) -> Router { + let scale_out_endpoints_enabled = state.api_server_options.enable_scale_out; build_router_with_options( state, server_dev_mode_enabled(), runtime_lora_updating_enabled(), - scale_out_endpoints_enabled(), + scale_out_endpoints_enabled, ) } @@ -121,10 +102,7 @@ fn build_router_with_options( if scale_out_endpoints_enabled { router = router.route("/inference/v1/generate", post(inference::generate)); } else { - info!( - "scale-out endpoints are disabled; set \ - VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 to enable them" - ); + info!("scale-out endpoints are disabled; pass --enable-scale-out to enable them"); } if runtime_lora_updating_enabled { diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 2851ec1e88..5ee2b51f3f 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -54,8 +54,7 @@ use zeromq::{DealerSocket, PushSocket, ZmqMessage}; use super::{ build_router, build_router_with_dev_mode, build_router_with_dev_mode_and_lora, - build_router_with_scale_out_endpoints, parse_scale_out_endpoints_flag, - render::build_router as build_render_router, + build_router_with_scale_out_endpoints, render::build_router as build_render_router, }; use crate::config::{ApiServerOptions, CorsConfig, LoraModulePath}; use crate::lora::LoadLoraError; @@ -731,23 +730,36 @@ async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router { ) } -#[test] -fn scale_out_flag_accepts_unset_empty_zero_one_and_whitespace() { - assert_eq!(parse_scale_out_endpoints_flag(None), Ok(false)); - assert_eq!(parse_scale_out_endpoints_flag(Some("")), Ok(false)); - assert_eq!(parse_scale_out_endpoints_flag(Some("0")), Ok(false)); - assert_eq!(parse_scale_out_endpoints_flag(Some("1")), Ok(true)); - assert_eq!(parse_scale_out_endpoints_flag(Some(" 1 ")), Ok(true)); -} +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn scale_out_generate_route_follows_api_server_options() { + let (chat, _engine_task) = test_models_with_engine_outputs_and_backend( + b"engine-scale-out-api-server-options", + default_stream_output_specs(), + Arc::new(FakeChatBackend::new()), + ) + .await; + let state = Arc::new( + AppState::new(vec!["model".to_string()], chat).with_api_server_options(ApiServerOptions { + enable_scale_out: true, + ..Default::default() + }), + ); + let mut app = build_router(state); -#[test] -fn scale_out_flag_rejects_values_other_than_zero_or_one() { - for value in ["-1", "2", "01", "+1", "invalid", " "] { - assert_eq!( - parse_scale_out_endpoints_flag(Some(value)), - Err("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1".to_string()) - ); - } + let response = app + .call( + Request::builder() + .method("POST") + .uri("/inference/v1/generate") + .header("content-type", "application/json") + .body(Body::from("{}")) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_ne!(response.status(), StatusCode::NOT_FOUND); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py index a90e93b6b1..94324dbed7 100644 --- a/tests/entrypoints/cohere/test_api_router.py +++ b/tests/entrypoints/cohere/test_api_router.py @@ -58,7 +58,6 @@ def _enable_cohere_api(monkeypatch): flag inside the test body. """ monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "1") - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1") # ---------------------------------------------------------------------- @@ -128,15 +127,19 @@ def _generate_request() -> GenerateRequest: ) -def _build_app(handler: _Handler | None) -> FastAPI: +def _build_app(handler: _Handler | None, *, enable_scale_out: bool = True) -> FastAPI: app = FastAPI() + app.state.args = Namespace(enable_scale_out=enable_scale_out, tokens_only=False) attach_router(app) app.state.cohere_serving_chat_v2 = handler return app -def _build_render_app(chat_handler, render_handler) -> FastAPI: +def _build_render_app( + chat_handler, render_handler, *, enable_scale_out: bool = True +) -> FastAPI: app = FastAPI() + app.state.args = Namespace(enable_scale_out=enable_scale_out, tokens_only=False) attach_router(app) app.state.cohere_serving_chat_v2 = chat_handler app.state.serving_render = render_handler @@ -151,11 +154,13 @@ def _build_app_with_vllm_handlers(handler: _Handler | None) -> FastAPI: into the ``CohereError`` wire shape. """ app = FastAPI() - attach_router(app) - app.state.cohere_serving_chat_v2 = handler # ``validation_exception_handler`` reads ``req.app.state.args``; the # real cli builds this via argparse. - app.state.args = Namespace(log_error_stack=False) + app.state.args = Namespace( + log_error_stack=False, enable_scale_out=True, tokens_only=False + ) + attach_router(app) + app.state.cohere_serving_chat_v2 = handler app.exception_handler(RequestValidationError)(validation_exception_handler) app.exception_handler(HTTPException)(http_exception_handler) app.exception_handler(VLLMError)(vllm_error_handler) @@ -402,10 +407,8 @@ class TestRenderEndpoint: paths = [getattr(r, "path", None) for r in app.routes] assert "/cohere/v2/chat/render" in paths - def test_route_not_registered_when_flag_unset(self, monkeypatch): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - app = FastAPI() - attach_router(app) + def test_route_not_registered_when_flag_unset(self): + app = _build_app(handler=None, enable_scale_out=False) paths = [getattr(r, "path", None) for r in app.routes] assert "/cohere/v2/chat/render" not in paths assert "/cohere/v2/chat" in paths diff --git a/tests/entrypoints/cohere/test_chat_v2.py b/tests/entrypoints/cohere/test_chat_v2.py index 36acecb622..4f07d3d873 100644 --- a/tests/entrypoints/cohere/test_chat_v2.py +++ b/tests/entrypoints/cohere/test_chat_v2.py @@ -57,6 +57,8 @@ def server(): # the conversation to surface a thinking block (SmolLM2 doesn't # emit Cohere-style reasoning tokens out of the box). "--no-cohere-is-reasoning-model", + # Registers /cohere/v2/chat/render alongside /cohere/v2/chat. + "--enable-scale-out", ] # Cap the CPU KV-cache pool the vLLM CPU backend reserves at @@ -71,7 +73,6 @@ def server(): env_dict = { "VLLM_CPU_KVCACHE_SPACE": "1", "VLLM_ENABLE_COHERE_API": "1", - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1", } with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server: diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index 5b313464da..adece09c0a 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -34,6 +34,7 @@ ARGS = [ "hermes", "--reasoning-parser", "deepseek_r1", + "--enable-scale-out", ] TOOLS = [ @@ -54,11 +55,7 @@ FORCE_WEATHER_TOOL = {"type": "function", "function": {"name": "get_weather"}} @pytest.fixture(scope="module") def server(): - with RemoteOpenAIServer( - MODEL, - ARGS, - env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}, - ) as remote_server: + with RemoteOpenAIServer(MODEL, ARGS) as remote_server: yield remote_server diff --git a/tests/entrypoints/scale_out/render/test_render_multimodal.py b/tests/entrypoints/scale_out/render/test_render_multimodal.py index 70bbc95016..c6a6931948 100644 --- a/tests/entrypoints/scale_out/render/test_render_multimodal.py +++ b/tests/entrypoints/scale_out/render/test_render_multimodal.py @@ -27,15 +27,10 @@ def vision_server(): "1", "--limit-mm-per-prompt.video", "0", + "--enable-scale-out", ] - env_overrides = {"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"} - - with RemoteOpenAIServer( - VISION_MODEL_NAME, - args, - env_dict=env_overrides, - ) as remote_server: + with RemoteOpenAIServer(VISION_MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/scale_out/test_factories.py b/tests/entrypoints/scale_out/test_factories.py index 9811bca8a2..6d8335395e 100644 --- a/tests/entrypoints/scale_out/test_factories.py +++ b/tests/entrypoints/scale_out/test_factories.py @@ -3,7 +3,6 @@ from argparse import Namespace -import pytest from fastapi import FastAPI from vllm.entrypoints.launchers.api_server.routers import register_api_routers @@ -22,69 +21,49 @@ SCALE_OUT_PATHS = RENDER_PATHS | {GENERATE_PATH} def registered_paths( - supported_tasks: tuple[SupportedTask, ...], *, tokens_only: bool = False + supported_tasks: tuple[SupportedTask, ...], + *, + tokens_only: bool = False, + enable_scale_out: bool = False, ) -> set[str]: app = FastAPI() - args = Namespace(tokens_only=tokens_only, enable_fault_tolerance=False) + args = Namespace( + tokens_only=tokens_only, + enable_scale_out=enable_scale_out, + enable_fault_tolerance=False, + ) app.state.args = args register_api_routers(args, app, supported_tasks) return {route.path for route in app.routes} -def test_scale_out_routes_are_disabled_by_default(monkeypatch): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - +def test_scale_out_routes_are_disabled_by_default(): paths = registered_paths(("generate",)) assert SCALE_OUT_PATHS.isdisjoint(paths) -def test_disabled_scale_out_routes_are_logged(monkeypatch, caplog): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - +def test_disabled_scale_out_routes_are_logged(caplog): with caplog.at_level("INFO", logger="vllm.entrypoints.scale_out.factories"): registered_paths(("generate",)) - assert any( - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1" in record.message - for record in caplog.records - ) + assert any("--enable-scale-out" in record.message for record in caplog.records) -def test_scale_out_routes_can_be_enabled_for_generate_server(monkeypatch): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1") - - paths = registered_paths(("generate",)) +def test_scale_out_routes_can_be_enabled_for_generate_server(): + paths = registered_paths(("generate",), enable_scale_out=True) assert paths >= SCALE_OUT_PATHS -def test_render_routes_remain_enabled_for_render_server(monkeypatch): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - +def test_render_routes_remain_enabled_for_render_server(): paths = registered_paths(("render",)) assert paths >= RENDER_PATHS assert GENERATE_PATH not in paths -def test_render_server_rejects_explicitly_disabled_scale_out_routes(monkeypatch): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0") - - with pytest.raises(ValueError, match="VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0"): - registered_paths(("render",)) - - -def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset(monkeypatch): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - +def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset(): paths = registered_paths(("generate",), tokens_only=True) assert paths >= {GENERATE_PATH, "/abort_requests"} - - -def test_tokens_only_mode_rejects_explicitly_disabled_scale_out_routes(monkeypatch): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0") - - with pytest.raises(ValueError, match="--tokens-only"): - registered_paths(("generate",), tokens_only=True) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py index 963e416e7e..e7527a071b 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_return_routed_experts.py @@ -34,12 +34,9 @@ def server(): "--enable-return-routed-experts", "--hf-overrides", '{"sliding_window": null}', + "--enable-scale-out", ] - with RemoteOpenAIServer( - MODEL_NAME, - args, - env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}, - ) as remote_server: + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py index 9491623fd1..892d3c8285 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py @@ -37,13 +37,10 @@ def server(): "4096", "--enforce-eager", "--no-enable-prefix-caching", + "--enable-scale-out", ] - with RemoteOpenAIServer( - MODEL_NAME, - args, - env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}, - ) as remote_server: + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index fb741df061..1c349e8d67 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -69,6 +69,7 @@ def server(request): # Option B: deterministic prefix caching # "--enable-prefix-caching", # "--deterministic-prefix-caching", + "--enable-scale-out", ] extra_args = getattr(request, "param", None) @@ -79,11 +80,7 @@ def server(request): else [str(extra_args)] ) - with RemoteOpenAIServer( - MODEL_NAME, - args, - env_dict={"VLLM_ENABLE_SCALE_OUT_ENDPOINTS": "1"}, - ) as remote_server: + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: yield remote_server diff --git a/tests/test_envs.py b/tests/test_envs.py index aaa53d2a54..1ef40cde0c 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -43,41 +43,6 @@ def test_api_key_is_not_compile_factor(monkeypatch: pytest.MonkeyPatch): assert "VLLM_API_KEY" not in envs.compile_factors() -def test_scale_out_endpoints_flag_is_runtime_only(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1") - - assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is True - assert "VLLM_ENABLE_SCALE_OUT_ENDPOINTS" not in envs.compile_factors() - - -def test_scale_out_endpoints_flag_distinguishes_unset_from_disabled( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.delenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", raising=False) - assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is None - - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "0") - assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is False - - -def test_scale_out_endpoints_flag_treats_empty_as_unset( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "") - - assert envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS is None - - -@pytest.mark.parametrize("value", ["-1", "2", "01", "+1", "invalid", " "]) -def test_scale_out_endpoints_flag_rejects_values_other_than_zero_or_one( - monkeypatch: pytest.MonkeyPatch, value: str -): - monkeypatch.setenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", value) - - with pytest.raises(ValueError, match="must be 0 or 1"): - _ = envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS - - def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index 445964e6aa..eaac36f133 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -17,9 +17,10 @@ log) and vLLM continues to boot normally. Even when the SDK is installed, :func:`attach_router` also requires ``VLLM_ENABLE_COHERE_API=1`` in the environment before it will expose -the chat route. The render route additionally requires -``VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1``. These gates keep unrelated -deployments from accidentally exposing the APIs. +the chat route. The render route additionally requires scale-out +endpoints to be enabled, i.e. ``--enable-scale-out`` or +``--tokens-only``. These gates keep unrelated deployments from +accidentally exposing the APIs. Note: the handlers must live at module scope (not inside ``attach_router``) so that FastAPI's ``typing.get_type_hints`` resolves @@ -306,6 +307,7 @@ def attach_router(app: FastAPI) -> None: ) return app.include_router(router) - if envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS: + args = getattr(app.state, "args", None) + if getattr(args, "enable_scale_out", False) or getattr(args, "tokens_only", False): app.include_router(render_router) app.add_middleware(CohereErrorEnvelopeMiddleware) diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 0da96b43e1..89af22cdc8 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -181,6 +181,13 @@ class BaseFrontendArgs: If set to True, only enable the Tokens In<>Out endpoint. This is intended for use in a Disaggregated Everything setup. """ + enable_scale_out: bool = False + """ + If set to True, register the scale-out endpoints (`/render`, `/derender`, + and `/inference/v1/generate`) on `vllm serve`. Has no effect on + `vllm launch render` or `vllm serve --tokens-only`, which always register + their required endpoints regardless of this flag. + """ fingerprint_mode: Literal["full", "hash", "custom", "none"] = "full" """Controls the ``system_fingerprint`` field on responses. diff --git a/vllm/entrypoints/scale_out/factories.py b/vllm/entrypoints/scale_out/factories.py index 30c329c51f..bd48c709dd 100644 --- a/vllm/entrypoints/scale_out/factories.py +++ b/vllm/entrypoints/scale_out/factories.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING from fastapi import FastAPI -from vllm import envs from vllm.engine.protocol import EngineClient from vllm.logger import init_logger from vllm.tasks import SupportedTask @@ -67,22 +66,15 @@ def register_scale_out_api_routers( app: FastAPI, supported_tasks: tuple["SupportedTask", ...], ): - dedicated_mode = None - if "render" in supported_tasks: - dedicated_mode = "`vllm launch render`" - elif getattr(app.state.args, "tokens_only", False): - dedicated_mode = "`--tokens-only`" - - enabled = envs.VLLM_ENABLE_SCALE_OUT_ENDPOINTS - if dedicated_mode is not None and enabled is False: - raise ValueError( - "`VLLM_ENABLE_SCALE_OUT_ENDPOINTS=0` conflicts with " - f"{dedicated_mode}; unset it or set it to 1" - ) - if dedicated_mode is None and not enabled: + args = getattr(app.state, "args", None) + # `vllm launch render` and `vllm serve --tokens-only` are dedicated to + # serving these endpoints, so they always register them regardless of + # `--enable-scale-out`. + dedicated_mode = "render" in supported_tasks or getattr(args, "tokens_only", False) + enabled = dedicated_mode or getattr(args, "enable_scale_out", False) + if not enabled: logger.info( - "Scale-out endpoints are disabled. Set " - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 to enable them." + "Scale-out endpoints are disabled. Set --enable-scale-out to enable them." ) return diff --git a/vllm/envs.py b/vllm/envs.py index 0269a96409..6158f86f89 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -261,7 +261,6 @@ if TYPE_CHECKING: VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True VLLM_ENABLE_RESPONSES_API_STORE: bool = False VLLM_ENABLE_COHERE_API: bool = False - VLLM_ENABLE_SCALE_OUT_ENDPOINTS: bool | None = None VLLM_HAS_FLASHINFER_CUBIN: bool = False VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True @@ -352,15 +351,6 @@ def maybe_convert_bool(value: str | None) -> bool | None: return bool(int(value)) -def maybe_convert_scale_out_endpoints(value: str | None) -> bool | None: - if value is None: - return None - normalized = value.strip() - if normalized not in ("0", "1"): - raise ValueError("VLLM_ENABLE_SCALE_OUT_ENDPOINTS must be 0 or 1") - return maybe_convert_bool(normalized) - - def maybe_convert_json_str_or_file(value: str | None) -> dict[str, Any] | None: if value is None: return None @@ -1858,12 +1848,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ENABLE_COHERE_API": lambda: bool( int(os.getenv("VLLM_ENABLE_COHERE_API", "0")) ), - # If set to 1, expose the scale-out endpoints on `vllm serve`. - # The dedicated `vllm launch render` server exposes render and derender - # endpoints when this variable is unset or set to 1. - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS": lambda: maybe_convert_scale_out_endpoints( - os.getenv("VLLM_ENABLE_SCALE_OUT_ENDPOINTS") or None - ), # If set, use the fp8 mfma in rocm paged attention. "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": lambda: bool( int(os.getenv("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0")) @@ -2289,7 +2273,6 @@ def compile_factors() -> dict[str, object]: "VLLM_CONFIG_ROOT", "LD_LIBRARY_PATH", "VLLM_SERVER_DEV_MODE", - "VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "VLLM_DP_MASTER_IP", "VLLM_DP_MASTER_PORT", "VLLM_NIXL_SIDE_CHANNEL_HOST", From 1678b396270406c27fcab8f5b86b21fd305ac605 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Mon, 14 Sep 2026 18:05:13 +0800 Subject: [PATCH 017/139] [Bugfix][EPD] Preserve media processing options in encoder requests (#56786) Signed-off-by: jiangkuaixue123 --- .../disaggregated_encoder/disagg_epd_proxy.py | 11 ++++- .../ec_connector/unit/test_epd_proxy_retry.py | 41 ++++++++++++++++--- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 3912c1c67a..830da9b5d6 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -278,7 +278,6 @@ async def fanout_encoder_primer( item_transfer_ids[idx] = transfer_id encoder_req = { - # You *may* need to keep additional fields "model": orig_request.get("model"), "messages": [ { @@ -292,6 +291,14 @@ async def fanout_encoder_primer( # once the prompt is encoded and its embeddings are published. "stream": False, } + for key in ( + "mm_processor_kwargs", + "media_io_kwargs", + "priority", + "session_id", + ): + if key in orig_request: + encoder_req[key] = orig_request[key] if consumer_zmq is not None: # No mm_hash here on purpose. The encoder's own # `mm_features[i].identifier` is derived from the uuid *and* the @@ -378,7 +385,7 @@ async def fanout_encoder_primer( # connector's own handle on the published embedding (for NIXL, # peer_host/peer_port/size_bytes). The decoder's connector # looks it up by mm_hash on the request, so carry it through. - ec_params[item_uuids.get(idx, ec_mm_hash)] = reported + ec_params[ec_mm_hash] = reported if NO_REWRITE and consumer_zmq is not None: ec_params.setdefault("ec_items", []).append( {"mm_hash": ec_mm_hash, "transfer_id": item_transfer_ids[idx]} diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index 9a92b1fdbc..1a8c877c72 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -116,7 +116,9 @@ def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch first, _, _ = asyncio.run(proxy.prepare_for_decode(body, *args)) reported = first["ec_transfer_params"] - assert [handle] == [value for key, value in reported.items() if key != "ec_items"] + assert {key: value for key, value in reported.items() if key != "ec_items"} == { + "encoder-side-hash": handle + } assert "ec_transfer_params" not in body # Attempt 2's encode reports nothing: decode must be told nothing. @@ -154,10 +156,22 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize( + "request_options", + [ + {}, + { + "mm_processor_kwargs": {"max_pixels": 262144}, + "media_io_kwargs": {"image": {"image_mode": "RGB"}}, + "priority": -3, + "session_id": "session-123", + }, + ], +) async def test_http_roundtrip_preserves_payload_and_response_bytes( - proxy, monkeypatch, stream, prefill + proxy, monkeypatch, stream, prefill, request_options ): - """Exercise real HTTP hops, rewrite and a decode retry without model servers.""" + """Preserve media and scheduling context across HTTP hops and retries.""" seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []} prefix = b'data: {"content":"' # Split a multibyte character at the old 1024-byte forwarding boundary. @@ -174,10 +188,13 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( seen[stage].append(body) if stage == "encode": item = body["messages"][0]["content"][0] + mm_hash = item["uuid"] + if body.get("mm_processor_kwargs"): + mm_hash += "-processed" return web.json_response( { "ec_transfer_params": { - item["uuid"]: { + mm_hash: { "metadata": {"image_grid_thw": [[1, 2, 2]]}, "peer_port": len(seen[stage]), } @@ -230,6 +247,7 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( "max_tokens": 32, "seed": 42, "structured_outputs": {"choice": ["A", "B"]}, + **request_options, } await proxy.on_startup() try: @@ -247,6 +265,15 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( assert len(seen["encode"]) == 4 assert len(seen["decode"]) == 2 + for requests in seen.values(): + for forwarded in requests: + for key in ( + "mm_processor_kwargs", + "media_io_kwargs", + "priority", + "session_id", + ): + assert forwarded.get(key) == body.get(key) final = seen["decode"][-1] for key in [ "model", @@ -265,7 +292,11 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( "image_embeds": {"image_grid_thw": [1, 2, 2]}, "uuid": proxy.content_uuid(item), } - assert final["ec_transfer_params"][proxy.content_uuid(item)]["peer_port"] in (3, 4) + mm_hash = proxy.content_uuid(item) + if request_options: + mm_hash += "-processed" + assert set(final["ec_transfer_params"]) - {"ec_items"} == {mm_hash} + assert final["ec_transfer_params"][mm_hash]["peer_port"] in (3, 4) if prefill: assert final["kv_transfer_params"] == {"remote_block_ids": [2]} From ea723c81c3ea26425cb69503a5d5e90822a04a45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Mon, 14 Sep 2026 12:22:18 +0200 Subject: [PATCH 018/139] [Security] Cap Qwen-VL video sampling knobs (#56729) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba Co-authored-by: Cursor --- vllm/multimodal/video.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 4142febc3f..c739191979 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -357,6 +357,9 @@ class PyNvVideoCodecVideoBackend(VideoBackend): video_processor=("Qwen3VLVideoProcessor", "Cosmos3EdgeVideoProcessor"), ) class Qwen3VLVideoBackend(VideoBackend): + _MAX_FRAMES: ClassVar[int] = 768 + _MAX_FPS: ClassVar[int] = 30 + @classmethod def compute_frames_index_to_sample( cls, @@ -366,10 +369,10 @@ class Qwen3VLVideoBackend(VideoBackend): ) -> list[int]: total_frames_num = source.total_frames_num original_fps = source.original_fps - fps = target.fps + fps = min(target.fps, cls._MAX_FPS) max_frame_idx = source.total_frames_num - 1 min_frames = kwargs.get("min_frames", 4) - max_frames = kwargs.get("max_frames", 768) + max_frames = min(kwargs.get("max_frames", cls._MAX_FRAMES), cls._MAX_FRAMES) # Refer to: # https://github.com/huggingface/transformers/blob/v5.9.0/src/transformers/models/qwen3_vl/video_processing_qwen3_vl.py#L119-L125 @@ -420,6 +423,9 @@ class Qwen2VLVideoBackend(VideoBackend): clip); it is clamped to the last valid frame. """ + _MAX_FRAMES: ClassVar[int] = 768 + _MAX_FPS: ClassVar[int] = 30 + @classmethod def compute_frames_index_to_sample( cls, @@ -433,7 +439,7 @@ class Qwen2VLVideoBackend(VideoBackend): original_fps = source.original_fps temporal_patch_size = kwargs.get("temporal_patch_size", 2) min_frames = kwargs.get("min_frames", 4) - max_frames = kwargs.get("max_frames", 768) + max_frames = min(kwargs.get("max_frames", cls._MAX_FRAMES), cls._MAX_FRAMES) # vLLM reports original_fps == 0 for clips with unknown/variable fps # (VFR, malformed, streaming); fail loudly instead of dividing by zero. @@ -447,7 +453,7 @@ class Qwen2VLVideoBackend(VideoBackend): math.floor(min(max_frames, total_frames_num) / temporal_patch_size) * temporal_patch_size ) - n = total_frames_num / original_fps * target.fps + n = total_frames_num / original_fps * min(target.fps, cls._MAX_FPS) n = min(max(n, min_frames), max_frames, total_frames_num) n = math.floor(n / temporal_patch_size) * temporal_patch_size From 0ca7aef4e341cb4f4d7233291c304a168f4384aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 14 Sep 2026 12:37:12 +0200 Subject: [PATCH 019/139] [Bugfix][NixlPush] Guard _remote_agents read in _do_send_reg_notif (X1) (#56317) Signed-off-by: NickLucche --- .../unit/test_nixl_push_connector.py | 104 ++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 5 +- .../kv_connector/v1/nixl/push_worker.py | 4 +- 3 files changed, 111 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 1cdafcfe48..2c7f0cb866 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -362,6 +362,7 @@ class _StubWriterWorker(NixlPushConnectorWorker): w.world_size = 1 w.engine_id = "test-decode-engine" w._remote_agents = {} + w._handshake_lock = threading.RLock() w._physical_blocks_per_logical_kv_block = 1 w._uses_region_group_mapping = False w.region_group_ids = [0] @@ -418,6 +419,84 @@ def _registration_data( } +class TestPushWriterRegSend: + """``_do_send_reg_notif``: the PUSH_REG goes to every handshaked agent of + the P engine, and the agent map is read under ``_handshake_lock`` so the + send can't race the handshake done-callback's insert or a concurrent + stale-engine eviction (X1).""" + + def _reg_send_worker(self) -> tuple[_StubWriterWorker, list]: + w = _StubWriterWorker.fresh() + w.nixl_wrapper = MagicMock() + failures: list = [] + w._handle_failed_transfer = lambda *a: failures.append(a) + return w, failures + + def test_sends_to_all_agents_with_prefix_and_payload(self): + w, failures = self._reg_send_worker() + w._remote_agents["prefill-engine"] = {(0, 0): "agent-0", (0, 1): "agent-1"} + rd = _registration_data("req-1") + + w._do_send_reg_notif("req-1", rd) + + assert failures == [] + calls = w.nixl_wrapper.send_notif.call_args_list + assert len(calls) == 2 + assert {c.args[0] for c in calls} == {"agent-0", "agent-1"} + for call in calls: + notif = call.kwargs["notif_msg"] + assert notif.startswith(PUSH_REG_NOTIF_PREFIX) + # (Tuples decode back as lists, so compare identity fields.) + decoded = msgspec.msgpack.decode(notif[len(PUSH_REG_NOTIF_PREFIX) :]) + assert decoded["request_id"] == "req-1" + assert decoded["decode_engine_id"] == rd["decode_engine_id"] + assert decoded["remote_engine_id"] == rd["remote_engine_id"] + assert decoded["local_block_ids"] == [ + list(group) for group in rd["local_block_ids"] + ] + + def test_fails_request_when_engine_not_handshaked(self): + w, failures = self._reg_send_worker() + + w._do_send_reg_notif("req-1", _registration_data("req-1")) + + assert failures == [("req-1", None)] + w.nixl_wrapper.send_notif.assert_not_called() + + def test_agent_read_is_serialized_with_handshake_lock(self): + """While another thread holds ``_handshake_lock`` (as the handshake + done-callback does while writing ``_remote_agents``), the writer's + check-then-send must not run ahead of the critical section.""" + w, failures = self._reg_send_worker() + w._remote_agents["prefill-engine"] = {(0, 0): "agent-0"} + rd = _registration_data("req-1") + + w._handshake_lock.acquire() + done = threading.Event() + errors: list[BaseException] = [] + + def _send(): + try: + w._do_send_reg_notif("req-1", rd) + except BaseException as e: # pragma: no cover + errors.append(e) + finally: + done.set() + + t = threading.Thread(target=_send) + t.start() + assert not done.wait(timeout=0.2) + w.nixl_wrapper.send_notif.assert_not_called() + assert failures == [] + + w._handshake_lock.release() + assert done.wait(timeout=2.0) + t.join(timeout=2.0) + assert errors == [] + assert failures == [] + assert w.nixl_wrapper.send_notif.call_count == 1 + + class TestPushWriterMatching: def test_handle_push_reg_matches_existing_finished_blocks(self): """PUSH_REG arrives second (P finished first): match + fire.""" @@ -727,6 +806,31 @@ def test_stale_engine_evicted_on_push(): w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old") +def test_cleanup_remote_engine_pops_under_handshake_lock(): + """Regression (X1): eviction must not remove ``_remote_agents`` entries + while another thread is inside the lock's critical section (e.g. the + push writer reading the map to send a registration).""" + w = _eviction_worker(engine_ttl=30.0) + w._remote_agents["D-old"] = {(0, 0): "agent-D-old"} + + w._handshake_lock.acquire() + done = threading.Event() + + def _cleanup(): + w._cleanup_remote_engine("D-old") + done.set() + + t = threading.Thread(target=_cleanup) + t.start() + assert "D-old" in w._remote_agents + + w._handshake_lock.release() + assert done.wait(timeout=2.0) + t.join(timeout=2.0) + assert "D-old" not in w._remote_agents + w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old") + + class TestPushWriterNotifs: def test_get_new_notifs_processes_forwarded_completion_notif(self): """Non-PUSH_REG notifs forwarded by the writer thread are drained diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 4edfaefa2f..58ab5850ea 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -3264,7 +3264,10 @@ class NixlBaseConnectorWorker: # Notif-only engines (push-mode D side) have no descriptor state. for handle in self.dst_xfer_side_handles.pop(engine_id, {}).values(): self.nixl_wrapper.release_dlist_handle(handle) - for agent_name in self._remote_agents.pop(engine_id).values(): + # Pop under the handshake lock; NIXL teardown stays outside it. + with self._handshake_lock: + agents = self._remote_agents.pop(engine_id) + for agent_name in agents.values(): self.nixl_wrapper.remove_remote_agent(agent_name) self.kv_caches_base_addr.pop(engine_id, None) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index e506aff1f4..603321260d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -354,7 +354,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None: engine_id = reg_data["remote_engine_id"] notif_msg = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(reg_data) - agents = self._remote_agents.get(engine_id) + # _remote_agents is mutated on other threads; snapshot under the lock. + with self._handshake_lock: + agents = dict(self._remote_agents.get(engine_id) or {}) if not agents: logger.error( "No remote agents for engine %s; cannot send registration for %s", From 7b1ea3f524445be0545ed962cb06ce21a93fcc33 Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:37:59 +0800 Subject: [PATCH 020/139] [Bugfix][Rust Frontend][Multimodal] Align DeepSeek V4.1 and Kimi K3 media with rendered placeholders (#56366) Co-authored-by: Bugen Zhao Co-authored-by: Codex Signed-off-by: reidliu41 Signed-off-by: Bugen Zhao --- rust/src/chat/src/lib.rs | 10 +- rust/src/chat/src/multimodal.rs | 205 +++++++++++++----- rust/src/chat/src/renderer/deepseek.rs | 65 +++++- .../src/chat/src/renderer/deepseek_v32/mod.rs | 1 + rust/src/chat/src/renderer/deepseek_v4/mod.rs | 1 + .../src/chat/src/renderer/deepseek_v41/mod.rs | 5 +- .../chat/src/renderer/deepseek_v41/tests.rs | 78 ++++++- rust/src/chat/src/renderer/harmony/mod.rs | 1 + rust/src/chat/src/renderer/hf/mod.rs | 1 + rust/src/chat/src/renderer/inkling/mod.rs | 1 + .../src/chat/src/renderer/kimi_k3/encoding.rs | 51 +++-- rust/src/chat/src/renderer/kimi_k3/mod.rs | 5 +- rust/src/chat/src/renderer/kimi_k3/tests.rs | 59 +++++ rust/src/chat/src/renderer/mod.rs | 13 ++ rust/src/chat/tests/chat.rs | 1 + rust/src/server/src/grpc/tests.rs | 1 + .../server/src/routes/http_client_tests.rs | 1 + rust/src/server/src/routes/tests.rs | 1 + 18 files changed, 415 insertions(+), 85 deletions(-) diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 1b5d5aea6f..9c59d14028 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -34,8 +34,8 @@ pub use parser::{ParserSelection, validate_parser_overrides}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DeepSeekV41ChatRenderer, - DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, RenderedPrompt, - RendererSelection, + DynChatRenderer, HarmonyChatRenderer, InklingChatRenderer, KimiK3ChatRenderer, MediaPartSource, + RenderedPrompt, RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -112,6 +112,10 @@ impl ChatRequestProcessor { request: &ChatRequest, rendered: RenderedPrompt, ) -> Result<(Prompt, Option)> { + let media_is_empty = match &rendered.media_order { + Some(media_order) => media_order.is_empty(), + None => !request.has_multimodal(), + }; match self.model_dtype { Some(model_dtype) => { multimodal::finalize_rendered_prompt( @@ -122,7 +126,7 @@ impl ChatRequestProcessor { ) .await } - None if !request.has_multimodal() => Ok((rendered.prompt, None)), + None if media_is_empty => Ok((rendered.prompt, None)), None => Err(Error::UnsupportedMultimodalRenderer), } } diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index 674979e6ce..a5dd3ba8ea 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -16,7 +16,7 @@ use std::fs; use std::path::Path; use std::sync::{Arc, LazyLock}; -use itertools::izip; +use itertools::{Either, izip}; use llm_multimodal::{ AsyncMultiModalTracker, AudioClip, AudioPreProcessor, EncoderFieldLayouts, FieldLayout, ImageFrame, MediaConnector, MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata, @@ -33,7 +33,7 @@ use vllm_text::Prompt; use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; use crate::error::{Error, Result, bail_multimodal, multimodal}; -use crate::renderer::RenderedPrompt; +use crate::renderer::{MediaPartSource, RenderedPrompt}; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; mod audio; @@ -549,7 +549,8 @@ pub(crate) async fn finalize_rendered_prompt( info: Option<&MultimodalModelInfo>, model_dtype: ModelDtype, ) -> Result<(Prompt, Option)> { - if !request.has_multimodal() { + let media_parts = extract_media_parts(request, rendered.media_order.as_deref())?; + if media_parts.is_empty() { return Ok((rendered.prompt, None)); } let info = info.ok_or(Error::UnsupportedMultimodalRenderer)?; @@ -561,63 +562,92 @@ pub(crate) async fn finalize_rendered_prompt( .map_err(|error| multimodal!("{error}"))?, Prompt::TokenIds(token_ids) => token_ids, }; - let media_parts = extract_media_parts(request)?; let prepared = info.prepare_multimodal(media_parts, &mut prompt_token_ids, model_dtype).await?; Ok((Prompt::TokenIds(prompt_token_ids), Some(prepared))) } -/// Extract media parts from chat messages in message/content order. -/// -/// Assistant history is skipped because generated assistant blocks are already -/// represented as text for prompt rendering in this crate. -fn extract_media_parts(request: &ChatRequest) -> Result> { - let mut all_parts = Vec::new(); - for message in &request.messages { - let content = match message { - ChatMessage::System { content } - | ChatMessage::Developer { content, .. } - | ChatMessage::User { content } - | ChatMessage::ToolResponse { content, .. } => content, - ChatMessage::Assistant { .. } => continue, - }; - let ChatContent::Parts(parts) = content else { - continue; - }; - for part in parts { - match part { - ChatContentPart::Text { .. } => {} - ChatContentPart::ImageUrl { - image_url, - detail, - uuid, - } => all_parts.push(MediaContentPart::ImageUrl { - url: image_url.clone(), - detail: *detail, - uuid: uuid.clone(), - }), - ChatContentPart::VideoUrl { video_url, uuid } => { - all_parts.push(MediaContentPart::VideoUrl { - url: video_url.clone(), - uuid: uuid.clone(), - }) - } - ChatContentPart::InputAudio { data, format, uuid } => { - all_parts.push(MediaContentPart::AudioUrl { - url: input_audio_data_url(data, format.as_deref())?, - uuid: uuid.clone(), - }) - } - ChatContentPart::AudioUrl { audio_url, uuid } => { - all_parts.push(MediaContentPart::AudioUrl { - url: audio_url.clone(), - uuid: uuid.clone(), - }) +/// Resolve media parts in the placeholder order reported by the renderer. +fn extract_media_parts( + request: &ChatRequest, + media_order: Option<&[MediaPartSource]>, +) -> Result> { + let parts = match media_order { + // Resolve exactly the sources selected by the renderer, in placeholder order. + Some(media_order) => Either::Left(media_order.iter().map(|source| -> Result<_> { + let message = request.messages.get(source.message_index).ok_or_else(|| { + multimodal!( + "renderer reported missing chat message {} for multimodal content", + source.message_index + ) + })?; + let content = match message { + ChatMessage::System { content } + | ChatMessage::Developer { content, .. } + | ChatMessage::User { content } + | ChatMessage::ToolResponse { content, .. } => content, + ChatMessage::Assistant { .. } => { + bail_multimodal!("renderer reported multimodal assistant content") } + }; + let ChatContent::Parts(parts) = content else { + bail_multimodal!("renderer reported multimodal content in a text-only chat message") + }; + let part = parts.get(source.content_part_index).ok_or_else(|| { + multimodal!( + "renderer reported missing content part {} in chat message {}", + source.content_part_index, + source.message_index + ) + })?; + Ok(part) + })), + // Fall back to all media in request message and content-part order (e.g. Jinja). + None => Either::Right(request.messages.iter().flat_map(|message| { + let content = match message { + ChatMessage::System { content } + | ChatMessage::Developer { content, .. } + | ChatMessage::User { content } + | ChatMessage::ToolResponse { content, .. } => Some(content), + ChatMessage::Assistant { .. } => None, + }; + match content { + Some(ChatContent::Parts(parts)) => parts.as_slice(), + _ => &[], } - } - } - Ok(all_parts) + .iter() + .filter(|part| part.is_multimodal()) + .map(Ok) + })), + }; + parts + .map(|part| match part? { + ChatContentPart::Text { .. } => { + bail_multimodal!("renderer reported a text part as multimodal content") + } + ChatContentPart::ImageUrl { + image_url, + detail, + uuid, + } => Ok(MediaContentPart::ImageUrl { + url: image_url.clone(), + detail: *detail, + uuid: uuid.clone(), + }), + ChatContentPart::VideoUrl { video_url, uuid } => Ok(MediaContentPart::VideoUrl { + url: video_url.clone(), + uuid: uuid.clone(), + }), + ChatContentPart::InputAudio { data, format, uuid } => Ok(MediaContentPart::AudioUrl { + url: input_audio_data_url(data, format.as_deref())?, + uuid: uuid.clone(), + }), + ChatContentPart::AudioUrl { audio_url, uuid } => Ok(MediaContentPart::AudioUrl { + url: audio_url.clone(), + uuid: uuid.clone(), + }), + }) + .collect() } /// Wrap OpenAI base64 audio in a data URL consumed by `MediaConnector`. @@ -1043,6 +1073,75 @@ mod tests { assert!(input_audio_data_url("AAAA", Some("flac")).is_err()); } + #[test] + fn extract_media_parts_uses_request_order_or_explicit_selection() { + let request = ChatRequest { + messages: vec![ + ChatMessage::user(vec![ + ChatContentPart::text("before image"), + ChatContentPart::ImageUrl { + image_url: "data:image/png;base64,image-b".to_string(), + detail: None, + uuid: Some("image-b".to_string()), + }, + ]), + ChatMessage::user(vec![ChatContentPart::ImageUrl { + image_url: "data:image/png;base64,image-a".to_string(), + detail: None, + uuid: Some("image-a".to_string()), + }]), + ChatMessage::assistant_text("answer"), + ChatMessage::user("follow-up"), + ], + ..ChatRequest::for_test() + }; + let media_order = [ + MediaPartSource { + message_index: 1, + content_part_index: 0, + }, + MediaPartSource { + message_index: 0, + content_part_index: 1, + }, + ]; + + let selections = [None, Some(&media_order[..]), Some(&[][..])]; + let uuids = selections.map(|order| { + extract_media_parts(&request, order) + .unwrap() + .into_iter() + .map(|part| match part { + MediaContentPart::ImageUrl { uuid, .. } => uuid, + _ => panic!("expected image URL"), + }) + .collect::>() + }); + + expect_test::expect![[r#" + [ + [ + Some( + "image-b", + ), + Some( + "image-a", + ), + ], + [ + Some( + "image-a", + ), + Some( + "image-b", + ), + ], + [], + ] + "#]] + .assert_debug_eq(&uuids); + } + fn image_url_part() -> MediaContentPart { MediaContentPart::ImageUrl { url: "https://example.com/image.png".to_string(), diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs index 5f900ab88b..5925c5fd16 100644 --- a/rust/src/chat/src/renderer/deepseek.rs +++ b/rust/src/chat/src/renderer/deepseek.rs @@ -16,6 +16,7 @@ use serde_json_fmt::JsonFormat; use llm_multimodal::DEEPSEEK_V41_IMAGE_PLACEHOLDER; +use super::MediaPartSource; use crate::error::{Error, Result}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ReasoningEffort, @@ -87,6 +88,14 @@ struct RenderedToolSchema<'a> { /// Render one chat request into the final prompt string. pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Result { + render_request_with_media_order(request, dialect).map(|(prompt, _)| prompt) +} + +/// Render one chat request and record media in placeholder order. +pub(super) fn render_request_with_media_order( + request: &ChatRequest, + dialect: DsDialect, +) -> Result<(String, Vec)> { let (thinking_mode, reasoning_effort_prompt) = match dialect { DsDialect::V4 => { resolve_thinking_options(request).map(|(mode, prompt)| (mode, Cow::Borrowed(prompt)))? @@ -112,6 +121,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul dialect, ); let mut out = String::from(BOS_TOKEN); + let mut media_order = Vec::new(); if dialect == DsDialect::V41 && (thinking_mode == ThinkingMode::Thinking || synthetic_tool_system @@ -126,7 +136,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul let mut request_tools_attached = false; let mut render_index = 0isize; if synthetic_tool_system { - render_system_message(&mut out, None, request_tools, dialect)?; + render_system_message(&mut out, None, request_tools, dialect, &mut media_order)?; request_tools_attached = true; render_index += 1; } @@ -157,14 +167,22 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul } else { &[] }; - render_system_message(&mut out, Some(content), tools, dialect)?; + render_system_message( + &mut out, + Some((message_index, content)), + tools, + dialect, + &mut media_order, + )?; } ChatMessage::Developer { content, tools } => { render_developer_message( &mut out, + message_index, content, tools.as_deref().unwrap_or(&[]), dialect, + &mut media_order, )?; } ChatMessage::User { .. } | ChatMessage::ToolResponse { .. } => { @@ -174,6 +192,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul message_index, &last_tool_call_order, dialect, + &mut media_order, )?; } ChatMessage::Assistant { content } => { @@ -223,7 +242,7 @@ pub(super) fn render_request(request: &ChatRequest, dialect: DsDialect) -> Resul } } - Ok(out) + Ok((out, media_order)) } /// Resolve DeepSeek V4's thinking controls. Unlike the Python tokenizer @@ -488,12 +507,13 @@ fn render_tool_schema(out: &mut String, tool: &ChatTool) -> Result<()> { /// Render a system turn, optionally followed by the V4 tool preamble. fn render_system_message( out: &mut String, - content: Option<&ChatContent>, + content: Option<(usize, &ChatContent)>, tools: &[ChatTool], dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { - if let Some(content) = content { - write_chat_content(out, content, dialect)?; + if let Some((message_index, content)) = content { + write_chat_content(out, content, dialect, message_index, media_order)?; } if !tools.is_empty() { out.push_str("\n\n"); @@ -505,9 +525,11 @@ fn render_system_message( /// Developer messages are rendered as user-like turns with optional tools. fn render_developer_message( out: &mut String, + message_index: usize, content: &ChatContent, tools: &[ChatTool], dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { if content.is_empty() { return Err(Error::ChatTemplate( @@ -516,7 +538,7 @@ fn render_developer_message( } out.push_str(USER_SP_TOKEN); - write_chat_content(out, content, dialect)?; + write_chat_content(out, content, dialect, message_index, media_order)?; if !tools.is_empty() { out.push_str("\n\n"); render_tools(out, tools, dialect)?; @@ -531,6 +553,7 @@ fn render_user_content_block( message_index: usize, tool_call_order: &HashMap, dialect: DsDialect, + media_order: &mut Vec, ) -> Result<()> { let (block_start, block_end) = user_content_block_bounds(messages, message_index); let mut sorted_tool_indices = @@ -542,7 +565,9 @@ fn render_user_content_block( out.push_str("\n\n"); } match &messages[message_index] { - ChatMessage::User { content } => write_chat_content(out, content, dialect)?, + ChatMessage::User { content } => { + write_chat_content(out, content, dialect, message_index, media_order)?; + } ChatMessage::ToolResponse { .. } => { let sorted_index = sorted_tool_indices .next() @@ -550,7 +575,7 @@ fn render_user_content_block( let ChatMessage::ToolResponse { content, .. } = &messages[sorted_index] else { unreachable!("sorted tool response index should reference a tool message"); }; - write_tool_result(out, content, dialect)?; + write_tool_result(out, content, dialect, sorted_index, media_order)?; } _ => unreachable!("user content block should only contain user content messages"), } @@ -597,9 +622,15 @@ fn sorted_tool_response_indices( } /// Render one tool response payload inside a V4 `` block. -fn write_tool_result(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> { +fn write_tool_result( + out: &mut String, + content: &ChatContent, + dialect: DsDialect, + message_index: usize, + media_order: &mut Vec, +) -> Result<()> { out.push_str(""); - write_chat_content(out, content, dialect)?; + write_chat_content(out, content, dialect, message_index, media_order)?; out.push_str(""); Ok(()) } @@ -727,7 +758,13 @@ fn encode_arguments_to_dsml( /// The V4.1 dialect inlines the image placeholder at each image part's /// position unconditionally (matching the Python encoding's /// `IMAGE_PLACEHOLDER`); other dialects reject multimodal parts. -fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialect) -> Result<()> { +fn write_chat_content( + out: &mut String, + content: &ChatContent, + dialect: DsDialect, + message_index: usize, + media_order: &mut Vec, +) -> Result<()> { match content { ChatContent::Text(text) => out.push_str(text), ChatContent::Parts(parts) => { @@ -738,6 +775,10 @@ fn write_chat_content(out: &mut String, content: &ChatContent, dialect: DsDialec match part { ChatContentPart::ImageUrl { .. } if dialect == DsDialect::V41 => { out.push_str(DEEPSEEK_V41_IMAGE_PLACEHOLDER); + media_order.push(MediaPartSource { + message_index, + content_part_index: index, + }); } _ => out.push_str(part.as_text()?), } diff --git a/rust/src/chat/src/renderer/deepseek_v32/mod.rs b/rust/src/chat/src/renderer/deepseek_v32/mod.rs index 1a41e14f8f..03c2765d19 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/mod.rs @@ -26,6 +26,7 @@ impl ChatRenderer for DeepSeekV32ChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(encoding::render_request(request)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v4/mod.rs b/rust/src/chat/src/renderer/deepseek_v4/mod.rs index 155a6ee140..54b111e505 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/mod.rs @@ -24,6 +24,7 @@ impl ChatRenderer for DeepSeekV4ChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V4)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v41/mod.rs b/rust/src/chat/src/renderer/deepseek_v41/mod.rs index 4e5ddb3725..7737f7ebda 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/mod.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/mod.rs @@ -21,9 +21,12 @@ impl DeepSeekV41ChatRenderer { impl ChatRenderer for DeepSeekV41ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let (prompt, media_order) = + deepseek::render_request_with_media_order(request, DsDialect::V41)?; Ok(RenderedPrompt { - prompt: Prompt::Text(deepseek::render_request(request, DsDialect::V41)?), + prompt: Prompt::Text(prompt), + media_order: Some(media_order), effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/deepseek_v41/tests.rs b/rust/src/chat/src/renderer/deepseek_v41/tests.rs index d7b0c538bb..3babb5c269 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/tests.rs @@ -8,7 +8,7 @@ use serde_json::{Value, json}; use super::DeepSeekV41ChatRenderer; use crate::ChatRenderer; -use crate::event::AssistantContentBlock; +use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort}; @@ -83,6 +83,82 @@ fn inlines_image_placeholder_at_image_part_positions() { ); } +#[test] +fn media_order_follows_reordered_tool_results() { + let renderer = DeepSeekV41ChatRenderer::new(); + let request = ChatRequest { + messages: vec![ + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-a".to_string(), + name: "tool-a".to_string(), + arguments: "{}".to_string(), + }), + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-b".to_string(), + name: "tool-b".to_string(), + arguments: "{}".to_string(), + }), + ]), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-b"), image_part("image-b")], + "call-b", + ), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-a"), image_part("image-a")], + "call-a", + ), + ], + ..ChatRequest::for_test() + }; + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + + assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap()); + expect![[r#" + [ + MediaPartSource { + message_index: 2, + content_part_index: 1, + }, + MediaPartSource { + message_index: 1, + content_part_index: 1, + }, + ] + "#]] + .assert_debug_eq(&media_order); +} + +#[test] +fn media_order_excludes_dropped_historical_developer_content() { + let renderer = DeepSeekV41ChatRenderer::new(); + let request = ChatRequest { + messages: vec![ + ChatMessage::developer(vec![image_part("historical-image")], None), + ChatMessage::user("question"), + ], + ..ChatRequest::for_test() + }; + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let prompt = rendered.prompt.into_text().unwrap(); + + assert!(!prompt.contains("<｜deepseek_image｜>")); + assert!(media_order.is_empty()); +} + +fn image_part(uuid: &str) -> ChatContentPart { + ChatContentPart::ImageUrl { + image_url: format!("data:image/png;base64,{uuid}"), + detail: None, + uuid: Some(uuid.to_string()), + } +} + #[test] fn maps_reasoning_effort_to_reference_numeric_budget() { for (effort, budget) in [ diff --git a/rust/src/chat/src/renderer/harmony/mod.rs b/rust/src/chat/src/renderer/harmony/mod.rs index b07d8458d7..71431cc92b 100644 --- a/rust/src/chat/src/renderer/harmony/mod.rs +++ b/rust/src/chat/src/renderer/harmony/mod.rs @@ -116,6 +116,7 @@ impl ChatRenderer for HarmonyChatRenderer { fn render(&self, request: &ChatRequest) -> Result { Ok(RenderedPrompt { prompt: Prompt::TokenIds(self.render_token_ids(request)?), + media_order: None, effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 5e7ba75e0a..ba2b76d100 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -217,6 +217,7 @@ impl HfChatRenderer { Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs, }) } diff --git a/rust/src/chat/src/renderer/inkling/mod.rs b/rust/src/chat/src/renderer/inkling/mod.rs index a33e3414f2..e8afe66396 100644 --- a/rust/src/chat/src/renderer/inkling/mod.rs +++ b/rust/src/chat/src/renderer/inkling/mod.rs @@ -346,6 +346,7 @@ impl ChatRenderer for InklingChatRenderer { Ok(RenderedPrompt { prompt: Prompt::TokenIds(out), + media_order: None, effective_template_kwargs, }) } diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs index 392b6b642d..88b4e3ae61 100644 --- a/rust/src/chat/src/renderer/kimi_k3/encoding.rs +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -10,6 +10,7 @@ use std::collections::HashMap; use serde_json::{Map, Value, json}; use vllm_tokenizer::Tokenizer; +use super::super::MediaPartSource; use crate::error::{Error, Result}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, @@ -29,6 +30,7 @@ const VALID_THINKING_EFFORTS: &[&str] = &["low", "high", "max"]; pub(super) struct K3TokenWriter<'a> { tokenizer: &'a dyn Tokenizer, token_ids: Vec, + media_order: Vec, } impl<'a> K3TokenWriter<'a> { @@ -36,6 +38,7 @@ impl<'a> K3TokenWriter<'a> { Self { tokenizer, token_ids: Vec::new(), + media_order: Vec::new(), } } @@ -55,13 +58,25 @@ impl<'a> K3TokenWriter<'a> { Ok(()) } - pub(super) fn finish(self) -> Vec { - self.token_ids + fn image(&mut self, message_index: usize, content_part_index: usize) -> Result<()> { + self.control(IMAGE_PLACEHOLDER)?; + self.media_order.push(MediaPartSource { + message_index, + content_part_index, + }); + Ok(()) + } + + pub(super) fn finish(self) -> (Vec, Vec) { + (self.token_ids, self.media_order) } } -/// Render and tokenize one chat request using K3's segment-aware contract. -pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) -> Result> { +/// Render and tokenize one chat request while recording media in placeholder order. +pub(super) fn render_request_with_media_order( + request: &ChatRequest, + tokenizer: &dyn Tokenizer, +) -> Result<(Vec, Vec)> { let thinking = thinking_enabled(request)?; let thinking_effort = thinking.then(|| thinking_effort(request)).transpose()?; let tools = request_tools(request); @@ -126,7 +141,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - resolved.sort_by_key(|item| (item.0, item.1)); } - for (position, _, content, name) in resolved { + for (position, message_index, content, name) in resolved { let tool_name = name.as_deref().ok_or_else(|| { Error::ChatTemplate( "Kimi K3 tool messages need a resolvable tool name: \ @@ -135,7 +150,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - .to_string(), ) })?; - write_tool_message(out, tool_name, position, &content)?; + write_tool_message(out, tool_name, position, message_index, &content)?; } Ok(()) }; @@ -161,14 +176,14 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - } if !local_tools.is_empty() => { write_tool_declare(&mut out, local_tools, true)?; if !content_is_empty(content) { - write_role_message(&mut out, "system", None, content)?; + write_role_message(&mut out, "system", None, message_index, content)?; } } ChatMessage::System { content } | ChatMessage::Developer { content, .. } => { - write_role_message(&mut out, "system", None, content)?; + write_role_message(&mut out, "system", None, message_index, content)?; } ChatMessage::User { content } => { - write_role_message(&mut out, "user", None, content)?; + write_role_message(&mut out, "user", None, message_index, content)?; } ChatMessage::Assistant { content } => { tool_call_id_index.clear(); @@ -364,6 +379,7 @@ fn write_role_message( out: &mut K3TokenWriter<'_>, role: &str, name: Option<&str>, + message_index: usize, content: &ChatContent, ) -> Result<()> { let mut attrs = vec![("role", role.to_string())]; @@ -372,7 +388,7 @@ fn write_role_message( } let attr_refs: Vec<(&str, &str)> = attrs.iter().map(|(k, v)| (*k, v.as_str())).collect(); write_open_tag(out, "message", &attr_refs)?; - write_content(out, content)?; + write_content(out, message_index, content)?; write_close_tag(out, "message")?; out.control(END_OF_MSG) } @@ -381,6 +397,7 @@ fn write_tool_message( out: &mut K3TokenWriter<'_>, tool_name: &str, index: usize, + message_index: usize, content: &ChatContent, ) -> Result<()> { let index_str = index.to_string(); @@ -389,7 +406,7 @@ fn write_tool_message( "message", &[("role", "tool"), ("tool", tool_name), ("index", &index_str)], )?; - write_content(out, content)?; + write_content(out, message_index, content)?; write_close_tag(out, "message")?; out.control(END_OF_MSG) } @@ -471,14 +488,20 @@ fn write_assistant_tool_call( write_close_tag(out, "call") } -fn write_content(out: &mut K3TokenWriter<'_>, content: &ChatContent) -> Result<()> { +fn write_content( + out: &mut K3TokenWriter<'_>, + message_index: usize, + content: &ChatContent, +) -> Result<()> { match content { ChatContent::Text(text) => write_text_with_images(out, text), ChatContent::Parts(parts) => { - for part in parts { + for (content_part_index, part) in parts.iter().enumerate() { match part { ChatContentPart::Text { text } => write_text_with_images(out, text)?, - ChatContentPart::ImageUrl { .. } => out.control(IMAGE_PLACEHOLDER)?, + ChatContentPart::ImageUrl { .. } => { + out.image(message_index, content_part_index)?; + } ChatContentPart::VideoUrl { .. } => { return Err(Error::UnsupportedMultimodalContent("video_url")); } diff --git a/rust/src/chat/src/renderer/kimi_k3/mod.rs b/rust/src/chat/src/renderer/kimi_k3/mod.rs index 09bf356f5d..90750eeac4 100644 --- a/rust/src/chat/src/renderer/kimi_k3/mod.rs +++ b/rust/src/chat/src/renderer/kimi_k3/mod.rs @@ -30,9 +30,12 @@ impl KimiK3ChatRenderer { impl ChatRenderer for KimiK3ChatRenderer { fn render(&self, request: &ChatRequest) -> Result { request.validate()?; + let (token_ids, media_order) = + encoding::render_request_with_media_order(request, self.tokenizer.as_ref())?; Ok(RenderedPrompt { - prompt: Prompt::TokenIds(encoding::render_request(request, self.tokenizer.as_ref())?), + prompt: Prompt::TokenIds(token_ids), + media_order: Some(media_order), effective_template_kwargs: request_template_kwargs(request), }) } diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs index 1d9f984dbb..d0ecbe44d3 100644 --- a/rust/src/chat/src/renderer/kimi_k3/tests.rs +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -284,6 +284,65 @@ fn partial_tool_results_keep_assistant_call_position() { ); } +#[test] +fn media_order_follows_reordered_tool_results() { + let mut request = crate::request::ChatRequest::for_test(); + request.messages = vec![ + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-a".to_string(), + name: "tool-a".to_string(), + arguments: "{}".to_string(), + }), + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call-b".to_string(), + name: "tool-b".to_string(), + arguments: "{}".to_string(), + }), + ]), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-b"), image_part("image-b")], + "call-b", + ), + ChatMessage::tool_response( + vec![ChatContentPart::text("result-a"), image_part("image-a")], + "call-a", + ), + ]; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + let renderer = KimiK3ChatRenderer::new(Arc::new(test_tokenizer())); + + let rendered = renderer.render(&request).unwrap(); + let media_order = rendered.media_order.unwrap(); + let Prompt::TokenIds(token_ids) = rendered.prompt else { + panic!("kimi k3 renderer should return token IDs") + }; + let prompt = test_tokenizer().decode(&token_ids, false).unwrap(); + + assert!(prompt.find("result-a").unwrap() < prompt.find("result-b").unwrap()); + expect![[r#" + [ + MediaPartSource { + message_index: 2, + content_part_index: 1, + }, + MediaPartSource { + message_index: 1, + content_part_index: 1, + }, + ] + "#]] + .assert_debug_eq(&media_order); +} + +fn image_part(uuid: &str) -> ChatContentPart { + ChatContentPart::ImageUrl { + image_url: format!("data:image/png;base64,{uuid}"), + detail: None, + uuid: Some(uuid.to_string()), + } +} + #[test] fn defaults_thinking_effort_to_max() { let rendered = render_request(&crate::request::ChatRequest::for_test()); diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index a4649a34f7..bc1831c375 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -30,11 +30,24 @@ pub use inkling::InklingChatRenderer; pub use kimi_k3::KimiK3ChatRenderer; pub use selection::RendererSelection; +/// Location of one multimodal content part in the source chat request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MediaPartSource { + /// Zero-based message index in the source request. + pub message_index: usize, + /// Zero-based content-part index within the source message. + pub content_part_index: usize, +} + /// Rendered chat prompt submitted to the text backend. #[derive(Debug, Clone, PartialEq)] pub struct RenderedPrompt { /// The rendered prompt, either as text or already tokenized. pub prompt: Prompt, + /// Media sources in rendered placeholder order. + /// `None` uses the message and content-part order in the request (e.g. Jinja). + /// `Some` uses exactly the listed sources; an empty list omits all media. + pub media_order: Option>, /// Effective chat-template kwargs visible to the renderer after applying /// server defaults, request overrides, and typed reasoning controls. pub effective_template_kwargs: HashMap, diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index ab99b4d8ba..ed96134cdc 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -250,6 +250,7 @@ impl ChatRenderer for FakeChatBackend { Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: request.chat_options.template_kwargs.clone(), }) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index f13dc11b27..f5f74d557e 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -252,6 +252,7 @@ impl ChatRenderer for FakeTextBackend { fn render(&self, _request: &ChatRequest) -> vllm_chat::Result { Ok(RenderedPrompt { prompt: Prompt::Text(String::new()), + media_order: None, effective_template_kwargs: Default::default(), }) } diff --git a/rust/src/server/src/routes/http_client_tests.rs b/rust/src/server/src/routes/http_client_tests.rs index 216670b695..a0e0872e64 100644 --- a/rust/src/server/src/routes/http_client_tests.rs +++ b/rust/src/server/src/routes/http_client_tests.rs @@ -186,6 +186,7 @@ impl ChatRenderer for FakeChatBackend { } Ok(RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: Default::default(), }) } diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 5ee2b51f3f..6fbe22e501 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -540,6 +540,7 @@ impl ChatRenderer for FakeChatBackend { } Ok(vllm_chat::RenderedPrompt { prompt: Prompt::Text(prompt), + media_order: None, effective_template_kwargs: Default::default(), }) } From e82794f48ff0812b01543f7c19c34eb0c5b3b1ee Mon Sep 17 00:00:00 2001 From: Xiaozhen Bi Date: Mon, 14 Sep 2026 03:48:53 -0700 Subject: [PATCH 021/139] [CI][Test] Mock CPU backend block sizes in kv_connector unit conftest (#56671) Signed-off-by: Xiaozhen Bi Co-authored-by: Claude Opus 5 --- tests/v1/kv_connector/unit/conftest.py | 27 ++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 tests/v1/kv_connector/unit/conftest.py diff --git a/tests/v1/kv_connector/unit/conftest.py b/tests/v1/kv_connector/unit/conftest.py new file mode 100644 index 0000000000..58c6e632d3 --- /dev/null +++ b/tests/v1/kv_connector/unit/conftest.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + +from vllm.platforms import current_platform + + +@pytest.fixture(autouse=True) +def _mock_cpu_backend_block_sizes(monkeypatch): + """Mock the CPU attention backend's kernel block-size declaration. + + Connector unit tests exercise transfer logic, never the attention + kernel, and hard-code block_size=16. On CPU-only machines + `CPUAttentionBackend` declares `MultipleOf(32)` (#54042), so every + worker construction would fail in `select_common_block_size` before + reaching the code under test. + """ + if not current_platform.is_cpu(): + return + from vllm.v1.attention.backend import MultipleOf + from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend + + monkeypatch.setattr( + CPUAttentionBackend, + "get_supported_kernel_block_sizes", + staticmethod(lambda: [MultipleOf(16)]), + ) From a6c5d6d0fcd7dd5b49e3146991c201b16caebef2 Mon Sep 17 00:00:00 2001 From: Shanshan Shen <467638484@qq.com> Date: Mon, 14 Sep 2026 19:21:30 +0800 Subject: [PATCH 022/139] [ROCm][Perf] Enable CSA multi-stream overlap for DeepSeek-V4 (#51794) Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> --- vllm/models/deepseek_v4/amd/model.py | 11 +- vllm/models/deepseek_v4/amd/rocm.py | 194 ++++++++++++++++++++++ vllm/models/deepseek_v4/attention.py | 31 ++-- vllm/models/deepseek_v4/cpu/cpu_sparse.py | 3 +- 4 files changed, 214 insertions(+), 25 deletions(-) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 349ff52679..c1a6903937 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -957,16 +957,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): self.hc_dim = self.hc_mult * config.hidden_size self.rms_norm_eps = config.rms_norm_eps - # Three aux streams: one per non-default input GEMM in - # DeepseekV4Attention._run_parallel_input_projections - # (compressor kv_score, indexer.weights_proj, indexer.compressor - # kv_score). fused_wqa_wkv stays on the default stream. - # Disable them on ROCm because of hang issues. - aux_stream_list = ( - None - if current_platform.is_rocm() - else [torch.cuda.Stream() for _ in range(3)] - ) + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] self.device = current_platform.device_type # Reserved topk indices buffer for all Indexer layers to reuse. diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index b39ad953e7..3cc9827c1f 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -23,6 +23,7 @@ from vllm.models.deepseek_v4.sparse_mla import ( from vllm.platforms import current_platform from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton +from vllm.utils.multi_stream_utils import execute_in_parallel from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) @@ -571,6 +572,199 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): self.register_buffer("_fused_compressor_weight", None, persistent=False) self._fused_compressor_split_sizes: tuple[int, int] | None = None + if self.indexer is None: + # Only enable multi-stream overlap for CSA layer now. + self.aux_stream_list = None + else: + # Disable indexer inner overlap. + self.indexer.aux_stream = None + + def _enable_csa_multi_stream(self) -> bool: + """All CSA multi-stream gates: env var, streams, and capture region. + + Dict metadata marks piecewise cudagraph, whose eager breaks rebuild + the attention inputs on the owning stream. Forking side streams + there would rely on runtime HIP event sync, which is unreliable in + this overlap on ROCm (event waits can hang), so multi-stream only + runs where the fork/join becomes static graph edges: inside capture, + or with non-dict metadata (full cudagraph or the profile run), which + has no eager breaks. + """ + attn_metadata = get_forward_context().attn_metadata + return self.aux_stream_list is not None and ( + torch.cuda.is_current_stream_capturing() + or not isinstance(attn_metadata, dict) + ) + + def _run_sequential_pipeline( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Disable ROCm streams when the current execution region cannot overlap.""" + aux_streams = self.aux_stream_list + self.aux_stream_list = None + try: + qr_kv, kv_score, indexer_kv_score, indexer_weights = ( + self._run_parallel_input_projections(hidden_states) + ) + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) + self._prepare_and_attn_fn( + hidden_states, + qr, + kv, + qr_scale, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + finally: + self.aux_stream_list = aux_streams + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-allocate attention output with FlashMLA-padded head count. + # The op writes into `o_padded`; we slice to n_local_heads after. + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + if self._enable_csa_multi_stream(): + # The ROCm override consumes these sentinels inside the capture + # boundary, moving the stream fan-out ahead of the projections. + self._prepare_and_attn_fn( + hidden_states, + None, + None, + None, + None, + None, + None, + positions, + o_padded, + ) + else: + self._run_sequential_pipeline(hidden_states, positions, o_padded) + + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + def _prepare_and_attn( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor | None, + kv: torch.Tensor | None, + qr_scale: torch.Tensor | None, + kv_score: torch.Tensor | None, + indexer_kv_score: torch.Tensor | None, + indexer_weights: torch.Tensor | None, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Run the ROCm CSA fork/join inside the graph capture boundary.""" + aux_streams = self.aux_stream_list + # The sequential pipeline disables aux_stream_list before calling + # back with real projection inputs; aux_streams is None ends that + # recursion here. + if aux_streams is None: + saved_streams = self.aux_stream_list + self.aux_stream_list = None + try: + super()._prepare_and_attn( + hidden_states, + cast(torch.Tensor, qr), + cast(torch.Tensor, kv), + qr_scale, + cast(torch.Tensor, kv_score), + cast(torch.Tensor, indexer_kv_score), + cast(torch.Tensor, indexer_weights), + positions, + o_padded, + ) + finally: + self.aux_stream_list = saved_streams + return + + # Re-check: forward's gate ran inside a captured segment that + # _prepare_and_attn_eager (MRV1) then broke, making this region eager. + if not self._enable_csa_multi_stream(): + self._run_sequential_pipeline(hidden_states, positions, o_padded) + return + + indexer = self.indexer + compressor = self.compressor + assert indexer is not None and compressor is not None + + def default_chain(): + qr_kv = self._fused_wqa_wkv_gemm(hidden_states) + qr_out, qr_scale_out, kv_out = self._split_qkv_and_norm(qr_kv) + q = self._wq_b_proj(qr_out, qr_scale_out).view( + -1, self.n_local_heads, self.head_dim + ) + attn_metadata = get_forward_context().attn_metadata + q = self._fused_qnorm_rope_kv_insert(q, kv_out, positions, attn_metadata) + return q, qr_out, qr_scale_out, kv_out + + def main_compressor_chain() -> None: + score = torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + compressor(score, positions, self.rotary_emb) + + def indexer_compressor_chain() -> None: + score = torch.mm( + hidden_states, + indexer.compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + indexer.compressor(score, positions, self.indexer_rotary_emb) + + (q, qr_out, qr_scale_out, kv_out), _ = execute_in_parallel( + default_chain, + [main_compressor_chain, indexer_compressor_chain], + self.ln_events[0], + self.ln_events[1:3], + aux_streams[:2], + enable=True, + ) + + indexer_weights_out, _ = indexer.weights_proj(hidden_states) + # The indexer compressor already ran on aux stream 1; build queries only. + index_q, index_q_scale, weights = indexer( + hidden_states, + qr_out, + None, + indexer_weights_out, + positions, + self.indexer_rotary_emb, + qr_scale_out, + skip_compressor=True, + ) + self._sparse_indexer_and_attn( + hidden_states, + index_q, + index_q_scale, + weights, + q, + kv_out, + positions, + o_padded, + ) + @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: return num_heads diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index d32a48cbe4..9c77980965 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -318,7 +318,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # graph and MRV1 produces garbage (#51430). self._prepare_and_attn_fn = self._prepare_and_attn_eager - # Will be None on ROCm for now. self.aux_stream_list = aux_stream_list # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins @@ -584,7 +583,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # Keep Q projection and KV insertion on the default stream. The indexer # and MLA compressor use aux streams 0 and 1; aux 2 is internal to the - # indexer. ROCm runs the same work sequentially without aux streams. + # indexer. if indexer is not None: assert compressor is not None q, (indexer_inputs, _) = execute_in_parallel( @@ -671,7 +670,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs # on aux streams 0..2 when their owning module exists. ln_events[0] # is the fan-out start event; ln_events[1..3] are per-aux done events. - # On ROCm, aux_streams is None and execute_in_parallel runs serially. aux_fns: list[Callable[[], Any] | None] = [None, None, None] if self.compressor is not None: @@ -1055,11 +1053,12 @@ class DeepseekV4Indexer(nn.Module): self, hidden_states: torch.Tensor, qr: torch.Tensor, - compressed_kv_score: torch.Tensor, + compressed_kv_score: torch.Tensor | None, indexer_weights: torch.Tensor, positions: torch.Tensor, rotary_emb: nn.Module, qr_scale: torch.Tensor | None = None, + skip_compressor: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: compressor = self.compressor @@ -1072,7 +1071,8 @@ class DeepseekV4Indexer(nn.Module): ): # candidates num smaller than topk, every candidate is selected # but we still need to build k cache - compressor(compressed_kv_score, positions, rotary_emb) + if not skip_compressor: + compressor(compressed_kv_score, positions, rotary_emb) assert self.topk_indices_buffer is not None num_tokens = ( indexer_metadata.num_decode_tokens @@ -1102,15 +1102,18 @@ class DeepseekV4Indexer(nn.Module): use_fp4=self.use_fp4_kv, ) - # compressor returns None and writes K to the indexer KV cache; the - # join orders that write before indexer_op (skip_k_cache_insert=True). - (q_quant, weights), _ = maybe_execute_in_parallel( - wq_b_and_q_quant, - lambda: compressor(compressed_kv_score, positions, rotary_emb), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) + if not skip_compressor: + # compressor returns None and writes K to the indexer KV cache; the + # join orders that write before indexer_op (skip_k_cache_insert=True). + (q_quant, weights), _ = maybe_execute_in_parallel( + wq_b_and_q_quant, + lambda: compressor(compressed_kv_score, positions, rotary_emb), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + else: + q_quant, weights = wq_b_and_q_quant() if isinstance(q_quant, tuple): q, q_scale = q_quant else: diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py index d51e589652..a2c3b00493 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_sparse.py +++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py @@ -198,11 +198,12 @@ class DeepseekV4CPUIndexer(DeepseekV4Indexer): self, hidden_states: torch.Tensor, qr: torch.Tensor, - compressed_kv_score: torch.Tensor, + compressed_kv_score: torch.Tensor | None, indexer_weights: torch.Tensor, positions: torch.Tensor, rotary_emb: nn.Module, qr_scale: torch.Tensor | None = None, + skip_compressor: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: """CPU override: no aux streams, so wq_b_and_q_quant and the compressor run straight-line instead of through From a5f6f61a8a8eb9d8e839cf7e322b23cc8ae406e5 Mon Sep 17 00:00:00 2001 From: Ganesh R Date: Mon, 14 Sep 2026 17:21:38 +0530 Subject: [PATCH 023/139] [CI][CPU] Add speculative-decoding coverage to CPU CI (#54934) Signed-off-by: R Signed-off-by: jiang1.li Co-authored-by: Cursor Co-authored-by: jiang1.li --- .buildkite/hardware_tests/cpu.yaml | 46 +++++ .../scripts/hardware_ci/run-cpu-test.sh | 7 +- docs/features/README.md | 2 +- tests/v1/e2e/test_cpu_spec_decode.py | 167 ++++++++++++++++++ 4 files changed, 219 insertions(+), 3 deletions(-) create mode 100644 tests/v1/e2e/test_cpu_spec_decode.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 769e3d4adc..b50a082727 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -52,6 +52,52 @@ steps: tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp" parallelism: 2 +# Speculative decoding under a CPU platform. The unit tests already pass on +# CPU but no CPU job runs them; the e2e tests gate the engine path itself, +# which unit tests cannot reach. +- label: CPU-Spec Decode Tests + key: cpu-spec-decode-tests + depends_on: [] + soft_fail: true + device: intel_cpu + no_plugin: true + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/cpu/ + - vllm/v1/worker/cpu_model_runner.py + - vllm/v1/worker/gpu/spec_decode/ + - vllm/v1/attention/backends/cpu_attn.py + - vllm/config/speculative.py + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/model_executor/models/llama_eagle3.py + - vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py + - vllm/model_executor/layers/mamba/mamba_mixer2.py + - vllm/model_executor/layers/mamba/ops/cpu/ + - tests/v1/e2e/test_cpu_spec_decode.py + - tests/v1/spec_decode/ + - tests/v1/worker/ + commands: + - | + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " + pytest -x -v -s \ + tests/v1/worker/test_gpu_autoregressive_speculator.py \ + tests/v1/worker/test_gpu_extract_hidden_states_speculator.py \ + tests/v1/worker/test_gpu_batch_ordering.py \ + tests/v1/worker/test_gpu_rejection_sampler_chunking.py \ + tests/v1/worker/test_eagle3_aux_hidden_states_pp.py \ + tests/v1/worker/test_spec_decode_embed_sharing_pp.py \ + tests/v1/spec_decode/test_adaptive_verification.py \ + tests/v1/spec_decode/test_dflash2.py \ + tests/v1/spec_decode/test_dflash_causality.py \ + tests/v1/spec_decode/test_eagle_draft_attn_metadata.py \ + tests/v1/spec_decode/test_draft_attention_backend_override.py \ + tests/v1/spec_decode/test_draft_moe_backend_override.py \ + tests/v1/spec_decode/test_ngram.py \ + tests/v1/spec_decode/test_llm_base_proposer_sampling.py + unset VLLM_CPU_CI_ENV + pytest -x -v -s tests/v1/e2e/test_cpu_spec_decode.py" + # Note: SDE can't be downloaded from CI host because of AWS WAF # - label: CPU-Compatibility Tests # depends_on: [] diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index e47cc22c8f..74c7749d50 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -83,8 +83,11 @@ rm -f "$build_log" # Run the image, setting --shm-size=4g for tensor parallel. Default to # HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network; -# retry once online if the cache is missing something. -OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)' +# retry once online if the cache is missing something. vllm's get_config() +# wraps the raw huggingface_hub offline-mode errors in a generic ValueError +# (see transformers_utils/config.py), so match that message too or the +# fallback never triggers for a config-lookup cache miss. +OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified' run_test() { local hf_offline=$1 docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \ diff --git a/docs/features/README.md b/docs/features/README.md index 28362f4014..1df8960ee3 100644 --- a/docs/features/README.md +++ b/docs/features/README.md @@ -64,7 +64,7 @@ th:not(:first-child) { | [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | -| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | +| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) | | [pooling](../models/pooling_models/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | enc-dec | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | diff --git a/tests/v1/e2e/test_cpu_spec_decode.py b/tests/v1/e2e/test_cpu_spec_decode.py new file mode 100644 index 0000000000..2d34c11656 --- /dev/null +++ b/tests/v1/e2e/test_cpu_spec_decode.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU speculative-decoding end-to-end correctness.""" + +import os + +import pytest + +from tests.models.utils import check_logprobs_close +from tests.v1.e2e.spec_decode.utils import compute_acceptance_len +from vllm import LLM, SamplingParams +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON + +pytestmark = pytest.mark.cpu_model + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +# Bound the KV cache so the run does not scale with host memory. +os.environ.setdefault("VLLM_CPU_KVCACHE_SPACE", "1") + +MODEL = "Qwen/Qwen3-0.6B" +NUM_LOGPROBS = 5 +SP = SamplingParams(max_tokens=48, temperature=0, logprobs=NUM_LOGPROBS) + +# Repetitive prompts so the ngram proposer finds matches and actually drafts; +# without hits the run is indistinguishable from plain decoding. +PROMPTS = [ + "The capital of France is Paris. The capital of Italy is Rome. " + "The capital of Spain is Madrid. The capital of France is", + "def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - b\n\n" + "def mul(a, b):\n return", + "one two three four five one two three four five one two three four", +] + + +def _make_llm(**overrides) -> LLM: + base = dict( + model=MODEL, + dtype="bfloat16", + max_model_len=1024, + disable_log_stats=False, + ) + base.update(overrides) + return LLM(**base) + + +def _tuples(outputs) -> list[tuple[list[int], str, object]]: + return [ + (list(o.outputs[0].token_ids), o.outputs[0].text, o.outputs[0].logprobs) + for o in outputs + ] + + +@pytest.fixture(scope="module") +def baseline_refs(): + llm = _make_llm() + refs = _tuples(llm.generate(PROMPTS, SP)) + del llm + return refs + + +def _spec_metric(llm: LLM, name: str) -> float: + for metric in llm.get_metrics(): + if metric.name == name: + return float(getattr(metric, "value", 0.0)) + return 0.0 + + +def test_ngram_spec_decode_matches_baseline(baseline_refs): + """Greedy output must be unchanged by speculative decoding. + + The drafted-token check keeps the test from passing vacuously when the + proposer silently produces nothing. + """ + llm = _make_llm( + speculative_config={ + "method": "ngram", + "prompt_lookup_max": 5, + "prompt_lookup_min": 3, + "num_speculative_tokens": 3, + }, + ) + got = _tuples(llm.generate(PROMPTS, SP)) + drafted = _spec_metric(llm, "vllm:spec_decode_num_draft_tokens") + accepted = _spec_metric(llm, "vllm:spec_decode_num_accepted_tokens") + del llm + + assert drafted > 0, "ngram proposer drafted no tokens; the test would be vacuous" + assert accepted > 0, f"no draft tokens accepted out of {drafted} drafted" + + check_logprobs_close( + outputs_0_lst=baseline_refs, + outputs_1_lst=got, + name_0="no_spec_decode", + name_1="ngram_spec_decode", + ) + + +SMOKE_PROMPTS = [ + "The capital of France is", + "2 + 2 equals", + "In one word, the color of the sky is", + "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:", +] +SMOKE_SP = SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True) + +# Methods routed to the V2 runner, using the same target/draft pairs as the GPU +# suite. MTP needs no separate draft checkpoint: its speculator layer is +# trained jointly with, and shipped inside, the target model's own weights. +SMOKE_CONFIGS = [ + pytest.param( + "meta-llama/Llama-3.2-1B-Instruct", + { + "method": "eagle3", + "model": "nm-testing/Llama3_2_1B_speculator.eagle3", + "num_speculative_tokens": 3, + }, + id="eagle3", + ), + pytest.param( + "Qwen/Qwen3-4B", + { + "method": "dflash", + "model": "z-lab/Qwen3-4B-DFlash-b16", + # The checkpoint allows 16, but accepted tokens plateau around 64 + # regardless, so a larger draft only burns CPU on rejected tokens. + "num_speculative_tokens": 4, + }, + id="dflash", + ), + pytest.param( + "Qwen/Qwen3.5-0.8B-Base", + { + "method": "mtp", + "num_speculative_tokens": 3, + }, + id="mtp", + ), +] + +# Measured 1.85 (eagle3), 1.99 (dflash), 2.29 (mtp) on CPU; this floor leaves +# room for numeric drift while still failing if speculation stops paying off. +MIN_ACCEPTANCE_LEN = 1.5 + + +@pytest.mark.skipif( + not HAS_TRITON, reason="the V2 runner needs triton-cpu, which CI builds" +) +@pytest.mark.parametrize("model,speculative_config", SMOKE_CONFIGS) +def test_v2_speculator_smoke(model: str, speculative_config: dict): + """Gates that drafts are accepted often enough for speculation to pay off.""" + llm = _make_llm(model=model, speculative_config=speculative_config) + try: + assert llm.llm_engine.vllm_config.use_v2_model_runner, ( + f"{speculative_config['method']} fell back to the V1 runner; " + "this test is meant to cover the V2 path" + ) + llm.generate(SMOKE_PROMPTS, SMOKE_SP) + acceptance = compute_acceptance_len(llm.get_metrics()) + assert acceptance >= MIN_ACCEPTANCE_LEN, ( + f"acceptance length {acceptance:.3f} below {MIN_ACCEPTANCE_LEN}; " + "drafts are being rejected far more than expected" + ) + finally: + del llm From c612e2bff0d045a5b2304371793985c330fba436 Mon Sep 17 00:00:00 2001 From: AlexHuang Date: Mon, 14 Sep 2026 20:02:45 +0800 Subject: [PATCH 024/139] [Bugfix][KV Offload] Reuse in-flight async lookup probes (#55823) Signed-off-by: Alex Co-authored-by: OpenAI Codex Co-authored-by: Or Ozeri --- .../kv_offload/tiering/test_async_lookup.py | 76 +++++++++++++++++-- vllm/v1/kv_offload/tiering/async_lookup.py | 57 ++++++++++---- 2 files changed, 112 insertions(+), 21 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py index ff5cbcd5dd..3649c8df04 100644 --- a/tests/v1/kv_offload/tiering/test_async_lookup.py +++ b/tests/v1/kv_offload/tiering/test_async_lookup.py @@ -8,7 +8,7 @@ from collections.abc import Iterable import pytest from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key -from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager +from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager, LookupPhase def _key(i: int) -> OffloadKey: @@ -46,10 +46,13 @@ class TestAsyncLookupManager: def test_found_key_returns_true(self): mgr = InMemoryLookupManager(existing_keys={_key(1)}) assert mgr.lookup(_key(1), _ctx()) is None + assert mgr._lookup_state[_key(1)].phase is LookupPhase.PENDING mgr.flush() + assert mgr._lookup_state[_key(1)].phase is LookupPhase.IN_FLIGHT mgr._results_ready.wait() mgr._results_ready.clear() assert mgr.lookup(_key(1), _ctx()) is True + assert mgr._lookup_state[_key(1)].phase is LookupPhase.RESOLVED mgr.shutdown() def test_not_found_key_returns_false(self): @@ -106,6 +109,68 @@ class TestAsyncLookupManager: assert _key(1) not in mgr._lookup_state mgr.shutdown() + def test_cleanup_reuses_in_flight_probe(self, monkeypatch: pytest.MonkeyPatch): + """A replacement request shares the in-flight probe and its verdict.""" + key = _key(1) + mgr = InMemoryLookupManager(existing_keys={key}) + ctx_b = _ctx("req_b") + probe_started = threading.Event() + release_probe = threading.Event() + batch_lookup = mgr.batch_lookup + + def blocking_lookup(keys, req_context): + probe_started.set() + if not release_probe.wait(timeout=5): + raise TimeoutError("Test did not release the backend probe") + return batch_lookup(keys, req_context) + + monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup) + try: + assert mgr.lookup(key, _ctx("req_a")) is None + mgr.flush() + assert probe_started.wait(timeout=5) + assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT + mgr.cleanup("req_a") + assert mgr.lookup(key, ctx_b) is None + assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT + mgr.flush() + + release_probe.set() + batch = mgr._pending_results.get(timeout=5) + mgr._pending_results.put(batch) + replacement_result = mgr.lookup(key, ctx_b) + finally: + release_probe.set() + mgr.shutdown() + + assert mgr.batch_lookup_calls == 1 + assert replacement_result is True + + @pytest.mark.parametrize( + "reclaim_at_shutdown", [False, True], ids=["flush", "shutdown"] + ) + def test_unclaimed_probe_reclaimed_without_lookup(self, reclaim_at_shutdown: bool): + """A completed orphan is released by flush or shutdown without a lookup.""" + key = _key(1) + mgr = InMemoryLookupManager(existing_keys={key}) + try: + mgr.lookup(key, _ctx("req_a")) + mgr.flush() + mgr.cleanup("req_a") + assert key in mgr._lookup_state + assert not mgr._req_keys + + if not reclaim_at_shutdown: + batch = mgr._pending_results.get(timeout=5) + mgr._pending_results.put(batch) + mgr.flush() + assert key not in mgr._lookup_state + finally: + mgr.shutdown() + + assert not mgr._lookup_state + assert mgr._pending_results.empty() + def test_stale_result_ignored_after_cleanup_and_key_reuse(self): key = _key(1) mgr = InMemoryLookupManager() @@ -119,11 +184,12 @@ class TestAsyncLookupManager: generation = mgr._lookup_state[key].generation assert generation != stale_generation + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + current_result = mgr._pending_results.get(timeout=5) mgr._pending_results.put([(key, stale_generation, True)]) - mgr.drain_results() - assert mgr.lookup(key, ctx_b) is None - - mgr._pending_results.put([(key, generation, False)]) + mgr._pending_results.put(current_result) mgr.drain_results() assert mgr.lookup(key, ctx_b) is False mgr.shutdown() diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index f111f233b6..c08c63d018 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -27,8 +27,10 @@ lookup() accumulates new keys in _lookup_batch without touching the queue. flush() is called once per step from the tier's on_schedule_end(), posting the entire batch as a single queue item so the background thread sees one batch per step. -drain_results() is called before any lookup() calls in the same step, so -lookup() is a pure OrderedDict operation. +Results are drained on the first lookup after each flush, at flush(), and +after worker shutdown. In-flight lookups with no remaining request references +are retained until their results are drained, allowing new requests to share +the same probe. """ import queue @@ -36,6 +38,7 @@ import threading from abc import ABC, abstractmethod from collections.abc import Collection, Iterable from dataclasses import dataclass, field +from enum import Enum, auto from vllm.logger import init_logger from vllm.v1.kv_offload.base import OffloadKey, ReqContext @@ -43,10 +46,21 @@ from vllm.v1.kv_offload.base import OffloadKey, ReqContext logger = init_logger(__name__) +class LookupPhase(Enum): + """Lifecycle phase of a lookup probe.""" + + PENDING = auto() # Accumulated in _lookup_batch, but not yet submitted. + IN_FLIGHT = auto() # Submitted to the worker, but not yet resolved. + RESOLVED = auto() # A final verdict has been applied to the state. + + @dataclass(slots=True) class LookupState: generation: int - result: bool | None = None # True (found), False (not found), None + phase: LookupPhase = LookupPhase.PENDING + # None while pending/in flight; True if the key exists; False if absent or + # explicitly marked missing after a failed load. + result: bool | None = None request_ids: set[str] = field(default_factory=set) # requests asking for the lookup @@ -155,24 +169,28 @@ class AsyncLookupManager(ABC): Called once per step from on_schedule_end() after all lookup() calls are done. The worker receives the full batch and processes it during the model-execution window, maximising time available before the next - step's drain_results(). Safe to call with an empty batch (no-op). + step's drain_results(). Also drains completed lookups when there + are no new keys to submit. """ + self.drain_results() self._need_to_drain = True batch = self._lookup_batch self._lookup_batch = [] - batch = [ - (key, req_context, generation) - for key, req_context, generation in batch - if (state := self._lookup_state.get(key)) is not None - and state.generation == generation - ] - if batch: - self._lookup_queue.put(batch) + in_flight_batch = [] + for key, req_context, generation in batch: + state = self._lookup_state.get(key) + if state is None or state.generation != generation: + continue + assert state.phase is LookupPhase.PENDING + state.phase = LookupPhase.IN_FLIGHT + in_flight_batch.append((key, req_context, generation)) + if in_flight_batch: + self._lookup_queue.put(in_flight_batch) def drain_results(self) -> None: """Apply pending worker results to _lookup_state. - Called from lookup() before checking state. + Called from lookup(), flush(), and shutdown() on the scheduler thread. """ while True: try: @@ -183,6 +201,10 @@ class AsyncLookupManager(ABC): state = self._lookup_state.get(key) if state is None or state.generation != generation: continue + if not state.request_ids: + del self._lookup_state[key] + continue + assert state.phase is LookupPhase.IN_FLIGHT # Each lookup generation is enqueued exactly once. A matching # generation must not receive a second result; stale # generations were discarded above. @@ -192,6 +214,7 @@ class AsyncLookupManager(ABC): "failed-load livelock" ) state.result = result + state.phase = LookupPhase.RESOLVED def mark_miss(self, keys: Collection[OffloadKey]) -> None: """Force the cached verdict for ``keys`` to False after a failed load, so @@ -201,9 +224,10 @@ class AsyncLookupManager(ABC): state = self._lookup_state.get(key) if state is not None: state.result = False + state.phase = LookupPhase.RESOLVED def cleanup(self, req_id: str) -> None: - """Remove entries no longer needed by any active request. + """Release request references, retaining in-flight lookups. Called from the tier's on_request_finished(). Uses the reverse index to visit only keys associated with this request. @@ -211,13 +235,14 @@ class AsyncLookupManager(ABC): for key in self._req_keys.pop(req_id, ()): state = self._lookup_state[key] state.request_ids.discard(req_id) - if not state.request_ids: + if not state.request_ids and state.phase is not LookupPhase.IN_FLIGHT: del self._lookup_state[key] def shutdown(self) -> None: - """Stop the worker thread.""" + """Stop the worker thread and drain completed lookups.""" self._lookup_queue.put(None) # unblock _worker from _lookup_queue.get() self._thread.join() + self.drain_results() # ------------------------------------------------------------------ # Internal helpers From 3f55ad2f073a7d1f573db5a5a3fe3dc9c688ddee Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Mon, 14 Sep 2026 20:16:25 +0800 Subject: [PATCH 025/139] [Qwen3.8-Flash-Next] Fuse PLE residual and QSA output gate (#55309) Signed-off-by: Canlin Signed-off-by: zjy0516 Co-authored-by: zjy0516 Co-authored-by: Kimi Code --- tests/models/qwen4_exp/test_ple.py | 15 ++++-- tests/models/qwen4_exp/test_qsa_reference.py | 3 ++ vllm/models/qwen4_exp/nvidia/model.py | 2 +- vllm/models/qwen4_exp/nvidia/ops/ple.py | 17 ++++++- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 48 ++++++++++++++++++++ vllm/models/qwen4_exp/nvidia/ple_layer.py | 30 ++++++++++-- vllm/models/qwen4_exp/nvidia/qsa.py | 8 +++- 7 files changed, 111 insertions(+), 12 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 020e1cc8fc..e7d0e72d8a 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -1610,6 +1610,12 @@ def test_fused_conv_correctness( dtype=torch.bfloat16, generator=rng, ) + outer_residual = torch.randn( + inputs.shape, + device=device, + dtype=torch.bfloat16, + generator=rng, + ) null_state = conv_state[NULL_BLOCK_ID].clone() residual_kernel = residual.clone() residual_reference = residual.clone() @@ -1617,6 +1623,7 @@ def test_fused_conv_correctness( module._short_conv_dilated_dispatch( inputs=inputs, residual=residual_kernel, + outer_residual=outer_residual, metadata=metadata, conv_state=conv_state, conv_weights=weights, @@ -1630,15 +1637,17 @@ def test_fused_conv_correctness( conv_state_len=module.conv_state_len, dilation=module.short_conv_dilation, ) + residual_reference = outer_residual + residual_reference - torch.testing.assert_close( - residual_kernel.float(), residual_reference.float(), atol=3e-2, rtol=3e-2 + assert torch.equal( + residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens] ) assert torch.equal(conv_state, state_reference) assert torch.equal(conv_state[NULL_BLOCK_ID], null_state) if case.graph_padding: assert torch.equal( - residual_kernel[num_real_tokens:], residual[num_real_tokens:] + residual_kernel[num_real_tokens:], + (outer_residual + residual)[num_real_tokens:], ) diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index 9259972ac9..16294a30d7 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -944,6 +944,7 @@ def test_qsa_sparse_paged_attention_correctness( q = torch.randn( num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16 ) + output_gate = torch.randn_like(q) kv_cache = torch.randn( num_cache_blocks, page_size, @@ -1023,6 +1024,7 @@ def test_qsa_sparse_paged_attention_correctness( block_table, token_to_req, use_prefill_config=use_prefill_config, + output_gate=output_gate, ) expected = _qsa_sparse_paged_attention_reference( q, @@ -1033,6 +1035,7 @@ def test_qsa_sparse_paged_attention_correctness( token_to_req, scale, ) + expected = expected * torch.sigmoid(output_gate) torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 28c3ea69d1..54ee6f02b4 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -297,7 +297,7 @@ class Qwen4ExpDecoderLayer(nn.Module): if input_ids is None or query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") - hidden_states = hidden_states + self.ple( + hidden_states = self.ple( hidden_states, input_ids, query_start_loc, diff --git a/vllm/models/qwen4_exp/nvidia/ops/ple.py b/vllm/models/qwen4_exp/nvidia/ops/ple.py index d76e2705a7..bc4c99957d 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/ple.py +++ b/vllm/models/qwen4_exp/nvidia/ops/ple.py @@ -322,6 +322,7 @@ def _ple_conv_kernel( state_ptr, w_ptr, residual_ptr, + outer_residual_ptr, state_idx_ptr, qsl_ptr, num_acc_ptr, @@ -442,11 +443,21 @@ def _ple_conv_kernel( mask=c_mask, other=0.0, ) + # Preserve the original eager operation boundaries: short convolution is + # first accumulated into the BF16/FP16 PLE output, then the outer residual + # is added to that rounded value. + ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) + outer_residual = tl.load( + outer_residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) if launch_pdl: tl.extra.cuda.gdc_launch_dependents() tl.store( residual_ptr + output_t * C + c_offs, - residual + conv_output, + ple_output, mask=c_mask, ) @@ -564,6 +575,7 @@ def ple_conv( conv_state: torch.Tensor, conv_weights: torch.Tensor, state_indices: torch.Tensor, + outer_residual: torch.Tensor, *, mode: Literal["decode", "spec", "prefill"], dilation: int, @@ -573,7 +585,7 @@ def ple_conv( spec_query_len: int = 1, token_indices: torch.Tensor | None = None, ) -> None: - """Add short-convolution output to ``residual`` and update its state.""" + """Add short convolution and the outer residual; update state.""" BLOCK_C = 512 kernel_spec_query_len = spec_query_len if mode == "spec" else 1 T, C = inputs.shape @@ -624,6 +636,7 @@ def ple_conv( conv_state, conv_weights, residual, + outer_residual, state_indices, query_start_loc, num_accepted_tokens, diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 709acbde00..d3986ba38a 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -24,6 +24,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, stride_q_row, stride_q_head, stride_k_block, @@ -36,6 +37,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( stride_table_req, stride_output_row, stride_output_head, + stride_output_gate_row, + stride_output_gate_head, num_rows, num_cache_blocks, num_requests, @@ -151,6 +154,18 @@ def _qsa_sparse_paged_gqa_splitk_kernel( ) output_mask = head_offsets[:, None] < GROUP_SIZE if NUM_SPLITS == 1: + # Preserve the unfused path's BF16 attention-output rounding before + # applying the gate in FP32. + normalized_output = normalized_output.to(output_ptr.dtype.element_ty) + output_gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + (first_head + head_offsets[:, None]) * stride_output_gate_head + + dim_offsets[None, :], + mask=output_mask, + other=0.0, + ).to(tl.float32) + normalized_output = normalized_output.to(tl.float32) * tl.sigmoid(output_gate) tl.store( output_ptr + row * stride_output_row @@ -186,8 +201,11 @@ def _qsa_merge_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, stride_output_row, stride_output_head, + stride_output_gate_row, + stride_output_gate_head, num_rows, HEAD_DIM: tl.constexpr, NUM_QUERY_HEADS: tl.constexpr, @@ -219,6 +237,16 @@ def _qsa_merge_splitk_kernel( ) merged = tl.sum(partial_output * weights[:, None], axis=0) merged = tl.where(denominator > 0, merged / denominator, 0.0) + # Preserve the unfused path's BF16 attention-output rounding before + # applying the gate in FP32. + merged = merged.to(output_ptr.dtype.element_ty) + output_gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + head * stride_output_gate_head + + dim_offsets + ).to(tl.float32) + merged = merged.to(tl.float32) * tl.sigmoid(output_gate) tl.store( output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, merged, @@ -455,6 +483,8 @@ def qsa_sparse_paged_attention( token_to_req: torch.Tensor, use_prefill_config: bool, out: torch.Tensor | None = None, + *, + output_gate: torch.Tensor, ) -> torch.Tensor: """Run sparse GQA directly over paged BF16 K/V caches. @@ -496,6 +526,8 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse output must match its query") assert out.dtype == q.dtype and out.device == q.device assert out.stride(2) == 1 + assert output_gate.is_contiguous() + output_gate_view = output_gate.view_as(q) if not q.shape[0]: return out @@ -533,6 +565,7 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + output_gate_view, q.stride(0), q.stride(1), k_cache.stride(0), @@ -545,6 +578,8 @@ def qsa_sparse_paged_attention( block_table.stride(0), out.stride(0), out.stride(1), + output_gate_view.stride(0), + output_gate_view.stride(1), q.shape[0], k_cache.shape[0], block_table.shape[0], @@ -568,8 +603,11 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + output_gate_view, out.stride(0), out.stride(1), + output_gate_view.stride(0), + output_gate_view.stride(1), q.shape[0], HEAD_DIM=q.shape[2], NUM_QUERY_HEADS=q.shape[1], @@ -632,6 +670,10 @@ def warmup_qsa_sparse_paged_attention( output_ptr = TritonWarmupTensor( torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) ) + # The output gate is mandatory at runtime; warm the gated specialization. + output_gate_ptr = TritonWarmupTensor( + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) + ) head_stride = head_dim row_stride = num_query_heads * head_dim num_cache_blocks = triton_scalar_specialization_rep(kv_cache.shape[0]) @@ -659,6 +701,7 @@ def warmup_qsa_sparse_paged_attention( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, row_stride, head_stride, key_cache.stride(0), @@ -671,6 +714,8 @@ def warmup_qsa_sparse_paged_attention( block_table.stride(0), row_stride, head_stride, + row_stride, + head_stride, num_rows, num_cache_blocks, num_requests, @@ -693,6 +738,9 @@ def warmup_qsa_sparse_paged_attention( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, + row_stride, + head_stride, row_stride, head_stride, num_rows, diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 26622357ca..2ec8790e8d 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -195,6 +195,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self, inputs: torch.Tensor, residual: torch.Tensor, + outer_residual: torch.Tensor, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, @@ -209,6 +210,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): has_non_spec = has_prefill or has_decode inputs = inputs[: metadata.num_actual_tokens] residual = residual[: metadata.num_actual_tokens] + outer_residual = outer_residual[: metadata.num_actual_tokens] spec_token_indices = None non_spec_token_indices = None @@ -235,6 +237,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): conv_state=conv_state, conv_weights=conv_weights, state_indices=spec_state_indices, + outer_residual=outer_residual, mode="spec", dilation=self.short_conv_dilation, query_start_loc=query_start_loc, @@ -259,11 +262,17 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): residual_d, residual_p = torch.split( residual, [num_decode_tokens, num_prefill_tokens], dim=0 ) + outer_residual_d, outer_residual_p = torch.split( + outer_residual, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) token_indices_d = None token_indices_p = None else: inputs_d = inputs_p = inputs residual_d = residual_p = residual + outer_residual_d = outer_residual_p = outer_residual token_indices_d, token_indices_p = torch.split( non_spec_token_indices, [num_decode_tokens, num_prefill_tokens], @@ -277,6 +286,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices_d, + outer_residual=outer_residual_d, mode="decode", dilation=self.short_conv_dilation, has_initial_states=metadata.has_initial_states_d, @@ -298,6 +308,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices_p, + outer_residual=outer_residual_p, mode="prefill", dilation=self.short_conv_dilation, query_start_loc=query_start_loc, @@ -316,6 +327,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices[:num_decode_rows], + outer_residual=outer_residual, mode="decode", dilation=self.short_conv_dilation, has_initial_states=metadata.has_initial_states_d, @@ -324,12 +336,18 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): # State routing consumes the current request metadata on every replay. @eager_break_during_capture - def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None: + def _short_conv( + self, + inputs: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, + ) -> None: forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata - # Profiling omits all metadata or this Mamba entry. The residual - # already contains the gated output, so short convolution is a no-op. + # Profiling omits all metadata or this Mamba entry. Short convolution + # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: + residual.add_(outer_residual) return if not isinstance(attn_metadata, dict): @@ -340,6 +358,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: + residual.add_(outer_residual) return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( @@ -367,6 +386,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self._short_conv_dilated_dispatch( inputs=inputs, residual=residual, + outer_residual=outer_residual, metadata=layer_attn_metadata, conv_state=conv_state, conv_weights=conv_weights.to(dtype=inputs.dtype), @@ -404,7 +424,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self.norm_conv.weight, self.norm_key.eps, ) - self._short_conv(conv_input, gated_output) + # State routing depends on runtime request metadata and remains outside + # the graph; short convolution accumulates into gated_output. + self._short_conv(conv_input, gated_output, hidden_states) return gated_output diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 90449d5706..1c2808a6c4 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -121,6 +121,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): output: torch.Tensor, token_to_req: torch.Tensor, use_prefill_config: bool, + output_gate: torch.Tensor, output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: @@ -157,6 +158,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): token_to_req, use_prefill_config, output[:num_tokens], + output_gate=output_gate[:num_tokens], ) return output @@ -353,6 +355,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, + output_gate: torch.Tensor, ) -> None: metadata = get_forward_context().attn_metadata if isinstance(metadata, list): @@ -396,6 +399,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): output, token_to_req=side_metadata.token_to_req, use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len, + output_gate=output_gate, ) def forward( @@ -405,6 +409,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v, gate = self._project_qkv_gate(qkv, positions) + assert gate is not None num_tokens = hidden_states.shape[0] query = q.view(num_tokens, self.num_heads, self.head_dim) key = k.view(num_tokens, self.num_kv_heads, self.head_dim) @@ -419,10 +424,9 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): key, value, attn_output, + gate, ) flat_output = attn_output.view(num_tokens, -1) - if gate is not None: - flat_output = flat_output * torch.sigmoid(gate) output, _ = self.o_proj(flat_output) return output From 1d0d1081c409176a0e7c24742f7c60b9943657aa Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:52:56 +0200 Subject: [PATCH 026/139] [ROCm][Connector] SWA+HMA-support in MoRI-IO connector (Gemma4) (#53721) Signed-off-by: simondanielsson --- .../unit/test_moriio_connector.py | 232 +++++++++++++++++- .../kv_connector/unit/test_moriio_routing.py | 7 +- .../kv_connector/v1/moriio/moriio_common.py | 7 +- .../v1/moriio/moriio_connector.py | 208 +++++++++++++--- 4 files changed, 408 insertions(+), 46 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index d169779c5d..836f9407bc 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -39,10 +39,12 @@ from vllm.utils.network_utils import ( make_zmq_path, ) from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + SlidingWindowSpec, compute_layer_kv_cache_shape_bytes, ) @@ -338,7 +340,7 @@ def test_write_mode_saves_local_block_ids(): req_meta = kv_connector_metadata.reqs_to_save[request_id] for block_id, block in zip( - req_meta.local_block_ids, + req_meta.local_block_ids[0], scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[ request_id ], @@ -393,7 +395,7 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids(): req_meta = kv_connector_metadata.reqs_to_save[request_id] for block_id, block in zip( - req_meta.local_block_ids, + req_meta.local_block_ids[0], scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[ request_id ], @@ -432,7 +434,7 @@ def test_read_mode_loads_remote_block_ids(): ].req_to_blocks[request_id] # Set remote block ids to be fetched. - request.kv_transfer_params["remote_block_ids"] = block_list + request.kv_transfer_params["remote_block_ids"] = [block_list] # Remote Prefill, triggers MoRIIOConnectorMetadata. @@ -457,7 +459,7 @@ def test_read_mode_loads_remote_block_ids(): req_meta = kv_connector_metadata.reqs_to_recv[request_id] for block_id, block in zip( - req_meta.local_block_ids, + req_meta.local_block_ids[0], scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[ request_id ], @@ -699,3 +701,225 @@ def test_resolve_host_ip_prefers_extra_config(): fallback = get_ip() assert resolve_host_ip({}) == fallback assert resolve_host_ip({"host_ip": ""}) == fallback + + +def _make_hybrid_kv_cache_config() -> KVCacheConfig: + """One full-attention group + one sliding-window group (Gemma-like).""" + full_spec = FullAttentionSpec( + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 + ) + sw_spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + sliding_window=32, + ) + num_blocks = 2 + page = full_spec.page_size_bytes + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=2 * num_blocks * page, + layers=["full0", "sw0"], + layer_stride=num_blocks * page, + block_stride=page, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["full0"], kv_cache_spec=full_spec), + KVCacheGroupSpec(layer_names=["sw0"], kv_cache_spec=sw_spec), + ], + ) + + +def _read_scheduler( + kv_cache_config: KVCacheConfig, disable_hma: bool = False +) -> MoRIIOConnectorScheduler: + vllm_config = create_vllm_config(role="kv_producer", read_mode=True) + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hma + with set_current_vllm_config(vllm_config): + connector = MoRIIOConnector( + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config + ) + assert connector.connector_scheduler is not None + return connector.connector_scheduler + + +def test_hma_blocks_per_sw_two_groups(): + """A Full + SlidingWindow config runs with HMA and computes block budgets + correctly""" + scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) + assert scheduler._is_hma_required is True + # cdiv(32, 16) + 1 == 3 for the sliding-window group, 0 for full attention. + assert scheduler.blocks_per_sw == [0, 3] + + +@pytest.mark.parametrize( + "swa_enabled, disable_hma, expected_is_hma", + [ + (True, False, True), # sliding-window group present, HMA enabled + (True, True, False), # sliding-window group present but HMA disabled + (False, False, False), # full-attention only, HMA not needed + ], +) +def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma): + """_is_hma_required tracks both the KV cache groups and the + --disable-hybrid-kv-cache-manager flag. When HMA is off, + get_exchange_clipped_blocks must be a no-op.""" + config = ( + _make_hybrid_kv_cache_config() if swa_enabled else _make_test_kv_cache_config() + ) + scheduler = _read_scheduler(config, disable_hma=disable_hma) + assert scheduler._is_hma_required is expected_is_hma + if not expected_is_hma: + blocks = [[1, 2, 3, 4, 5]] + assert scheduler.get_exchange_clipped_blocks(blocks) == blocks + + +def test_non_sliding_window_hybrid_is_rejected(): + """A hybrid group that is not sliding-window (e.g. chunked-local + attention) must fail closed rather than be silently mistransferred.""" + full_spec = FullAttentionSpec( + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16 + ) + local_spec = ChunkedLocalAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + attention_chunk_size=32, + ) + num_blocks = 2 + page = full_spec.page_size_bytes + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=2 * num_blocks * page, + layers=["full0", "local0"], + layer_stride=num_blocks * page, + block_stride=page, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["full0"], kv_cache_spec=full_spec), + KVCacheGroupSpec(layer_names=["local0"], kv_cache_spec=local_spec), + ], + ) + with pytest.raises(NotImplementedError, match="sliding-window hybrid"): + _read_scheduler(config) + + +def test_token_count_basis_uses_full_attention_group(): + """Chunked-prefill token counting must use an unclipped full-attention + group (blocks_per_sw == 0), not a clipped sliding-window group.""" + scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) + # Group 0 is the full-attention group + assert scheduler._full_attn_group_idx == 0 + assert scheduler._full_attn_block_size == 16 + + +def test_get_exchange_clipped_blocks_clips_only_sw_group(): + """get_exchange_clipped_blocks keeps the full attn group intact and clips the + sliding-window group to its window tail.""" + scheduler = _read_scheduler(_make_hybrid_kv_cache_config()) + full = [10, 11, 12, 13, 14] + sw = [20, 21, 22, 23, 24] + clipped = scheduler.get_exchange_clipped_blocks([full, sw]) + assert clipped[0] == full + assert clipped[1] == [22, 23, 24] + + +def test_metadata_hma_block_ids_preserved_per_group(): + """add_new_req stores per-group (BlockIds) block lists unchanged for both + read and write, so the hybrid group structure is retained.""" + metadata = MoRIIOConnectorMetadata() + + # Assume: + # - Full-attention group (6 blocks) + + # - sliding-window group already clipped to its window tail (3 blocks). + fa_blocks = [0, 1, 2, 3, 4, 5] + sw_blocks = [10, 11, 12] + local_block_ids = [fa_blocks, sw_blocks] + remote_block_ids = [[100, 101, 102, 103, 104, 105], [200, 201, 202]] + + base_params = { + "remote_engine_id": "remote-engine", + "remote_host": "127.0.0.1", + "remote_handshake_port": 6301, + "remote_notify_port": 61005, + "remote_block_ids": remote_block_ids, + } + + # Read mode: both local and remote ids stay per-group. + metadata.add_new_req( + request_id="recv-req", + local_block_ids=local_block_ids, + kv_transfer_params={**base_params, "transfer_id": "recv-req"}, + ) + recv_meta = metadata.reqs_to_recv["recv-req"] + assert recv_meta.local_block_ids == [fa_blocks, sw_blocks] + assert recv_meta.remote_block_ids == remote_block_ids + + # Write mode: the decode peer allocates its own blocks, so #remote_block_ids may + # be empty, Local group structure is kept. + metadata.add_new_req( + request_id="save-req", + local_block_ids=local_block_ids, + kv_transfer_params={ + **base_params, + "transfer_id": "save-req", + "remote_block_ids": [], + }, + write_mode=True, + ) + save_meta = metadata.reqs_to_save["save-req"] + assert save_meta.local_block_ids == [fa_blocks, sw_blocks] + assert save_meta.remote_block_ids == [] + + +def test_single_group_path_unchanged(): + """A full attn config does not enable HMA and never clips blocks.""" + scheduler = _read_scheduler(_make_test_kv_cache_config()) + assert scheduler._is_hma_required is False + assert scheduler.blocks_per_sw == [0] + blocks = [[1, 2, 3, 4, 5]] + assert scheduler.get_exchange_clipped_blocks(blocks) == blocks + + +def test_hybrid_write_mode_rejected(): + """Hybrid KV cache groups are unsupported in WRITE mode and fail closed.""" + vllm_config = create_vllm_config(role="kv_producer", read_mode=False) + with ( + set_current_vllm_config(vllm_config), + pytest.raises(NotImplementedError), + ): + MoRIIOConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + _make_hybrid_kv_cache_config(), + ) + + +def test_worker_layer_to_group_routing(mock_parallel_groups): + """The worker maps every layer to its KV cache group correctly.""" + vllm_config = create_vllm_config(role="kv_consumer", read_mode=True) + # Building the worker directly bypasses MoRIIOConnector._set_port_defaults, + # so provide the ports manually. + vllm_config.kv_transfer_config.kv_connector_extra_config.update( + { + "http_port": 12346, + "handshake_port": 12347, + "notify_port": 12348, + } + ) + with set_current_vllm_config(vllm_config): + worker = FakeMoRIIOConnectorWorker( + vllm_config, + "engine0", + hand_shake_latency=0, + kv_cache_config=_make_hybrid_kv_cache_config(), + ) + assert worker.layer_to_group == {"full0": 0, "sw0": 1} diff --git a/tests/v1/kv_connector/unit/test_moriio_routing.py b/tests/v1/kv_connector/unit/test_moriio_routing.py index a9a3521beb..fc8b323218 100644 --- a/tests/v1/kv_connector/unit/test_moriio_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_routing.py @@ -20,12 +20,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import from vllm.v1.request import RequestStatus _request_finished = MoRIIOConnectorScheduler.request_finished +_clip_blocks = MoRIIOConnectorScheduler.get_exchange_clipped_blocks def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int): """Minimal prefill-side (producer) scheduler stand-in for request_finished.""" - return SimpleNamespace( + s = SimpleNamespace( is_producer=True, + _is_hma_required=False, + blocks_per_sw=[0], _global_dp_rank=global_dp_rank, engine_id="engine-abc", host_ip="10.0.0.1", @@ -39,6 +42,8 @@ def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int): ) ), ) + s.get_exchange_clipped_blocks = _clip_blocks.__get__(s) + return s def _req( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index a3010203ee..7d06b925ba 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -14,6 +14,7 @@ import torch import zmq from vllm.config import KVTransferConfig, VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, ) @@ -444,8 +445,8 @@ class ReqMeta: """Metadata for a single request.""" transfer_id: TransferId - local_block_ids: list[int] - remote_block_ids: list[int] + local_block_ids: BlockIds + remote_block_ids: BlockIds remote_host: str remote_port: int remote_handshake_port: int @@ -481,7 +482,7 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata): def add_new_req( self, request_id: ReqId, - local_block_ids: list[int], + local_block_ids: BlockIds, kv_transfer_params: dict[str, Any], write_mode=False, ): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 349cf50775..ecdd59fca3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -18,10 +18,12 @@ import torch import zmq from vllm.config import CUDAGraphMode, VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( ROLE, @@ -67,19 +69,24 @@ from vllm.distributed.parallel_state import ( ) from vllm.forward_context import ForwardContext, get_forward_context from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import ( make_zmq_path, make_zmq_socket, ) from vllm.v1.attention.selector import get_attn_backend from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + SlidingWindowSpec, +) from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus if TYPE_CHECKING: from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.kv_cache_manager import KVCacheBlocks - from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request logger = init_logger(__name__) @@ -189,7 +196,7 @@ def resolve_moriio_transfer_ack( return transfer_id -class MoRIIOConnector(KVConnectorBase_V1): +class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): def __init__( self, vllm_config: VllmConfig, @@ -225,7 +232,7 @@ class MoRIIOConnector(KVConnectorBase_V1): ) if role == KVConnectorRole.SCHEDULER: self.connector_scheduler: MoRIIOConnectorScheduler | None = ( - MoRIIOConnectorScheduler(vllm_config, self.engine_id) + MoRIIOConnectorScheduler(vllm_config, self.engine_id, kv_cache_config) ) self.connector_worker: MoRIIOConnectorWorker | None = None elif role == KVConnectorRole.WORKER: @@ -283,6 +290,14 @@ class MoRIIOConnector(KVConnectorBase_V1): self, request: "Request", block_ids: list[int], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, (block_ids,)) + + def request_finished_all_groups( + self, + request: "Request", + block_ids: BlockIds, ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None return self.connector_scheduler.request_finished(request, block_ids) @@ -368,7 +383,12 @@ class MoRIIOConnector(KVConnectorBase_V1): class MoRIIOConnectorScheduler: """Implementation of Scheduler side methods""" - def __init__(self, vllm_config: VllmConfig, engine_id: str): + def __init__( + self, + vllm_config: VllmConfig, + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): self.vllm_config = vllm_config assert vllm_config.kv_transfer_config is not None, ( @@ -378,6 +398,66 @@ class MoRIIOConnectorScheduler: self.block_size = vllm_config.cache_config.block_size self.engine_id: EngineId = engine_id self.mode = get_moriio_mode(self.kv_transfer_config) + + self._is_hma_required = ( + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + and any( + not isinstance(g.kv_cache_spec, FullAttentionSpec) + for g in kv_cache_config.kv_cache_groups + ) + ) + if self._is_hma_required: + # TODO(simondanielsson): support non-sliding window hybrids + unsupported = [ + type(g.kv_cache_spec).__name__ + for g in kv_cache_config.kv_cache_groups + if not isinstance( + g.kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec) + ) + ] + if unsupported: + raise NotImplementedError( + "MoRIIO only supports sliding-window hybrid models (e.g. " + "Gemma) under the hybrid KV cache manager. Unsupported KV " + f"cache group spec(s): {sorted(set(unsupported))}. Pass " + "--disable-hybrid-kv-cache-manager to run without HMA." + ) + if self.mode == MoRIIOMode.WRITE: + # TODO(simondanielsson): support HMA in WRITE mode + raise NotImplementedError( + "MoRIIO WRITE mode does not support hybrid KV cache groups " + "(sliding-window attention). Use READ mode, or pass " + "--disable-hybrid-kv-cache-manager." + ) + + sw_sizes_tokens: list[tuple[int, int]] = [ + (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) + if isinstance(g.kv_cache_spec, SlidingWindowSpec) + else (0, self.block_size) + for g in kv_cache_config.kv_cache_groups + ] + # add 1 to conservatively account for boundary overlap eg window isn't fully + # aligned with blocks. + self.blocks_per_sw = [ + cdiv(n_tokens, block_size) + 1 if n_tokens else 0 + for n_tokens, block_size in sw_sizes_tokens + ] + # In WRITE mode and chunked prefill, we perform the write after the last chunk. + # Hence we need to track once we are in the last chunk. We do this by computing + # len(block_ids[g]) * block_size[g] for a full attn group g, as hybrid groups + # may be clipped and hence not reflect the full sequence length. Here we store + # g and block_size[g] for later. + self._full_attn_group_idx = 0 + self._full_attn_block_size = self.block_size + for gi, group in enumerate(kv_cache_config.kv_cache_groups): + is_full_attn = self.blocks_per_sw[gi] == 0 + if is_full_attn: + self._full_attn_group_idx = gi + self._full_attn_block_size = getattr( + group.kv_cache_spec, "block_size", self.block_size + ) + break + self.host_ip = resolve_host_ip( self.kv_transfer_config.kv_connector_extra_config ) @@ -409,14 +489,14 @@ class MoRIIOConnectorScheduler: # Requests that need to start recv/send. # New requests are added by update_state_after_alloc in # the scheduler. Used to make metadata passed to Worker. - self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {} - self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {} + self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {} + self._reqs_need_save: dict[ReqId, tuple[Request, BlockIds]] = {} # Snapshot of kv_transfer_params for chunked prefill recovery. self._req_kv_params: dict[ReqId, dict] = {} # For chunked prefill, we perform layer-wise access within the final chunk. # TODO: Perform transfer at end chunk. - self._reqs_need_pending_save: dict[ReqId, tuple[Request, list[int]]] = {} + self._reqs_need_pending_save: dict[ReqId, tuple[Request, BlockIds]] = {} if self.is_producer: set_role(ROLE.PRODUCER) @@ -442,6 +522,23 @@ class MoRIIOConnectorScheduler: self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} + def get_exchange_clipped_blocks( + self, + block_ids: BlockIds, + ) -> BlockIds: + """Clip the block list based on the layer type. + + Behavior: + - sliding window: keep only the last window size blocks + - full attention: keep all blocks + """ + if len(block_ids) == 0 or not self._is_hma_required: + return [list(blocks) for blocks in block_ids] + return [ + blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks + for i, blocks in enumerate(block_ids) + ] + def map_request_id(self, request_id: ReqId, transfer_id: TransferId): self.transfer_id_to_request_id[transfer_id] = request_id self.request_id_to_transfer_id[request_id] = transfer_id @@ -642,7 +739,7 @@ class MoRIIOConnectorScheduler: request_id = request.request_id self.map_request_id(request_id, transfer_id) if params.get("do_remote_decode"): - local_block_ids = blocks.get_block_ids()[0] + local_block_ids: BlockIds = blocks.get_block_ids() self._reqs_need_save[request.request_id] = (request, local_block_ids) # Snapshot params now so chunked-prefill build_connector_meta # can recover them on the final chunk even if the live @@ -657,17 +754,22 @@ class MoRIIOConnectorScheduler: if "remote_engine_id" in params: if num_external_tokens > 0: # Get unhashed blocks to pull from remote. - local_block_ids = blocks.get_block_ids()[0] - assert len(local_block_ids) <= len(remote_block_ids) - if len(local_block_ids) != len(remote_block_ids): - local_block_ids = remote_block_ids[ - -len(local_block_ids) : - ] + local_block_ids = self.get_exchange_clipped_blocks( + blocks.get_block_ids() + ) + # Partial cache hits can lead to fewer local blocks vs + # remote blocks, but never more + assert all( + len(local) <= len(remote) + for local, remote in zip( + local_block_ids, remote_block_ids, strict=True + ) + ) else: # If remote_blocks and num_external_tokens = 0, we have # a full prefix cache hit on the D worker. We need to call # send_notify in _read_blocks to free the memory on the P. - local_block_ids = [] + local_block_ids = [[] for _ in remote_block_ids] self._reqs_need_recv[request.request_id] = ( request, @@ -818,8 +920,6 @@ class MoRIIOConnectorScheduler: new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i] if new_block_ids is not None: - block_ids = new_block_ids[0] - # TODO : hybrid attn, etc # A non-disagg request (no kv_transfer_params, e.g. smoke # test) is never registered in _reqs_need_pending_save; # indexing it unconditionally would KeyError and crash the @@ -827,12 +927,16 @@ class MoRIIOConnectorScheduler: if req_id not in self._reqs_need_pending_save: continue req, existing_blocks = self._reqs_need_pending_save[req_id] - updated_blocks = list(existing_blocks) + (block_ids) + updated_blocks = [ + existing_blocks[g] + new_block_ids[g] + for g in range(len(new_block_ids)) + ] self._reqs_need_pending_save[req_id] = (req, updated_blocks) - if ( - len(self._reqs_need_pending_save[req_id][1]) * self.block_size - >= req.num_prompt_tokens - ): + saved_tokens = ( + len(updated_blocks[self._full_attn_group_idx]) + * self._full_attn_block_size + ) + if saved_tokens >= req.num_prompt_tokens: # Final chunk: live kv_transfer_params may be cleared, # so prefer the snapshot from update_state_after_alloc. kv_params = self._req_kv_params.pop( @@ -857,7 +961,10 @@ class MoRIIOConnectorScheduler: for req_id, (req, block_ids) in self._reqs_need_save.items(): kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {}) - if req.num_prompt_tokens > len(block_ids) * self.block_size: + tokens_covered = ( + len(block_ids[self._full_attn_group_idx]) * self._full_attn_block_size + ) + if req.num_prompt_tokens > tokens_covered: # not last chunk prefill self._reqs_need_pending_save[req_id] = (req, block_ids) continue @@ -898,7 +1005,7 @@ class MoRIIOConnectorScheduler: def request_finished( self, request: "Request", - block_ids: list[int], + block_ids: BlockIds, ) -> tuple[bool, dict[str, Any] | None]: """ Once a request is finished, determine whether request blocks @@ -936,7 +1043,8 @@ class MoRIIOConnectorScheduler: if self.mode == MoRIIOMode.WRITE: self._release_write_prefill_blocks(request.request_id, params) else: - self._reqs_need_recv[request.request_id] = (request, []) + empty_groups: BlockIds = [[] for _ in self.blocks_per_sw] + self._reqs_need_recv[request.request_id] = (request, empty_groups) params["do_remote_prefill"] = False return False, None @@ -946,10 +1054,9 @@ class MoRIIOConnectorScheduler: ): return False, None - # computed_block_ids = block_ids if all_full else block_ids[:-1] - computed_block_ids = block_ids - # If prompt < block_size, no xfer so free blocks immediately. - delay_free_blocks = len(computed_block_ids) > 0 + # On producer + computed_block_ids = self.get_exchange_clipped_blocks(block_ids) + delay_free_blocks = any(len(group) > 0 for group in computed_block_ids) if delay_free_blocks: # Prefill request on remote. It will be read from D upon completion @@ -1102,6 +1209,12 @@ class MoRIIOConnectorWorker: self.kv_transfer_config = vllm_config.kv_transfer_config self.is_producer = self.kv_transfer_config.is_kv_producer self.layer_to_spec = build_layer_to_spec(kv_cache_config) + # we transfer by layer, but store by kv cache group + self.layer_to_group: dict[str, int] = { + layer_name: group_idx + for group_idx, group in enumerate(kv_cache_config.kv_cache_groups) + for layer_name in group.layer_names + } if self.is_producer: set_role(ROLE.PRODUCER) @@ -1789,8 +1902,6 @@ class MoRIIOConnectorWorker: self.block_size, ) self.block_size = first_geometry.block_size - # TODO(tms): self.block_len needs to be per-layer for sliding window, - # hybrid attn, etc # block size in bytes self.block_len = first_geometry.block_len self.kv_cache_shape = first_kv_cache.shape @@ -1803,14 +1914,30 @@ class MoRIIOConnectorWorker: self.layer_base_addr_index = {} base_addr_idx = 0 + hma_enabled = ( + not self.vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + ) for layer_name in kv_caches: self.layer_base_addr_index[layer_name] = base_addr_idx geometry = self._get_layer_transfer_geometry(layer_name) - if geometry.block_size != self.block_size: + if hma_enabled: + if geometry.block_len != self.block_len: + raise ValueError( + "MoRIIO KV cache block length mismatch for layer " + f"{layer_name}: {geometry.block_len} != {self.block_len}" + ) + elif geometry.block_size != self.block_size: raise ValueError( "MoRIIO KV cache block size mismatch for layer " f"{layer_name}: {geometry.block_size} != {self.block_size}" ) + # num_blocks is advertised as a single scalar to the peer, so must it + # be uniform + if geometry.num_blocks != self.num_blocks: + raise ValueError( + "MoRIIO KV cache num_blocks mismatch for layer " + f"{layer_name}: {geometry.num_blocks} != {self.num_blocks}" + ) self.block_lens[layer_name] = geometry.block_len for cache, region_len in self._iter_layer_registration_regions(layer_name): base_addr = cache.data_ptr() @@ -2455,12 +2582,16 @@ class MoRIIOConnectorWorker: self.remote_dp_size_local = int(meta.remote_dp_size_local) else: self.remote_dp_size_local = int(meta.remote_dp_size) + # WRITE does not support HMA, so unwrap blocks into flat lists. + # remote_block_ids can itself be empty so need to be careful when unwrapping. + local_block_ids = meta.local_block_ids[0] + remote_block_ids = meta.remote_block_ids[0] if meta.remote_block_ids else [] self.schedule_write_blocks( request_id=req_id, transfer_id=meta.transfer_id, dst_engine_id=meta.remote_engine_id, - local_block_ids=meta.local_block_ids, - remote_block_ids=meta.remote_block_ids, + local_block_ids=local_block_ids, + remote_block_ids=remote_block_ids, layer_name=layer_name, kv_layer=kv_layer, remote_notify_port=meta.remote_notify_port, @@ -2607,8 +2738,8 @@ class MoRIIOConnectorWorker: def _read_blocks( self, - local_block_ids: list[int], - remote_block_ids: list[int], + local_block_ids: BlockIds, + remote_block_ids: BlockIds, dst_engine_id: str, request_id: str, transfer_id: str, @@ -2653,10 +2784,11 @@ class MoRIIOConnectorWorker: sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( layer_name ) + group_idx = self.layer_to_group[layer_name] offs = self._compute_block_transfer_offsets( layer_name, - local_block_ids, - remote_block_ids, + local_block_ids[group_idx], + remote_block_ids[group_idx], remote_moriio_meta, remote_tp_size=remote_tp_size, remote_engine_id=remote_dp_engine_id, From e0c04c7b4dddceb25eb9d23d87c3914bfc2ad3f9 Mon Sep 17 00:00:00 2001 From: Ronald Date: Mon, 14 Sep 2026 21:05:47 +0800 Subject: [PATCH 027/139] [Bugfix][Models] Fix OpenPangu sleep mode with static sinks (#53696) Signed-off-by: Ronald1995 --- .../layers/attention/static_sink_attention.py | 2 -- vllm/model_executor/models/openpangu.py | 18 +++++++++++------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/attention/static_sink_attention.py b/vllm/model_executor/layers/attention/static_sink_attention.py index d27ccc728a..d321d675a5 100644 --- a/vllm/model_executor/layers/attention/static_sink_attention.py +++ b/vllm/model_executor/layers/attention/static_sink_attention.py @@ -152,8 +152,6 @@ class StaticSinkAttention(Attention, CustomOp): attn_backend=attn_backend, **kwargs, ) - CustomOp.__init__(self) - self.sink_len = sink_len self.sink_populated = False self.sink_key = None diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 60b5b301d5..b8937f2e09 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -677,14 +677,18 @@ class OpenPanguSinkAttention(nn.Module): }, ) else: - self.param_sink_value = torch.zeros( - ( - self.param_sink_number, - self.num_kv_heads, - self.v_channels, + self.register_buffer( + "param_sink_value", + torch.zeros( + ( + self.param_sink_number, + self.num_kv_heads, + self.v_channels, + ), + device=current_platform.current_device(), + dtype=config.torch_dtype, ), - device=current_platform.current_device(), - dtype=config.torch_dtype, + persistent=False, ) # To enable dummy run with out weight self.post_weight_load() From 47fbd36e6dc908277e386424beeb8647355c1c8b Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Mon, 14 Sep 2026 21:06:26 +0800 Subject: [PATCH 028/139] [Bugfix] Trim stale consequence claims from unannotated-eagle warning (#56791) Signed-off-by: zjy0516 Co-authored-by: Kimi Code --- tests/v1/core/test_kv_cache_utils.py | 21 +++++++++++++++++---- vllm/v1/core/kv_cache_utils.py | 14 ++------------ 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index dbddc556c9..365338ad4f 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -4135,9 +4135,9 @@ def test_draft_group_not_annotated_without_spec_decode(): def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): - # No group carries the draft marker, so every consumer falls back to - # flagging all groups -- including Mamba ones, which then can never report - # a hit. That is silent today; it must at least be visible. + # No group carries the draft marker, so consumers fall back to + # conservative behavior that silently breaks reuse for Mamba groups. + # That must at least be visible. groups = get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) ) @@ -4146,7 +4146,20 @@ def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): assert "no KV cache group could be identified as the draft model's" in ( caplog_vllm.text ) - assert "Mamba groups" in caplog_vllm.text + + +def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): + # Pure-attention models degrade gracefully under the consumers' + # conservative fallback (a one-block hit drop at most), so the warning + # stays silent to avoid noise on every unannotated EAGLE deployment. + specs = { + "target.attn.0": new_mla_spec(block_size=64), + "target.attn.1": new_mla_spec(block_size=64), + } + groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" not in caplog_vllm.text def test_no_warning_when_draft_group_is_identified(caplog_vllm): diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 996cce8c43..4e4f34bd52 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2182,12 +2182,7 @@ def _warn_if_unannotated_eagle_mamba( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], ) -> None: - """Warn when the flag-all eagle fallback will silently disable reuse. - - With no group annotated, consumers flag every group as a draft group. That - widens a Mamba group's required lookup window to two consecutive chunks, - which align-mode checkpointing never produces, so reuse drops to zero with - no error and no metric to show it. + """Warn when no KV cache group could be identified as the draft model's. Args: vllm_config: Config supplying the speculative method, if any. @@ -2210,13 +2205,8 @@ def _warn_if_unannotated_eagle_mamba( return logger.warning( "Speculative decoding (method=%s) is enabled but no KV cache group " - "could be identified as the draft model's, so every group -- " - "including Mamba groups %s -- will be treated as a draft group. A " - "Mamba group cannot satisfy the widened lookup window that implies, " - "so prefix-cache reuse across requests will be disabled and any " - "external KV offload tier will store without ever serving a hit.", + "could be identified as the draft model's.", spec_config.method, - mamba_groups, ) From 6623fe5b4ec783e7264f140e0caea20bdc45a341 Mon Sep 17 00:00:00 2001 From: arkohut <39525455+arkohut@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:18:55 +0800 Subject: [PATCH 029/139] [Bugfix] MiniCPM-V 4.6: fix ViT self-attn qkv weight loading (#49417) Signed-off-by: arkohut <39525455+arkohut@users.noreply.github.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/model_executor/models/minicpmv4_6.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index e43f13f79f..d8e10d3498 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -674,14 +674,6 @@ class MiniCPMV4_6ProcessingInfo(MiniCPMVProcessingInfo): class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - } - ) - def __init__( self, config, @@ -729,10 +721,6 @@ class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module): out, _ = self.out_proj(attn_out) return out - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - class MiniCPMV4_6ViTWindowAttentionMerger(nn.Module): def __init__( @@ -971,6 +959,11 @@ class MiniCPMV4_6ForConditionalGeneration( supports_encoder_tp_data = True hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + }, orig_to_new_prefix={ # transformers v5.7+ uses `vision_tower` and nests `vit_merger` # inside it. Order matters: more specific prefix must come first. @@ -982,7 +975,7 @@ class MiniCPMV4_6ForConditionalGeneration( "model.language_model.": "language_model.model.", "lm_head.": "language_model.lm_head.", "mtp.": None, - } + }, ) packed_modules_mapping = { From 9d4d9aa5bced3fc84c9e7c89b1257f206fbff6e2 Mon Sep 17 00:00:00 2001 From: linitra24 Date: Mon, 14 Sep 2026 21:19:15 +0800 Subject: [PATCH 030/139] [LoRA][Refactor] Unify multimodal LoRA token count hooks (#55071) Signed-off-by: linitra24 --- .../processing/test_llava_next_video.py | 37 ++++++++------ vllm/model_executor/models/blip2.py | 35 ++++++-------- vllm/model_executor/models/dots_ocr.py | 16 ++++--- vllm/model_executor/models/gemma3_mm.py | 48 +++++-------------- vllm/model_executor/models/glm4_1v.py | 20 ++++---- vllm/model_executor/models/h2ovl.py | 27 ++++++----- vllm/model_executor/models/idefics3.py | 23 ++++----- vllm/model_executor/models/interfaces.py | 27 ++++------- vllm/model_executor/models/internvl.py | 26 +++++----- vllm/model_executor/models/keye.py | 16 ++++--- vllm/model_executor/models/lfm2_vl.py | 18 +++---- vllm/model_executor/models/llava.py | 23 ++++----- .../model_executor/models/llava_next_video.py | 40 ++++++++-------- vllm/model_executor/models/mllama4.py | 28 +++++------ vllm/model_executor/models/moondream3.py | 15 ++++-- vllm/model_executor/models/openvla.py | 20 +++++--- vllm/model_executor/models/paligemma.py | 15 ++++-- vllm/model_executor/models/pixtral.py | 20 ++++---- vllm/model_executor/models/qwen2_5_vl.py | 23 ++++----- vllm/model_executor/models/qwen2_vl.py | 23 ++++----- vllm/model_executor/models/qwen3_asr.py | 16 ++++--- vllm/model_executor/models/qwen3_vl.py | 22 ++++----- vllm/model_executor/models/ultravox.py | 33 +++++++------ 23 files changed, 276 insertions(+), 295 deletions(-) diff --git a/tests/models/multimodal/processing/test_llava_next_video.py b/tests/models/multimodal/processing/test_llava_next_video.py index 2dcbb7d4b5..c4fdecd21a 100644 --- a/tests/models/multimodal/processing/test_llava_next_video.py +++ b/tests/models/multimodal/processing/test_llava_next_video.py @@ -14,7 +14,7 @@ from ...utils import build_model_context class _StubModel: - """Carries only the two attributes the token helpers read from + """Carries only the attributes the token-count method reads from `self`, so the real methods can be exercised without constructing the full `nn.Module` (vision tower, language model, etc.).""" @@ -22,18 +22,15 @@ class _StubModel: pooled_grid_length: int -get_num_mm_encoder_tokens = ( - LlavaNextVideoForConditionalGeneration.get_num_mm_encoder_tokens -) -get_num_mm_connector_tokens = ( - LlavaNextVideoForConditionalGeneration.get_num_mm_connector_tokens +get_mm_lora_token_counts = ( + LlavaNextVideoForConditionalGeneration.get_mm_lora_token_counts ) @pytest.mark.parametrize("model_id", ["llava-hf/LLaVA-NeXT-Video-7B-hf"]) def test_num_mm_tokens_match_real_config(model_id): """The stored grid lengths must match what `__init__` derives from - the real HF config, and the two helpers must invert each other's + the real HF config, and the LoRA token counts must preserve the frame-level scaling exactly.""" ctx = build_model_context(model_id, limit_mm_per_prompt={"video": 1}) config = ctx.model_config.hf_config @@ -49,10 +46,13 @@ def test_num_mm_tokens_match_real_config(model_id): for num_frames in (1, 2, 8, 16, 32): num_video_tokens = num_frames * pooled_grid_length**2 - encoder_tokens = get_num_mm_encoder_tokens(stub, num_video_tokens) + encoder_tokens, connector_tokens = get_mm_lora_token_counts( + stub, + modality="video", + mm_kwargs=None, + num_mm_embeds=num_video_tokens, + ) assert encoder_tokens == num_frames * patch_grid_length**2 - - connector_tokens = get_num_mm_connector_tokens(stub, encoder_tokens) assert connector_tokens == num_video_tokens @@ -72,10 +72,13 @@ def test_num_mm_tokens_roundtrip(patch_grid_length, pooled_grid_length, num_fram num_video_tokens = num_frames * pooled_grid_length**2 - encoder_tokens = get_num_mm_encoder_tokens(stub, num_video_tokens) + encoder_tokens, connector_tokens = get_mm_lora_token_counts( + stub, + modality="video", + mm_kwargs=None, + num_mm_embeds=num_video_tokens, + ) assert encoder_tokens == num_frames * patch_grid_length**2 - - connector_tokens = get_num_mm_connector_tokens(stub, encoder_tokens) assert connector_tokens == num_video_tokens @@ -84,5 +87,9 @@ def test_num_mm_tokens_zero(): stub.patch_grid_length = 24 stub.pooled_grid_length = 12 - assert get_num_mm_encoder_tokens(stub, 0) == 0 - assert get_num_mm_connector_tokens(stub, 0) == 0 + assert get_mm_lora_token_counts( + stub, + modality="video", + mm_kwargs=None, + num_mm_embeds=0, + ) == (0, 0) diff --git a/vllm/model_executor/models/blip2.py b/vllm/model_executor/models/blip2.py index 091001fc01..de74d1f5ac 100644 --- a/vllm/model_executor/models/blip2.py +++ b/vllm/model_executor/models/blip2.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import MultiModalDataItems @@ -687,28 +688,22 @@ class Blip2ForConditionalGeneration( tower_model="vision_model", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - if num_image_tokens <= 0: - return 0 - assert num_image_tokens % self.config.num_query_tokens == 0, ( + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + if num_mm_embeds <= 0: + return 0, 0 + assert num_mm_embeds % self.config.num_query_tokens == 0, ( "The number of image tokens must be a multiple of " "the number of query tokens." ) - num_images = num_image_tokens / self.config.num_query_tokens - return num_images * self._vision_tokens_per_image - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - if num_vision_tokens <= 0: - return 0 - assert num_vision_tokens % self._vision_tokens_per_image == 0, ( - "The number of vision tokens must be a multiple of " - "the number of tokens per image." + num_images = num_mm_embeds // self.config.num_query_tokens + return ( + num_images * self._vision_tokens_per_image, + num_images * self.config.num_query_tokens, ) - num_images = num_vision_tokens / self._vision_tokens_per_image - return num_images * self.config.num_query_tokens diff --git a/vllm/model_executor/models/dots_ocr.py b/vllm/model_executor/models/dots_ocr.py index e91d6ce94c..c939c2bef5 100644 --- a/vllm/model_executor/models/dots_ocr.py +++ b/vllm/model_executor/models/dots_ocr.py @@ -54,6 +54,7 @@ from vllm.model_executor.models.utils import ( ) from vllm.model_executor.models.vision import get_vit_attn_backend from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalKwargsItem from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.dotsocr import DotsOCRConfig, DotsVisionConfig from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -718,13 +719,16 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA return image_embeds.split(sizes) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs merge_size = self.vision_tower.spatial_merge_size - return num_image_tokens * (merge_size**2) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - merge_size = self.vision_tower.spatial_merge_size - return num_vision_tokens // (merge_size**2) + return num_mm_embeds * (merge_size**2), num_mm_embeds def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: image_input = self._parse_and_validate_image_input(**kwargs) diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index b45327fbeb..abe1f5f1bd 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -19,6 +19,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems @@ -686,43 +687,16 @@ class Gemma3ForConditionalGeneration( tower_model="vision_tower", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - """ - Calculate the number of tokens output by the vision encoder. - - The vision encoder processes images into patch embeddings. For Gemma3, - the relationship between prompt placeholder tokens and actual vision - encoder output tokens depends on the patch grid size. - - Args: - num_image_tokens: Number of image placeholder tokens in the prompt - (typically mm_tokens_per_image per image) - - Returns: - Number of tokens output by the vision encoder - """ - # For Gemma3, the vision encoder outputs tokens_per_side x tokens_per_side - # tokens per image. Since num_image_tokens represents the number of - # connector output tokens (mm_tokens_per_image = 256), and tokens_per_side - # is sqrt(256) = 16, we need to account for the token expansion. - # Based on empirical testing, the multiplier of 16 works correctly. - return num_image_tokens * 16 - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - """ - Calculate the number of tokens output by the multimodal connector. - - The connector applies projection and normalization but maintains the - token count for Gemma3. - - Args: - num_vision_tokens: Number of tokens from vision encoder - - Returns: - Number of tokens after connector processing - """ - # The Gemma3 connector maintains a 1:1 token mapping - return num_vision_tokens + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + tower_tokens = num_mm_embeds * 16 + return tower_tokens, tower_tokens def get_encoder_cudagraph_config(self): from vllm.v1.worker.encoder_cudagraph_defs import ( diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 163115ff3d..730a015cfe 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -83,6 +83,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, VideoItem, ) @@ -2378,19 +2379,16 @@ class Glm4vForConditionalGeneration( tower_model="visual.", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs merge_size = self.config.vision_config.spatial_merge_size - return num_image_tokens * (merge_size**2) - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - merge_size = self.config.vision_config.spatial_merge_size - return num_vision_tokens // (merge_size**2) + return num_mm_embeds * (merge_size**2), num_mm_embeds @MULTIMODAL_REGISTRY.register_processor( diff --git a/vllm/model_executor/models/h2ovl.py b/vllm/model_executor/models/h2ovl.py index 5c4f984c13..5be2d7b46d 100644 --- a/vllm/model_executor/models/h2ovl.py +++ b/vllm/model_executor/models/h2ovl.py @@ -14,7 +14,7 @@ from transformers import PretrainedConfig from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import BatchedTensorInputs +from vllm.multimodal.inputs import BatchedTensorInputs, MultiModalKwargsItem from vllm.multimodal.parse import ( ImageEmbeddingItems, ImageProcessorItems, @@ -178,16 +178,17 @@ class H2OVLChatModel(InternVLChatModel): prefix=prefix, ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - if num_image_tokens <= 0 or self.num_image_token <= 0: - return 0 + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + if num_mm_embeds <= 0 or self.num_image_token <= 0: + return 0, 0 - num_patches = num_image_tokens // self.num_image_token - return num_patches * (self.patch_tokens + 1) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - if num_vision_tokens <= 0 or self.num_image_token <= 0: - return 0 - - num_patches = num_vision_tokens // (self.patch_tokens + 1) - return num_patches * self.num_image_token + num_patches = num_mm_embeds // self.num_image_token + tower_tokens = num_patches * (self.patch_tokens + 1) + return tower_tokens, num_patches * self.num_image_token diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index baaa4106cb..4ac8eb5eeb 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -39,6 +39,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems @@ -837,20 +838,14 @@ class Idefics3ForConditionalGeneration( tower_model="model.vision_model", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs hf_config = self.config scale_factor = hf_config.scale_factor - - return num_image_tokens * scale_factor**2 - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - hf_config = self.config - scale_factor = hf_config.scale_factor - - return num_vision_tokens // scale_factor**2 + return num_mm_embeds * scale_factor**2, num_mm_embeds diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 254f2c635b..2e79001add 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -31,7 +31,7 @@ import torch import torch.nn as nn from torch import Tensor from transformers.models.whisper.tokenization_whisper import LANGUAGES -from typing_extensions import Self, TypeIs +from typing_extensions import Self, TypeIs, deprecated from vllm.logger import init_logger from vllm.model_executor.layers.quantization import QuantizationConfig @@ -420,23 +420,16 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol): yield - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - """ - Implement this function to enable LoRA support - for the tower module of the multi-modal model. - Given the number of image tokens, output the number of - multi-modal encoder tokens. - """ - ... + @deprecated( + "get_num_mm_encoder_tokens is deprecated; use get_mm_lora_token_counts instead." + ) + def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: ... - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - """ - Implement this function to enable LoRA support - for the connector module of the multi-modal model. - Given the number of vision tokens, output the number of - multi-modal connector tokens. - """ - ... + @deprecated( + "get_num_mm_connector_tokens is deprecated; use " + "get_mm_lora_token_counts instead." + ) + def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: ... def get_mm_lora_token_counts( self, diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index 1fcb6b8681..c57359b7b5 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -29,6 +29,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( BatchedTensorInputs, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -914,19 +915,20 @@ class InternVLChatModel( tower_model="vision_model", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - if num_image_tokens <= 0 or self.num_image_token <= 0: - return 0 + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + if num_mm_embeds <= 0 or self.num_image_token <= 0: + return 0, 0 - num_patches = num_image_tokens // self.num_image_token - return num_patches * (self.patch_tokens + 1) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - if num_vision_tokens <= 0 or self.num_image_token <= 0: - return 0 - - num_patches = num_vision_tokens // (self.patch_tokens + 1) - return num_patches * self.num_image_token + num_patches = num_mm_embeds // self.num_image_token + tower_tokens = num_patches * (self.patch_tokens + 1) + return tower_tokens, num_patches * self.num_image_token # -- SupportsEncoderCudaGraph protocol methods -- diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 53cda0ed17..d117b91d53 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -40,6 +40,7 @@ from vllm.multimodal.inputs import ( ImageItem, MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, VideoItem, ) @@ -1495,13 +1496,16 @@ class BaseKeyeModule(nn.Module, SupportsMultiModal): tower_model="visual.", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs merge_size = self.config.vision_config.spatial_merge_size - return num_image_tokens * merge_size**2 - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - merge_size = self.config.vision_config.spatial_merge_size - return num_vision_tokens // merge_size**2 + return num_mm_embeds * merge_size**2, num_mm_embeds @MULTIMODAL_REGISTRY.register_processor( diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index a8b57e7c28..15872a1f3f 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -35,6 +35,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems @@ -1284,12 +1285,13 @@ class Lfm2VLForConditionalGeneration( tower_model="vision_tower", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs downsample_factor = self.config.downsample_factor - - return num_image_tokens * downsample_factor**2 - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - downsample_factor = self.config.downsample_factor - - return num_vision_tokens // downsample_factor**2 + return num_mm_embeds * downsample_factor**2, num_mm_embeds diff --git a/vllm/model_executor/models/llava.py b/vllm/model_executor/models/llava.py index 1f1f45bfea..a32e0e6c63 100644 --- a/vllm/model_executor/models/llava.py +++ b/vllm/model_executor/models/llava.py @@ -28,6 +28,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import BaseMultiModalProcessorCache from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -739,18 +740,12 @@ class LlavaForConditionalGeneration( tower_model="vision_tower", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: - # LLaVA's vision encoder outputs one token per patch without - # spatial merging or pixel shuffle - return num_image_tokens - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - # LLaVA's MLP projector outputs the same number of tokens - # as it receives from the vision encoder (1:1 mapping) - return num_vision_tokens + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + return num_mm_embeds, num_mm_embeds diff --git a/vllm/model_executor/models/llava_next_video.py b/vllm/model_executor/models/llava_next_video.py index 038da17b75..96ebdaea8e 100644 --- a/vllm/model_executor/models/llava_next_video.py +++ b/vllm/model_executor/models/llava_next_video.py @@ -17,6 +17,7 @@ from vllm.model_executor.models.clip import CLIPVisionModel from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -494,31 +495,30 @@ class LlavaNextVideoForConditionalGeneration( tower_model="vision_tower", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_video_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + # Invert the spatial pooling done by `vision_resampler`: each frame # contributes `pooled_grid_length ** 2` tokens after the language # model's placeholder count, but `patch_grid_length ** 2` tokens # when it leaves the vision encoder. pooled_tokens_per_frame = self.pooled_grid_length**2 - if num_video_tokens <= 0 or pooled_tokens_per_frame <= 0: - return 0 - - num_frames = num_video_tokens // pooled_tokens_per_frame - return num_frames * (self.patch_grid_length**2) - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - # The projector is length-preserving; it runs after - # `vision_resampler` has already pooled each frame down to - # `pooled_grid_length ** 2` tokens. patch_tokens_per_frame = self.patch_grid_length**2 - if num_vision_tokens <= 0 or patch_tokens_per_frame <= 0: - return 0 + if ( + num_mm_embeds <= 0 + or pooled_tokens_per_frame <= 0 + or patch_tokens_per_frame <= 0 + ): + return 0, 0 - num_frames = num_vision_tokens // patch_tokens_per_frame - return num_frames * (self.pooled_grid_length**2) + num_frames = num_mm_embeds // pooled_tokens_per_frame + return ( + num_frames * patch_tokens_per_frame, + num_frames * pooled_tokens_per_frame, + ) diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 5a6e654417..0096eb4241 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -60,6 +60,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems @@ -1281,21 +1282,18 @@ class Llama4ForConditionalGeneration( tower_model="vision_model.", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs vision_config = self.config.vision_config patches_per_chunk = Mllama4ProcessingInfo.get_patch_per_chunk(vision_config) - if num_image_tokens <= 0 or patches_per_chunk <= 0: - return 0 + if num_mm_embeds <= 0 or patches_per_chunk <= 0: + return 0, 0 raw_patches = (vision_config.image_size // vision_config.patch_size) ** 2 - num_chunks = num_image_tokens // patches_per_chunk - # Encoder processes raw_patches + 1 (CLS) per chunk - return num_chunks * (raw_patches + 1) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - vision_config = self.config.vision_config - raw_patches = (vision_config.image_size // vision_config.patch_size) ** 2 - if num_vision_tokens <= 0: - return 0 - num_chunks = num_vision_tokens // (raw_patches + 1) - patches_per_chunk = Mllama4ProcessingInfo.get_patch_per_chunk(vision_config) - return num_chunks * patches_per_chunk + num_chunks = num_mm_embeds // patches_per_chunk + return num_chunks * (raw_patches + 1), num_chunks * patches_per_chunk diff --git a/vllm/model_executor/models/moondream3.py b/vllm/model_executor/models/moondream3.py index 818015e1ba..47a240b589 100644 --- a/vllm/model_executor/models/moondream3.py +++ b/vllm/model_executor/models/moondream3.py @@ -47,6 +47,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ImageSize, MultiModalDataItems @@ -1104,11 +1105,15 @@ class Moondream3ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): def get_language_model(self) -> nn.Module: return self.text - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - return num_image_tokens - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - return num_vision_tokens + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + return num_mm_embeds, num_mm_embeds def _split_pixel_values( self, diff --git a/vllm/model_executor/models/openvla.py b/vllm/model_executor/models/openvla.py index a676b3447c..ecbfb8a54c 100644 --- a/vllm/model_executor/models/openvla.py +++ b/vllm/model_executor/models/openvla.py @@ -24,7 +24,11 @@ from vllm.model_executor.models.interfaces import ( SupportsPP, ) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItem, + MultiModalKwargsItems, +) from vllm.multimodal.parse import ( ImageEmbeddingItems, ImageProcessorItems, @@ -487,11 +491,15 @@ class OpenVLAForActionPrediction(nn.Module, SupportsMultiModal, SupportsPP): tower_model="vision_backbone", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - return num_image_tokens - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - return num_vision_tokens + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + return num_mm_embeds, num_mm_embeds def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def maybe_rename_vision_weights( diff --git a/vllm/model_executor/models/paligemma.py b/vllm/model_executor/models/paligemma.py index 2fa498ee26..8e3524eeec 100644 --- a/vllm/model_executor/models/paligemma.py +++ b/vllm/model_executor/models/paligemma.py @@ -14,6 +14,7 @@ from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -401,8 +402,12 @@ class PaliGemmaForConditionalGeneration( tower_model="vision_tower", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: - return num_image_tokens - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - return num_vision_tokens + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs + return num_mm_embeds, num_mm_embeds diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index b6848b09de..1485b90e39 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -42,6 +42,7 @@ from vllm.model_executor.models.utils import WeightsMapper from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsOptionalItems, NestedTensors, ) @@ -623,17 +624,18 @@ class PixtralForConditionalGeneration( tower_model="vision_encoder", ) - def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs if getattr(self, "patch_merger", None) is None: - return num_image_tokens + return num_mm_embeds, num_mm_embeds merge_size = self.vision_args.spatial_merge_size - return num_image_tokens * (merge_size**2) - - def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: - if getattr(self, "patch_merger", None) is None: - return num_vision_tokens - merge_size = self.vision_args.spatial_merge_size - return num_vision_tokens // (merge_size**2) + return num_mm_embeds * (merge_size**2), num_mm_embeds # Vision encoder diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 61d0153f86..6cfacc1421 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -70,6 +70,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import MultiModalDataItems @@ -2050,21 +2051,15 @@ class Qwen2_5_VLForConditionalGeneration( tower_model="visual.", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs hf_config = self.config vision_config = hf_config.vision_config merge_size = vision_config.spatial_merge_size - - return num_image_tokens * merge_size**2 - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - hf_config = self.config - vision_config = hf_config.vision_config - merge_size = vision_config.spatial_merge_size - return num_vision_tokens // merge_size**2 + return num_mm_embeds * merge_size**2, num_mm_embeds diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index 80abdc62b5..c3761f8893 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -66,6 +66,7 @@ from vllm.multimodal.inputs import ( ImageItem, MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, VideoItem, ) @@ -1747,21 +1748,15 @@ class Qwen2VLForConditionalGeneration( tower_model="visual.", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs hf_config = self.config vision_config = hf_config.vision_config merge_size = vision_config.spatial_merge_size - - return num_image_tokens * merge_size**2 - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - hf_config = self.config - vision_config = hf_config.vision_config - merge_size = vision_config.spatial_merge_size - return num_vision_tokens // merge_size**2 + return num_mm_embeds * merge_size**2, num_mm_embeds diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index e98f37ba37..3bea939340 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -67,6 +67,7 @@ from vllm.multimodal.inputs import ( AudioItem, MultiModalFeatureSpec, MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, ) from vllm.multimodal.parse import ( @@ -625,16 +626,19 @@ class Qwen3ASRForConditionalGeneration( tower_model=["audio_tower."], ) - def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int: - """Return the number of tokens processed by the audio tower encoder. - - Required for LoRA support on the tower module. - """ + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs # For Qwen3-ASR, the audio tower produces one embedding per audio # placeholder token inserted into the prompt (no additional # merge/downsample step like vision towers). Therefore, the encoder # token budget is identity. - return num_audio_tokens + return num_mm_embeds, None @classmethod def get_speech_to_text_config( diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index d99cfe7aa2..644fd134ad 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -3077,24 +3077,18 @@ class Qwen3VLForConditionalGeneration( tower_model="visual.", ) - def get_num_mm_encoder_tokens( + def get_mm_lora_token_counts( self, - num_image_tokens: int, - ) -> int: + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs hf_config = self.config vision_config = hf_config.vision_config merge_size = vision_config.spatial_merge_size - - return num_image_tokens * merge_size**2 - - def get_num_mm_connector_tokens( - self, - num_vision_tokens: int, - ) -> int: - hf_config = self.config - vision_config = hf_config.vision_config - merge_size = vision_config.spatial_merge_size - return num_vision_tokens // merge_size**2 + return num_mm_embeds * merge_size**2, num_mm_embeds @lru_cache diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index a394e5f02e..bffe94fdce 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -29,6 +29,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, NestedTensors, ) @@ -373,9 +374,9 @@ def _build_chunk_attn_metadata( valid/padding boundary (equivalent to the key-padding mask the HF implementation uses), while every row still flows through the (potentially LoRA-wrapped) linears, keeping the per-chunk token counts - constant as required by `get_num_mm_encoder_tokens` / - `get_num_mm_connector_tokens`. Queries at padding positions produce - (garbage) outputs, which are trimmed by `audio_token_len` downstream. + constant as required by `get_mm_lora_token_counts`. Queries at padding + positions produce (garbage) outputs, which are trimmed by `audio_token_len` + downstream. """ batch_size = feature_lens.shape[0] starts = np.arange(batch_size, dtype=np.int64) * seq_len @@ -684,7 +685,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA): # LoRA on the tower/connector requires per-item token counts that are # predictable from the placeholder count alone (see - # `get_num_mm_encoder_tokens`), so pad every audio chunk's mel input + # `get_mm_lora_token_counts`), so pad every audio chunk's mel input # to the tower's full context instead of the batch's max length. self.pad_audio_to_max_context = bool( lora_config is not None and lora_config.enable_tower_connector_lora @@ -769,22 +770,26 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA): self.config.audio_config.max_source_positions / self.config.stack_factor ) - def get_num_mm_encoder_tokens(self, num_audio_tokens: int) -> int: + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + del modality, mm_kwargs # With `pad_audio_to_max_context` (enforced when tower/connector LoRA # is enabled), the tower's LoRA-wrapped linears always run on # `max_source_positions` conv-downsampled tokens per chunk, regardless # of the valid frame count. - num_chunks = math.ceil(num_audio_tokens / self._get_max_tokens_per_chunk()) - return num_chunks * self.config.audio_config.max_source_positions - - def get_num_mm_connector_tokens(self, num_encoder_tokens: int) -> int: + num_chunks = math.ceil(num_mm_embeds / self._get_max_tokens_per_chunk()) + tower_tokens = num_chunks * self.config.audio_config.max_source_positions # The connector runs on the frame-stacked tower output. Stacking pads # each chunk to a multiple of `stack_factor`, so this is # ceil(max_source_positions / stack_factor) tokens per chunk (188 for - # whisper's 1500), not `num_encoder_tokens // stack_factor` (187). - max_source_positions = self.config.audio_config.max_source_positions - num_chunks = num_encoder_tokens // max_source_positions - return num_chunks * self._get_max_tokens_per_chunk() + # whisper's 1500), not `tower_tokens // stack_factor` (187). + connector_tokens = num_chunks * self._get_max_tokens_per_chunk() + return tower_tokens, connector_tokens def _audio_features_to_embeddings( self, @@ -866,7 +871,7 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA): # Pad every chunk to the tower's full context so the number of # tokens processed by the tower/connector per chunk is constant, # keeping the LoRA token mappings computed by - # `get_num_mm_encoder_tokens`/`get_num_mm_connector_tokens` exact. + # `get_mm_lora_token_counts` exact. # Attention to the extra padding is masked via `audio_lens`. # Over-long inputs are not trimmed here; the tower raises on them. pad_len = self.audio_tower.max_context_length - audio_features.shape[-1] From 767d1c4d473a3e81429563aac995f22f7151c325 Mon Sep 17 00:00:00 2001 From: Russell Bryant Date: Mon, 14 Sep 2026 09:44:50 -0400 Subject: [PATCH 031/139] [Docs] Move russellb to emeritus committer (#56467) Signed-off-by: Russell Bryant Co-authored-by: Claude Opus 5 --- .github/CODEOWNERS | 16 ++++++++-------- docs/contributing/vulnerability_management.md | 2 -- docs/governance/committers.md | 9 +++++---- 3 files changed, 13 insertions(+), 14 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 10fd125768..aa5b2a9fc9 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -34,9 +34,9 @@ # Entrypoints /vllm/entrypoints/anthropic @mgoin @DarkLight1337 -/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb +/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 /vllm/entrypoints/mcp @heheda12345 -/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb +/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 /vllm/entrypoints/speech_to_text/realtime @njhill /vllm/entrypoints/speech_to_text @NickLucche /vllm/entrypoints/pooling @noooop @@ -73,7 +73,7 @@ /vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /vllm/v1/sample @22quinn @houseroad @njhill /vllm/v1/spec_decode @benchislett @luccafong @MatthewBonanni -/vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett @arpera +/vllm/v1/structured_output @mgoin @aarnphm @benchislett @arpera /vllm/v1/kv_cache_interface.py @heheda12345 @ivanium /vllm/v1/kv_offload @ApostaC @orozery /vllm/v1/simple_kv_offload @ivanium @@ -107,8 +107,8 @@ /tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan /tests/multimodal/media @Isotr0py /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas -/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm -/tests/v1/structured_output @mgoin @russellb @aarnphm @arpera +/tests/entrypoints/llm/test_struct_output_generate.py @mgoin @aarnphm @arpera +/tests/v1/structured_output @mgoin @aarnphm @arpera /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery @ivanium /tests/weight_loading @mgoin @youkaichao @yewentao256 /tests/lora @jeejeelee @@ -232,6 +232,6 @@ mkdocs.yaml @hmellor /docs/training/weight_transfer @aoshen02 # Security guide and policies -/docs/usage/security.md @russellb @jperezdealgaba -/SECURITY.md @russellb @jperezdealgaba -/docs/contributing/vulnerability_management.md @russellb @jperezdealgaba +/docs/usage/security.md @jperezdealgaba +/SECURITY.md @jperezdealgaba +/docs/contributing/vulnerability_management.md @jperezdealgaba diff --git a/docs/contributing/vulnerability_management.md b/docs/contributing/vulnerability_management.md index 8dd908a7ce..a73d451a62 100644 --- a/docs/contributing/vulnerability_management.md +++ b/docs/contributing/vulnerability_management.md @@ -33,9 +33,7 @@ on GitHub. However, if you need to contact the VMT directly for an urgent issue, you may contact the following individuals: - Simon Mo - -- Russell Bryant - - Juan Pérez de Algaba - -- Huzaifa Sidhpurwala - ## Slack Discussion diff --git a/docs/governance/committers.md b/docs/governance/committers.md index 0ca6e24ded..bd36759122 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -32,6 +32,7 @@ Sorted alphabetically by GitHub handle: - [@ivanium](https://github.com/ivanium): KV Connector and offload - [@jeejeelee](https://github.com/jeejeelee): LoRA, new model support - [@jikunshang](https://github.com/jikunshang): Intel CPU/XPU integration +- [@jperezdealgaba](https://github.com/jperezdealgaba): Security - [@khluu](https://github.com/khluu): CI infrastructure - [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance - [@markmc](https://github.com/markmc): Observability @@ -44,7 +45,6 @@ Sorted alphabetically by GitHub handle: - [@pavanimajety](https://github.com/pavanimajety): NVIDIA GPU integration - [@ProExpertProg](https://github.com/ProExpertProg): Compilation, startup UX - [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg -- [@russellb](https://github.com/russellb): Structured output, engine core, security - [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser - [@shen-shanshan](https://github.com/shen-shanshan): AMD GPU / ROCm integration, multimodality, ViT CUDA graph - [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community @@ -82,6 +82,7 @@ Committers who have contributed to vLLM significantly in the past (thank you!) b - [@pcmoritz](https://github.com/pcmoritz): MoE - [@rkooo567](https://github.com/rkooo567): Chunked prefill - [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support +- [@russellb](https://github.com/russellb): Structured output, engine core, security - [@sighingnow](https://github.com/sighingnow): Qwen models, new model support - [@sroy745](https://github.com/sroy745): Speculative decoding - [@Yard1](https://github.com/Yard1): kernels and performance @@ -100,7 +101,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - KV Cache Manager: memory management layer within scheduler maintaining KV cache logical block data - @heheda12345, @WoosukKwon - AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints - - @robertgshaw2-redhat, @njhill, @russellb + - @robertgshaw2-redhat, @njhill - ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni - KV Connector: Connector interface and implementation for KV cache offload and transfer @@ -153,7 +154,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP. - @WoosukKwon, @benchislett, @MatthewBonanni - Structured Output: The structured output implementation - - @russellb, @aarnphm + - @aarnphm - RL: The RL related features such as collective rpc, sleep mode, etc. - @youkaichao, @zhuohan123 - LoRA: @jeejeelee @@ -167,7 +168,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Documentation: @hmellor, @DarkLight1337, @simon-mo - Benchmarks: @ywang96, @simon-mo - CI, Build, Release Process: @khluu, @njhill, @simon-mo, @Harry-Chen, @vadiklyutiy -- Security: @russellb +- Security: @jperezdealgaba ### External Kernels Integration From c676e4930bb3c87b97790520e7979d43887ca5dd Mon Sep 17 00:00:00 2001 From: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:07:07 +0200 Subject: [PATCH 032/139] [Spec Decode] Fix Qwen3 DSpark d2t requirement for padded-vocab drafts (#55133) Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Signed-off-by: Orestis Zambounis Co-authored-by: Benjamin Chislett Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5 --- tests/model_executor/test_qwen3_omni.py | 46 ++++++++++++++++------ vllm/model_executor/models/qwen3_dspark.py | 7 +++- 2 files changed, 38 insertions(+), 15 deletions(-) diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 92d527e625..42accbef85 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -293,6 +293,25 @@ def test_qwen3_omni_text_model_collects_post_deepstack_aux_hidden_states(): torch.testing.assert_close(aux_hidden_states[0], torch.tensor([[15.0]])) +def _dspark_vocab_stub(input_vocab_size: int, draft_vocab_size: int): + """A Qwen3 DSpark model with only what `load_weights` reads.""" + from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM + + model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM) + nn.Module.__init__(model) + object.__setattr__( + model, + "config", + SimpleNamespace( + vocab_size=input_vocab_size, + draft_vocab_size=draft_vocab_size, + ), + ) + object.__setattr__(model, "target_vocab_size", 100) + object.__setattr__(model, "model", Mock(confidence_head=None)) + return model + + @pytest.mark.skip_global_cleanup @pytest.mark.parametrize( ("input_vocab_size", "draft_vocab_size", "weights", "error"), @@ -311,24 +330,25 @@ def test_qwen3_omni_text_model_collects_post_deepstack_aux_hidden_states(): def test_qwen3_dspark_rejects_incomplete_vocab_weights( input_vocab_size, draft_vocab_size, weights, error ): - from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM - - model = Qwen3DSparkForCausalLM.__new__(Qwen3DSparkForCausalLM) - nn.Module.__init__(model) - object.__setattr__( - model, - "config", - SimpleNamespace( - vocab_size=input_vocab_size, - draft_vocab_size=draft_vocab_size, - ), - ) - object.__setattr__(model, "target_vocab_size", 100) + model = _dspark_vocab_stub(input_vocab_size, draft_vocab_size) with pytest.raises(ValueError, match=error): model.load_weights(weights) +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("draft_vocab_size", [100, 101]) +def test_qwen3_dspark_accepts_draft_vocab_at_least_target(draft_vocab_size): + """A draft vocab >= the target's maps identically, so needs no lm_head/d2t. + + Guards the padded case (draft > target): the draft's unembedding is padded + to its physical size, which is not a reduced vocabulary. + """ + model = _dspark_vocab_stub(100, draft_vocab_size) + + model.load_weights([]) + + @pytest.mark.skip_global_cleanup def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary(): from vllm.v1.worker.gpu.spec_decode.dspark import utils as dspark_utils diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index c2c252280d..7318d225f8 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -254,7 +254,10 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): self.config.draft_vocab_size, scale=logit_scale ) self.target_vocab_size = vllm_config.model_config.get_vocab_size() - if self.config.draft_vocab_size != self.target_vocab_size: + # Only a pruned draft (subset of the target vocab) needs a d2t mapping to + # scatter ids up. A draft >= the target vocab (e.g. padded to the target's + # physical embedding size) maps identically, so no d2t is needed. + if self.config.draft_vocab_size < self.target_vocab_size: self.draft_id_to_target_id = nn.Parameter( torch.zeros(self.config.draft_vocab_size, dtype=torch.long), requires_grad=False, @@ -334,7 +337,7 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): # embed_tokens / lm_head are optional; when omitted they are shared from # the target by load_dspark_model, so skip the unloaded params here. uses_expanded_input_vocab = self.config.vocab_size > self.target_vocab_size - uses_reduced_vocab = self.config.draft_vocab_size != self.target_vocab_size + uses_reduced_vocab = self.config.draft_vocab_size < self.target_vocab_size if uses_expanded_input_vocab and not includes_embed_tokens: raise ValueError( "Qwen3 DSpark checkpoints whose input vocab_size is larger than " From ff5f6d41b1ccf1e305e425db41bc7f52dbce8621 Mon Sep 17 00:00:00 2001 From: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:07:09 +0200 Subject: [PATCH 033/139] [Bugfix] Scale KV page size for hidden states extraction with TP (#50894) Signed-off-by: Claude Signed-off-by: Orestis Zambounis Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Signed-off-by: Speculators Bot Signed-off-by: Speculators Bot Co-authored-by: Orestis Zambounis Co-authored-by: Claude Opus 4.6 Co-authored-by: Benjamin Chislett Co-authored-by: Speculators Bot Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/core/test_kv_cache_utils.py | 39 ++++++++++++++++++++++++++++ vllm/v1/core/kv_cache_utils.py | 34 ++++++++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 365338ad4f..c6354cdc0f 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3208,6 +3208,45 @@ def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): assert "attention compute is unchanged" in caplog_vllm.text +def test_hidden_states_with_tp_scales_page_size(): + """When TP shrinks KV pages below the hidden-state per-token cost, + get_kv_cache_groups must scale up target block sizes so that the + common page accommodates the unsharded hidden states.""" + # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. + # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) + # = 16 * 2 * 64 * 2 = 4096 bytes. + kv_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.bfloat16, + ) + # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) + # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. + hs_spec = HiddenStateCacheSpec( + block_size=16, + num_kv_heads=6, + head_size=512, + dtype=torch.bfloat16, + ) + specs = { + "target.0.attn": kv_spec, + "target.1.attn": kv_spec, + "cache_only_layers.48": hs_spec, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + + # The hidden-state layer should be present and no assertion should fire. + all_layers = {name for g in groups for name in g.layer_names} + assert "cache_only_layers.48" in all_layers + + # The target group block sizes must have been scaled up. + for g in groups: + if "cache_only_layers.48" not in g.layer_names: + assert g.kv_cache_spec.block_size > kv_spec.block_size + + def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 4e4f34bd52..b47df61add 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2217,6 +2217,38 @@ def _largest_divisor_at_most(value: int, limit: int) -> int: return 1 +def _ensure_min_page_size( + groups: list[KVCacheGroupSpec], + common_page: int, + hidden_specs: dict[str, HiddenStateCacheSpec], +) -> tuple[list[KVCacheGroupSpec], int]: + """Scale up group block sizes so the common page is at least as large as + the biggest hidden-state per-token cost. + + This protects the hidden state extraction feature, where we store hidden states in + the KV cache, and return them to clients using a custom hidden state connector. + + Returns the (possibly scaled) groups and updated common page size. + """ + max_per_token = max( + spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + for spec in hidden_specs.values() + ) + if max_per_token <= common_page: + return groups, common_page + + scale = cdiv(max_per_token, common_page) + common_page *= scale + scaled: list[KVCacheGroupSpec] = [] + for g in groups: + s = g.kv_cache_spec + kw: dict[str, int] = {"block_size": s.block_size * scale} + if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: + kw["page_size_padded"] = s.page_size_padded * scale + scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) + return scaled, common_page + + def get_kv_cache_groups( vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec], @@ -2288,6 +2320,8 @@ def get_kv_cache_groups( # Add hidden-state layers back with page aligned to the common page. if hidden_specs: common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) + # TP may shrink the common page below hidden-state per-token cost. + groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) for name, spec in hidden_specs.items(): per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) From 4be3dcf0fc7a9086d978eaea3c5a4d78fb98e44a Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 14 Sep 2026 10:20:46 -0400 Subject: [PATCH 034/139] [PCP][DCP] Declare FlashMLASparse MTP support at CP interleave > 1 (#56722) Co-authored-by: Claude Opus 5 (1M context) --- vllm/v1/attention/backends/mla/flashmla_sparse.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 433cded567..d8c4368401 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -634,6 +634,7 @@ class FlashMLASparseMetadataBuilder( class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): can_return_lse_for_decode: bool = True supports_dcp: bool = True + supports_mtp_with_cp_non_trivial_interleave_size: bool = True @staticmethod def _compute_fp8_decode_padded_heads(num_heads: int) -> int: From b3124a8237f21fcd3a4510002a5e2925b0dbedfe Mon Sep 17 00:00:00 2001 From: zql <37731799+zqlcode@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:06:01 +0800 Subject: [PATCH 035/139] [Model] Add support for Nanbeige4.2 (transformers backend) (#56071) Signed-off-by: zqlcode Signed-off-by: zql <37731799+zqlcode@users.noreply.github.com> Co-authored-by: zqlcode Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/models/supported_models.md | 1 + tests/models/registry.py | 4 ++++ vllm/model_executor/models/registry.py | 1 + 3 files changed, 6 insertions(+) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 81f4fb7df1..d55348ad72 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -490,6 +490,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor | `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | | | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ | | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ | +| `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ | | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index 874fa9c11c..cb39ed7ce4 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -461,6 +461,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { trust_remote_code=True, is_available_online=False, ), + "NanbeigeForCausalLM": _HfExamplesInfo( + "Nanbeige/Nanbeige4.2-3B", + trust_remote_code=True, + ), "OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"), "Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"), "Olmo3ForCausalLM": _HfExamplesInfo("allenai/Olmo-3-7B-Instruct"), diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 1a5f8399b6..6fc9f9a080 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -713,6 +713,7 @@ _TRANSFORMERS_SUPPORTED_MODELS = { "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"), + "NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"), "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"), "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"), "Olmo3ForCausalLM": ("transformers", "TransformersForCausalLM"), From 1713a094915d643c4ffaa25f421353108ee2049f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 14 Sep 2026 17:17:02 +0200 Subject: [PATCH 036/139] [Rust Frontend] Forward per-request watermarking controls (#56338) Co-authored-by: Codex Signed-off-by: NickLucche --- docs/features/watermarking.md | 7 ++- rust/proto/inference.proto | 3 ++ .../src/protocol/sampling.rs | 9 ++++ .../engine-core-client/src/tests/client.rs | 2 + .../src/tests/python_compat.py | 2 + rust/src/server/src/grpc/convert.rs | 22 ++++++++++ .../src/routes/inference/generate/convert.rs | 24 +++++++++++ .../routes/openai/chat_completions/convert.rs | 1 + .../routes/openai/chat_completions/types.rs | 5 +++ .../src/routes/openai/completions/convert.rs | 1 + .../src/routes/openai/completions/types.rs | 4 ++ rust/src/server/src/routes/tests.rs | 43 +++++++++++++++++++ rust/src/text/src/lower.rs | 8 ++++ rust/src/text/src/request.rs | 3 ++ 14 files changed, 133 insertions(+), 1 deletion(-) diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index f9ceda0547..54c7e58ce7 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -29,7 +29,12 @@ from vllm import SamplingParams sampling_params = SamplingParams(watermarking=False) ``` -The OpenAI-compatible APIs accept the same `watermarking: false` request field. +The Python OpenAI-compatible APIs and the Rust chat/completions APIs accept the +same `watermarking: false` request field. The Rust token API accepts it in +`sampling_params`, and the Rust gRPC `GenerateRequest` accepts +the optional `watermarking` field. Omitting the field keeps watermarking enabled +when the engine has a watermark configuration. + Deployments that require watermarking must restrict this field to trusted callers, or strip and validate it at the ingress boundary, so untrusted clients cannot opt out. diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 858d8a7a7e..46313f915a 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -53,6 +53,9 @@ message GenerateRequest { // Loaded LoRA name; empty selects the base model. string lora_name = 15; + + // Apply the engine's configured watermark. Defaults to true when omitted. + optional bool watermarking = 16; } message RandomSampling { diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index 19c579ac33..92f354be6f 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -20,6 +20,10 @@ fn default_temperature() -> f32 { 1.0 } +fn default_watermarking() -> bool { + true +} + fn default_max_tokens() -> u32 { 16 } @@ -68,6 +72,9 @@ pub struct EngineCoreSamplingParams { /// greedy sampling. #[serde(default = "default_temperature")] pub temperature: f32, + /// Whether to apply the engine's configured watermark to this request. + #[serde(default = "default_watermarking")] + pub watermarking: bool, /// Cumulative probability threshold for nucleus sampling. #[serde(default = "default_top_p")] pub top_p: f32, @@ -155,6 +162,7 @@ impl EngineCoreSamplingParams { pub fn for_test() -> Self { Self { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, @@ -229,6 +237,7 @@ mod tests { // Omitted fields -> Python defaults. assert_eq!(sampling.temperature, 1.0); + assert!(sampling.watermarking); assert_eq!(sampling.top_p, 1.0); assert_eq!(sampling.top_k, 0); assert_eq!(sampling.seed, None); diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 13525375f8..27fbbe9abd 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -151,6 +151,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest { prompt_token_ids: Some(vec![11, 22]), sampling_params: Some(EngineCoreSamplingParams { temperature: 0.8, + watermarking: false, top_p: 0.9, top_k: 8, max_tokens: 32, @@ -2662,6 +2663,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { sampling, EngineCoreSamplingParams { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index fd6cb0c86b..20cec81808 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -34,6 +34,7 @@ class FinishReason(IntEnum): # Mirror of real SamplingParams; omit_defaults makes fixtures match real maps. class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True): temperature: float = 1.0 + watermarking: bool = True top_p: float = 1.0 top_k: int = 0 seed: int | None = None @@ -123,6 +124,7 @@ request = EngineCoreRequest( mm_features=None, sampling_params=EngineCoreSamplingParams( temperature=0.8, + watermarking=False, top_p=0.9, top_k=8, seed=None, diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index aea04495eb..5f6ce5d42b 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -65,6 +65,7 @@ pub fn to_text_request( let mut sampling_params = build_sampling_params(req.temperature, sampling, decoding, stopping, response)?; + sampling_params.watermarking = req.watermarking.unwrap_or(true); // Thread KVCacheParameters → SamplingParams fields. if let Some(kv) = kv { @@ -521,6 +522,7 @@ impl ResponseOpts { #[cfg(test)] mod tests { + use prost::Message as _; use vllm_engine_core_client::protocol::output::StopReason; use vllm_text::{ FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params, @@ -539,6 +541,26 @@ mod tests { } } + #[test] + fn watermarking_defaults_and_opt_out_survive_protobuf_conversion() { + for watermarking in [None, Some(true), Some(false)] { + let request = pb::GenerateRequest { + watermarking, + ..base_request() + }; + let encoded = request.encode_to_vec(); + for stream in [false, true] { + let decoded = pb::GenerateRequest::decode(encoded.as_slice()).unwrap(); + let text = to_text_request(decoded, stream, &["test-model".to_string()]) + .expect("convert request"); + assert_eq!( + text.sampling_params.watermarking, + watermarking.unwrap_or(true) + ); + } + } + } + #[test] fn temperature_propagates_from_top_level_request_field() { let req = pb::GenerateRequest { diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index cb06cfe33c..08c4ce18be 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -164,6 +164,30 @@ mod tests { ); } + #[test] + fn prepare_generate_request_preserves_watermarking_defaults_and_opt_out() { + for watermarking in [None, Some(true), Some(false)] { + let mut body = json!({ + "token_ids": [11, 22], + "sampling_params": {} + }); + if let Some(watermarking) = watermarking { + body["sampling_params"]["watermarking"] = json!(watermarking); + } + let prepared = prepare_generate_request( + serde_json::from_value(body).expect("parse request"), + &served(&["test-model"]), + ResolvedRequestContext::default(), + None, + ) + .expect("prepare request"); + assert_eq!( + prepared.text_request.sampling_params.watermarking, + watermarking.unwrap_or(true) + ); + } + } + #[test] fn prepare_generate_request_forwards_thinking_token_budget() { let request: GenerateRequest = serde_json::from_value(json!({ diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 2c64ffbef4..51b0fb68ad 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -150,6 +150,7 @@ pub(super) fn prepare_chat_request( messages, sampling_params: SamplingParams { temperature: request.temperature, + watermarking: request.watermarking, top_p: request.top_p, top_k: request.top_k, seed: request.seed, diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index e4da86c22d..a371462694 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -116,6 +116,10 @@ pub struct ChatCompletionRequest { pub user: Option, // -------- vLLM Sampling Parameters -------- + /// Whether to apply the engine's configured watermark to this request. + #[serde(default = "default_true")] + pub watermarking: bool, + /// Use beam search instead of sampling #[serde(default)] pub use_beam_search: bool, @@ -282,6 +286,7 @@ impl Default for ChatCompletionRequest { include_reasoning: true, parallel_tool_calls: None, user: None, + watermarking: true, use_beam_search: false, top_k: None, min_p: None, diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index d944dca657..a4e8b6d3e4 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -115,6 +115,7 @@ pub(super) fn prepare_completion_request( mm_features: None, sampling_params: SamplingParams { temperature: request.temperature, + watermarking: request.watermarking, top_p: request.top_p, top_k: request.top_k, seed: request.seed, diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 8d93a89ed4..2a69013706 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -90,6 +90,10 @@ pub struct CompletionRequest { pub user: Option, // -------- vLLM Sampling Parameters -------- + /// Whether to apply the engine's configured watermark to this request. + #[serde(default = "default_true")] + pub watermarking: bool, + /// Options for streaming response pub stream_options: Option, diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 6fbe22e501..4731958e2f 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -2736,6 +2736,49 @@ async fn non_stream_chat_returns_json_response() { assert!(!message.contains_key("tool_calls"), "{json}"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn openai_watermarking_defaults_and_opt_out_reach_engine_core() { + for endpoint in ["/v1/chat/completions", "/v1/completions"] { + for stream in [false, true] { + for watermarking in [None, Some(true), Some(false)] { + let (mut app, engine_task) = test_app_with_backend_and_engine_request_check( + Arc::new(FakeChatBackend::new()), + move |request| { + let params = request.sampling_params.as_ref().expect("sampling params"); + assert_eq!(params.watermarking, watermarking.unwrap_or(true)); + }, + ) + .await; + + let mut body = json!({"stream": stream, "max_tokens": 2}); + if endpoint == "/v1/chat/completions" { + body["messages"] = json!([{"role": "user", "content": "hi"}]); + } else { + body["prompt"] = json!("hi"); + } + if let Some(watermarking) = watermarking { + body["watermarking"] = json!(watermarking); + } + let response = app + .call( + Request::builder() + .method("POST") + .uri(endpoint) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .expect("build request"), + ) + .await + .expect("call app"); + assert_eq!(response.status(), StatusCode::OK); + to_bytes(response.into_body(), usize::MAX).await.expect("drain response"); + engine_task.await.expect("mock engine task"); + } + } + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn non_stream_chat_image_url_reaches_engine_mm_features() { diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index bff4a6d48b..466b9d5b31 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -92,6 +92,7 @@ pub fn lower_sampling_params( ) -> Result { let SamplingParams { temperature, + watermarking, top_p, top_k, seed, @@ -164,6 +165,7 @@ pub fn lower_sampling_params( let params = EngineCoreSamplingParams { temperature, + watermarking, top_p, top_k, seed, @@ -632,6 +634,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, @@ -686,6 +689,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, @@ -854,6 +858,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 0.6, + watermarking: true, top_p: 0.95, top_k: 20, seed: None, @@ -918,6 +923,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 1.0, + watermarking: true, top_p: 1.0, top_k: 0, seed: None, @@ -990,6 +996,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 0.2, + watermarking: true, top_p: 0.3, top_k: 4, seed: None, @@ -1240,6 +1247,7 @@ mod tests { expect_test::expect![[r#" EngineCoreSamplingParams { temperature: 0.8, + watermarking: true, top_p: 0.9, top_k: 12, seed: None, diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 182e20c660..f0dd2fb1f8 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -69,6 +69,8 @@ pub struct SamplingParams { /// Controls randomness. Lower values are more deterministic; zero means /// greedy sampling. `None` means no explicit user override. pub temperature: Option, + /// Whether to apply the engine's configured watermark to this request. + pub watermarking: bool, /// Cumulative probability threshold for nucleus sampling. pub top_p: Option, /// Maximum number of top tokens to consider. `Some(0)` means all tokens. @@ -144,6 +146,7 @@ impl Default for SamplingParams { fn default() -> Self { Self { temperature: None, + watermarking: true, top_p: None, top_k: None, seed: None, From b7e8dd8f372c1324c376ba5dc03d92264eae8f24 Mon Sep 17 00:00:00 2001 From: Tomas Ruiz Date: Mon, 14 Sep 2026 17:29:13 +0200 Subject: [PATCH 037/139] [watermarking hardening] Add e2e test and basic GSM8K quality tests (#56309) Signed-off-by: Tomas Ruiz Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test_areas/lm_eval.yaml | 16 +++++++ .../gsm8k/configs/Qwen3.5-4B-watermark.yaml | 22 ++++++++++ .../evals/gsm8k/configs/models-watermark.txt | 1 + tests/watermarking/test_watermarking_e2e.py | 42 +++++++++++++++++++ 4 files changed, 81 insertions(+) create mode 100644 tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml create mode 100644 tests/evals/gsm8k/configs/models-watermark.txt create mode 100644 tests/watermarking/test_watermarking_e2e.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 3c65611966..d9afefc82e 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -30,6 +30,22 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py +- label: ":nvidia: (H200 MIG 35GB) LM Eval Watermarking" + device: h200_35gb + key: lm-eval-watermarking + # Recorded runs took 80–140 seconds + timeout_in_minutes: 10 + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/config/watermarking.py + - vllm/sampling_params.py + - vllm/v1/watermarking/ + - vllm/v1/worker/gpu/ + autorun_on_main: true + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt + # - label: LM Eval Large Models (4xA100) # key: lm-eval-large-models-4xa100 # device: a100 diff --git a/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml b/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml new file mode 100644 index 0000000000..bdd42a0f1a --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3.5-4B-watermark.yaml @@ -0,0 +1,22 @@ +model_name: "Qwen/Qwen3.5-4B" +# full GSM8K set is 1319 questions +num_questions: 1319 +num_fewshot: 5 +max_concurrency: 128 +max_tokens: 1024 +use_chat_completions: true +temperature: 0.7 +seed: 42 + +# Calibrated by running non-watermaked 13 times +# the threshold is the average accuracy, and the tolerance is set to 3 sigma (standard deviations) +accuracy_threshold: 0.885 +tolerance: 0.015 + +server_args: >- + --language-model-only + --max-model-len 8192 + --max-num-seqs 128 + --gpu-memory-utilization 0.4 + --default-chat-template-kwargs '{"enable_thinking": false}' + --watermark-config '{"key": 42, "context_width": 4}' diff --git a/tests/evals/gsm8k/configs/models-watermark.txt b/tests/evals/gsm8k/configs/models-watermark.txt new file mode 100644 index 0000000000..a761b41c1b --- /dev/null +++ b/tests/evals/gsm8k/configs/models-watermark.txt @@ -0,0 +1 @@ +Qwen3.5-4B-watermark.yaml diff --git a/tests/watermarking/test_watermarking_e2e.py b/tests/watermarking/test_watermarking_e2e.py new file mode 100644 index 0000000000..c73d48f6b6 --- /dev/null +++ b/tests/watermarking/test_watermarking_e2e.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm import SamplingParams +from vllm.platforms import current_platform +from vllm.v1.watermarking import GumbelWatermarkDetector + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_llm_generated_sequence_is_watermarked(vllm_runner): + """Engine-level watermarking must produce tokens the detector recognizes.""" + watermark_config = {"key": 42, "context_width": 4} + detector = GumbelWatermarkDetector(**watermark_config) + + runner = vllm_runner( + "facebook/opt-125m", + watermark_config=watermark_config, + seed=0, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.2, + ) + request = "Tell me a story about an explorer who discovers a mysterious island." + params_use_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=True) + params_no_wm = SamplingParams(temperature=1.0, max_tokens=256, watermarking=False) + + with runner: + output_use_wm = runner.llm.generate(request, params_use_wm) + output_no_wm = runner.llm.generate(request, params_no_wm) + + result_use_wm = detector.detect(list(output_use_wm[0].outputs[0].token_ids)) + result_no_wm = detector.detect(list(output_no_wm[0].outputs[0].token_ids)) + + assert result_use_wm.is_watermarked, result_use_wm + recorded_p_value = 9.83e-20 + assert result_use_wm.p_value < recorded_p_value * 10 # tolerance + + assert not result_no_wm.is_watermarked, result_no_wm From 435c96f9dbdd29258cb8e0f433c5b54a00cf6b16 Mon Sep 17 00:00:00 2001 From: Thillai Chithambaram <79466435+thillai-c@users.noreply.github.com> Date: Mon, 14 Sep 2026 12:04:55 -0400 Subject: [PATCH 038/139] [V1][Metrics] Support Sliding Window Attention (SWA) and hybrid layers in MFU/MBU estimation (#55624) Signed-off-by: Thillai Chithambaram Co-authored-by: Mark McLoughlin --- tests/v1/metrics/test_perf_metrics.py | 374 +++++++++++++++++++++++++- vllm/v1/metrics/perf.py | 314 +++++++++++++++++---- 2 files changed, 640 insertions(+), 48 deletions(-) diff --git a/tests/v1/metrics/test_perf_metrics.py b/tests/v1/metrics/test_perf_metrics.py index ab30f1bb9e..c94ec1b701 100644 --- a/tests/v1/metrics/test_perf_metrics.py +++ b/tests/v1/metrics/test_perf_metrics.py @@ -8,6 +8,7 @@ import types from types import SimpleNamespace import pytest +from transformers import Gemma2Config, MistralConfig from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config from transformers.models.llama4.configuration_llama4 import ( Llama4Config, @@ -31,6 +32,7 @@ from vllm.v1.metrics.perf import ( MLAAttentionMetrics, ModelMetrics, ParsedArgs, + SlidingWindowAttentionParser, UnembedMetrics, ) @@ -38,7 +40,8 @@ from vllm.v1.metrics.perf import ( class MockModelConfig: """Mock ModelConfig that implements the getter methods used by parsers.""" - def __init__(self, hf_config, dtype): + def __init__(self, hf_config, dtype, disable_sliding_window=False): + """Initialize MockModelConfig with HF config and options.""" self.hf_config = hf_config self.hf_text_config = get_hf_text_config(hf_config) convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get( @@ -49,8 +52,10 @@ class MockModelConfig: ).convert() self.dtype = dtype self.is_attention_free = False + self.disable_sliding_window = disable_sliding_window def __getattr__(self, name): + """Delegate attribute lookups to ModelConfig.""" # 1. Check if ModelConfig actually has this attribute if not hasattr(ModelConfig, name): raise AttributeError( @@ -86,9 +91,13 @@ def create_mock_vllm_config( tensor_parallel_size=1, pipeline_parallel_size=1, enable_expert_parallel=False, + disable_sliding_window=False, ) -> SimpleNamespace: + """Create a mock VllmConfig object for unit tests.""" vllm_config = SimpleNamespace() - vllm_config.model_config = MockModelConfig(hf_config, model_dtype) + vllm_config.model_config = MockModelConfig( + hf_config, model_dtype, disable_sliding_window=disable_sliding_window + ) vllm_config.cache_config = SimpleNamespace() vllm_config.cache_config.cache_dtype = cache_dtype @@ -1336,3 +1345,364 @@ def test_mla_attention_scaling_with_layers(): assert double_metrics.get_num_flops(ctx) == 2 * base_metrics.get_num_flops(ctx) assert double_metrics.get_read_bytes(ctx) == 2 * base_metrics.get_read_bytes(ctx) assert double_metrics.get_write_bytes(ctx) == 2 * base_metrics.get_write_bytes(ctx) + + +#### Sliding Window Attention Tests #### + + +def test_sliding_window_parser_no_swa(): + """Test parser defaults to no SWA when sliding_window is unset.""" + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_hidden_layers=12, + ) + vllm_config = create_mock_vllm_config(hf_config) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 12 + result = parser.parse(args, vllm_config) + + assert result.sliding_window is None + assert result.num_swa_layers == 0 + + +def test_sliding_window_parser_uniform_swa(): + """Test SlidingWindowAttentionParser for full SWA model (e.g. Mistral).""" + hf_config = MistralConfig( + hidden_size=4096, + num_attention_heads=32, + num_key_value_heads=8, + num_hidden_layers=32, + sliding_window=4096, + ) + vllm_config = create_mock_vllm_config(hf_config) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 32 + result = parser.parse(args, vllm_config) + + assert result.sliding_window == 4096 + assert result.num_swa_layers == 32 + + +def test_sliding_window_parser_explicit_layer_types(): + """Test SlidingWindowAttentionParser with explicit layer_types list.""" + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_hidden_layers=6, + ) + hf_config.sliding_window = 2048 + hf_config.layer_types = [ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ] + vllm_config = create_mock_vllm_config(hf_config) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 6 + result = parser.parse(args, vllm_config) + + assert result.sliding_window == 2048 + assert result.num_swa_layers == 3 + + +def test_sliding_window_parser_gemma2(): + """Test SlidingWindowAttentionParser detects alternating SWA for Gemma 2.""" + hf_config = Gemma2Config( + hidden_size=2304, + num_attention_heads=8, + num_key_value_heads=4, + num_hidden_layers=26, + sliding_window=4096, + ) + vllm_config = create_mock_vllm_config(hf_config) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 26 + result = parser.parse(args, vllm_config) + + assert result.sliding_window == 4096 + # Gemma 2 alternates: even layers are sliding window (26 layers -> 13 SWA) + assert result.num_swa_layers == 13 + + +def test_sliding_window_parser_disabled(): + """Test SlidingWindowAttentionParser respects disable_sliding_window.""" + hf_config = MistralConfig( + hidden_size=4096, + num_attention_heads=32, + num_key_value_heads=8, + num_hidden_layers=32, + sliding_window=4096, + ) + vllm_config = create_mock_vllm_config(hf_config, disable_sliding_window=True) + parser = SlidingWindowAttentionParser() + args = ParsedArgs() + args.num_hidden_layers = 32 + result = parser.parse(args, vllm_config) + + assert result.sliding_window is None + assert result.num_swa_layers == 0 + + +def test_attention_metrics_sliding_window_flops_bounds(): + """Test that SWA clamps attention core FLOPs when context_len > sliding_window.""" + sw = 1024 + dense_config = MistralConfig( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=8, + sliding_window=None, + ) + swa_config = MistralConfig( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=8, + sliding_window=sw, + ) + + dense_vllm = create_mock_vllm_config(dense_config) + swa_vllm = create_mock_vllm_config(swa_config) + + dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm) + swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm) + + # 1. When context <= sliding_window: FLOPs should be identical + short_ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=512, is_prefill=False + ) + assert swa_metrics.get_num_flops(short_ctx) == dense_metrics.get_num_flops( + short_ctx + ) + + # 2. When context > sliding_window: FLOPs for SWA must be strictly smaller + long_ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=8192, is_prefill=False + ) + dense_flops = dense_metrics.get_num_flops_breakdown(long_ctx) + swa_flops = swa_metrics.get_num_flops_breakdown(long_ctx) + + # Projections (qkv_proj, out_proj) are unaffected by sliding window + assert swa_flops["qkv_proj"] == dense_flops["qkv_proj"] + assert swa_flops["out_proj"] == dense_flops["out_proj"] + + # Attention core FLOPs are clamped: 8192 vs 1024 (8x difference) + assert swa_flops["attn_qk"] < dense_flops["attn_qk"] + assert swa_flops["attn_av"] < dense_flops["attn_av"] + assert swa_flops["attn_qk"] * 8 == dense_flops["attn_qk"] + assert swa_flops["attn_av"] * 8 == dense_flops["attn_av"] + + +def test_attention_metrics_sliding_window_read_bytes(): + """Test that SWA bounds decode KV-cache read bytes to sliding_window.""" + sw = 2048 + dense_config = MistralConfig( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=4, + num_hidden_layers=4, + sliding_window=None, + ) + swa_config = MistralConfig( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=4, + num_hidden_layers=4, + sliding_window=sw, + ) + + dense_vllm = create_mock_vllm_config(dense_config) + swa_vllm = create_mock_vllm_config(swa_config) + + dense_metrics = AttentionMetrics.from_vllm_config(dense_vllm) + swa_metrics = AttentionMetrics.from_vllm_config(swa_vllm) + + # Long decode context: 8192 tokens + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=8192, is_prefill=False + ) + dense_reads = dense_metrics.get_read_bytes_breakdown(ctx) + swa_reads = swa_metrics.get_read_bytes_breakdown(ctx) + + # Attention input read traffic is lower for SWA + assert swa_reads["attn_input"] < dense_reads["attn_input"] + + # Weight reads and projection inputs are identical + assert swa_reads["qkv_weight"] == dense_reads["qkv_weight"] + assert swa_reads["out_weight"] == dense_reads["out_weight"] + assert swa_reads["qkv_input"] == dense_reads["qkv_input"] + + +def test_attention_metrics_hybrid_alternating_layers(): + """Test hybrid model where half layers are full attention and half are SWA.""" + sw = 1024 + hybrid_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=4, + ) + hybrid_config.sliding_window = sw + hybrid_config.layer_types = [ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ] + + vllm_config = create_mock_vllm_config(hybrid_config) + metrics = AttentionMetrics.from_vllm_config(vllm_config) + + assert metrics.num_swa_layers == 2 + assert metrics.sliding_window == sw + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=4096, is_prefill=False + ) + flops = metrics.get_num_flops_breakdown(ctx) + + # 2 full layers attend to 4096; 2 SWA layers attend to 1024 + # Expected TC total = 2 * 4096 + 2 * 1024 = 8192 + 2048 = 10240 + # Formula: 2 * q * d * (L_full * 4096 + L_swa * 1024) + q = metrics.num_attention_heads + d = metrics.head_dim + expected_core_flops = 2 * q * d * (2 * 4096 + 2 * 1024) + assert flops["attn_qk"] == expected_core_flops + assert flops["attn_av"] == expected_core_flops + + +def test_attention_metrics_sliding_window_per_gpu(): + """Test that SWA metrics scale properly with tensor and pipeline parallelism.""" + sw = 1024 + hf_config = MistralConfig( + hidden_size=4096, + num_attention_heads=32, + num_key_value_heads=8, + num_hidden_layers=16, + sliding_window=sw, + ) + single_vllm = create_mock_vllm_config(hf_config) + tp_vllm = create_mock_vllm_config(hf_config, tensor_parallel_size=2) + pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2) + + single_metrics = AttentionMetrics.from_vllm_config(single_vllm) + tp_metrics = AttentionMetrics.from_vllm_config(tp_vllm) + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=4096, is_prefill=False + ) + + # Both TP=2 and PP=2 should halve the per-GPU attention core FLOPs + single_flops = single_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + tp_flops = tp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + + assert tp_flops["attn_qk"] == single_flops["attn_qk"] // 2 + assert pp_flops["attn_qk"] == single_flops["attn_qk"] // 2 + + +def test_attention_metrics_hybrid_sliding_window_per_gpu(): + """Test that hybrid SWA models preserve L_full + L_swa == L under PP.""" + sw = 1024 + hybrid_config = Gemma2Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=30, + sliding_window=sw, + ) + pp_vllm = create_mock_vllm_config(hybrid_config, pipeline_parallel_size=2) + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=4096, is_prefill=False + ) + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + pp_read_bytes = pp_metrics.get_read_bytes_breakdown(ctx, per_gpu=True) + + # Under PP=2, each global count is halved: 30 layers -> 15 per GPU, and the + # 15 SWA / 15 full split becomes 7.5 each. Gemma 2 alternates, so an even + # split is the correct average stage; deriving it after reducing L instead + # produced a lopsided 8/7. + q = pp_metrics.num_attention_heads + d = pp_metrics.head_dim + expected_pp_attn = int(2 * q * d * (7.5 * 4096 + 7.5 * sw)) + assert pp_flops["attn_qk"] == expected_pp_attn + assert pp_flops["attn_av"] == expected_pp_attn + assert "attn_input" in pp_read_bytes + + +def test_attention_metrics_minority_swa_survives_pipeline_parallelism(): + """A minority SWA sublayer must not round away under PP. + + Regression test for a 6-layer model with a single sliding-window layer at + pp_size=2: deriving the split after reducing L gave + round(3 * 1 / 6) == 0, billing all three per-GPU layers as full attention + and dropping the sliding-window contribution entirely. + """ + sw = 512 + context_len = 4096 + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=6, + ) + hf_config.sliding_window = sw + hf_config.layer_types = ["sliding_attention"] + ["full_attention"] * 5 + + pp_vllm = create_mock_vllm_config(hf_config, pipeline_parallel_size=2) + pp_metrics = AttentionMetrics.from_vllm_config(pp_vllm) + assert pp_metrics.num_swa_layers == 1 + + ctx = ExecutionContext.from_single_request( + num_tokens=1, context_len=context_len, is_prefill=False + ) + pp_flops = pp_metrics.get_num_flops_breakdown(ctx, per_gpu=True) + + # 5 full and 1 SWA layer, each halved by PP=2: 2.5 full and 0.5 SWA. + q = pp_metrics.num_attention_heads + d = pp_metrics.head_dim + expected = int(2 * q * d * (2.5 * context_len + 0.5 * sw)) + assert pp_flops["attn_qk"] == expected + assert pp_flops["attn_av"] == expected + + # The sliding-window layer must change the result, i.e. it is not being + # billed at the full context length. + all_full = int(2 * q * d * 3 * context_len) + assert pp_flops["attn_qk"] < all_full + + +def test_attention_metrics_per_gpu_layers_sum_to_whole_model(): + """Per-GPU counts times pp_size must account for every layer. + + Flooring the layer count per stage silently drops the remainder, so a model + whose layers do not divide evenly across stages under-reports its FLOPs. + """ + hf_config = Qwen3Config( + hidden_size=2048, + num_attention_heads=16, + num_key_value_heads=16, + num_hidden_layers=7, + ) + pp_size = 2 + vllm_config = create_mock_vllm_config(hf_config, pipeline_parallel_size=pp_size) + metrics = AttentionMetrics.from_vllm_config(vllm_config) + + ctx = ExecutionContext.from_single_request( + num_tokens=8, context_len=1024, is_prefill=True + ) + per_gpu = metrics.get_num_flops_breakdown(ctx, per_gpu=True) + whole = metrics.get_num_flops_breakdown(ctx, per_gpu=False) + + for key, value in whole.items(): + assert per_gpu[key] * pp_size == pytest.approx(value, rel=1e-9) diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index 75524ef128..ab90ffcbe1 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -10,7 +10,7 @@ import json import time from abc import ABC, abstractmethod from collections.abc import Iterable -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from typing import Any, Protocol import prometheus_client @@ -124,8 +124,12 @@ class ExecutionContext: decode_context_len: int = 0 # sum of context_len for decode requests decode_token_context_product: int = 0 # sum of (num_tokens * context_len) + # Per-request records: (num_tokens, context_len, is_prefill) + requests: list[tuple[int, int, bool]] = field(default_factory=list) + def add(self, num_tokens: int, context_len: int, is_prefill: bool) -> None: """Add a single request's statistics to this batch context.""" + self.requests.append((num_tokens, context_len, is_prefill)) if is_prefill: self.num_prefill_requests += 1 self.prefill_num_tokens += num_tokens @@ -141,9 +145,40 @@ class ExecutionContext: """Total number of tokens across all requests in the batch.""" return self.prefill_num_tokens + self.decode_num_tokens - def total_token_context_product(self) -> int: - """Total sum of (num_tokens * context_len) across all requests.""" - return self.prefill_token_context_product + self.decode_token_context_product + def total_token_context_product(self, sliding_window: int | None = None) -> int: + """Total sum of (num_tokens * context_len) across all requests. + + If sliding_window is provided, context length per request is clamped to + sliding_window. + """ + if sliding_window is None: + return ( + self.prefill_token_context_product + self.decode_token_context_product + ) + return sum( + num_tokens * min(context_len, sliding_window) + for num_tokens, context_len, _ in self.requests + ) + + def prefill_context_len_for_window(self, sliding_window: int | None = None) -> int: + """Sum of context_len for prefill requests, clamped to sliding_window if set.""" + if sliding_window is None: + return self.prefill_context_len + return sum( + min(context_len, sliding_window) + for _, context_len, is_prefill in self.requests + if is_prefill + ) + + def decode_context_len_for_window(self, sliding_window: int | None = None) -> int: + """Sum of context_len for decode requests, clamped to sliding_window if set.""" + if sliding_window is None: + return self.decode_context_len + return sum( + min(context_len, sliding_window) + for _, context_len, is_prefill in self.requests + if not is_prefill + ) def num_logits_tokens(self) -> int: """Number of tokens that require logits computation (unembedding). @@ -214,6 +249,7 @@ class ParserChain: self.parsers.append(parser) def parse(self, vllm_config: VllmConfig) -> ParsedArgs: + """Apply all parsers in sequence and return the accumulated ParsedArgs.""" args = ParsedArgs() for parser in self.parsers: args = parser.parse(args, vllm_config) @@ -234,7 +270,9 @@ class ComponentMetrics(BaseModel, ABC): @classmethod @abstractmethod - def component_type(cls) -> str: ... + def component_type(cls) -> str: + """Return the unique string key identifying this component in the registry.""" + ... @classmethod @abstractmethod @@ -273,17 +311,23 @@ class ComponentMetrics(BaseModel, ABC): @abstractmethod def get_num_flops_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True - ) -> dict[str, int]: ... + ) -> dict[str, int]: + """Return per-operation FLOPs breakdown for this component.""" + ... @abstractmethod def get_read_bytes_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True - ) -> dict[str, int]: ... + ) -> dict[str, int]: + """Return per-operation read-memory-traffic breakdown for this component.""" + ... @abstractmethod def get_write_bytes_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True - ) -> dict[str, int]: ... + ) -> dict[str, int]: + """Return per-operation write-memory-traffic breakdown for this component.""" + ... def get_num_flops(self, ctx: ExecutionContext, per_gpu: bool = True) -> int: return sum(self.get_num_flops_breakdown(ctx, per_gpu).values()) @@ -306,6 +350,7 @@ class BaseConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Parse base model dimensions, dtype byte sizes, and parallelism config.""" model_config = vllm_config.model_config args.vocab_size = model_config.get_vocab_size() @@ -357,6 +402,7 @@ class BaseAttentionConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Parse KV head count, head dimension, and KV-cache dtype byte size.""" model_config = vllm_config.model_config args.num_key_value_heads = model_config.get_total_num_kv_heads() @@ -371,6 +417,54 @@ class BaseAttentionConfigParser(Parser): return args +class SlidingWindowAttentionParser(Parser): + """ + Parses sliding window attention configuration and layer breakdown. + Provides: sliding_window, num_swa_layers + """ + + def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Determine sliding_window size and count of SWA layers. + + Detects per-layer type lists (e.g. Llama-4), Gemma 2 alternating + pattern, and uniform SWA models (Mistral, Qwen2.5). + """ + model_config = vllm_config.model_config + + sliding_window: int | None = None + if not getattr(model_config, "disable_sliding_window", False): + sliding_window = model_config.get_sliding_window() + if sliding_window is None and hasattr(vllm_config, "cache_config"): + sliding_window = getattr( + vllm_config.cache_config, "sliding_window", None + ) + + args.sliding_window = sliding_window + + num_layers = args.num_hidden_layers + if sliding_window is not None and sliding_window > 0: + hf_text_config = getattr(model_config, "hf_text_config", None) + layer_types = getattr(hf_text_config, "layer_types", None) + if layer_types is not None and isinstance(layer_types, list): + num_swa = sum( + 1 for lt in layer_types[:num_layers] if "sliding" in lt.lower() + ) + args.num_swa_layers = num_swa + elif ( + getattr(getattr(model_config, "hf_config", None), "model_type", None) + == "gemma2" + ): + # Gemma 2 alternates sliding window: even layers (0, 2, ...) are sliding + args.num_swa_layers = (num_layers + 1) // 2 + else: + # Full SWA model (e.g. Mistral, Qwen2.5 with sliding_window) + args.num_swa_layers = num_layers + else: + args.num_swa_layers = 0 + + return args + + class AttentionQuantizationConfigParser(Parser): """ Parses quantization configuration for attention layers. @@ -378,6 +472,7 @@ class AttentionQuantizationConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Override weight_byte_size based on the active quantization method.""" cfg = vllm_config.quant_config if cfg is None: @@ -401,6 +496,7 @@ class AttentionDetectionParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Raise InvalidComponent for MLA models so AttentionMetrics is skipped.""" if vllm_config.model_config.is_deepseek_mla: raise InvalidComponent( "Model uses MLA attention; use MLAAttentionMetrics instead" @@ -425,53 +521,129 @@ class AttentionMetrics(ComponentMetrics): # From BaseConfig Parser, overridden by AttentionQuantizationConfigParser weight_byte_size: int | float = Field(..., gt=0) - # TODO: discern cases where we have mixture of different attention layer types - # such as SWA, MLA, etc. + # From SlidingWindowAttentionParser + sliding_window: int | None = Field(None) + num_swa_layers: int = Field(0, ge=0) @classmethod def component_type(cls) -> str: + """Return the component registry key for standard (non-MLA) attention.""" return "attn" @classmethod def get_parser(cls) -> ParserChain: + """Return the parser chain for AttentionMetrics. + + Parsers run in order: MLA detection guard, base model config, + attention-specific config, sliding-window breakdown, quantization + weight-size override. + """ return ParserChain( AttentionDetectionParser(), BaseConfigParser(), BaseAttentionConfigParser(), + SlidingWindowAttentionParser(), AttentionQuantizationConfigParser(), ) + def _layer_counts(self, per_gpu: bool) -> tuple[float, float, float]: + """Return ``(L, L_full, L_swa)``, the layer counts to bill for. + + Per-GPU counts divide each global count by ``pp_size`` independently. + Deriving the sublayer split after reducing ``L`` instead rounds a + minority sublayer away: a 6-layer model with 1 sliding-window layer at + ``pp_size=2`` yields ``round(3 * 1 / 6) == 0``, billing all three layers + as full attention. Dividing each count separately keeps + ``L_full + L_swa == L`` and keeps the per-GPU values summing back to the + whole model across stages. + + The counts are fractional because these metrics are built from config + alone (see ``from_vllm_config``); no pipeline rank is available, so a + per-GPU figure is the average stage rather than a specific one. A model + whose layers do not divide evenly across stages has no single integer + answer here. + """ + L = float(self.num_hidden_layers) + L_swa = float(self.num_swa_layers) + L_full = L - L_swa + + if per_gpu: + L /= self.pp_size + L_swa /= self.pp_size + L_full /= self.pp_size + + return L, L_full, L_swa + def get_num_flops_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True ) -> dict[str, int]: - L, D, q, kv, d = ( - self.num_hidden_layers, + """Compute FLOPs breakdown for attention layers. + + Accounts for hybrid models by splitting total layers into full-attention + (L_full) and sliding-window (L_swa) sublayers. See ``_layer_counts`` for + how those are scaled under pipeline parallelism. + + Args: + ctx: Execution context describing the current batch. + per_gpu: If True, scale counts down by tensor/pipeline parallelism. + + Returns: + Dict with keys: qkv_proj, out_proj, attn_qk, attn_av. + """ + D, q, kv, d = ( self.hidden_size, self.num_attention_heads, self.num_key_value_heads, self.head_dim, ) T = ctx.total_num_tokens() - TC = ctx.total_token_context_product() + TC_full = ctx.total_token_context_product() + + L, L_full, L_swa = self._layer_counts(per_gpu) if per_gpu: - L //= self.pp_size # tensor parallel along heads q = max(1, q // self.tp_size) kv = max(1, kv // self.tp_size) - return { - "qkv_proj": 2 * T * D * (q + 2 * kv) * d * L, - "attn_qk": 2 * q * TC * d * L, - "attn_av": 2 * q * TC * d * L, - "out_proj": 2 * T * D * q * d * L, + flops = { + "qkv_proj": int(2 * T * D * (q + 2 * kv) * d * L), + "out_proj": int(2 * T * D * q * d * L), } + attn_qk = 0.0 + attn_av = 0.0 + if L_full > 0: + attn_qk += 2 * q * TC_full * d * L_full + attn_av += 2 * q * TC_full * d * L_full + + if L_swa > 0 and self.sliding_window is not None: + TC_swa = ctx.total_token_context_product(self.sliding_window) + attn_qk += 2 * q * TC_swa * d * L_swa + attn_av += 2 * q * TC_swa * d * L_swa + + flops["attn_qk"] = int(attn_qk) + flops["attn_av"] = int(attn_av) + + return flops + def get_read_bytes_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True ) -> dict[str, int]: - L, D, q, kv, d = ( - self.num_hidden_layers, + """Compute read-memory-traffic breakdown for attention layers. + + Splits layers into full-attention and sliding-window sublayers using the + same counts as get_num_flops_breakdown; see ``_layer_counts``. + + Args: + ctx: Execution context describing the current batch. + per_gpu: If True, scale counts down by tensor/pipeline parallelism. + + Returns: + Dict with keys: qkv_input, qkv_weight, attn_input (conditional), + out_input, out_weight. + """ + D, q, kv, d = ( self.hidden_size, self.num_attention_heads, self.num_key_value_heads, @@ -479,35 +651,63 @@ class AttentionMetrics(ComponentMetrics): ) T = ctx.total_num_tokens() + L, L_full, L_swa = self._layer_counts(per_gpu) + if per_gpu: - L //= self.pp_size # tensor parallel along heads q = max(1, q // self.tp_size) kv = max(1, kv // self.tp_size) read_bytes = {} - read_bytes["qkv_input"] = T * D * self.activation_byte_size * L + read_bytes["qkv_input"] = int(T * D * self.activation_byte_size * L) read_bytes["qkv_weight"] = int(D * (q + 2 * kv) * d * self.weight_byte_size * L) - # Attention input reads differ between prefill and decode - # Prefill: read Q, K, V activations (all in activation_byte_size) - if ctx.prefill_num_tokens > 0: - read_bytes["attn_input"] = ( - (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv) - * d - * self.activation_byte_size - * L - ) + attn_input = 0.0 - # Decode: read Q activations + read K, V from cache (in cache_byte_size) - if ctx.decode_num_tokens > 0: - read_bytes["attn_input"] = read_bytes.get("attn_input", 0) + ( - ctx.decode_num_tokens * q * d * self.activation_byte_size * L - + 2 * ctx.decode_context_len * kv * d * self.cache_byte_size * L - ) + # Full attention layers + if L_full > 0: + if ctx.prefill_num_tokens > 0: + attn_input += ( + (ctx.prefill_num_tokens * q + 2 * ctx.prefill_context_len * kv) + * d + * self.activation_byte_size + * L_full + ) + if ctx.decode_num_tokens > 0: + attn_input += ( + ctx.decode_num_tokens * q * d * self.activation_byte_size * L_full + + 2 + * ctx.decode_context_len + * kv + * d + * self.cache_byte_size + * L_full + ) - read_bytes["out_input"] = T * q * d * self.activation_byte_size * L + # Sliding window attention layers + if L_swa > 0: + sw = self.sliding_window + prefill_sw_len = ctx.prefill_context_len_for_window(sw) + decode_sw_len = ctx.decode_context_len_for_window(sw) + + if ctx.prefill_num_tokens > 0: + attn_input += ( + (ctx.prefill_num_tokens * q + 2 * prefill_sw_len * kv) + * d + * self.activation_byte_size + * L_swa + ) + if ctx.decode_num_tokens > 0: + attn_input += ( + ctx.decode_num_tokens * q * d * self.activation_byte_size * L_swa + + 2 * decode_sw_len * kv * d * self.cache_byte_size * L_swa + ) + + if ctx.prefill_num_tokens > 0 or ctx.decode_num_tokens > 0: + read_bytes["attn_input"] = int(attn_input) + + read_bytes["out_input"] = int(T * q * d * self.activation_byte_size * L) read_bytes["out_weight"] = int(q * d * D * self.weight_byte_size * L) return read_bytes @@ -515,9 +715,13 @@ class AttentionMetrics(ComponentMetrics): def get_write_bytes_breakdown( self, ctx: ExecutionContext, per_gpu: bool = True ) -> dict[str, int]: - """Calculate write memory traffic for attention layers.""" - L, D, q, kv, d = ( - self.num_hidden_layers, + """Calculate write memory traffic for attention layers. + + Writes do not depend on the attention window, so only the total layer + count is needed; it comes from ``_layer_counts`` so all three breakdowns + bill the same number of layers. + """ + D, q, kv, d = ( self.hidden_size, self.num_attention_heads, self.num_key_value_heads, @@ -525,16 +729,17 @@ class AttentionMetrics(ComponentMetrics): ) T = ctx.total_num_tokens() + L, _, _ = self._layer_counts(per_gpu) + if per_gpu: - L //= self.pp_size # tensor parallel along heads q = max(1, q // self.tp_size) kv = max(1, kv // self.tp_size) return { - "qkv_output": T * (q + 2 * kv) * d * self.activation_byte_size * L, - "kv_cache": 2 * T * kv * d * self.cache_byte_size * L, - "out_output": T * D * self.activation_byte_size * L, + "qkv_output": int(T * (q + 2 * kv) * d * self.activation_byte_size * L), + "kv_cache": int(2 * T * kv * d * self.cache_byte_size * L), + "out_output": int(T * D * self.activation_byte_size * L), } @@ -549,6 +754,7 @@ class MLADetectionParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Raise InvalidComponent if the model does not use MLA attention.""" if not vllm_config.model_config.is_deepseek_mla: raise InvalidComponent("Model does not use MLA attention") return args @@ -562,6 +768,7 @@ class MLAConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Parse MLA-specific compression dimensions and KV-cache dtype byte size.""" model_config = vllm_config.model_config cfg = model_config.hf_text_config @@ -616,10 +823,16 @@ class MLAAttentionMetrics(ComponentMetrics): @classmethod def component_type(cls) -> str: + """Return the component registry key for MLA attention.""" return "mla_attn" @classmethod def get_parser(cls) -> ParserChain: + """Return the parser chain for MLAAttentionMetrics. + + Parsers run in order: MLA detection guard, base model config, + MLA-specific config, quantization weight-size override. + """ return ParserChain( MLADetectionParser(), BaseConfigParser(), @@ -923,6 +1136,7 @@ class FfnQuantizationConfigParser(Parser): """ def parse(self, args: ParsedArgs, vllm_config: VllmConfig) -> ParsedArgs: + """Override weight_byte_size based on the active FFN quantization method.""" cfg = vllm_config.quant_config if cfg is None: @@ -980,10 +1194,16 @@ class FfnMetrics(ComponentMetrics): @classmethod def component_type(cls) -> str: + """Return the component registry key for FFN layers.""" return "ffn" @classmethod def get_parser(cls) -> ParserChain: + """Return the parser chain for FfnMetrics. + + Parsers run in order: base model config, FFN parallelism, base FFN + config, MoE interleave step, MoE layer frequency, quantization override. + """ return ParserChain( BaseConfigParser(), FfnParallelParser(), @@ -1207,10 +1427,12 @@ class UnembedMetrics(ComponentMetrics): @classmethod def component_type(cls) -> str: + """Return the component registry key for the unembedding (LM head) layer.""" return "unembed" @classmethod def get_parser(cls) -> ParserChain: + """Return the parser chain for UnembedMetrics (base model config only).""" return ParserChain( BaseConfigParser(), ) From 03dc26e639e1927cda51d7b62639a11a354e1c0d Mon Sep 17 00:00:00 2001 From: Samuel Nordmann Date: Mon, 14 Sep 2026 19:26:15 +0200 Subject: [PATCH 039/139] [Perf][Kernel][Quantization] Fuse ReLU2 with static FP8 activation quantization (#53793) Signed-off-by: mgoin Signed-off-by: ElizaWszola Signed-off-by: snordmann Co-authored-by: mgoin Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: ElizaWszola Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../test_silu_mul_quant_manual_fusion.py | 11 +- .../fusion/test_quant_activation_contract.py | 5 +- tests/kernels/test_fused_quant_activation.py | 30 ++- tests/kernels/test_relu2_fp8_quant.py | 229 ++++++++++++++++++ tests/lora/test_layers.py | 10 + .../test_nemotron_h_quantization.py | 27 +++ vllm/lora/layers/base_linear.py | 4 + .../layers/fusion/fused_act_quant.py | 66 ++++- .../layers/fusion/quant_activation.py | 22 +- .../layers/fusion/relu2_fp8_quant.py | 76 ++++++ vllm/model_executor/layers/linear.py | 2 + vllm/model_executor/models/nemotron_h.py | 9 +- 12 files changed, 459 insertions(+), 32 deletions(-) create mode 100644 tests/kernels/test_relu2_fp8_quant.py create mode 100644 vllm/model_executor/layers/fusion/relu2_fp8_quant.py diff --git a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py index 5ee749659b..ee94960844 100644 --- a/tests/compile/passes/test_silu_mul_quant_manual_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_manual_fusion.py @@ -34,6 +34,7 @@ from vllm.model_executor.layers.fusion.fused_act_quant import ( from vllm.model_executor.layers.fusion.quant_activation import ( QuantizedActivation, expose_input_quant_key, + get_input_quant_key, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic128Sym, @@ -45,7 +46,7 @@ from vllm.platforms import current_platform # Mock linear layer for testing fusion paths that don't have real kernel support class MockLinearForFusion(torch.nn.Module): - """Mock linear layer that exposes input_quant_key for fusion testing.""" + """Mock linear layer that exposes an input key for fusion testing.""" def __init__( self, @@ -55,7 +56,7 @@ class MockLinearForFusion(torch.nn.Module): input_global_scale_inv=None, ): super().__init__() - self.input_quant_key = quant_key + self._input_quant_key = quant_key if input_scale is not None: self.input_scale = input_scale if input_global_scale is not None: @@ -118,7 +119,7 @@ def test_manual_fusion_fp8_static_with_linear( # Enable fusion expose_input_quant_key(fp8_linear, fp8_linear.kernel) - if not hasattr(fp8_linear, "input_quant_key"): + if get_input_quant_key(fp8_linear) is None: pytest.skip( f"Kernel {force_kernel.__name__} doesn't support input_quant_key" ) @@ -315,7 +316,7 @@ def test_manual_fusion_nvfp4_dynamic( envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm" ) def test_manual_fusion_fallback_no_key(dtype: torch.dtype): - """Test that maybe_fused_act_quant falls back when no input_quant_key.""" + """Test that maybe_fused_act_quant falls back without an input key.""" torch.set_default_device("cuda") torch.set_default_dtype(dtype) @@ -327,7 +328,7 @@ def test_manual_fusion_fallback_no_key(dtype: torch.dtype): with set_current_vllm_config(config): silu_and_mul = SiluAndMul() - # Linear without input_quant_key attribute + # Linear without an input quantization key mock_linear = torch.nn.Linear(128, 128) result = maybe_fused_act_quant(silu_and_mul, x, mock_linear) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index f698a272f5..ab12393692 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -42,6 +42,7 @@ from vllm.model_executor.layers.fusion.quant_activation import ( QuantizedActivation, as_quantized_activation, expose_input_quant_key, + get_input_quant_key, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, @@ -118,13 +119,13 @@ def test_bridge_marks_supporting_and_skips_others(): supported = _probe(FlashInferCutlassNvFp4LinearKernel) layer = torch.nn.Module() expose_input_quant_key(layer, supported) - assert layer.input_quant_key == kNvfp4Dynamic + assert get_input_quant_key(layer) == kNvfp4Dynamic unsupported = _probe(FlashInferTrtllmNvFp4LinearKernel) assert unsupported.input_quant_key() is None layer = torch.nn.Module() expose_input_quant_key(layer, unsupported) - assert not hasattr(layer, "input_quant_key") + assert get_input_quant_key(layer) is None def test_as_quantized_activation_validates_key(): diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py index 267484e433..321cb6978e 100644 --- a/tests/kernels/test_fused_quant_activation.py +++ b/tests/kernels/test_fused_quant_activation.py @@ -91,7 +91,7 @@ class MockLinearFp8Static(torch.nn.Module): def __init__(self, input_scale: torch.Tensor): super().__init__() - self.input_quant_key = kFp8StaticTensorSym + self._input_quant_key = kFp8StaticTensorSym self.input_scale = input_scale @@ -100,15 +100,21 @@ class MockLinearFp8Dynamic128(torch.nn.Module): def __init__(self): super().__init__() - self.input_quant_key = kFp8Dynamic128Sym + self._input_quant_key = kFp8Dynamic128Sym class MockLinearNoQuant(torch.nn.Module): - """Mock linear layer with no input_quant_key (no fusion).""" + """Mock linear layer with no input quantization key.""" pass +class MockLinearRequiresUnquantized(MockLinearFp8Static): + """Mock consumer that needs the original activation for another branch.""" + + requires_unquantized_input = True + + @pytest.mark.parametrize("num_tokens", [1, 16, 128]) @pytest.mark.parametrize("hidden_size", [128, 512, 1024]) @pytest.mark.parametrize("dtype", DTYPES) @@ -199,7 +205,7 @@ def test_maybe_fused_act_quant_fallback( hidden_size: int, dtype: torch.dtype, ) -> None: - """Test maybe_fused_act_quant falls back when no input_quant_key.""" + """Test maybe_fused_act_quant falls back without an input quantization key.""" device = "cuda:0" torch.set_default_device(device) @@ -216,3 +222,19 @@ def test_maybe_fused_act_quant_fallback( ref_out = act_fn(x) torch.testing.assert_close(result, ref_out) + + +@torch.inference_mode() +def test_maybe_fused_act_quant_preserves_required_unquantized_input( + default_vllm_config, +) -> None: + device = "cuda:0" + act_fn = SiluAndMul() + scale = torch.tensor([0.5], device=device, dtype=torch.float32) + linear = MockLinearRequiresUnquantized(scale) + x = torch.randn(17, 256, dtype=torch.bfloat16, device=device) + + result = maybe_fused_act_quant(act_fn, x, linear) + + assert isinstance(result, torch.Tensor) + torch.testing.assert_close(result, act_fn(x)) diff --git a/tests/kernels/test_relu2_fp8_quant.py b/tests/kernels/test_relu2_fp8_quant.py new file mode 100644 index 0000000000..cb605914c7 --- /dev/null +++ b/tests/kernels/test_relu2_fp8_quant.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config +from vllm.config.vllm import OptimizationLevel +from vllm.model_executor.layers.activation import ReLUSquaredActivation +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + kFp8StaticTensorSym, +) +from vllm.platforms import current_platform + +requires_sm90 = pytest.mark.skipif( + not (current_platform.is_cuda() and current_platform.has_device_capability(90)), + reason="This test requires SM90 or newer", +) + + +class _StaticFp8Linear(torch.nn.Module): + def __init__(self, input_scale: torch.Tensor) -> None: + super().__init__() + self._input_quant_key = kFp8StaticTensorSym + self.input_scale = input_scale + + +def _assert_fp8_bitwise_equal(actual: torch.Tensor, expected: torch.Tensor) -> None: + assert actual.dtype == expected.dtype == current_platform.fp8_dtype() + torch.testing.assert_close( + actual.view(torch.uint8), expected.view(torch.uint8), rtol=0, atol=0 + ) + + +def _assert_quantized_activation( + result: QuantizedActivation, x: torch.Tensor, scale: torch.Tensor +) -> None: + assert result.scale is scale + assert result.orig_dtype == x.dtype + assert result.orig_shape == x.shape + assert result.quant_key == kFp8StaticTensorSym + + +@pytest.fixture(scope="module") +def o2_relu2_fp8_ops(): + config = VllmConfig( + optimization_level=OptimizationLevel.O2, + compilation_config=CompilationConfig(custom_ops=["none"]), + ) + with set_current_vllm_config(config): + yield ( + ReLUSquaredActivation(), + QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR), + ) + + +@requires_sm90 +@pytest.mark.parametrize("shape", [(1, 1), (1, 16), (17, 5120), (128, 10240)]) +@torch.inference_mode() +def test_relu2_static_fp8_quant(o2_relu2_fp8_ops, shape) -> None: + torch.manual_seed(0) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.1495535671710968, device="cuda", dtype=torch.float32) + + relu2, quant_fp8 = o2_relu2_fp8_ops + expected, _ = quant_fp8(relu2(x), scale) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + _assert_fp8_bitwise_equal(result.data, expected) + + +@requires_sm90 +@pytest.mark.parametrize( + "enabled_ops", + [(), ("relu2",), ("quant_fp8",), ("relu2", "quant_fp8")], +) +@torch.inference_mode() +def test_relu2_static_fp8_quant_ignores_custom_op_dispatch(enabled_ops) -> None: + custom_ops = ["none", *(f"+{op}" for op in enabled_ops)] + config = VllmConfig( + optimization_level=OptimizationLevel.O2, + compilation_config=CompilationConfig(custom_ops=custom_ops), + ) + with set_current_vllm_config(config): + torch.manual_seed(0) + relu2 = ReLUSquaredActivation() + quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32) + + expected, _ = quant_fp8(relu2(x), scale) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + _assert_fp8_bitwise_equal(result.data, expected) + + +@requires_sm90 +@torch.inference_mode() +def test_relu2_static_fp8_quant_without_compilation() -> None: + config = VllmConfig(optimization_level=OptimizationLevel.O0) + with set_current_vllm_config(config): + torch.manual_seed(0) + relu2 = ReLUSquaredActivation() + quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32) + + expected, _ = quant_fp8(relu2(x), scale) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + _assert_fp8_bitwise_equal(result.data, expected) + + +@requires_sm90 +@pytest.mark.parametrize( + "scale_value", + [ + 0.00435965, + 0.007149832788854837, + 0.03125, + 0.0714285746216774, + 0.1495535671710968, + 0.1517857164144516, + 6.03571415, + ], +) +@torch.inference_mode() +def test_relu2_static_fp8_quant_all_bf16_values(o2_relu2_fp8_ops, scale_value) -> None: + bits = torch.arange(1 << 16, dtype=torch.int32).to(torch.uint16) + x = bits.view(torch.bfloat16).reshape(1, -1).to("cuda") + scale = torch.tensor(scale_value, device="cuda", dtype=torch.float32) + + relu2, quant_fp8 = o2_relu2_fp8_ops + expected, _ = quant_fp8(relu2(x), scale) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + _assert_fp8_bitwise_equal(result.data, expected) + + +@requires_sm90 +@torch.inference_mode() +def test_relu2_static_fp8_quant_empty(o2_relu2_fp8_ops) -> None: + x = torch.empty(0, device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32) + + relu2, _ = o2_relu2_fp8_ops + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, QuantizedActivation) + _assert_quantized_activation(result, x, scale) + assert result.data.shape == x.shape + assert result.data.dtype == current_platform.fp8_dtype() + + +@requires_sm90 +@pytest.mark.parametrize("unsupported", ["dtype", "layout", "scale"]) +@torch.inference_mode() +def test_relu2_static_fp8_quant_falls_back(o2_relu2_fp8_ops, unsupported: str) -> None: + x = torch.randn((17, 32), device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32) + if unsupported == "dtype": + x = x.to(torch.float16) + elif unsupported == "layout": + x = x.T + else: + scale = scale.to(torch.float16) + + relu2, _ = o2_relu2_fp8_ops + expected = relu2(x) + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale)) + + assert isinstance(result, torch.Tensor) + assert not isinstance(result, QuantizedActivation) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +@requires_sm90 +@torch.inference_mode() +def test_relu2_static_fp8_quant_torch_compile_dynamic( + o2_relu2_fp8_ops, +) -> None: + torch.manual_seed(1) + scale = torch.tensor(0.00435965, device="cuda", dtype=torch.float32) + relu2, quant_fp8 = o2_relu2_fp8_ops + linear = _StaticFp8Linear(scale) + + def fused(x: torch.Tensor) -> torch.Tensor: + result = maybe_fused_act_quant(relu2, x, linear) + assert isinstance(result, QuantizedActivation) + return result.data + + compile_count = 0 + + def counting_inductor_backend(gm, example_inputs): + nonlocal compile_count + compile_count += 1 + return torch._inductor.compile(gm, example_inputs) + + torch._dynamo.reset() + try: + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16) + torch._dynamo.mark_dynamic(x, 0) + compiled_fused = torch.compile( + fused, backend=counting_inductor_backend, fullgraph=True + ) + + for candidate in ( + x, + torch.randn((31, 5120), device="cuda", dtype=torch.bfloat16), + torch.randn((63, 5120), device="cuda", dtype=torch.bfloat16), + ): + expected, _ = quant_fp8(relu2(candidate), scale) + _assert_fp8_bitwise_equal(compiled_fused(candidate), expected) + + assert compile_count == 1 + finally: + torch._dynamo.reset() diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index bce91510d4..2dab86bc0d 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -31,6 +31,9 @@ from vllm.lora.layers import ( ) from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights from vllm.lora.punica_wrapper import get_punica_wrapper +from vllm.model_executor.layers.fusion.quant_activation import ( + get_input_quant_key, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, @@ -56,6 +59,13 @@ TOLERANCES = { torch.bfloat16: (3e-2, 2e-2), } + +def test_lora_linear_requires_unquantized_input() -> None: + layer = RowParallelLinearWithLoRA.__new__(RowParallelLinearWithLoRA) + layer._input_quant_key = object() + assert get_input_quant_key(layer) is None + + pytestmark = [ pytest.mark.skipif( not ( diff --git a/tests/model_executor/test_nemotron_h_quantization.py b/tests/model_executor/test_nemotron_h_quantization.py index 81a3fb5b13..e4e1c87db4 100644 --- a/tests/model_executor/test_nemotron_h_quantization.py +++ b/tests/model_executor/test_nemotron_h_quantization.py @@ -3,6 +3,8 @@ from unittest.mock import Mock, patch +import torch + def test_nemotron_h_lm_head_receives_quant_config(): from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM @@ -32,3 +34,28 @@ def test_nemotron_h_lm_head_receives_quant_config(): MockLMHead.assert_called_once() call_kwargs = MockLMHead.call_args.kwargs assert call_kwargs["quant_config"] is mock_quant_config + + +def test_relu2_fp8_fusion_uses_registry(): + from vllm.model_executor.models.nemotron_h import NemotronHMLP + + projected = torch.empty((1, 1), dtype=torch.bfloat16) + fused = Mock() + act_fn = Mock() + down_proj = Mock(side_effect=lambda x: (x, None)) + + mlp = NemotronHMLP.__new__(NemotronHMLP) + torch.nn.Module.__init__(mlp) + mlp.up_proj = Mock(return_value=(projected, None)) + mlp.down_proj = down_proj + mlp.act_fn = act_fn + + with patch( + "vllm.model_executor.models.nemotron_h.maybe_fused_act_quant", + return_value=fused, + ) as maybe_fused: + result = mlp(Mock()) + + maybe_fused.assert_called_once_with(act_fn, projected, down_proj) + act_fn.assert_not_called() + assert result is fused diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py index 5c8e829b29..df7cfe7122 100644 --- a/vllm/lora/layers/base_linear.py +++ b/vllm/lora/layers/base_linear.py @@ -68,6 +68,10 @@ if envs.VLLM_LORA_ENABLE_DUAL_STREAM: class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): + # The adapter branch consumes the original activation, independently of + # whether the wrapped base layer can consume a pre-quantized activation. + requires_unquantized_input = True + def __init__(self, base_layer: LinearBase): super().__init__() diff --git a/vllm/model_executor/layers/fusion/fused_act_quant.py b/vllm/model_executor/layers/fusion/fused_act_quant.py index 9fe253f2d9..e18a35a70d 100644 --- a/vllm/model_executor/layers/fusion/fused_act_quant.py +++ b/vllm/model_executor/layers/fusion/fused_act_quant.py @@ -5,21 +5,27 @@ Producer side of the QuantizedActivation contract for activation layers. Given an activation module and the downstream linear it feeds, fuse the activation with that linear's input quantization into a single kernel when the -linear advertises a consumable input_quant_key (see quant_activation.py). +linear advertises a consumable input quantization key (see quant_activation.py). Falls back to the plain activation when nothing matches, so a model forward can always call maybe_fused_act_quant unconditionally. This is the manual-fusion counterpart to ActivationQuantFusionPass: when fusion -fires here the silu_and_mul pattern is already consumed, so the compiler pass -finds nothing to rewrite and the two never double-fuse. +fires here, the activation and quantization are already consumed, so a compiler +pass cannot fuse the same boundary again. """ from collections.abc import Callable import torch -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.activation import ReLUSquaredActivation, SiluAndMul +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) +from vllm.model_executor.layers.fusion.relu2_fp8_quant import ( + relu_squared_static_fp8_quant, +) from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -139,12 +145,46 @@ def _silu_and_mul_nvfp4_dynamic( ) -# (activation module type, consumer input_quant_key) -> fused producer. -# Mirrors ActivationQuantFusionPass.FUSED_OPS; add a row to migrate a scheme. +# (activation module type, consumer input quantization key) -> fused producer. +# Add a row for each supported manual producer. _FUSED_ACT_QUANT: dict[tuple[type, QuantKey], Callable] = { (SiluAndMul, kFp8StaticTensorSym): _silu_and_mul_fp8_static, } + +def _relu_squared_static_fp8_quant_supported( + act_fn: torch.nn.Module, + x: torch.Tensor, + linear: LinearBase, +) -> bool: + """Return whether the ReLU2 static-FP8 producer can consume this input.""" + scale = getattr(linear, "input_scale", None) + return ( + isinstance(act_fn, ReLUSquaredActivation) + and x.is_cuda + and x.dtype == torch.bfloat16 + and x.is_contiguous() + and isinstance(scale, torch.Tensor) + and scale.dtype == torch.float32 + and scale.device == x.device + and scale.numel() == 1 + ) + + +# Optional per-entry predicates preserve the unconditional fallback contract +# for producers whose supported inputs are narrower than their registry key. +_FUSED_ACT_QUANT_SUPPORT: dict[ + tuple[type, QuantKey], Callable[[torch.nn.Module, torch.Tensor, LinearBase], bool] +] = {} + +# This path is validated on SM90 and newer CUDA devices. +if current_platform.is_cuda() and current_platform.has_device_capability(90): + _RELU2_STATIC_FP8_KEY = (ReLUSquaredActivation, kFp8StaticTensorSym) + _FUSED_ACT_QUANT[_RELU2_STATIC_FP8_KEY] = relu_squared_static_fp8_quant + _FUSED_ACT_QUANT_SUPPORT[_RELU2_STATIC_FP8_KEY] = ( + _relu_squared_static_fp8_quant_supported + ) + # Add CUDA-specific entries for dynamic block quantization if current_platform.is_cuda_alike(): _FUSED_ACT_QUANT[(SiluAndMul, kFp8Dynamic128Sym)] = _silu_and_mul_fp8_dynamic_128 @@ -161,12 +201,14 @@ def maybe_fused_act_quant( ) -> "torch.Tensor | QuantizedActivation": """Apply act_fn, fusing the downstream linear's input quant when possible. - Returns a QuantizedActivation when a fused kernel matches - (act_fn, linear.input_quant_key), else the plain activated tensor. + Returns a QuantizedActivation when a fused kernel matches the activation and + the consumer's effective input quantization key, else the plain activation. """ - key = getattr(linear, "input_quant_key", None) + key = get_input_quant_key(linear) if key is not None: - producer = _FUSED_ACT_QUANT.get((type(act_fn), key)) - if producer is not None: + registry_key = (type(act_fn), key) + producer = _FUSED_ACT_QUANT.get(registry_key) + support = _FUSED_ACT_QUANT_SUPPORT.get(registry_key) + if producer is not None and (support is None or support(act_fn, x, linear)): return producer(x, linear) return act_fn(x) diff --git a/vllm/model_executor/layers/fusion/quant_activation.py b/vllm/model_executor/layers/fusion/quant_activation.py index bd6cab668e..d69415448f 100644 --- a/vllm/model_executor/layers/fusion/quant_activation.py +++ b/vllm/model_executor/layers/fusion/quant_activation.py @@ -5,7 +5,9 @@ A QuantizedActivation is a pre-quantized activation produced by a fused kernel and consumed directly by a linear layer, letting the layer skip its own input quantization. A linear advertises the key its kernel can consume via expose_input_quant_key; the kernel validates and reads the activation via -as_quantized_activation. +as_quantized_activation. Producers query the effective capability through +get_input_quant_key, which hides the key when another consumer branch needs the +original activation. """ from dataclasses import dataclass, replace @@ -46,12 +48,11 @@ class QuantizedActivation: def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None: - """Advertise the kernel's pre-quantized input key on the layer, if any. + """Store the kernel's pre-quantized input key on the layer, if any. - This is the bridge from a kernel's input_quant_key() to the - layer.input_quant_key attribute that fusion call sites read. The attribute - is left unset when the kernel quantizes its own input, so non-supporting - backends never receive a QuantizedActivation. + This is the bridge from a kernel's input_quant_key() to the layer capability + that fusion call sites read through get_input_quant_key. The raw key is left + unset when the kernel quantizes its own input. TODO(mgoin): Producers also need the consumer's quantization scales (e.g. static input scale, global scale). Expose those here as well so producers @@ -59,7 +60,14 @@ def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None: """ key = kernel.input_quant_key() if key is not None: - layer.input_quant_key = key + layer._input_quant_key = key + + +def get_input_quant_key(layer: torch.nn.Module) -> QuantKey | None: + """Return the pre-quantized input key the complete layer can consume.""" + if getattr(layer, "requires_unquantized_input", False): + return None + return getattr(layer, "_input_quant_key", None) def as_quantized_activation( diff --git a/vllm/model_executor/layers/fusion/relu2_fp8_quant.py b/vllm/model_executor/layers/fusion/relu2_fp8_quant.py new file mode 100644 index 0000000000..4dc0fed4cc --- /dev/null +++ b/vllm/model_executor/layers/fusion/relu2_fp8_quant.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, + kFp8StaticTensorSym, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +_FP8_MIN, _FP8_MAX = get_fp8_min_max() + + +@triton.jit +def _relu_squared_static_fp8_quant_kernel( + x_ptr, + scale_ptr, + output_ptr, + num_elements, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + input_is_nan = x != x + relu = tl.maximum(x, 0.0) + # Preserve the standalone ReLU2 kernel's BF16 output boundary. + activated = (relu * relu).to(tl.bfloat16).to(tl.float32) + # Match O2's Inductor-compiled QuantFP8 reciprocal-then-multiply. + quantized = activated * (1.0 / tl.load(scale_ptr)) + quantized = tl.maximum(tl.minimum(quantized, FP8_MAX), FP8_MIN) + # Match the O2 native chain by preserving NaN through the FP8 conversion. + quantized = tl.where(input_is_nan, x, quantized) + tl.store(output_ptr + offsets, quantized, mask=mask) + + +def relu_squared_static_fp8_quant( + x: torch.Tensor, linear: LinearBase +) -> QuantizedActivation: + """BF16 ReLU2 followed by static per-tensor FP8 quantization.""" + assert x.dtype == torch.bfloat16 + assert x.is_contiguous() + scale = linear.input_scale + assert scale.dtype == torch.float32 + assert x.device == scale.device + assert scale.numel() == 1 + + output = torch.empty_like(x, dtype=current_platform.fp8_dtype()) + if x.numel() != 0: + block_size = min(triton.next_power_of_2(x.shape[-1]), 2048) + num_warps = min(max(block_size // 256, 1), 4) + grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),) + _relu_squared_static_fp8_quant_kernel[grid]( + x, + scale, + output, + x.numel(), + FP8_MIN=_FP8_MIN, + FP8_MAX=_FP8_MAX, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + + return QuantizedActivation( + data=output, + scale=scale, + orig_dtype=x.dtype, + orig_shape=x.shape, + quant_key=kFp8StaticTensorSym, + ) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 212dd95375..95dc852fb3 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -252,6 +252,8 @@ class LinearBase(PluggableLayer): disable_tp: If true, tensor parallelism will be disabled for this layer. """ + requires_unquantized_input = False + def __init__( self, input_size: int, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index e65e51f3bb..c1bc4c6823 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -25,7 +25,11 @@ import torch from torch import nn from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, ModelConfig, VllmConfig +from vllm.config import ( + CacheConfig, + ModelConfig, + VllmConfig, +) from vllm.config.parallel import ParallelConfig from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size from vllm.distributed.communication_op import tensor_model_parallel_all_gather @@ -38,6 +42,7 @@ from vllm.model_executor.layers.fused_moe import ( activation_without_mul, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -118,7 +123,7 @@ class NemotronHMLP(nn.Module): def forward(self, x: torch.Tensor): x, _ = self.up_proj(x) - x = self.act_fn(x) + x = maybe_fused_act_quant(self.act_fn, x, self.down_proj) x, _ = self.down_proj(x) return x From e85c8826ce2a810367e8a70eedb987e31b140800 Mon Sep 17 00:00:00 2001 From: "Zhiyao(Ian) Yu" <105995086+Ianniu123@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:00:20 -0400 Subject: [PATCH 040/139] [Bugfix] Make KV cache and MFU log lines backend-neutral (#55650) Signed-off-by: Ianniu <985683179yzy@gmail.com> Co-authored-by: Claude Opus 5 (1M context) --- vllm/v1/core/kv_cache_utils.py | 4 +++- vllm/v1/metrics/loggers.py | 8 ++++++-- vllm/v1/metrics/perf.py | 2 +- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index b47df61add..28478b69f5 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2391,9 +2391,11 @@ def update_kv_cache_capacity( vllm_config.cache_config.kv_cache_size_tokens = num_tokens vllm_config.cache_config.kv_cache_max_concurrency = max_concurrency max_model_len = vllm_config.model_config.max_model_len + device_type = vllm_config.device_config.device_type logger.info_once( - "GPU KV cache size: %s tokens, " + "%s KV cache size: %s tokens, " "Maximum concurrency for %s tokens per request: %.2fx", + "GPU" if device_type == "cuda" else device_type.upper(), f"{num_tokens:,}", f"{max_model_len:,}", max_concurrency, diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 9e3f9c63a6..3a8adb6e32 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -101,6 +101,8 @@ class LoggingStatLogger(StatLoggerBase): def __init__(self, vllm_config: VllmConfig, engine_index: int = 0): self.engine_index = engine_index self.vllm_config = vllm_config + device_type = vllm_config.device_config.device_type + self.kv_cache_device = "GPU" if device_type == "cuda" else device_type.upper() self._reset(time.monotonic()) self.last_scheduler_stats = SchedulerStats() @@ -184,7 +186,7 @@ class LoggingStatLogger(StatLoggerBase): "%sIteration(%d): %d context requests, %d context tokens, " "%d generation requests, %d generation tokens, " "iteration elapsed time: %.2f ms%s, " - "GPU KV cache usage: %.1f%%%s", + "%s KV cache usage: %.1f%%%s", self._log_prefix_for_engine(engine_idx), details.iteration_index, details.num_ctx_requests, @@ -193,6 +195,7 @@ class LoggingStatLogger(StatLoggerBase): details.num_generation_tokens, details.elapsed_ms, " (dummy)" if details.is_dummy else "", + self.kv_cache_device, scheduler_stats.kv_cache_usage * 100, encoder_msg, ) @@ -287,12 +290,13 @@ class LoggingStatLogger(StatLoggerBase): log_parts.extend( [ - "GPU KV cache usage: %.1f%%", + "%s KV cache usage: %.1f%%", "Prefix cache hit rate: %.1f%%", ] ) log_args.extend( [ + self.kv_cache_device, self.last_scheduler_stats.kv_cache_usage * 100, self.prefix_caching_metrics.hit_rate * 100, ] diff --git a/vllm/v1/metrics/perf.py b/vllm/v1/metrics/perf.py index ab90ffcbe1..ebb676e66d 100644 --- a/vllm/v1/metrics/perf.py +++ b/vllm/v1/metrics/perf.py @@ -1749,7 +1749,7 @@ class PerfMetricsLogging: ) log_fn( - "%sMFU: %.1f TF/s/GPU %.1f GB/s/GPU", + "%sMFU: %.1f TF/s/device %.1f GB/s/device", log_prefix, avg_tflops_per_gpu, avg_gbps_per_gpu, From 56ca9904ce76975fe3aad28d1af7af961a07f054 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 11:37:08 -0700 Subject: [PATCH 041/139] [CI] Initialize ubatch runner in cudagraph unit test (#56808) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/v1/cudagraph/test_breakable_cudagraph.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index 0e10e4a51f..e792d68d89 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -43,6 +43,7 @@ def test_piecewise_capture_builds_fresh_metadata_for_both_passes(): manager._capture_descs = {CUDAGraphMode.PIECEWISE: [desc]} manager._graphs_captured = False manager.use_breakable_cg = True + manager.ubatch_runner = None create_calls = [] forward_calls = [] From 5372e72a9884c141874465f0eb8054a4a81da6bb Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 14 Sep 2026 11:57:37 -0700 Subject: [PATCH 042/139] [Perf][DSv4.1] Fold the mHC post block into the delayed pre projection (#56633) Signed-off-by: Yongye Zhu --- tests/distributed/test_engram_dp_shard.py | 7 +- tests/kernels/test_mhc_jit_warmup.py | 49 ++- tests/kernels/test_mhc_kernels.py | 248 ++++++++++++++- vllm/model_executor/kernels/mhc/tilelang.py | 285 +++++++++++++++++- .../kernels/mhc/tilelang_kernels.py | 140 ++++++++- vllm/model_executor/kernels/mhc/warmup.py | 22 +- vllm/models/deepseek_v41/nvidia/dspark.py | 2 +- vllm/models/deepseek_v41/nvidia/model.py | 178 +++++++---- 8 files changed, 838 insertions(+), 93 deletions(-) diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 35ffc659d4..c2fa0f857f 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -491,7 +491,9 @@ def _check_dummy_hash_model_forward( state.forward = lambda *args: ids.unsqueeze(1) class Decoder(SimpleNamespace): - def __call__(self, hidden, positions, input_ids, *args): + # Engram hashes and the mask stay the last positional arguments; + # **kwargs absorbs the decoder's keyword-only flags. + def __call__(self, hidden, positions, input_ids, *args, **kwargs): hashes, keep = args[-2:] if dp_shared_memory and dp_rank == 1: assert hashes is None and keep is None @@ -504,7 +506,8 @@ def _check_dummy_hash_model_forward( if dp_rank == 1: assert torch.all(hashes == engram_ops.DEAD_ID) output = engram.embed(hashes[:, 0]) - return output, None, None, None, None + # Trailing None is previous_aux; this stub captures no aux states. + return output, None, None, None, None, None model = SimpleNamespace( use_mega_moe=False, diff --git a/tests/kernels/test_mhc_jit_warmup.py b/tests/kernels/test_mhc_jit_warmup.py index 28e746c092..512b3e6ba5 100644 --- a/tests/kernels/test_mhc_jit_warmup.py +++ b/tests/kernels/test_mhc_jit_warmup.py @@ -17,6 +17,8 @@ from vllm.model_executor.kernels.mhc.tilelang_kernels import ( MhcFusedTileLangKernel, MhcPostTileLangKernel, MhcPreBigFuseTileLangKernel, + mhc_fused_post_pre_split_config, + require_fused_post_pre_config, ) from vllm.model_executor.warmup import jit_warmup_tilelang_helper @@ -113,17 +115,32 @@ def test_mhc_pre_big_fuse_dispatch_matches_legacy_runtime_config( @pytest.mark.parametrize( - ("num_tokens", "hidden_size", "expected_n_splits", "expected_tile_n"), - [(4, 4096, 8, 2), (4, 8192, 4, 2), (8, 4096, 4, 3)], + ("num_tokens", "hidden_size", "expected_tile_n", "expected_n_splits"), + [ + (1, 4096, 2, 8), + (4, 8192, 2, 8), + (8, 5120, 2, 8), + (16, 5120, 6, 8), + (32, 7168, 6, 8), + ], ) -def test_mhc_fused_dispatch_matches_legacy_runtime_config( +def test_mhc_fused_dispatch_matches_the_launch_config( num_tokens: int, hidden_size: int, - expected_n_splits: int, expected_tile_n: int, + expected_n_splits: int, ) -> None: - kernel = MhcFusedTileLangKernel() + """The compile key must be the config the launch will actually use. + Both read mhc_fused_post_pre_split_config, so this pins the tuned bands + and guards against the two drifting apart again. + """ + kernel = MhcFusedTileLangKernel() + tile_n, n_splits, n_thr = require_fused_post_pre_config( + num_tokens, hidden_size, hc_mult=4 + ) + + assert (tile_n, n_splits) == (expected_tile_n, expected_n_splits) assert kernel.dispatch( num_tokens=num_tokens, hidden_size=hidden_size, @@ -131,11 +148,20 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config( ) == kernel.CompileKey( hidden_size=hidden_size, hc_mult=4, - n_splits=expected_n_splits, - tile_n=expected_tile_n, + n_splits=n_splits, + tile_n=tile_n, + n_thr=n_thr, ) +def test_mhc_fused_declines_shapes_it_cannot_tile() -> None: + """Above the token cutoff, and for a hidden size the block cannot split.""" + assert mhc_fused_post_pre_split_config(33, 5120, 4) is None + assert mhc_fused_post_pre_split_config(1, 5137, 4) is None + with pytest.raises(ValueError, match="does not cover num_tokens"): + require_fused_post_pre_config(33, 5120, 4) + + @pytest.mark.parametrize( ("kernel", "compile_key"), [ @@ -235,8 +261,9 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config( MhcFusedTileLangKernel.CompileKey( hidden_size=4096, hc_mult=4, - n_splits=4, - tile_n=3, + n_splits=8, + tile_n=6, + n_thr=128, ), ), ( @@ -246,6 +273,7 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config( hc_mult=4, n_splits=8, tile_n=2, + n_thr=128, ), ), ( @@ -253,8 +281,9 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config( MhcFusedTileLangKernel.CompileKey( hidden_size=8192, hc_mult=4, - n_splits=4, + n_splits=8, tile_n=2, + n_thr=128, ), ), ( diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 46f311dfe6..5049815380 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -15,8 +15,12 @@ from vllm.model_executor.kernels.mhc.tilelang import ( ) from vllm.model_executor.kernels.mhc.tilelang_kernels import ( _HC_PRENORM_GEMM_TILELANG_KERNEL, + mhc_fused_post_pre_split_config, +) +from vllm.model_executor.kernels.mhc.torch import ( + mhc_post_torch, + mhc_pre_delayed_torch, ) -from vllm.model_executor.kernels.mhc.torch import mhc_pre_delayed_torch from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton from vllm.model_executor.layers.mhc import ( HAS_AITER_MHC, @@ -115,9 +119,9 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): carried_mixes = [] def make_layer(mix): - def forward(hidden, positions, ids, pre, post, res, residual): + def forward(hidden, positions, ids, pre, post, res, residual, **kwargs): carried_mixes.append(pre) - return hidden, None, None, None, mix + return hidden, None, None, None, mix, None return forward @@ -235,6 +239,171 @@ def test_deepseek_v41_mhc_pre_delayed( torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2) +def mhc_fused_post_pre_delayed_ref( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + *mix_args, + pre_mix: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Project the post result before it is rounded to BF16. + + The fused kernel keeps the updated residual streams in registers, so its + projection sees the FP32 post result while the streams it stores are the + BF16 rounding of it. + """ + post = torch.einsum("tij,tih->tjh", comb_res_mix, residual.float()) + post = post + post_layer_mix * x.float().unsqueeze(-2) + return mhc_pre_delayed_torch( + post.bfloat16(), *mix_args, pre_mix=pre_mix, x=post.flatten(1) + ) + + +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("num_tokens", [0, 1, 8, 128]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) +@pytest.mark.parametrize("carried", [False, True]) +def test_deepseek_v41_mhc_fused_post_pre_delayed(num_tokens, hidden_size, carried): + """Fold the post block into the delayed pre without changing its outputs.""" + set_random_seed(0) + hc_mult = 4 + mix_size = hc_mult * (hc_mult + 2) + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) + residual = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + post_layer_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE) + comb_res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE) + fn = torch.randn(mix_size, hc_mult * hidden_size, device=DEVICE) * 0.02 + scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE) + base = torch.randn(mix_size, device=DEVICE) + pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) if carried else None + weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE) + weight.uniform_(0.5, 1.5) + mix_args = (fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20) + + residual_cur, post, comb, layer_input, next_pre, aux = ( + torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6 + ) + ) + assert aux.shape == (0, hidden_size) + assert [ + residual_cur.shape, + post.shape, + comb.shape, + layer_input.shape, + next_pre.shape, + ] == [ + (num_tokens, hc_mult, hidden_size), + (num_tokens, hc_mult, 1), + (num_tokens, hc_mult, hc_mult), + (num_tokens, hidden_size), + (num_tokens, hc_mult), + ] + if num_tokens == 0: + return + + # The post mapping and the collapse are what the unfused kernels produce. + residual_ref = torch.ops.vllm.mhc_post_tilelang( + x, residual, post_layer_mix, comb_res_mix + ) + layer_input_ref = mhc_pre_delayed_tilelang( + residual_ref, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6 + )[2] + torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=0) + torch.testing.assert_close(layer_input, layer_input_ref, atol=0, rtol=0) + + post_ref, comb_ref, _, next_pre_ref = mhc_fused_post_pre_delayed_ref( + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix=pre_mix + ) + # Above the fused kernel's token count the projection falls back to a + # split-k TF32 GEMM over the stored BF16 streams. + tol = ( + 1e-5 + if mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) is not None + else 5e-3 + ) + for actual, expected in ( + (post, post_ref), + (comb, comb_ref), + (next_pre, next_pre_ref), + ): + torch.testing.assert_close(actual, expected, atol=tol, rtol=tol) + + # Folding the stream mean into the collapse must not change it, or anything + # else the op returns. + captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + *mix_args, + pre_mix, + weight, + 1e-6, + True, + ) + torch.testing.assert_close(captured[5], residual_cur.mean(dim=1), atol=0, rtol=0) + unchanged = (residual_cur, post, comb, layer_input, next_pre) + for with_aux, without in zip(captured[:5], unchanged, strict=True): + torch.testing.assert_close(with_aux, without, atol=0, rtol=0) + + # The unnormalized epilogue has its own fused aux path. + bare = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, None, 1e-6, True + ) + torch.testing.assert_close(bare[5], bare[0].mean(dim=1), atol=0, rtol=0) + + # Decode replays this op from a captured graph. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang( + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6 + ) + graph.replay() + eager = (residual_cur, post, comb, layer_input, next_pre, aux) + for replayed, expected in zip(captured, eager, strict=True): + torch.testing.assert_close(replayed, expected, atol=0, rtol=0) + + +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("carried", [False, True]) +@pytest.mark.parametrize("capture_aux", [False, True]) +def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_aux): + set_random_seed(0) + x = torch.randn(2, 5120, dtype=torch.bfloat16, device=DEVICE) + residual = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE) + post_layer_mix = torch.rand(2, 4, 1, device=DEVICE) + comb_res_mix = torch.rand(2, 4, 4, device=DEVICE) + fn = torch.randn(24, 20480, device=DEVICE) * 0.02 + pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None + scale = torch.ones(3, device=DEVICE) + base = torch.zeros(24, device=DEVICE) + torch.library.opcheck( + torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang.default, + ( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + scale, + base, + 1e-20, + 1e-6, + 1e-6, + 2.0, + 20, + pre_mix, + None, + 1e-6, + capture_aux, + ), + ) + + @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"]) def test_deepseek_v41_decoder_mixes_match_torch( @@ -295,15 +464,86 @@ def test_deepseek_v41_decoder_mixes_match_torch( post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs) return post, res, decoder.attn_norm(collapsed), pre + def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, **kw): + residual = mhc_post_torch(x, residual, post_mix, res_mix) + aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0) + return residual, *reference(residual, *args, **kw), aux + monkeypatch.setattr( "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang", reference, ) + monkeypatch.setattr( + "vllm.models.deepseek_v41.nvidia.model.mhc_fused_post_pre_delayed_tilelang", + fused_reference, + ) expected = decoder(x, positions, None, **kwargs) - for result, ref in zip(actual, expected, strict=True): + assert actual[-1] is None and expected[-1] is None + for result, ref in zip(actual[:-1], expected[:-1], strict=True): torch.testing.assert_close(result, ref, atol=2e-2, rtol=1e-2) +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("entry", ["fused", "engram"]) +def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_config): + """Read the aux hidden state back out of the seam that already computed it. + + The draft model's input is the mean over hc streams of the previous + layer's post. That post now runs inside this layer's fused pre, so the + captured value must equal a standalone ``mhc_post_tilelang`` on the same + inputs -- and Engram must not be folded into it, since the injection + happens after the aux consumers read the stream. + """ + set_random_seed(0) + decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer) + nn.Module.__init__(decoder) + decoder.hc_mult = 4 + decoder.hc_sinkhorn_iters = 20 + decoder.hc_eps = 1e-6 + decoder.rms_norm_eps = 1e-6 + decoder.hc_post_alpha = 2.0 + decoder.use_sequence_parallel = False + decoder.engram = None + from vllm.model_executor.layers.layernorm import RMSNorm + + decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to( + device=DEVICE, dtype=torch.bfloat16 + ) + decoder.attn = lambda positions, x, _: x * 0.5 + decoder.ffn = lambda x, input_ids: x * 0.25 + with torch.device(DEVICE): + decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02 + decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02 + decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3) + decoder.hc_attn_base = torch.randn(24) + decoder.hc_ffn_base = torch.randn(24) + x = torch.randn(3, 5120, dtype=torch.bfloat16) + kwargs = dict( + pre_mix=torch.rand(3, 4), + residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16), + post_mix=torch.rand(3, 4, 1), + res_mix=torch.rand(3, 4, 4), + ) + if entry == "engram": + + class FakeEngram(nn.Module): + layer_hash_index = 0 + + def forward(self, residual, hashes, mask): + return residual + 0.125 + + decoder.engram = FakeEngram() + kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32) + positions = torch.arange(3) + + expected = torch.ops.vllm.mhc_post_tilelang( + x, kwargs["residual"], kwargs["post_mix"], kwargs["res_mix"] + ).mean(dim=1) + previous_aux = decoder(x, positions, None, **kwargs, capture_previous_aux=True)[-1] + torch.testing.assert_close(previous_aux, expected, atol=0, rtol=0) + assert decoder(x, positions, None, **kwargs)[-1] is None + + @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize("carried", [False, True]) def test_mhc_pre_delayed_custom_op_supports_compile(carried): diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 97ca23b541..cea186e49e 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -194,6 +194,7 @@ def mhc_pre_delayed_tilelang( norm_weight, pre_mix if pre_mix is not None else post, next_pre_mix, + layer_input, hidden_size=hidden_size, rms_eps=rms_eps, hc_pre_eps=hc_pre_eps, @@ -218,6 +219,7 @@ def mhc_pre_delayed_tilelang( layer_input, pre_mix if pre_mix is not None else post, next_pre_mix, + layer_input, hidden_size, rms_eps, hc_pre_eps, @@ -263,6 +265,278 @@ def _mhc_pre_delayed_tilelang_fake( ) +def mhc_fused_post_pre_delayed_tilelang( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Run one mHC post block followed by the next delayed mHC pre block. + + Within the fused kernel's token range the post mapping is folded into the + pre-norm GEMM, so the updated residual streams feed the projection from + registers instead of a second pass over global memory. Above it this runs + the same post kernel and split-k GEMM as the unfused pair. + + Args: + x: BF16 sublayer output of shape (tokens, hidden_size). + residual: BF16 residual streams of shape (tokens, hc_mult, hidden_size). + post_layer_mix: FP32 post coefficients of shape (tokens, hc_mult, 1). + comb_res_mix: FP32 residual coefficients, (tokens, hc_mult, hc_mult). + fn: FP32 projection of shape (hc_mult * (hc_mult + 2), input_size). + hc_scale: FP32 scales of shape (3,). + hc_base: FP32 bias of shape (hc_mult * (hc_mult + 2),). + rms_eps: RMS normalization epsilon. + hc_pre_eps: Pre-mix epsilon. + hc_sinkhorn_eps: Sinkhorn epsilon. + hc_post_mult_value: Post-mix multiplier. + sinkhorn_repeat: Number of Sinkhorn iterations. + pre_mix: FP32 coefficients from the previous sublayer, or None to + select residual stream zero. + norm_weight: Optional BF16 RMSNorm weight for the collapsed input. + norm_eps: RMSNorm epsilon for the collapsed input. + capture_aux: Also return the mean over the post-mapped streams, which + draft models consume as the target's hidden state. It is folded + into the collapse, which already reads those streams. + + Returns: + The post-mapped residual streams, the post and residual coefficients, + the optionally normalized BF16 layer input, the next FP32 pre-mix, and + the BF16 stream mean (empty unless capture_aux), with shapes + (tokens, hc_mult, hidden_size), (tokens, hc_mult, 1), + (tokens, hc_mult, hc_mult), (tokens, hidden_size), (tokens, hc_mult), + and (tokens, hidden_size). + """ + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + _HC_PRENORM_GEMM_TILELANG_KERNEL, + _MHC_FUSED_TILELANG_KERNEL, + _MHC_POST_TILELANG_KERNEL, + mhc_fused_post_pre_split_config, + mhc_pre_big_fuse_tilelang, + ) + from vllm.model_executor.kernels.mhc.warmup import ( + MHC_PRE_NORM_KERNEL, + compute_mhc_pre_num_splits, + ) + from vllm.utils.deep_gemm import ( + is_deep_gemm_supported, + tf32_hc_prenorm_gemm, + ) + + assert residual.ndim == 3 and residual.dtype == torch.bfloat16 + assert residual.is_contiguous() + num_tokens, hc_mult, hidden_size = residual.shape + input_size = hc_mult * hidden_size + mix_size = hc_mult * (hc_mult + 2) + assert x.shape == (num_tokens, hidden_size) and x.dtype == torch.bfloat16 + assert x.is_contiguous() + assert post_layer_mix.shape[:2] == (num_tokens, hc_mult) + assert post_layer_mix.dtype == torch.float32 and post_layer_mix.is_contiguous() + assert comb_res_mix.shape == (num_tokens, hc_mult, hc_mult) + assert comb_res_mix.dtype == torch.float32 and comb_res_mix.is_contiguous() + assert fn.shape == (mix_size, input_size) and fn.dtype == torch.float32 + assert hc_scale.shape == (3,) and hc_scale.dtype == torch.float32 + assert hc_base.shape == (mix_size,) and hc_base.dtype == torch.float32 + if pre_mix is not None: + assert pre_mix.shape == (num_tokens, hc_mult) + assert pre_mix.dtype == torch.float32 and pre_mix.is_contiguous() + + next_pre_mix = torch.empty( + num_tokens, hc_mult, dtype=torch.float32, device=residual.device + ) + post = torch.empty_like(next_pre_mix) + comb = torch.empty( + num_tokens, hc_mult * hc_mult, dtype=torch.float32, device=residual.device + ) + layer_input = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device + ) + aux = torch.empty( + num_tokens if capture_aux else 0, + hidden_size, + dtype=torch.bfloat16, + device=residual.device, + ) + if num_tokens == 0: + return ( + torch.empty_like(residual), + post.unsqueeze(-1), + comb.view(num_tokens, hc_mult, hc_mult), + layer_input, + next_pre_mix, + aux, + ) + + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) + if fused_config is not None: + mixes, sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL( + comb_res_mix, + residual, + post_layer_mix.view(num_tokens, hc_mult), + x, + fn.view(mix_size, hc_mult, hidden_size), + hc_mult, + hidden_size, + mix_size, + ) + else: + residual_cur = torch.empty_like(residual) + _MHC_POST_TILELANG_KERNEL( + comb_res_mix, + residual, + post_layer_mix.view(num_tokens, hc_mult), + x, + residual_cur, + hc_mult, + hidden_size, + ) + # The delayed epilogue is compiled per bucketed split count, so the + # projection has to use the same bucket rather than the raw estimate. + use_deep_gemm = is_deep_gemm_supported() + n_splits = ( + compute_mhc_pre_num_splits(input_size, num_tokens) if use_deep_gemm else 1 + ) + mixes = torch.empty( + n_splits, num_tokens, mix_size, dtype=torch.float32, device=residual.device + ) + sqrsum = torch.empty( + n_splits, num_tokens, dtype=torch.float32, device=residual.device + ) + residual_cur_2d = residual_cur.view(num_tokens, input_size) + if use_deep_gemm: + tf32_hc_prenorm_gemm(residual_cur_2d, fn, mixes, sqrsum, n_splits) + else: + _HC_PRENORM_GEMM_TILELANG_KERNEL( + residual_cur_2d, + fn, + mixes, + sqrsum, + input_size, + 1, + ) + + outputs = ( + residual_cur, + post.unsqueeze(-1), + comb.view(num_tokens, hc_mult, hc_mult), + layer_input, + next_pre_mix, + aux, + ) + if norm_weight is not None: + assert norm_weight.shape == (hidden_size,) + assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous() + MHC_PRE_NORM_KERNEL( + mixes, + sqrsum, + hc_scale, + hc_base, + residual_cur, + post, + comb, + layer_input, + norm_weight, + pre_mix if pre_mix is not None else post, + next_pre_mix, + aux if capture_aux else layer_input, + hidden_size=hidden_size, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + norm_eps=norm_eps, + hc_mult=hc_mult, + use_pre_mix_in=pre_mix is not None, + save_pre_mix=True, + rms_numel=input_size, + write_aux=capture_aux, + ) + return outputs + mhc_pre_big_fuse_tilelang( + mixes, + sqrsum, + hc_scale, + hc_base, + residual_cur, + post, + comb, + layer_input, + pre_mix if pre_mix is not None else post, + next_pre_mix, + aux if capture_aux else layer_input, + hidden_size, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + mixes.shape[0], + hc_mult, + use_pre_mix_in=pre_mix is not None, + save_pre_mix=True, + rms_numel=input_size, + write_aux=capture_aux, + ) + return outputs + + +def _mhc_fused_post_pre_delayed_tilelang_fake( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + num_tokens, hc_mult, hidden_size = residual.shape + return ( + torch.empty_like(residual), + torch.empty( + num_tokens, hc_mult, 1, dtype=torch.float32, device=residual.device + ), + torch.empty( + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=residual.device + ), + torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device + ), + torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=residual.device), + torch.empty( + num_tokens if capture_aux else 0, + hidden_size, + dtype=torch.bfloat16, + device=residual.device, + ), + ) + + def mhc_pre_tilelang( residual: torch.Tensor, fn: torch.Tensor, @@ -574,6 +848,7 @@ def mhc_fused_post_pre_tilelang( _MHC_FUSED_TILELANG_KERNEL, _MHC_POST_TILELANG_KERNEL, _MHC_PRE_BIG_FUSE_TILELANG_KERNEL, + mhc_fused_post_pre_split_config, ) assert residual.dtype == torch.bfloat16 @@ -614,7 +889,7 @@ def mhc_fused_post_pre_tilelang( post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult) comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult) - use_small_fma = num_tokens <= 16 + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) post_mix_cur = torch.empty( num_tokens, @@ -635,7 +910,7 @@ def mhc_fused_post_pre_tilelang( device=residual.device, ) - if use_small_fma: + if fused_config is not None: gemm_out_mul, gemm_out_sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL( comb_res_mix_flat, residual_flat, @@ -815,6 +1090,12 @@ direct_register_custom_op( mutates_args=[], fake_impl=_mhc_pre_delayed_tilelang_fake, ) +direct_register_custom_op( + op_name="mhc_fused_post_pre_delayed_tilelang", + op_func=mhc_fused_post_pre_delayed_tilelang, + mutates_args=[], + fake_impl=_mhc_fused_post_pre_delayed_tilelang_fake, +) direct_register_custom_op( op_name="mhc_pre_tilelang", op_func=mhc_pre_tilelang, diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index b4d3190c0e..9e803d123e 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -37,8 +37,67 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: return split_k -def _mhc_fused_n_splits(num_tokens: int, hidden_size: int) -> int: - return 8 if num_tokens < 8 and hidden_size <= 4096 else 4 +# ``mhc_fused_tilelang`` splits the hidden size across a fixed thread block, so +# every (tile_n, n_thr) pair is one more kernel to compile. Tuned on GB300 +# (hc_mult 4, hidden_size 5120 and 7168) against the separate post + split-k +# GEMM it replaces; the bands are coarse because the loss against a per-shape +# optimum is ~1-2%, while one fixed config gives up to 18% at the top of the +# range. +_FUSED_POST_PRE_N_THR = 128 +_FUSED_POST_PRE_MAX_TOKENS = 32 + + +def mhc_fused_post_pre_split_config( + num_tokens: int, hidden_size: int, hc_mult: int +) -> tuple[int, int, int] | None: + """Pick ``(tile_n, n_splits, n_thr)`` for the fused post + pre-norm GEMM. + + Returns None when a separate post kernel followed by a split-k GEMM is the + better choice, or when the hidden size does not divide evenly across the + fused kernel's block. One source of truth for the dispatch decision, the + compile key and the launch, which must agree. + """ + if num_tokens > _FUSED_POST_PRE_MAX_TOKENS: + return None + n_thr = _FUSED_POST_PRE_N_THR + if hidden_size % n_thr: + return None + # The projection tiles supply the parallelism a small token count cannot, + # so the tile grows with the batch: 1.3x at one token and 1.03x at 32 + # against a split-k GEMM given the split its own estimator computes. + if num_tokens < 16: + tile_n, n_splits = 2, 8 + else: + tile_n, n_splits = 6, 8 + # The kernel drops projection tiles and k-slices it cannot fill evenly. + mix_size = hc_mult * (hc_mult + 2) + while tile_n > 1 and mix_size % tile_n: + tile_n -= 1 + while n_splits > 1 and hidden_size % (n_splits * n_thr): + n_splits //= 2 + return tile_n, n_splits, n_thr + + +def require_fused_post_pre_config( + num_tokens: int, hidden_size: int, hc_mult: int +) -> tuple[int, int, int]: + """The config for a shape the caller has already committed to fusing.""" + config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) + if config is None: + raise ValueError( + "the fused mHC post + pre-norm GEMM does not cover num_tokens=" + f"{num_tokens} at hidden_size={hidden_size}" + ) + return config + + +def mhc_fused_post_pre_splits(hidden_size: int, hc_mult: int) -> tuple[int, ...]: + """Every split-k factor the fused post + pre-norm GEMM path can pick.""" + configs = ( + mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) + for num_tokens in range(1, _FUSED_POST_PRE_MAX_TOKENS + 1) + ) + return tuple(sorted({config[1] for config in configs if config is not None})) @tilelang_jit @@ -53,6 +112,7 @@ def mhc_pre_big_fuse_tilelang( layer_input, pre_mix_in, pre_mix_out, + aux_out, hidden_size: int, rms_eps: float, hc_pre_eps: float, @@ -64,6 +124,7 @@ def mhc_pre_big_fuse_tilelang( use_pre_mix_in: bool = False, save_pre_mix: bool = False, rms_numel: int = 0, + write_aux: bool = False, ): """Fuse coefficient generation and residual collapse after the projection. @@ -88,6 +149,7 @@ def mhc_pre_big_fuse_tilelang( pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] + aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: if ENABLE_PDL: @@ -188,11 +250,24 @@ def mhc_pre_big_fuse_tilelang( ol = T.alloc_fragment(hidden_block, T.float32) T.clear(ol) + if write_aux: + # Aux consumers (draft models) take the plain stream mean + # of the same residual this collapse reads. + aux = T.alloc_fragment(hidden_block, T.float32) + T.clear(aux) for i_hc in T.serial(hc_mult): pre = pre_mix_shared[i_hc] for i1_h in T.Parallel(hidden_block): ol[i1_h] += pre * xl[i_hc, i1_h] + if write_aux: + aux[i1_h] += xl[i_hc, i1_h] + + if write_aux: + for i1_h in T.Parallel(hidden_block): + aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16( + aux[i1_h] / hc_mult + ) T.copy(ol, layer_input[i, i0_h * hidden_block]) @@ -216,6 +291,7 @@ def mhc_pre_big_fuse_with_norm_tilelang( norm_weight, pre_mix_in, pre_mix_out, + aux_out, hidden_size: int, rms_eps: float, hc_pre_eps: float, @@ -229,6 +305,7 @@ def mhc_pre_big_fuse_with_norm_tilelang( use_pre_mix_in: bool = False, save_pre_mix: bool = False, rms_numel: int = 0, + write_aux: bool = False, ): num_tokens = T.dynamic("num_tokens") hc_mult3 = hc_mult * (2 + hc_mult) @@ -250,6 +327,7 @@ def mhc_pre_big_fuse_with_norm_tilelang( pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] + aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: rms = T.alloc_fragment(1, T.float32) @@ -345,11 +423,24 @@ def mhc_pre_big_fuse_with_norm_tilelang( ol = T.alloc_fragment(hidden_block, T.float32) T.clear(ol) + if write_aux: + # Aux consumers (draft models) take the plain stream mean + # of the same residual this collapse reads. + aux = T.alloc_fragment(hidden_block, T.float32) + T.clear(aux) for i_hc in T.serial(hc_mult): pre = pre_mix_shared[i_hc] for i1_h in T.Parallel(hidden_block): ol[i1_h] += pre * xl[i_hc, i1_h] + if write_aux: + aux[i1_h] += xl[i_hc, i1_h] + + if write_aux: + for i1_h in T.Parallel(hidden_block): + aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16( + aux[i1_h] / hc_mult + ) if save_pre_mix: # Keep the BF16 boundary before the delayed input RMSNorm. @@ -1211,9 +1302,17 @@ class MhcPreBigFuseTileLangKernel( middle_args = (residual_out, *output_args, norm_weight) elif compile_key.use_norm_weight: assert norm_weight is not None - middle_args = (*output_args, norm_weight, post_mix, post_mix) + # layer_input stands in for the aux buffer on both epilogues: same + # shape and dtype, and write_aux is off here so it is never written. + middle_args = ( + *output_args, + norm_weight, + post_mix, + post_mix, + layer_input, + ) else: - middle_args = (*output_args, post_mix, post_mix) + middle_args = (*output_args, post_mix, post_mix, layer_input) norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else () return ( *common_args, @@ -1247,8 +1346,14 @@ class MhcPreBigFuseTileLangKernel( if use_pre_gemm_splits else n_splits ) - fused_n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) - actual_n_splits = fused_n_splits if use_fused_tilelang else pre_gemm_n_splits + # The epilogue reduces over whatever the fused kernel split the GEMM + # into, so both have to read the same config. + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) + actual_n_splits = ( + fused_config[1] + if use_fused_tilelang and fused_config is not None + else pre_gemm_n_splits + ) actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps actual_use_norm_weight = use_norm_weight or is_broadcast return self.CompileKey( @@ -1540,6 +1645,7 @@ class MhcFusedTileLangKernel( hc_mult: int n_splits: int tile_n: int + n_thr: int @staticmethod def kernel() -> Any: @@ -1552,14 +1658,15 @@ class MhcFusedTileLangKernel( hidden_size: int, hc_mult: int, ) -> CompileKey: - # TODO(gnovack): investigate autotuning these heuristics - tile_n = 2 if num_tokens < 8 else 3 - n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) + # Subscripts rather than an unpack: the warmup tracer parses this body + # and allows only plain assignments before the return. + config = require_fused_post_pre_config(num_tokens, hidden_size, hc_mult) return self.CompileKey( hidden_size=hidden_size, hc_mult=hc_mult, - n_splits=n_splits, - tile_n=tile_n, + n_splits=config[1], + tile_n=config[0], + n_thr=config[2], ) def get_warmup_keys( @@ -1574,11 +1681,11 @@ class MhcFusedTileLangKernel( num_tokens=WarmupIntRange(1, max_tokens + 1), hidden_size=hidden_size, hc_mult=hc_mult, - _when=lambda *, num_tokens: num_tokens <= 16, + _when=lambda *, num_tokens: num_tokens <= _FUSED_POST_PRE_MAX_TOKENS, ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: - num_tokens = 1 if compile_key.tile_n == 2 else 8 + num_tokens = 1 if compile_key.tile_n == 2 else 16 hidden_size = compile_key.hidden_size hc_mult = compile_key.hc_mult hc_mult3 = hc_mult * (2 + hc_mult) @@ -1617,8 +1724,9 @@ class MhcFusedTileLangKernel( hc_mult3: int, ) -> TileLangLaunchSpec: num_tokens = residual_in.shape[0] - tile_n = 2 if num_tokens < 8 else 3 - n_splits = _mhc_fused_n_splits(num_tokens, hidden_size) + tile_n, n_splits, n_thr = require_fused_post_pre_config( + num_tokens, hidden_size, hc_mult + ) yp_out = residual_in.new_empty( (n_splits, num_tokens, hc_mult3), dtype=torch.float32 ) @@ -1639,7 +1747,7 @@ class MhcFusedTileLangKernel( hidden_size, hc_mult3, ), - dict(tile_n=tile_n, split_k=n_splits), + dict(n_thr=n_thr, tile_n=tile_n), (yp_out, rp_out, residual_out), ) diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index 0abe1a2161..4e6a19781c 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/vllm/model_executor/kernels/mhc/warmup.py @@ -36,23 +36,33 @@ class MHCPreNormKernel(VllmJitKernel["MHCPreNormKernel.CompileKey"]): use_pre_mix_in: bool rms_numel: int save_pre_mix: bool = True + write_aux: bool = False kernel: Any = staticmethod(mhc_pre_big_fuse_with_norm_tilelang) def dispatch(self, *, n_splits, **fields) -> CompileKey: # type: ignore[override] return self.CompileKey(n_splits=n_splits, **fields) - def get_warmup_keys(self, *, max_tokens: int, **fields) -> list[CompileKey]: - # The split heuristic changes only at 64-token boundaries. - splits = ( - sorted( + def get_warmup_keys( + self, + *, + max_tokens: int, + extra_splits: Iterable[int] = (), + **fields, + ) -> list[CompileKey]: + # The split heuristic changes only at 64-token boundaries. Callers that + # reach this kernel through another GEMM (the fused post + pre-norm + # GEMM) pass the splits that GEMM picks in extra_splits. + splits = sorted( + set(extra_splits) + | ( { compute_mhc_pre_num_splits(fields["rms_numel"], num_tokens) for num_tokens in range(1, max_tokens + 1, 64) } + if is_deep_gemm_supported() + else {1} ) - if is_deep_gemm_supported() - else [1] ) return self._trace_dispatch(self.dispatch)(n_splits=splits, **fields) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 46fe2558ab..3c711681aa 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -207,7 +207,7 @@ class DSparkDeepseekV4Model(nn.Module): residual = post_mix = res_mix = pre_mix = None for layer in self.layers: - hidden_states, residual, post_mix, res_mix, pre_mix = layer( + hidden_states, residual, post_mix, res_mix, pre_mix, _ = layer( hidden_states, positions, input_ids, diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index df12a9d9bb..f3530544da 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -20,6 +20,7 @@ from vllm.distributed import ( from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_fused_post_pre_delayed_tilelang, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -214,17 +215,45 @@ class DeepseekV4DecoderLayer(nn.Module): self.hc_eps = config.hc_eps self.hc_post_alpha = 2.0 if vllm_config.kernel_config.enable_jit_warmup and current_platform.is_cuda(): + from vllm.model_executor.kernels.mhc.tilelang_kernels import ( + mhc_fused_post_pre_splits, + ) from vllm.model_executor.kernels.mhc.warmup import MHC_PRE_NORM_KERNEL max_tokens = vllm_config.scheduler_config.max_num_batched_tokens if self.use_sequence_parallel: tp_size = vllm_config.parallel_config.tensor_parallel_size max_tokens = (max_tokens + tp_size - 1) // tp_size - for input_size, use_pre_mix in ( - (self.hidden_size, False), - (self.hc_mult * self.hidden_size, False), - (self.hc_mult * self.hidden_size, True), - ): + # The epilogue compiles per projection width and per pre-mix mode. + # The first layer projects the broadcast embedding, so it reads one + # hidden_size-wide row and selects stream zero; every later sublayer + # projects the full hc stream, and collapses it with the pre-mix the + # previous sublayer carried in. Those later shapes also arrive + # through the fused post + pre-norm GEMM, which picks split-k + # factors the token sweep below never produces. + broadcast_embedding = { + "rms_numel": self.hidden_size, + "use_pre_mix_in": False, + "extra_splits": (), + } + hc_stream = { + "rms_numel": self.hc_mult * self.hidden_size, + "extra_splits": mhc_fused_post_pre_splits( + self.hidden_size, self.hc_mult + ), + } + variants = [ + broadcast_embedding, + {**hc_stream, "use_pre_mix_in": False}, + {**hc_stream, "use_pre_mix_in": True}, + ] + if vllm_config.speculative_config is not None: + # Draft setups read aux hidden states out of the same collapse, + # which is its own specialization of the kernel. + variants.append( + {**hc_stream, "use_pre_mix_in": True, "write_aux": True} + ) + for variant in variants: MHC_PRE_NORM_KERNEL.register_warmup( max_tokens=max_tokens, hidden_size=self.hidden_size, @@ -235,8 +264,7 @@ class DeepseekV4DecoderLayer(nn.Module): sinkhorn_repeat=self.hc_sinkhorn_iters, norm_eps=self.rms_norm_eps, hc_mult=self.hc_mult, - use_pre_mix_in=use_pre_mix, - rms_numel=input_size, + **variant, ) mix_hc = (2 + self.hc_mult) * self.hc_mult hc_dim = self.hc_mult * self.hidden_size @@ -295,7 +323,17 @@ class DeepseekV4DecoderLayer(nn.Module): residual: torch.Tensor | None = None, engram_hashes: torch.Tensor | None = None, engram_mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + *, + capture_previous_aux: bool = False, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + ]: + previous_aux: torch.Tensor | None = None # The reference collapses each sublayer's input with the *previous* # sublayer's pre-mix: attention uses the pre-mix carried in (identity # for the first layer), the FFN uses this layer's attention pre-mix. @@ -335,17 +373,19 @@ class DeepseekV4DecoderLayer(nn.Module): norm_weight=self.attn_norm.weight, norm_eps=self.attn_norm.variance_epsilon, ) - else: - residual = mhc_post_tilelang(x, residual, post_mix, res_mix) - if self.engram is not None and engram_hashes is not None: - # Engram injection happens between the previous sublayer's - # post and this block's pre, on the full hc stream, so the - # mix coefficients see the injected stream. - residual = self.engram( - residual, - engram_hashes[:, self.engram.layer_hash_index], - engram_mask, - ) + elif self.engram is not None and engram_hashes is not None: + # Engram injection happens between the previous sublayer's post + # and this block's pre, on the full hc stream, so the mix + # coefficients see the injected stream. The injection also keeps + # the post out of the pre-norm GEMM's fused prologue. + previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix) + if capture_previous_aux: + previous_aux = previous_post.mean(dim=1) + residual = self.engram( + previous_post, + engram_hashes[:, self.engram.layer_hash_index], + engram_mask, + ) post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang( residual, self.hc_attn_fn, @@ -360,6 +400,31 @@ class DeepseekV4DecoderLayer(nn.Module): norm_weight=self.attn_norm.weight, norm_eps=self.attn_norm.variance_epsilon, ) + else: + # The collapse already reads the post-mapped streams, so the mean + # aux consumers want comes out of the same kernel. + residual, post_mix, res_mix, x, attn_pre, aux = ( + mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + norm_weight=self.attn_norm.weight, + norm_eps=self.attn_norm.variance_epsilon, + capture_aux=capture_previous_aux, + ) + ) + if capture_previous_aux: + previous_aux = aux if self.use_sequence_parallel: x = sp_all_gather(x)[: positions.shape[0]] @@ -368,23 +433,27 @@ class DeepseekV4DecoderLayer(nn.Module): if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual = mhc_post_tilelang(x, residual, post_mix, res_mix) - post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_tilelang( - residual, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=attn_pre, - norm_weight=self.ffn_norm.weight, - norm_eps=self.ffn_norm.variance_epsilon, + residual, post_mix, res_mix, x, ffn_pre, _ = ( + mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=attn_pre, + norm_weight=self.ffn_norm.weight, + norm_eps=self.ffn_norm.variance_epsilon, + ) ) x = self.ffn(x, input_ids) - return x, residual, post_mix, res_mix, ffn_pre + return x, residual, post_mix, res_mix, ffn_pre, previous_aux class DeepseekV4Model(nn.Module, EagleModelMixin): @@ -631,13 +700,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): if not get_pp_group().is_first_rank: assert intermediate_tensors is not None pre_mix = intermediate_tensors["pre_mix"] - aux_hidden_states: list[torch.Tensor] = [] - final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call + # Every layer's post runs inside the next layer's fused pre, so aux + # hidden states are read back from there instead of recomputed. + aux_hidden_by_layer: dict[int, torch.Tensor] = {} for idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): - hidden_states, residual, post_mix, res_mix, pre_mix = layer( + hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = layer( hidden_states, positions, input_ids, @@ -647,25 +717,29 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): residual, engram_hashes, engram_mask, + capture_previous_aux=idx in self.aux_hidden_state_layers, ) - if idx + 1 in self.aux_hidden_state_layers: - # Reconstruct the aux hidden state for draft models - aux_recon = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix - ) - aux_hidden_state = aux_recon.mean(dim=1) + if previous_aux is not None: + # idx is the one-based id of the layer whose post this is. if self.use_sequence_parallel: - aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] - aux_hidden_states.append(aux_hidden_state) - final_aux_recon = aux_recon + previous_aux = sp_all_gather(previous_aux)[:full_num_tokens] + aux_hidden_by_layer[idx] = previous_aux if layer is not None: - # Reuse if the last layer was captured as an aux hidden state + # The last layer has no successor to fold its post into. + hidden_states = mhc_post_tilelang( + hidden_states, residual, post_mix, res_mix + ) if self.end_layer in self.aux_hidden_state_layers: - hidden_states = final_aux_recon - else: - hidden_states = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix - ) + final_aux = hidden_states.mean(dim=1) + if self.use_sequence_parallel: + final_aux = sp_all_gather(final_aux)[:full_num_tokens] + aux_hidden_by_layer[self.end_layer] = final_aux + + aux_hidden_states = [ + aux_hidden_by_layer[layer_id] + for layer_id in self.aux_hidden_state_layers + if layer_id in aux_hidden_by_layer + ] if not get_pp_group().is_last_rank: return IntermediateTensors( From 0a5747d410da9b08145dba2b6f3e6614bab35753 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=90=98=E5=A4=A9=E6=A5=BD?= Date: Mon, 14 Sep 2026 15:02:07 -0400 Subject: [PATCH 043/139] [Profiler] Add Proton CUDA graph attribution for MRV2 (#51084) Signed-off-by: Luossu Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- docs/contributing/profiling.md | 44 ++- tests/v1/worker/test_gpu_profiler.py | 400 +++++++++++++++++++++- vllm/config/profiler.py | 16 + vllm/config/vllm.py | 51 ++- vllm/profiler/wrapper.py | 149 +++++++- vllm/v1/worker/gpu/model_runner.py | 10 + vllm/v1/worker/gpu_worker.py | 29 +- vllm/v1/worker/mm_encoder_model_runner.py | 3 + 8 files changed, 657 insertions(+), 45 deletions(-) diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index e420a3aa68..880d065a2c 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -90,18 +90,18 @@ $ curl -X POST http://localhost:8000/stop_profile is Triton's GPU profiler. It can collect a low-overhead aggregate tree or a Chrome trace and works through the same vLLM profiling controls as the PyTorch and CUDA profilers. Proton currently supports NVIDIA GPUs through CUPTI and -requires eager execution. +supports CUDA graph attribution. -Start a server with a local output directory: +Start a server with a local output directory and graph attribution: ```bash vllm serve meta-llama/Llama-3.1-8B-Instruct \ - --enforce-eager \ --profiler-config '{ "profiler": "proton", "proton_profiler_dir": "./proton_profile", "proton_output_format": "hatchet", - "proton_hook": "triton" + "proton_hook": "triton", + "proton_graph_attribution": true }' ``` @@ -111,7 +111,7 @@ rank-qualified output name, such as `proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet`, so distributed workers, restarted servers, and repeated profiling runs do not overwrite one another. A `profile_prefix` is included when supplied. Each -profile is finalized by `/stop_profile` and is ready to inspect immediately. +profile is written by `/stop_profile` and is ready to inspect immediately. The Proton-specific options are: @@ -121,21 +121,39 @@ The Proton-specific options are: - `proton_mode`: an optional backend mode string - `proton_hook`: `triton` to record Triton launch metadata, or unset - `proton_output_format`: `hatchet`, `hatchet_msgpack`, `chrome_trace`, or unset +- `proton_graph_attribution`: observe CUDA graph capture for replay attribution; + disabled by default and requires `proton_data: "tree"` `hatchet` and `hatchet_msgpack` require `proton_data: "tree"`, while -`chrome_trace` requires `proton_data: "trace"`. CUDA graph profiling is not yet -supported, so `--enforce-eager` is required for every Proton mode. +`chrome_trace` requires `proton_data: "trace"`. Automatic backend selection is recommended. vLLM currently supports Proton's `cupti` backend on NVIDIA GPUs. ROCm support is not yet available. vLLM does not expose Proton's experimental instrumentation backend because current upstream -Triton builds can produce profiles without timing metrics. Backend-specific -modes, including `pcsampling`, can be selected with `proton_mode`. +can produce profiles without timing metrics. When `proton_graph_attribution` is +enabled, Proton observes vLLM's CUDA graph capture with the configured profiling +session active, then deactivates that same session until profiling starts. This +lets later profiles attribute replayed kernels without retaining model-startup +activity. Backend-specific modes can be selected with `proton_mode`; +`pcsampling` synchronizes the CUDA context and therefore requires +`--enforce-eager`. When CUDA graphs are enabled (including encoder graphs), +Proton requires `proton_graph_attribution: true` to collect replayed kernels. +For Chrome traces, disable CUDA graphs with `--enforce-eager`. -Triton 3.6 supports explicit `hatchet` and `chrome_trace` output. The -`hatchet_msgpack` format and `periodic_flushing` mode require Triton 3.7 or -newer; vLLM rejects these options with the detected version before starting a -session. +CUDA graph-attributed profiles support repeated `start_profile`/`stop_profile` +runs. Each stop flushes and writes one tree-data phase while preserving the +graph-aware session. Without `proton_graph_attribution`, each `stop_profile` +instead finalizes and writes an independent Proton session. + +CUDA graph attribution requires Triton 3.7 or newer. It uses the phase data API +to discard graph-capture activity and separate profiling runs. The `hatchet_msgpack` +output format and `periodic_flushing` mode also require Triton 3.7 or newer. +`periodic_flushing` cannot be combined with graph attribution because both +manage the session's data phases. Ordinary Proton profiling remains available +with Triton 3.6. + +Graph attribution retains capture metadata for the worker lifetime. With eager +execution or no graphs to capture, profiling uses ordinary independent sessions. Inspect tree profiles with: diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index 347b5b379a..f6263b68da 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -350,17 +350,24 @@ def test_profiler_entered_during_capture(): def make_proton(session_id: int | None = 7): + data = SimpleNamespace( + advance_phase=Mock(side_effect=range(1, 100)), + clear=Mock(), + get=Mock(return_value={"traceEvents": []}), + get_msgpack=Mock(return_value=b"profile"), + ) return SimpleNamespace( start=Mock(return_value=session_id), activate=Mock(), deactivate=Mock(), finalize=Mock(), scope=Mock(return_value=nullcontext()), + data=data, ) def make_proton_wrapper( - tmp_path, proton=None, triton_version="3.6.0", **config_overrides + tmp_path, proton=None, triton_version="3.7.0", **config_overrides ): proton = proton or make_proton() config = ProfilerConfig( @@ -415,6 +422,13 @@ class TestProtonConfig: }, "requires proton_data", ), + ( + { + "proton_data": "trace", + "proton_graph_attribution": True, + }, + "requires proton_data='tree'", + ), ], ) def test_rejects_invalid_option_combinations(self, tmp_path, options, message): @@ -450,6 +464,7 @@ class TestProtonConfig: ("proton_mode", "pcsampling"), ("proton_hook", "triton"), ("proton_output_format", "chrome_trace"), + ("proton_graph_attribution", True), ], ) def test_rejects_proton_options_for_other_profilers(self, field, value): @@ -478,15 +493,101 @@ class TestProtonConfig: compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE), ) - def test_rejects_proton_when_cuda_graphs_are_enabled(self, tmp_path): - with pytest.raises(ValueError, match="requires CUDA graphs to be disabled"): + @pytest.mark.parametrize("encoder_only", [False, True]) + @pytest.mark.parametrize("attribution", [False, True]) + def test_cuda_graphs_require_attribution(self, tmp_path, encoder_only, attribution): + # Encoder graphs are independent of the decoder's cudagraph_mode. + expected = ( + nullcontext() + if attribution + else pytest.raises( + ValueError, match="requires proton_graph_attribution=True" + ) + ) + with expected: + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=attribution, + ), + compilation_config=CompilationConfig( + cudagraph_mode=( + CUDAGraphMode.NONE if encoder_only else CUDAGraphMode.FULL + ), + cudagraph_mm_encoder=encoder_only, + ), + ) + + def test_validates_default_cuda_graph_mode_after_resolution(self, tmp_path): + with pytest.raises(ValueError, match="requires proton_graph_attribution=True"): VllmConfig( profiler_config=ProfilerConfig( profiler="proton", proton_profiler_dir=str(tmp_path) ), - compilation_config=CompilationConfig( - cudagraph_mode=CUDAGraphMode.PIECEWISE + ) + + @pytest.mark.parametrize( + "mode", + [ + "periodic_flushing", + "periodic_flushing:format=hatchet", + "PERIODIC_FLUSHING:format=hatchet", + ], + ) + def test_graph_attribution_rejects_periodic_flushing(self, tmp_path, mode): + # Reject before Proton's native phase manager can abort the worker. + with pytest.raises(ValueError, match="incompatible with periodic_flushing"): + ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=True, + proton_mode=mode, + ) + + @pytest.mark.parametrize("attribution", [False, True]) + @pytest.mark.parametrize( + "mode", ["pcsampling", "pcsampling:interval=100", "PcSampling:interval=100"] + ) + @pytest.mark.parametrize("encoder_only", [False, True]) + def test_pcsampling_requires_graphs_disabled( + self, tmp_path, attribution, mode, encoder_only + ): + with pytest.raises(ValueError, match="PC sampling requires CUDA graphs"): + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=attribution, + proton_mode=mode, ), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.NONE + if encoder_only + else CUDAGraphMode.FULL, + cudagraph_mm_encoder=encoder_only, + ), + ) + + def test_ordinary_proton_keeps_mrv1_graph_restriction(self, tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + with pytest.raises(ValueError, match="requires proton_graph_attribution=True"): + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", proton_profiler_dir=str(tmp_path) + ), + compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.FULL), + ) + + def test_graph_attribution_rejects_mrv1(self, tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + with pytest.raises(ValueError, match="requires the V2 model runner"): + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=True, + ) ) @@ -564,7 +665,20 @@ class TestProtonProfilerWrapper: ) def test_newer_features_reject_triton_3_6(self, tmp_path, option, value, feature): with pytest.raises(RuntimeError, match=feature): - make_proton_wrapper(tmp_path, **{option: value}) + make_proton_wrapper(tmp_path, triton_version="3.6.0", **{option: value}) + + @pytest.mark.parametrize("version", ["3.6.0", "unknown"]) + def test_graph_attribution_requires_phase_api(self, tmp_path, version): + with pytest.raises(RuntimeError, match="requires Triton >= 3.7"): + make_proton_wrapper( + tmp_path, triton_version=version, proton_graph_attribution=True + ) + + def test_ordinary_profiling_still_supports_triton_3_6(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, triton_version="3.6.0") + wrapper.start() + wrapper.stop() + proton.finalize.assert_called_once_with(session=7) @pytest.mark.parametrize( ("option", "value"), @@ -583,6 +697,125 @@ class TestProtonProfilerWrapper: with pytest.raises(RuntimeError, match="does not support selecting"): make_proton_wrapper(tmp_path, proton, proton_output_format="hatchet") + def test_cuda_graph_tree_phase_is_written_at_stop(self, tmp_path): + wrapper, proton = make_proton_wrapper( + tmp_path, + proton_context="python", + proton_graph_attribution=True, + ) + with wrapper.capture_cuda_graphs(): + proton.start.assert_called_once() + proton.deactivate.assert_not_called() + + capture_args = proton.start.call_args.kwargs + assert capture_args["context"] == "python" + assert capture_args["data"] == "tree" + proton.data.advance_phase.assert_called_once_with(7) + proton.deactivate.assert_called_once_with(session=7, flushing=True) + proton.data.clear.assert_called_once_with(7, 0) + + wrapper.start() + wrapper.stop() + wrapper.start() + wrapper.stop() + + assert proton.data.get.call_args_list == [call(7, 1), call(7, 2)] + proton.data.clear.assert_has_calls([call(7, 0), call(7, 1), call(7, 2)]) + output_names = sorted(tmp_path.glob("proton_rank_3_*.hatchet")) + assert len(output_names) == 2 + assert output_names[0].name.endswith("_run0.hatchet") + assert output_names[1].name.endswith("_run1.hatchet") + + def test_capture_is_noop_without_opt_in(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path) + with wrapper.capture_cuda_graphs(): + pass + proton.start.assert_not_called() + + @pytest.mark.parametrize("delay", [0, 2]) + def test_duplicate_start_preserves_output_prefix(self, tmp_path, delay): + wrapper, proton = make_proton_wrapper( + tmp_path, proton_graph_attribution=True, delay_iterations=delay + ) + with wrapper.capture_cuda_graphs(): + pass + wrapper.set_output_name("first") + wrapper.start() + wrapper.set_output_name("duplicate") + wrapper.start() + for _ in range(delay): + wrapper.step() + wrapper.stop() + assert len(list(tmp_path.glob("proton_first_*.hatchet"))) == 1 + assert not list(tmp_path.glob("proton_duplicate_*")) + wrapper.set_output_name("second") + wrapper.start() + for _ in range(delay): + wrapper.step() + wrapper.stop() + assert len(list(tmp_path.glob("proton_second_*.hatchet"))) == 1 + + def test_failed_export_clears_activity_and_allows_next_interval(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + with wrapper.capture_cuda_graphs(): + pass + proton.data.get.side_effect = [OSError("disk full"), []] + wrapper.start() + wrapper.stop() + proton.data.clear.assert_called_with(7, 1) + wrapper.start() + wrapper.stop() + proton.data.clear.assert_called_with(7, 2) + assert len(list(tmp_path.glob("*_run1.hatchet"))) == 1 + + def test_cuda_graph_context_deactivates_after_capture_error(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + + with ( + pytest.raises(RuntimeError, match="capture failed"), + wrapper.capture_cuda_graphs(), + ): + raise RuntimeError("capture failed") + + proton.deactivate.assert_called_once_with(session=7, flushing=True) + + def test_cuda_graph_capture_deactivates_when_phase_advance_fails(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + proton.data.advance_phase.side_effect = RuntimeError("advance failed") + + with ( + pytest.raises(RuntimeError, match="advance failed"), + wrapper.capture_cuda_graphs(), + ): + pass + + proton.deactivate.assert_called_once_with(session=7, flushing=True) + proton.data.clear.assert_not_called() + + def test_cuda_graph_stop_deactivates_when_phase_advance_fails(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + with wrapper.capture_cuda_graphs(): + pass + proton.data.clear.reset_mock() + proton.data.advance_phase.side_effect = RuntimeError("advance failed") + + wrapper._start() + with pytest.raises(RuntimeError, match="advance failed"): + wrapper._stop() + + proton.deactivate.assert_called_with(session=7, flushing=True) + proton.data.clear.assert_not_called() + + def test_shutdown_finalizes_cuda_graph_capture_session(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) + with wrapper.capture_cuda_graphs(): + pass + + wrapper.shutdown() + wrapper.shutdown() + + proton.finalize.assert_called_once_with(session=7) + def test_missing_proton_has_actionable_error(self, tmp_path): config = ProfilerConfig(profiler="proton", proton_profiler_dir=str(tmp_path)) with ( @@ -638,6 +871,7 @@ def test_gpu_worker_recreates_proton_profiler_for_each_run(): ), patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper, ): + wrapper.return_value.has_cuda_graph_session = False Worker.profile(worker, profile_prefix="first") Worker.profile(worker, is_start=False) Worker.profile(worker, profile_prefix="second") @@ -647,3 +881,157 @@ def test_gpu_worker_recreates_proton_profiler_for_each_run(): call(worker.profiler_config, worker_name="second_rank1"), ] assert wrapper.return_value.start.call_count == 2 + + +@_requires_cuda_for_proton +def test_gpu_worker_reuses_cuda_graph_proton_session(): + worker = MagicMock() + worker.rank = 1 + worker.profiler = MagicMock(spec=ProtonProfilerWrapper) + worker.profiler.has_cuda_graph_session = True + worker.profiler_config.profiler = "proton" + + with patch( + "vllm.distributed.utils.get_worker_rank_suffix", + return_value="rank1", + ): + Worker.profile(worker, profile_prefix="first") + Worker.profile(worker, is_start=False) + + worker.profiler.set_output_name.assert_called_once_with("first_rank1") + worker.profiler.start.assert_called_once_with() + worker.profiler.stop.assert_called_once_with() + + +@_requires_cuda_for_proton +@pytest.mark.parametrize("runner", ["attribution_off", "v1", "no_capture"]) +def test_proton_not_initialized_without_capture(runner): + worker = MagicMock() + worker.profiler = None + worker.profiler_config.profiler = "proton" + worker.profiler_config.proton_graph_attribution = runner != "attribution_off" + worker.use_v2_model_runner = runner != "v1" + worker.model_runner.needs_cudagraph_capture.return_value = runner != "no_capture" + + context = Worker._get_cudagraph_capture_context(worker) + + assert worker.profiler is None + with context: + pass + + +@_requires_cuda_for_proton +def test_proton_initializes_before_cuda_graph_capture(): + class FakeProtonProfiler: + def __init__(self, config, worker_name): + self.config = config + self.worker_name = worker_name + self.capture_context = nullcontext() + + def capture_cuda_graphs(self): + return self.capture_context + + worker = MagicMock() + worker.rank = 2 + worker.profiler = None + worker.profiler_config.profiler = "proton" + worker.profiler_config.proton_graph_attribution = True + worker.use_v2_model_runner = True + worker.model_runner.needs_cudagraph_capture.return_value = True + + with ( + patch( + "vllm.distributed.utils.get_worker_rank_suffix", + return_value="rank2", + ), + patch( + "vllm.profiler.wrapper.ProtonProfilerWrapper", + FakeProtonProfiler, + ), + ): + context = Worker._get_cudagraph_capture_context(worker) + + assert worker.profiler.config is worker.profiler_config + assert worker.profiler.worker_name == "rank2" + assert context is worker.profiler.capture_context + + +@_requires_cuda_for_proton +@pytest.mark.parametrize("context", ["shadow", "python"]) +@pytest.mark.parametrize("output_format", ["hatchet", "hatchet_msgpack"]) +def test_proton_cuda_graph_replay_attribution_on_gpu(tmp_path, context, output_format): + """Both intervals contain replay kernels, without capture-only activity.""" + import json + + import torch + import triton + import triton.profiler as proton + from packaging.version import Version + + if Version(triton.__version__) < Version("3.7"): + pytest.skip("Graph attribution requires Triton >= 3.7") + + wrapper = ProtonProfilerWrapper( + ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + proton_graph_attribution=True, + proton_context=context, + proton_output_format=output_format, + ), + worker_name="gpu", + ) + worker = SimpleNamespace( + profiler=wrapper, + use_v2_model_runner=True, + ) + x = torch.ones(1024, device="cuda") + graph = torch.cuda.CUDAGraph() + + def capture_only(): + x.add_(1) + + def captured_add(): + x.add_(1) + + def kernel_metrics(value): + if isinstance(value, list): + return [metric for child in value for metric in kernel_metrics(child)] + if isinstance(value, dict): + metrics = value.get("metrics", {}) + return ([metrics] if metrics.get("count", 0) else []) + [ + metric for child in value.values() for metric in kernel_metrics(child) + ] + return [] + + try: + with Worker._get_cudagraph_capture_context(worker): + with proton.scope("capture_only"): + capture_only() + torch.accelerator.synchronize() + with torch.cuda.graph(graph), proton.scope("captured_add"): + captured_add() + for run in range(2): + wrapper.start() + with wrapper.annotate_context_manager(f"replay_{run}"): + graph.replay() + wrapper.stop() + (path,) = tmp_path.glob(f"*_run{run}.{output_format}") + if output_format == "hatchet_msgpack": + import msgpack + + data = msgpack.unpackb(path.read_bytes()) + else: + data = json.loads(path.read_text()) + serialized = json.dumps(data) + assert "" in serialized + assert "captured_add" in serialized + assert "capture_only" not in serialized + assert f"replay_{1 - run}" not in serialized + metrics = kernel_metrics(data) + assert sum(metric["count"] for metric in metrics) == 1 + assert sum(metric["time (ns)"] for metric in metrics) > 0 + torch.testing.assert_close(x, torch.full_like(x, 4)) + finally: + wrapper.shutdown() + assert not list(tmp_path.glob(".proton_cuda_graph_session*")) diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index b3b232490f..535960a8d7 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -76,6 +76,10 @@ class ProfilerConfig: """Optional format passed to Proton when finalizing a profile. ``None`` uses the default format for ``proton_data``.""" + proton_graph_attribution: bool = False + """Observe CUDA graph capture so replayed kernels can be attributed. + Requires Triton >= 3.7 and ``proton_data='tree'``.""" + torch_profiler_with_stack: bool = True """If `True`, enables stack tracing in the torch profiler. Enabled by default as it is useful for debugging. Can be disabled via @@ -199,6 +203,7 @@ class ProfilerConfig: ("proton_mode", self.proton_mode, None), ("proton_hook", self.proton_hook, None), ("proton_output_format", self.proton_output_format, None), + ("proton_graph_attribution", self.proton_graph_attribution, False), ) if value != default ] @@ -220,6 +225,17 @@ class ProfilerConfig: if self.profiler == "proton": output_format = self.proton_output_format + if self.proton_graph_attribution and self.proton_data != "tree": + raise ValueError("proton_graph_attribution requires proton_data='tree'") + if ( + self.proton_graph_attribution + and self.proton_mode + and self.proton_mode.split(":", 1)[0].lower() == "periodic_flushing" + ): + raise ValueError( + "proton_graph_attribution is incompatible with periodic_flushing: " + "both manage Proton data phases." + ) if output_format == "chrome_trace" and self.proton_data != "trace": raise ValueError("chrome_trace output requires proton_data='trace'") if ( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 2fcc144af0..7d2baea26c 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1551,18 +1551,6 @@ class VllmConfig: self.compilation_config.mode = CompilationMode.NONE self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE - if self.profiler_config.profiler == "proton": - if not current_platform.is_cuda(): - raise ValueError( - "The Proton profiler currently supports NVIDIA CUDA only" - ) - if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: - raise ValueError( - "The Proton profiler requires CUDA graphs to be disabled. " - "Use --enforce-eager or set " - "--compilation-config.cudagraph_mode=none." - ) - if os.environ.get("TORCH_COMPILE_DISABLE") == "1": logger.warning_once( "TORCH_COMPILE_DISABLE is set, disabling torch.compile. " @@ -1860,6 +1848,7 @@ class VllmConfig: else: self._validate_v1_model_runner() + self._validate_profiler_config() self._validate_batch_sharded_sampling() self._validate_adaptive_verification() @@ -3042,6 +3031,44 @@ class VllmConfig: return unsupported + def _validate_profiler_config(self) -> None: + if self.profiler_config.profiler != "proton": + return + + from vllm.platforms import current_platform + + if not current_platform.is_cuda(): + raise ValueError("The Proton profiler currently supports NVIDIA CUDA only") + if ( + self.profiler_config.proton_graph_attribution + and not self.use_v2_model_runner + ): + raise ValueError( + "Proton CUDA graph attribution requires the V2 model runner." + ) + + has_cuda_graphs = ( + self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE + or ( + self.compilation_config.cudagraph_mm_encoder + and not (self.model_config and self.model_config.enforce_eager) + ) + ) + if not has_cuda_graphs: + return + mode = self.profiler_config.proton_mode + if mode and mode.split(":", 1)[0].lower() == "pcsampling": + raise ValueError( + "Proton PC sampling requires CUDA graphs to be disabled. " + "Use --enforce-eager." + ) + if not self.profiler_config.proton_graph_attribution: + raise ValueError( + "Proton profiling with CUDA graphs requires " + "proton_graph_attribution=True to capture replayed kernels. " + "Enable attribution or use --enforce-eager." + ) + def _validate_v2_model_runner(self) -> None: """Check for features not yet supported by the V2 model runner.""" if not HAS_TRITON: diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index fb89631986..abe8a34317 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -6,8 +6,9 @@ import inspect import json import os from abc import ABC, abstractmethod -from collections.abc import Callable -from contextlib import nullcontext +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress +from glob import glob from typing import Literal from uuid import uuid4 @@ -160,6 +161,15 @@ class WorkerProfiler(ABC): if self._running: self.stop() + @property + def has_cuda_graph_session(self) -> bool: + """Whether a capture-time session is retained to attribute replays.""" + return False + + def capture_cuda_graphs(self) -> AbstractContextManager[None]: + """Observe graph creation for backends that attribute replay activity.""" + return nullcontext() + def annotate_context_manager(self, name: str): """Return a context manager to annotate profiler traces.""" return nullcontext() @@ -387,6 +397,7 @@ class ProtonProfilerWrapper(WorkerProfiler): self._mode = profiler_config.proton_mode self._hook = profiler_config.proton_hook self._output_format = profiler_config.proton_output_format + self._graph_attribution = profiler_config.proton_graph_attribution self._triton_version_string = getattr(triton, "__version__", "unknown") try: self._triton_version = Version(self._triton_version_string) @@ -398,6 +409,13 @@ class ProtonProfilerWrapper(WorkerProfiler): # worker cannot overwrite profiles left by an earlier server process. self._instance_id = f"pid{os.getpid()}_{uuid4().hex}" self._run_id = 0 + self._graph_session = False + self._phase = 0 + self._active_output_path: str | None = None + self._session_storage_path = os.path.join( + self._output_dir, + f".proton_cuda_graph_session_{worker_name}_{self._instance_id}", + ) logger.info_once( "Proton profiling enabled. Output will be saved under: %s", @@ -412,6 +430,10 @@ class ProtonProfilerWrapper(WorkerProfiler): ) def _validate_capabilities(self) -> None: + if self._graph_attribution: + self._require_triton_version( + "CUDA graph attribution", _TRITON_PROTON_3_7_VERSION + ) if self._output_format is not None: parameters = inspect.signature(self._proton.finalize).parameters supports_output_format = "output_format" in parameters or any( @@ -428,7 +450,7 @@ class ProtonProfilerWrapper(WorkerProfiler): self._require_triton_version( "hatchet_msgpack output", _TRITON_PROTON_3_7_VERSION ) - if self._mode and self._mode.split(":", 1)[0] == "periodic_flushing": + if self._mode and self._mode.split(":", 1)[0].lower() == "periodic_flushing": self._require_triton_version( "periodic flushing", _TRITON_PROTON_3_7_VERSION ) @@ -447,28 +469,112 @@ class ProtonProfilerWrapper(WorkerProfiler): raise RuntimeError("Proton did not create a profiling session") return session_id + @property + def has_cuda_graph_session(self) -> bool: + return self._graph_session + + def set_output_name(self, worker_name: str) -> None: + """Set the next run's output name after startup graph capture.""" + if self._active: + return + self._output_path = os.path.join(self._output_dir, f"proton_{worker_name}") + + @override + @contextmanager + def capture_cuda_graphs(self) -> Iterator[None]: + """Keep a Proton session active while vLLM captures CUDA graphs.""" + if not self._graph_attribution: + yield + return + if self._session_id is None: + self._session_id = self._create_session(self._session_storage_path) + else: + self._proton.activate(session=self._session_id) + self._graph_session = True + + try: + yield + finally: + captured_phase = self._phase + try: + self._phase = self._proton.data.advance_phase(self._session_id) + finally: + self._proton.deactivate(session=self._session_id, flushing=True) + self._proton.data.clear(self._session_id, captured_phase) + @override def _start(self) -> None: - output_path = f"{self._output_path}_{self._instance_id}_run{self._run_id}" - self._session_id = self._create_session(output_path) + self._active_output_path = ( + f"{self._output_path}_{self._instance_id}_run{self._run_id}" + ) self._run_id += 1 + if self._graph_session: + assert self._session_id is not None + self._proton.activate(session=self._session_id) + else: + self._session_id = self._create_session(self._active_output_path) + + def _write_graph_phase(self, phase: int) -> None: + assert self._active_output_path is not None + output_format = self._output_format or "hatchet" + output_path = f"{self._active_output_path}.{output_format}" + if output_format == "hatchet_msgpack": + with open(output_path, "wb") as output_file: + output_file.write( + self._proton.data.get_msgpack(self._session_id, phase) + ) + else: + with open(output_path, "w", encoding="utf-8") as output_file: + json.dump(self._proton.data.get(self._session_id, phase), output_file) + + def _finalize_session(self, session_id: int) -> None: + if self._output_format is None: + self._proton.finalize(session=session_id) + else: + self._proton.finalize(session=session_id, output_format=self._output_format) @override def _stop(self) -> None: assert self._session_id is not None session_id = self._session_id + if self._graph_session: + completed_phase = self._phase + try: + self._phase = self._proton.data.advance_phase(session_id) + finally: + self._proton.deactivate(session=session_id, flushing=True) + try: + self._write_graph_phase(completed_phase) + finally: + self._proton.data.clear(session_id, completed_phase) + self._active_output_path = None + return + try: self._proton.deactivate(session=session_id) finally: try: - if self._output_format is None: - self._proton.finalize(session=session_id) - else: - self._proton.finalize( - session=session_id, output_format=self._output_format - ) + self._finalize_session(session_id) finally: self._session_id = None + self._active_output_path = None + + @override + def shutdown(self) -> None: + super().shutdown() + if self._graph_session and self._session_id is not None: + session_id = self._session_id + self._session_id = None + try: + self._finalize_session(session_id) + except Exception: + logger.exception( + "Failed to finalize Proton CUDA graph session during shutdown." + ) + finally: + for output_path in glob(f"{self._session_storage_path}.*"): + with suppress(FileNotFoundError): + os.remove(output_path) @override def annotate_context_manager(self, name: str): @@ -496,3 +602,24 @@ class CudaProfilerWrapper(WorkerProfiler): @override def annotate_context_manager(self, name: str): return torch.cuda.nvtx.range(name) + + +def create_graph_capture_profiler( + profiler_config: ProfilerConfig, global_rank: int +) -> WorkerProfiler | None: + """Create a profiler to observe CUDA graph capture, if configured. + + Applies only to backends that attribute graph replay activity which + need a session around capture (Proton with ``proton_graph_attribution``). + """ + if ( + profiler_config.profiler == "proton" + and profiler_config.proton_graph_attribution + ): + from vllm.distributed.utils import get_worker_rank_suffix + + return ProtonProfilerWrapper( + profiler_config, + worker_name=get_worker_rank_suffix(global_rank=global_rank), + ) + return None diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d66fc6c90f..4b1520f276 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -987,6 +987,16 @@ class GPUModelRunner(LoRAModelRunnerMixin): """Estimate the GPU memory required to capture CUDA graphs.""" return _profile_cudagraph_memory(self) + def needs_cudagraph_capture(self) -> bool: + """Whether capture_model() has any CUDA graphs left to capture.""" + return ( + self.model_state.supports_mm_inputs + and self.model_state.encoder_runner.has_cudagraph() + ) or ( + self.cudagraph_manager is not None + and self.cudagraph_manager.needs_capture() + ) + @torch.inference_mode() def capture_model(self, *, profile_only: bool = False) -> int: assert self.cudagraph_manager is not None diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 2e84fd511a..3f457b9b7b 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -9,7 +9,7 @@ from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import timedelta from types import NoneType -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np import regex as re @@ -64,6 +64,7 @@ from vllm.profiler.wrapper import ( CudaProfilerWrapper, ProtonProfilerWrapper, TorchProfilerWrapper, + create_graph_capture_profiler, ) from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask @@ -142,6 +143,7 @@ def maybe_rocm_profiling_fallback(profile_result: MemoryProfilingResult) -> int if TYPE_CHECKING: from vllm.device_allocator.sleep_mode_backend import SleepModeBackend from vllm.model_executor.model_loader.tensorizer import TensorizerConfig + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2 from vllm.v1.worker.gpu_model_runner import GPUModelRunner @@ -810,7 +812,8 @@ class Worker(WorkerBase): cuda_graph_memory_bytes = 0 if not self.model_config.enforce_eager: - cuda_graph_memory_bytes = self.model_runner.capture_model() + with self._get_cudagraph_capture_context(): + cuda_graph_memory_bytes = self.model_runner.capture_model() # Compare actual vs estimated CUDA graph memory (if we did profiling) if ( @@ -949,6 +952,21 @@ class Worker(WorkerBase): encoder=self.compilation_config.encoder_compilation_time, ) + def _get_cudagraph_capture_context(self) -> AbstractContextManager[None]: + """Let the configured profiler observe CUDA graph capture.""" + if not self.use_v2_model_runner: + return nullcontext() + if self.profiler is None: + model_runner = cast("GPUModelRunnerV2", self.model_runner) + if not model_runner.needs_cudagraph_capture(): + return nullcontext() + self.profiler = create_graph_capture_profiler( + self.profiler_config, global_rank=self.rank + ) + if self.profiler is None: + return nullcontext() + return self.profiler.capture_cuda_graphs() + def reset_mm_cache(self) -> None: self.model_runner.reset_mm_cache() @@ -1248,6 +1266,9 @@ class Worker(WorkerBase): else: trace_name = rank_suffix + if profiler_type == "proton" and self.profiler is not None: + self.profiler.set_output_name(trace_name) + # Create the profiler wrapper only on the first start call if self.profiler is None: if profiler_type == "torch": @@ -1284,7 +1305,9 @@ class Worker(WorkerBase): try: self.profiler.stop() finally: - if self.profiler_config.profiler == "proton": + if self.profiler_config.profiler == "proton" and not ( + self.profiler.has_cuda_graph_session + ): # Proton output names are fixed when the wrapper is constructed. # Recreate it so the next profile_prefix is honored. self.profiler = None diff --git a/vllm/v1/worker/mm_encoder_model_runner.py b/vllm/v1/worker/mm_encoder_model_runner.py index 8208a6ce3d..e19f1b332a 100644 --- a/vllm/v1/worker/mm_encoder_model_runner.py +++ b/vllm/v1/worker/mm_encoder_model_runner.py @@ -70,6 +70,9 @@ class MMEncoderModelRunner(GPUModelRunner): def capture_model(self, *, profile_only: bool = False) -> int: return 0 + def needs_cudagraph_capture(self) -> bool: + return False + def _dummy_run( self, *args: Any, **kwargs: Any ) -> tuple[torch.Tensor, torch.Tensor]: From 995e8581f462a13e32f30cfba946c63d48cf31d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Mon, 14 Sep 2026 21:16:39 +0200 Subject: [PATCH 044/139] [Model] Voxtral Realtime: add support for `CUDAGraphMode.FULL_DECODE_ONLY` (#51167) Signed-off-by: NickLucche Co-authored-by: Jorge Gallego Feliciano --- .../generation/test_voxtral_realtime.py | 106 ++++++++++------ .../model_executor/models/voxtral_realtime.py | 12 +- vllm/model_executor/models/whisper_causal.py | 118 +++++++++++++++--- 3 files changed, 174 insertions(+), 62 deletions(-) diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 0ef3c33aeb..1bd9f62303 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -14,7 +14,9 @@ from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy from vllm import LLM, SamplingParams from vllm.assets.audio import AudioAsset +from vllm.config import CUDAGraphMode from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.kv_cache_interface import SlidingWindowSpec @@ -130,6 +132,45 @@ async def async_engine(): wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) +def _build_inputs( + audio_assets: list[AudioAsset], tokenizer: MistralTokenizer +) -> tuple[list[dict], list[SamplingParams]]: + audio_config = tokenizer.instruct_tokenizer.tokenizer.audio + inputs = [] + sampling_params = [] + for audio_asset in audio_assets: + audio = Audio.from_file(audio_asset.get_local_path(), strict=False) + req = TranscriptionRequest( + audio=audio.to_base64(audio.format), + streaming=StreamingMode.OFFLINE, + language=None, + ) + tokenized = tokenizer.instruct_tokenizer.encode_transcription(req) + tokens = tokenized.tokens + audio_array = tokenized.audios[0].audio_array + + num_samples = audio_array.shape[0] + max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1 + sampling_params.append(SamplingParams(temperature=0.0, max_tokens=max_tokens)) + inputs.append( + { + "multi_modal_data": {"audio": [(audio_array, None)]}, + "prompt_token_ids": tokens, + } + ) + return inputs, sampling_params + + +def _assert_expected_text(outputs) -> None: + texts = _normalize([out.outputs[0].text for out in outputs]) + for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)): + assert got == expected, ( + f"Output mismatch at index {i}:\n" + f" got: {got!r}\n" + f" expected: {expected!r}" + ) + + def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch): monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") @@ -138,51 +179,38 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypa with vllm_runner(**vllm_kwargs) as vllm_model: assert_encoder_kv_cache_spec(vllm_model.llm) - audio_config = tokenizer.instruct_tokenizer.tokenizer.audio + inputs, sampling_params = _build_inputs(audio_assets, tokenizer) - def from_file(file_path: str): - audio = Audio.from_file(file_path, strict=False) - req = TranscriptionRequest( - audio=audio.to_base64(audio.format), - streaming=StreamingMode.OFFLINE, - language=None, - ) - tokenized = tokenizer.instruct_tokenizer.encode_transcription(req) + outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params) + _assert_expected_text(outputs) - return (tokenized.tokens, tokenized.audios[0].audio_array) - tokenized_list = [ - from_file(audio_asset.get_local_path()) for audio_asset in audio_assets - ] +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="Full cudagraphs require CUDA" +) +@pytest.mark.parametrize( + "cudagraph_mode", + [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], +) +def test_voxtral_realtime_cudagraph( + audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode +): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") - inputs = [] - sampling_params = [] + vllm_kwargs = {**ENGINE_CONFIG} + vllm_kwargs["model_name"] = vllm_kwargs.pop("model") + vllm_kwargs["enforce_eager"] = False + vllm_kwargs["compilation_config"] = { + "cudagraph_mode": cudagraph_mode, + # max_num_seqs is 4, so this is the only decode size worth capturing. + "cudagraph_capture_sizes": [4], + } - for tokens, audio_array in tokenized_list: - num_samples = audio_array.shape[0] - max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1 - sampling_params.append( - SamplingParams(temperature=0.0, max_tokens=max_tokens) - ) + with vllm_runner(**vllm_kwargs) as vllm_model: + inputs, sampling_params = _build_inputs(audio_assets, tokenizer) - input_dict = { - "multi_modal_data": {"audio": [(audio_array, None)]}, - "prompt_token_ids": tokens, - } - inputs.append(input_dict) - - outputs = vllm_model.llm.generate( - inputs, - sampling_params=sampling_params, - ) - - texts = _normalize([out.outputs[0].text for out in outputs]) - for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)): - assert got == expected, ( - f"Output mismatch at index {i}:\n" - f" got: {got!r}\n" - f" expected: {expected!r}" - ) + outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params) + _assert_expected_text(outputs) @pytest.mark.asyncio diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 738185a7ae..961aa3af9c 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -217,9 +217,15 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) - assert ( - not vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs() - ), "Voxtral realtime doesn't support full cudagraphs yet. Please use PIECEWISE." + # Full cudagraphs are supported for decode-only batches (the encoder's + # block-pooling attention builder is capture-safe for uniform + # single-token decode). Mixed/prefill batches are not, so block only + # pure FULL; FULL_DECODE_ONLY and FULL_AND_PIECEWISE are allowed. + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode + assert not cudagraph_mode.mixed_mode().has_full_cudagraphs(), ( + "Voxtral realtime supports full cudagraphs for decode-only batches. " + "Use cudagraph_mode=FULL_DECODE_ONLY or FULL_AND_PIECEWISE (not FULL)." + ) self.time_embedding: TimeEmbedding = TimeEmbedding( dim=self.config.text_config.hidden_size diff --git a/vllm/model_executor/models/whisper_causal.py b/vllm/model_executor/models/whisper_causal.py index 7004bba283..ec84cae8cf 100644 --- a/vllm/model_executor/models/whisper_causal.py +++ b/vllm/model_executor/models/whisper_causal.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import copy import functools import logging import math from dataclasses import replace from fractions import Fraction from functools import partial +from typing import ClassVar import torch import torch.nn.functional as F @@ -27,6 +27,7 @@ from vllm.model_executor.models.whisper import WhisperPosEmbedType from vllm.utils.math_utils import cdiv from vllm.v1.attention.backend import ( AttentionBackend, + AttentionCGSupport, AttentionMetadata, AttentionType, CommonAttentionMetadata, @@ -123,6 +124,16 @@ def create_whisper_attention_backend_with_block_pooling( underlying_impl = underlying_attn_backend.get_impl_cls() class WhisperCausalAttentionWithBlockPoolingBuilder(underlying_builder): # type: ignore + # Full cudagraphs only for uniform single-token decode: capture bakes in + # tensor addresses, so `build` writes metadata into persistent buffers + # (below). Prefill/mixed batches fall back to piecewise. + _cudagraph_support: ClassVar[AttentionCGSupport] = ( + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + ) + # Disable the update_block_table fast path: it would bypass block pooling + # by splicing the raw slot_mapping/block_table into cached metadata. + supports_update_block_table: bool = False + def __init__( self, kv_cache_spec: AttentionSpec, @@ -150,6 +161,38 @@ def create_whisper_attention_backend_with_block_pooling( # num_heads_q for the encoder is the same as num_kv_heads # (no GQA in whisper encoder) self.num_heads_q = kv_cache_spec.num_kv_heads + # Persistent capture-safe buffers for pooled metadata; lazily sized + # on first `build` to match the incoming tensors' dtype/device. + self._pool_buffers_ready = False + + def _maybe_init_pool_buffers( + self, common_attn_metadata: CommonAttentionMetadata + ) -> None: + if self._pool_buffers_ready: + return + device = common_attn_metadata.slot_mapping.device + sched = self.vllm_config.scheduler_config + self._qsl_buf = torch.empty( + sched.max_num_seqs + 1, + dtype=common_attn_metadata.query_start_loc.dtype, + device=device, + ) + self._seq_buf = torch.empty( + sched.max_num_seqs, + dtype=common_attn_metadata.seq_lens.dtype, + device=device, + ) + self._slot_buf = torch.empty( + sched.max_num_batched_tokens * block_pool_size, + dtype=common_attn_metadata.slot_mapping.dtype, + device=device, + ) + self._pool_arange = torch.arange( + block_pool_size, + dtype=common_attn_metadata.slot_mapping.dtype, + device=device, + ) + self._pool_buffers_ready = True def build( self, @@ -157,27 +200,62 @@ def create_whisper_attention_backend_with_block_pooling( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> AttentionMetadata: - new_common_attn_metadata = copy.deepcopy(common_attn_metadata) - new_common_attn_metadata.query_start_loc *= block_pool_size - new_common_attn_metadata.query_start_loc_cpu *= block_pool_size - new_common_attn_metadata.seq_lens *= block_pool_size - if new_common_attn_metadata.seq_lens_cpu_upper_bound is not None: - new_common_attn_metadata.seq_lens_cpu_upper_bound *= block_pool_size - new_common_attn_metadata.num_actual_tokens *= block_pool_size - new_common_attn_metadata.max_query_len *= block_pool_size - new_common_attn_metadata.max_seq_len *= block_pool_size - original_slot_mapping = common_attn_metadata.slot_mapping - common_prefix_len *= block_pool_size - new_common_attn_metadata.slot_mapping = ( - ( - original_slot_mapping.unsqueeze(1) * block_pool_size - + torch.arange(block_pool_size, device=original_slot_mapping.device) - ) - .flatten() - .clamp(min=-1) + cm = common_attn_metadata + self._maybe_init_pool_buffers(cm) + + # Scale query_start_loc / seq_lens into the persistent buffers. + n = cm.query_start_loc.numel() + assert n <= self._qsl_buf.numel(), ( + f"{n=} exceeds pooled query_start_loc buffer {self._qsl_buf.numel()}" + ) + qsl_buf = self._qsl_buf[:n] + qsl_buf.copy_(cm.query_start_loc) + qsl_buf.mul_(block_pool_size) + + r = cm.seq_lens.numel() + assert r <= self._seq_buf.numel(), ( + f"{r=} exceeds pooled seq_lens buffer {self._seq_buf.numel()}" + ) + seq_buf = self._seq_buf[:r] + seq_buf.copy_(cm.seq_lens) + seq_buf.mul_(block_pool_size) + + # Expand each slot into block_pool_size contiguous slots. Padding + # slots (-1) stay -1 after clamp so the KV write is skipped. + m = cm.slot_mapping.numel() + assert m * block_pool_size <= self._slot_buf.numel(), ( + f"pooled slot_mapping {m * block_pool_size} exceeds buffer " + f"{self._slot_buf.numel()}" + ) + slot_buf = self._slot_buf[: m * block_pool_size] + slot_buf.view(m, block_pool_size).copy_( + cm.slot_mapping.unsqueeze(1) * block_pool_size + self._pool_arange + ) + slot_buf.clamp_(min=-1) + + # CPU tensors aren't captured; fresh scaled copies avoid mutating + # the runner's shared buffers. + qsl_cpu = cm.query_start_loc_cpu * block_pool_size + seq_lens_cpu_upper_bound = ( + cm.seq_lens_cpu_upper_bound * block_pool_size + if cm.seq_lens_cpu_upper_bound is not None + else None + ) + + new_common_attn_metadata = cm.replace( + query_start_loc=qsl_buf, + query_start_loc_cpu=qsl_cpu, + seq_lens=seq_buf, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + num_actual_tokens=cm.num_actual_tokens * block_pool_size, + max_query_len=cm.max_query_len * block_pool_size, + max_seq_len=cm.max_seq_len * block_pool_size, + slot_mapping=slot_buf, ) return super().build( - common_prefix_len, new_common_attn_metadata, fast_build + common_prefix_len * block_pool_size, + new_common_attn_metadata, + fast_build, ) # NOTE: We need a custom impl so we can use the transformed slot_mapping From d4ee7fe7a98b3012f64309e18dd425e565069a01 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:24:18 +0300 Subject: [PATCH 045/139] [Quantization] Select linear backends per quantization (#51204) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- docs/features/quantization/README.md | 17 +++++ .../test_scaled_mm_kernel_selection.py | 70 +++++++++++++++++++ .../kernels/test_b12x_linear.py | 8 ++- vllm/config/kernel.py | 9 +++ .../model_executor/kernels/linear/__init__.py | 65 +++++++++++++---- .../kernels/linear/mxfp4/aiter.py | 2 +- 6 files changed, 154 insertions(+), 17 deletions(-) diff --git a/docs/features/quantization/README.md b/docs/features/quantization/README.md index 69ece36076..ba1a5b86fe 100644 --- a/docs/features/quantization/README.md +++ b/docs/features/quantization/README.md @@ -23,6 +23,23 @@ The following are the supported quantization formats for vLLM: - [TorchAO](torchao.md) - [FP8 ViT Encoder Attention](fp8_vit_attn.md) +## Selecting Linear Backends per Quantization + +`--linear-backend` selects one backend for all quantized linear layers. For +mixed-precision models that use more than one linear quantization scheme, use +`linear_backend_per_quant` to override the backend for individual schemes: + +```bash +vllm serve \ + --linear-backend cutlass \ + --kernel-config '{"linear_backend_per_quant":{"nvfp4_w4a16":"humming"}}' +``` + +Here, NVFP4 W4A16 linear layers use Humming, while all other quantized linear +layers use CUTLASS. Per-quantization overrides take precedence over +`--linear-backend`; schemes without an override continue to use the global +setting, including automatic selection when it is `auto`. + ## Supported Hardware The table below shows the compatibility of various quantization implementations with different hardware platforms in vLLM: diff --git a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py index bedebdb59b..8275714744 100644 --- a/tests/kernels/quantization/test_scaled_mm_kernel_selection.py +++ b/tests/kernels/quantization/test_scaled_mm_kernel_selection.py @@ -12,20 +12,90 @@ from unittest.mock import patch import pytest import torch +from vllm.config import KernelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.kernels.linear import ( AiterInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel, + HummingFP8ScaledMMLinearKernel, Int8ScaledMMLinearKernel, Int8ScaledMMLinearLayerConfig, ScaledMMLinearKernel, + _get_linear_backend, + _resolve_backend_kernels, + init_fp8_linear_kernel, init_int8_linear_kernel, register_linear_kernel, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, +) from vllm.platforms import PlatformEnum pytestmark = pytest.mark.cpu_test +def test_auto_linear_backend_is_unchanged_without_overrides(): + kernels = [object] + quantizations = ( + "fp8_block_w8a8", + "fp8_w8a8", + "int8_w8a8", + "mixed_precision", + "mxfp8", + "mxfp4", + "mxfp6", + "w8a16_fp8", + "nvfp4_w4a4", + "nvfp4_w4a16", + ) + config = VllmConfig(kernel_config=KernelConfig(linear_backend="auto")) + + with set_current_vllm_config(config): + for quantization in quantizations: + assert _get_linear_backend(quantization=quantization) == "auto" + assert ( + _resolve_backend_kernels(kernels, "test", quantization=quantization) + is kernels + ) + + +def test_linear_backend_override_is_quantization_specific(): + config = VllmConfig( + kernel_config=KernelConfig( + linear_backend="cutlass", + linear_backend_per_quant={"nvfp4_w4a16": "humming"}, + ) + ) + + with set_current_vllm_config(config): + assert _get_linear_backend(quantization="nvfp4_w4a16") == "humming" + assert _get_linear_backend(quantization="nvfp4_w4a4") == "cutlass" + assert _get_linear_backend(quantization="fp8_w8a8") == "cutlass" + + +@patch.object(HummingFP8ScaledMMLinearKernel, "is_supported", return_value=(True, None)) +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_fp8_linear_backend_override(platform_mock, _): + platform_mock._enum = PlatformEnum.CUDA + config = VllmConfig( + kernel_config=KernelConfig( + linear_backend="cutlass", + linear_backend_per_quant={"fp8_w8a8": "humming"}, + ) + ) + + with set_current_vllm_config(config): + kernel = init_fp8_linear_kernel( + activation_quant_key=kFp8StaticTensorSym, + weight_quant_key=kFp8StaticTensorSym, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + weight_shape=(128, 128), + ) + + assert type(kernel) is HummingFP8ScaledMMLinearKernel + + def test_is_supported_is_abstract(): """Test that is_supported() is properly defined as abstract.""" assert issubclass(ScaledMMLinearKernel, ABC) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 27368ead27..9d4f9beeb9 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -115,7 +115,9 @@ def test_b12x_backend_registration_priority_and_selection( assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after) monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") + monkeypatch.setattr( + linear_mod, "_get_linear_backend", lambda *, quantization: "b12x" + ) monkeypatch.setattr( kernel_cls, "is_supported", @@ -784,7 +786,9 @@ def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: import vllm.model_executor.kernels.linear as linear_mod monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") + monkeypatch.setattr( + linear_mod, "_get_linear_backend", lambda *, quantization: "b12x" + ) monkeypatch.setitem( linear_mod._POSSIBLE_NVFP4_KERNELS, PlatformEnum.CUDA, diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 4505e0cfd9..0ac5bf4391 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -326,6 +326,13 @@ class KernelConfig: - "xpu_woq": Use XPU kernels for weight-only quantization (e.g. W8A16) """ + linear_backend_per_quant: dict[str, LinearBackend] | None = Field( + default=None, min_length=1 + ) + """Backend overrides keyed by linear quantization scheme. Overrides take + precedence over ``linear_backend``; for example, + ``{"nvfp4_w4a16": "humming"}``.""" + @field_validator("moe_backend", mode="before") @classmethod def _normalize_moe_backend(cls, value: Any) -> Any: @@ -359,6 +366,8 @@ class KernelConfig: "enable_flashinfer_autotune", "ir_op_priority", # handled separately below } + if self.linear_backend_per_quant is None: + ignored_factors.add("linear_backend_per_quant") factors = get_hash_factors(self, ignored_factors) factors["ir_op_priority"] = self.ir_op_priority.compute_hash() return hash_factors(factors) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 1df73273ca..d3c41f97d8 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -235,14 +235,19 @@ from vllm.platforms import PlatformEnum, current_platform logger = init_logger(__name__) -def _get_linear_backend() -> str: +def _get_linear_backend(*, quantization: str) -> str: """Get the linear_backend setting from the current vllm config.""" from vllm.config import get_current_vllm_config_or_none - config = get_current_vllm_config_or_none() - if config is not None: - return config.kernel_config.linear_backend - return "auto" + if (config := get_current_vllm_config_or_none()) is None: + return "auto" + overrides = config.kernel_config.linear_backend_per_quant + if overrides is not None and (override := overrides.get(quantization)): + logger.info_once( + "Applied linear backend override for %r: %r", quantization, override + ) + return override + return config.kernel_config.linear_backend # Kernel classes covered by each --linear-backend value. @@ -370,6 +375,8 @@ def _filter_kernels_by_backend( def _resolve_backend_kernels( kernels: list[type], layer_desc: str, + *, + quantization: str, ) -> list[type]: """Apply --linear-backend filtering to one layer type's kernel list. @@ -380,7 +387,7 @@ def _resolve_backend_kernels( layer types (e.g. NVFP4 MoE projections next to FP8 attention projections). """ - linear_backend = _get_linear_backend() + linear_backend = _get_linear_backend(quantization=quantization) if linear_backend == "auto": return kernels filtered = _filter_kernels_by_backend(linear_backend, kernels) @@ -625,6 +632,8 @@ def choose_scaled_mm_linear_kernel( possible_kernels: dict[PlatformEnum, list[type[_KernelT]]], compute_capability: int | None = None, force_kernel: type[_KernelT] | None = None, + *, + quantization: str, ) -> type[_KernelT]: """ Choose a _KernelT that can implement the given config for the @@ -642,6 +651,7 @@ def choose_scaled_mm_linear_kernel( force_kernel (Optional[type[_KernelT]]): An Optional forced kernel to override the possible_kernels if it can be implemented. If None, it will only try the possible kernels. + quantization: Quantization scheme used to select a backend override. Raises: ValueError: If no kernel can implement the given config. @@ -668,7 +678,11 @@ def choose_scaled_mm_linear_kernel( platform_kernels = possible_kernels.get(current_platform._enum, []) # Apply --linear-backend filtering when set. - platform_kernels = _resolve_backend_kernels(platform_kernels, "scaled-mm") + platform_kernels = _resolve_backend_kernels( + platform_kernels, + "scaled-mm", + quantization=quantization, + ) for kernel in platform_kernels: is_supported_and_can_implement, failure_reason = ( @@ -705,6 +719,7 @@ def init_fp8_linear_kernel( kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_BLOCK_KERNELS, # type: ignore[misc] + quantization="fp8_block_w8a8", force_kernel=force_kernel, ) if module_name: @@ -736,6 +751,7 @@ def init_fp8_linear_kernel( kernel_type = choose_scaled_mm_linear_kernel( config=scaled_mm_linear_kernel_config, possible_kernels=_POSSIBLE_FP8_KERNELS, # type: ignore[arg-type] + quantization="fp8_w8a8", force_kernel=force_kernel, ) if module_name: @@ -772,6 +788,7 @@ def init_int8_linear_kernel( kernel_type = choose_scaled_mm_linear_kernel( config, _POSSIBLE_INT8_KERNELS, + quantization="int8_w8a8", ) logger.info_once( @@ -824,7 +841,11 @@ def choose_mp_linear_kernel( platform_kernels = _POSSIBLE_KERNELS.get(current_platform._enum, []) # Apply --linear-backend filtering when set. - platform_kernels = _resolve_backend_kernels(platform_kernels, "mixed-precision") + platform_kernels = _resolve_backend_kernels( + platform_kernels, + "mixed-precision", + quantization="mixed_precision", + ) failure_reasons = [] for kernel in platform_kernels: @@ -875,7 +896,11 @@ def init_mxfp8_linear_kernel(*, bmm_batch_size: int | None = None) -> Mxfp8Linea possible = list(_POSSIBLE_MXFP8_KERNELS.get(platform, [])) # Apply --linear-backend filtering when set. - possible = _resolve_backend_kernels(possible, "MXFP8") + possible = _resolve_backend_kernels( + possible, + "MXFP8", + quantization="mxfp8", + ) failure_reasons = [] for kernel_cls in possible: @@ -917,7 +942,11 @@ def init_mxfp4_linear_kernel( possible = list(_POSSIBLE_MXFP4_KERNELS.get(platform, [])) # Apply --linear-backend filtering when set. - possible = _resolve_backend_kernels(possible, "MXFP4") + possible = _resolve_backend_kernels( + possible, + "MXFP4", + quantization="mxfp4", + ) failure_reasons = [] for kernel_cls in possible: @@ -957,7 +986,7 @@ def init_mxfp6_linear_kernel( activation_quant_key=activation_quant_key, ) - linear_backend = _get_linear_backend() + linear_backend = _get_linear_backend(quantization="mxfp6") platform = current_platform._enum possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, [])) @@ -1017,7 +1046,10 @@ def init_wfp8_a16_linear_kernel( ) kernel_type = choose_scaled_mm_linear_kernel( - config, _POSSIBLE_WFP8A16_KERNELS, force_kernel=force_kernel + config, + _POSSIBLE_WFP8A16_KERNELS, + quantization="w8a16_fp8", + force_kernel=force_kernel, ) if module_name: @@ -1048,7 +1080,8 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: # batch-invariant CUTLASS implementation when available, otherwise fall # back to emulation. It overrides --linear-backend. force_kernel: type[NvFp4LinearKernel] | None = None - linear_backend = _get_linear_backend() + quantization = "nvfp4_w4a16" if use_a16 else "nvfp4_w4a4" + linear_backend = _get_linear_backend(quantization=quantization) if envs.VLLM_BATCH_INVARIANT: bi_supported, reason = CutlassNvFp4LinearKernel.is_supported() if bi_supported: @@ -1110,7 +1143,11 @@ def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel: possible = [kernel for kernel in possible if kernel in a16_kernels] # Apply --linear-backend filtering when set. - possible = _resolve_backend_kernels(possible, "NVFP4") + possible = _resolve_backend_kernels( + possible, + "NVFP4", + quantization=quantization, + ) failure_reasons = [] for kernel_cls in possible: diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py index 13fd5db275..420e7f9493 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -158,7 +158,7 @@ class AiterMxfp4LinearKernel(MxFp4LinearKernel): from vllm._aiter_ops import is_aiter_found_and_supported from vllm.model_executor.kernels.linear import _get_linear_backend - linear_backend = _get_linear_backend() + linear_backend = _get_linear_backend(quantization="mxfp4") if ( current_platform.is_rocm() From 7702ee87dba0d8eed7f201e77a0a6613aac738a8 Mon Sep 17 00:00:00 2001 From: sychen52 <41452870+sychen52@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:25:33 -0700 Subject: [PATCH 046/139] [Bugfix][DSA] Write nvfp4_ds_mla from the fused norm+rope kernel (#55538) Signed-off-by: Shiyang Chen --- .../nvfp4_ds_mla_cache_kernels.cu | 14 +- .../test_fused_deepseek_v32_norm_rope.py | 120 ++++++++++++++++++ vllm/models/deepseek_v32/attention.py | 4 +- vllm/models/deepseek_v32/common/kernels.py | 68 +++++++++- 4 files changed, 186 insertions(+), 20 deletions(-) diff --git a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu b/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu index c2faf7b760..b1f369e5ab 100644 --- a/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu +++ b/csrc/libtorch_stable/nvfp4_ds_mla_cache_kernels.cu @@ -109,9 +109,6 @@ __device__ __forceinline__ void nvfp4_unpack16_e2m1(uint64_t raw, #endif } -// Rounded UP to the next e4m3 value so amax / sf never exceeds the payload -// range: round-to-nearest can land up to 33% low near e4m3's 2^-9 floor, which -// would saturate the tile's largest values. __device__ __forceinline__ float nvfp4_tile_scale(const float* vals, float max_val, uint8_t* sf_out) { @@ -121,15 +118,8 @@ __device__ __forceinline__ float nvfp4_tile_scale(const float* vals, amax = fmaxf(amax, fabsf(vals[i])); } const float sf_f = fmaxf(amax / max_val, 0.001953125f); // 2^-9 - __nv_fp8_e4m3 sf8(sf_f); - uint8_t sf_bits = *reinterpret_cast(&sf8); - if (float(sf8) < sf_f && sf_bits < 0x7E) { - // Bump to the next e4m3 value (bit patterns of positive e4m3 values are - // monotonic; 0x7E = 448 is the max finite value). - sf_bits += 1; - sf8 = *reinterpret_cast(&sf_bits); - } - *sf_out = sf_bits; + const __nv_fp8_e4m3 sf8(sf_f); + *sf_out = *reinterpret_cast(&sf8); return float(sf8); } diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index fbe22af2f8..7f74cde05d 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -573,6 +573,126 @@ def test_fused_norm_rope_ds_mla(num_tokens: int): assert (topk == 7).all(), "topk buffer should be untouched (no indexer)" +E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + + +def quantize_to_e2m1(x: torch.Tensor) -> torch.Tensor: + """Round to nearest e2m1, saturating to +-6 like cvt.rn.satfinite.e2m1x2.f32.""" + mags = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x.device) + a = x.float().abs().clamp_max(6.0) + mids = (mags[:-1] + mags[1:]) / 2 + code = torch.bucketize(a, mids, right=True) + on_tie = (a.unsqueeze(-1) == mids).any(dim=-1) + tie_code = torch.bucketize(a, mids, right=False) + code = torch.where(on_tie, tie_code + (tie_code & 1), code) + return (torch.signbit(x).to(torch.uint8) << 3) | code.to(torch.uint8) + + +def nvfp4_sf_byte(s: torch.Tensor) -> torch.Tensor: + """Scale-factor byte permutation shared with FlashMLA: 8*(s&3) + (s>>2).""" + return 8 * (s % 4) + (s // 4) + + +@pytest.mark.skipif( + not current_platform.is_device_capability(100), + reason="nvfp4_ds_mla requires SM100 (Blackwell)", +) +@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) +def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int): + """nvfp4_ds_mla MLA cache layout (FlashMLA sparse, SM100 only). + + Per-token 352-byte entry: 256 B of 512 e2m1 NoPE packed 2/byte (low nibble + = even element) | 64 B unscaled e4m3 RoPE | 32 B byte-permuted e4m3 tile + scales, one per 16 NoPE elements. + """ + torch.manual_seed(5) + dev = "cuda" + max_pos = 8192 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos + + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16) + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16) + # Tile 0 is deliberately low-magnitude: amax/6 lands between two e4m3 + # subnormals (spaced a flat 2^-9 there), so round-to-nearest picks a scale + # BELOW amax/6 and the tile's peak saturates at +-6. Asserted below, so a + # future change to the scale rule cannot silently stop covering this. + # amax must land in (6, 9) * 2^-9 after rms_norm; 7.5 sits mid-band, with + # margin for the per-token RMS to vary. + kv_c[:, :16] = ( + 7.5 * (2.0**-9) * torch.linspace(0.2, 1.0, 16, device=dev, dtype=torch.bfloat16) + ) + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16) + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16) + kvw = torch.ones(KV_LORA, device=dev, dtype=torch.bfloat16) + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) + + bs = max_pos + mla_cache = torch.zeros(1, bs, 352, device=dev, dtype=torch.uint8) + slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) + + q_out = K.fused_norm_rope( + pos, + q_c, + qw, + EPS, + kv_c, + kvw, + EPS, + k_pe, + mla_cos_sin, + None, + None, + None, + EPS, + None, + topk, + slot_mapping=slot, + indexer_k_cache=None, + mla_kv_cache=mla_cache, + mla_kv_cache_dtype="nvfp4_ds_mla", + mla_k_scale=None, + has_indexer=False, + index_rope_interleave=False, + ) + + assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (nvfp4_ds_mla)") + + kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32 + kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64] + tiles = kv_ref.view(num_tokens, 32, 16) + amax = tiles.abs().amax(dim=-1) + scale_target = torch.clamp_min(amax / 6.0, 2.0**-9) + ref_scale = scale_target.to(FP8) # round-to-nearest + + # The crafted tile must actually exercise the saturating path. + assert (ref_scale[:, 0].float() < scale_target[:, 0]).all(), ( + "tile 0 should round its scale DOWN; adjust the crafted magnitude" + ) + assert (amax[:, 0] / ref_scale[:, 0].float() > 6.0).all(), ( + "tile 0 should saturate e2m1; adjust the crafted magnitude" + ) + + ref_codes = quantize_to_e2m1(tiles / ref_scale.float().unsqueeze(-1)) + ref_codes = ref_codes.reshape(num_tokens, KV_LORA) + + cache = mla_cache[0, :num_tokens] # [N, 352] uint8 + packed = cache[:, : KV_LORA // 2] + got_codes = torch.stack([packed & 0xF, packed >> 4], dim=-1) + got_codes = got_codes.reshape(num_tokens, KV_LORA) # low nibble = even elem + torch.testing.assert_close(got_codes, ref_codes, rtol=0, atol=0) + + got_rope = cache[:, KV_LORA // 2 : KV_LORA // 2 + ROPE_DIM].view(FP8) + assert_fp8(got_rope, kpe_ref.to(FP8), "nvfp4_ds_mla RoPE e4m3") + + perm = nvfp4_sf_byte(torch.arange(32, device=dev)) + got_scale = cache[:, KV_LORA // 2 + ROPE_DIM :].view(FP8)[:, perm] + torch.testing.assert_close(got_scale.float(), ref_scale.float(), rtol=0, atol=0) + + # No indexer on this call: top-k buffer must be untouched. + assert (topk == 7).all(), "topk buffer should be untouched (no indexer)" + + def test_fused_norm_rope_supports_large_token_count(): """Keep the token count off CUDA grid-y at its 65,536-block boundary.""" num_tokens = 65536 diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 00dd2e18a8..87cd618524 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -237,7 +237,9 @@ class DeepseekV32Attention(MLAAttention): fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype) self._fp8_query = fp8_attention and self.impl.supports_quant_query_input - self._fp8_kv_needs_view = fp8_attention and self.kv_cache_dtype != "fp8_ds_mla" + self._fp8_kv_needs_view = fp8_attention and not self.kv_cache_dtype.endswith( + "_ds_mla" + ) self._index_rope_interleave = getattr(config, "indexer_rope_interleave", False) diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index ef7fcde5d3..fd02a51076 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -150,6 +150,7 @@ def _fused_norm_rope_kernel( mla_cache_ds_scale_ptr, mla_cache_ds_rope_ptr, MLA_CACHE_DS_MLA: tl.constexpr, + MLA_CACHE_NVFP4: tl.constexpr, MLA_NUM_TILES: tl.constexpr, MLA_TILE_DIM: tl.constexpr, # Top k indices @@ -261,6 +262,56 @@ def _fused_norm_rope_kernel( mla_block_idx = slot_idx // mla_block_size mla_block_off = slot_idx % mla_block_size + if MLA_CACHE_NVFP4: + # nvfp4_ds_mla layout (352 B/token, KV_DIM == 512): + # [0, 256) 512 e2m1 NoPE, packed 2/byte (low nibble = even) + # [256, 320) 64 e4m3 RoPE, unscaled + # [320, 352) 32 e4m3 NoPE scales, byte-permuted + # Keep in lockstep with concat_and_cache_nvfp4_ds_mla (CUDA) and + # nvfp4_sf_byte() in FlashMLA: byte(s) = 8*(s&3) + (s>>2). + byte_base = ( + mla_block_idx * mla_cache_block_stride + + mla_block_off * mla_cache_entry_stride + ) + kv_2d = tl.reshape(kv_c.to(tl.float32), (MLA_NUM_TILES, MLA_TILE_DIM)) + amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True) + # sf = e4m3(max(amax/6, 2^-9)), round-to-nearest like + # cvt_warp_fp16_to_fp4. + sf8 = tl.maximum(amax * (1.0 / 6.0), 0.001953125).to(tl.float8e4nv) + q = kv_2d * (1.0 / sf8.to(tl.float32)) + qp = tl.reshape(q, (KV_DIM // 2, 2)) + sel = tl.arange(0, 2)[None, :] + even = tl.sum(tl.where(sel == 0, qp, 0.0), axis=1) + odd = tl.sum(tl.where(sel == 1, qp, 0.0), axis=1) + # Same instruction the CUDA writer uses, so the two are bit-exact + # by construction: dest byte = {high: odd, low: even}. + packed = tl.inline_asm_elementwise( + "{ .reg .b8 t; .reg .b32 w;" + " cvt.rn.satfinite.e2m1x2.f32 t, $2, $1;" + " mov.b32 w, {t, t, t, t}; and.b32 $0, w, 255; }", + "=r,f,f", + [even, odd], + dtype=tl.int32, + is_pure=True, + pack=1, + ).to(tl.uint8) + tl.store( + mla_cache_ptr + byte_base + tl.arange(0, KV_DIM // 2), + packed.to(tl.float8e4nv, bitcast=True), + ) + # RoPE: unscaled e4m3, interleaved like the fp8_ds_mla path + rope_base = mla_cache_ptr + byte_base + KV_DIM // 2 + tl.store(rope_base + dim_off * 2, r1.to(tl.float8e4nv)) + tl.store(rope_base + dim_off * 2 + 1, r2.to(tl.float8e4nv)) + # scales, permuted: byte(s) = 8*(s&3) + (s>>2) + sf_tile = tl.arange(0, MLA_NUM_TILES) + sf_perm = 8 * (sf_tile % 4) + (sf_tile // 4) + tl.store( + mla_cache_ptr + byte_base + (KV_DIM // 2 + 64) + sf_perm, + tl.reshape(sf8, (MLA_NUM_TILES,)), + ) + return + if MLA_CACHE_DS_MLA: # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element # tile of the NoPE is dynamically quantized to fp8 with its own @@ -491,7 +542,8 @@ def fused_norm_rope( idx_cache_stride = 0 # --- MLA KV cache setup --- - mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" + mla_cache_nvfp4 = mla_kv_cache_dtype == "nvfp4_ds_mla" + mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla" or mla_cache_nvfp4 mla_cache_fp8 = is_quantized_kv_cache(mla_kv_cache_dtype) and not mla_cache_ds_mla mla_num_tiles = 1 mla_ds_scale_view = torch.empty(0, dtype=torch.float32, device=device) @@ -499,11 +551,12 @@ def fused_norm_rope( if mla_kv_cache is not None: mla_block_size = mla_kv_cache.shape[1] if mla_cache_ds_mla: - # 656-byte custom layout addressed in bytes; mla_cache_ptr is the - # 1-byte fp8 view, so block/entry strides are byte offsets and the - # fp32/bf16 views share the same buffer. - assert kv_dim == 512, "fp8_ds_mla requires kv_lora_rank == 512" - mla_num_tiles = kv_dim // 128 + # Custom byte-addressed layout: 656 B/token for fp8_ds_mla, 352 B + # for nvfp4_ds_mla. mla_cache_ptr is the 1-byte fp8 view, so + # block/entry strides are byte offsets and the fp32/bf16 views + # (fp8_ds_mla only) share the same buffer. + assert kv_dim == 512, "ds_mla layouts require kv_lora_rank == 512" + mla_num_tiles = kv_dim // 16 if mla_cache_nvfp4 else kv_dim // 128 u8_cache = mla_kv_cache.view(torch.uint8) mla_block_stride = u8_cache.stride(0) mla_entry_stride = u8_cache.stride(1) @@ -596,7 +649,8 @@ def fused_norm_rope( mla_k_scale, mla_ds_scale_view, mla_ds_rope_view, - mla_cache_ds_mla, + mla_cache_ds_mla and not mla_cache_nvfp4, + mla_cache_nvfp4, mla_num_tiles, kv_dim // mla_num_tiles if mla_cache_ds_mla else 1, # Top k indices buffer From dabc4362b47ad2665b0802b0b42b13f660efe634 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:47:32 -0400 Subject: [PATCH 047/139] [ROCm][DSV4.1][Perf] Stride the DSA decode candidate mask over the live context (#56628) Signed-off-by: Fangzhou Ai Co-authored-by: Cursor --- .../test_rocm_aiter_candidate_mask.py | 226 ++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 120 +++++++++- 2 files changed, 344 insertions(+), 2 deletions(-) create mode 100644 tests/kernels/attention/test_rocm_aiter_candidate_mask.py diff --git a/tests/kernels/attention/test_rocm_aiter_candidate_mask.py b/tests/kernels/attention/test_rocm_aiter_candidate_mask.py new file mode 100644 index 0000000000..85e742a096 --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_candidate_mask.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm correctness tests for the strided DSA candidate mask. + +``_apply_candidate_mask_strided`` is a decode-path variant of the shared +``apply_candidate_mask``. It differs in two ways that need pinning down: + +* it launches a fixed number of programs that stride the column axis, rather + than one program per 1024-column tile, so at ``max_model_len`` widths the + grid no longer scales with a workspace that is mostly padding; +* it stops at each row's ``end`` instead of sanitizing the whole width. + +The second is only sound because every consumer bounds its scan by the same +ends -- ``top_k_per_row_decode`` takes ``rowEnd`` from ``seq_lens``. So the +contract is agreement with the shared kernel over ``[0, end)`` and nothing +beyond it, which is what these tests assert rather than comparing whole rows. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-specific tests" +) + +BLOCK_SIZE = 8 +TOPK_BLOCKS = 64 +# Above 128 tiles the grid saturates and programs start striding; below it the +# min() clamp applies instead. Both regimes need covering, so the widths here +# deliberately straddle 128 * 1024. +NARROW = 8192 +WIDE = 262144 + + +def _inputs(rows, width, ends, starts, *, block_size, col_stride, dtype): + """Logits, bounds and candidate blocks shared by both kernels.""" + torch.manual_seed(0) + # A column stride > 1 keeps the kernels honest about stride_col; the decode + # caller slices a workspace, so contiguous rows are not guaranteed. + logits = torch.randn(rows, width * col_stride, device="cuda")[:, ::col_stride] + logits[0, : min(16, width)] = 0.0 + if rows > 2: + logits[2, 5] = float("nan") + row_ke = torch.tensor(ends, device="cuda", dtype=dtype) + row_ks = torch.tensor(starts, device="cuda", dtype=dtype) if starts else None + nblocks = (width + block_size - 1) // block_size + candidates = torch.randint( + 0, nblocks, (rows, TOPK_BLOCKS), device="cuda", dtype=torch.int32 + ) + # -1 is the "no block" sentinel the selector emits for short rows. + candidates[:, ::7] = -1 + return logits, row_ks, row_ke, candidates + + +def _run_both(logits, row_ks, row_ke, candidates, block_size, row_repeat): + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask, + ) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _apply_candidate_mask_strided, + ) + + ref = logits.clone() + got = logits.clone() + apply_candidate_mask(ref, row_ks, row_ke, candidates, block_size, row_repeat) + _apply_candidate_mask_strided( + got, row_ks, row_ke, candidates, block_size, row_repeat + ) + return ref, got + + +def _assert_agrees_within_ends(ref, got, row_ke, row_repeat): + """Compare only the region a consumer can read.""" + for row in range(ref.shape[0]): + end = int(row_ke[row // row_repeat]) + if end <= 0: + continue + torch.testing.assert_close( + got[row, :end], + ref[row, :end], + rtol=0, + atol=0, + equal_nan=True, + msg=lambda m, r=row, e=end: f"row {r} differs within [0, {e}): {m}", + ) + + +@pytest.mark.parametrize("width", [NARROW, WIDE]) +@pytest.mark.parametrize("col_stride", [1, 2]) +def test_matches_shared_kernel_ragged_ends(width, col_stride): + """Ends on and around tile boundaries, plus the empty and full rows.""" + ends = [0, 1, 1023, 1024, 1025, width - 1, width, width // 2] + logits, row_ks, row_ke, candidates = _inputs( + len(ends), + width, + ends, + None, + block_size=BLOCK_SIZE, + col_stride=col_stride, + dtype=torch.int32, + ) + ref, got = _run_both(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1) + _assert_agrees_within_ends(ref, got, row_ke, 1) + + +@pytest.mark.parametrize("block_size", [8, 16]) +def test_matches_shared_kernel_with_starts(block_size): + """Non-zero starts: the decode caller passes zeros, prefill does not.""" + ends = [0, 64, 4096, 8192] + starts = [0, 7, 1024, 3] + logits, row_ks, row_ke, candidates = _inputs( + len(ends), + NARROW, + ends, + starts, + block_size=block_size, + col_stride=1, + dtype=torch.int64, + ) + ref, got = _run_both(logits, row_ks, row_ke, candidates, block_size, 1) + _assert_agrees_within_ends(ref, got, row_ke, 1) + + +def test_matches_shared_kernel_row_repeat(): + """Speculative decode: next_n logit rows share one bound.""" + row_repeat = 3 + ends = [0, 1025, WIDE] + logits, row_ks, row_ke, candidates = _inputs( + len(ends) * row_repeat, + WIDE, + ends, + None, + block_size=BLOCK_SIZE, + col_stride=1, + dtype=torch.int32, + ) + ref, got = _run_both(logits, row_ks, row_ke, candidates, BLOCK_SIZE, row_repeat) + _assert_agrees_within_ends(ref, got, row_ke, row_repeat) + + +def test_leaves_columns_past_end_untouched(): + """The deliberate difference from the shared kernel, pinned explicitly. + + Skipping the tail is where the speedup comes from, so a change that + quietly restores full-width sanitizing should fail here rather than just + get slower. Work stops at the tile boundary containing ``end``, not at + ``end`` itself, so that is the bound asserted. + """ + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _MASK_TILE, + _apply_candidate_mask_strided, + ) + + ends = [0, 1, 1025, WIDE // 2] + logits, row_ks, row_ke, candidates = _inputs( + len(ends), + WIDE, + ends, + None, + block_size=BLOCK_SIZE, + col_stride=1, + dtype=torch.int32, + ) + sentinel = 1.5 + logits.fill_(sentinel) + _apply_candidate_mask_strided(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1) + for row, end in enumerate(ends): + touched = min(-(-end // _MASK_TILE) * _MASK_TILE, WIDE) + tail = logits[row, touched:] + assert torch.equal(tail, torch.full_like(tail, sentinel)), ( + f"row {row} (end {end}) was written past column {touched}" + ) + # And the bound is tight: a row that ends early must not have cost a + # full-width pass, which is the whole point of the kernel. + assert touched <= end + _MASK_TILE + + +def test_cudagraph_replay_tracks_changing_ends(): + """The grid is static so a FULL capture stays valid when ends change. + + Capturing bakes the grid, so a data-dependent program count would freeze + whatever lengths were resident at capture. Replaying against longer and + shorter contexts than the captured ones is the property that matters. + """ + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + apply_candidate_mask, + ) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _apply_candidate_mask_strided, + ) + + rows = 4 + captured_ends = [1024, 4096, 8192, 2048] + logits, row_ks, row_ke, candidates = _inputs( + rows, + WIDE, + captured_ends, + None, + block_size=BLOCK_SIZE, + col_stride=1, + dtype=torch.int32, + ) + base = logits.clone() + + # Warm up outside the capture: Triton compiles on first call, and autotune + # launches would otherwise happen inside the graph. + _apply_candidate_mask_strided( + logits.clone(), row_ks, row_ke, candidates, BLOCK_SIZE, 1 + ) + torch.accelerator.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _apply_candidate_mask_strided(logits, row_ks, row_ke, candidates, BLOCK_SIZE, 1) + + for replay_ends in ([1, 131072, WIDE, 0], [WIDE, 1025, 3, 65536]): + row_ke.copy_(torch.tensor(replay_ends, device="cuda", dtype=row_ke.dtype)) + logits.copy_(base) + graph.replay() + torch.accelerator.synchronize() + + ref = base.clone() + apply_candidate_mask(ref, row_ks, row_ke, candidates, BLOCK_SIZE, 1) + _assert_agrees_within_ends(ref, logits, row_ke, 1) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 598baee0b0..44b6017a05 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -819,6 +819,123 @@ def rocm_fp8_mqa_logits( return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) +# Programs along the column axis of the decode mask grid. The shared kernel in +# model_executor/kernels/attention/dsa/candidate_blocks.py launches one program +# per 1024-column tile, which is 1024 of them per row at max_model_len 1048576 +# with only the first few doing work. This is a fixed count that strides +# instead, measured on gfx950; it is not a portable choice, which is why this +# variant lives here rather than replacing the shared one. +_MASK_GRID_COLS = 128 +# Columns each program handles per iteration. Work stops at the tile boundary +# containing a row's end rather than at the end itself; the overshoot is +# masked the same way the shared kernel masks it, so it costs a tile and +# changes nothing. +_MASK_TILE = 1024 + + +@triton.jit(do_not_specialize=["width", "nblocks"]) +def _mask_candidates_strided_kernel( + logits, + starts, + ends, + flags, + stride_row, + stride_col, + stride_start, + stride_end, + width, + nblocks, + BLOCK_SIZE: tl.constexpr, + HAS_STARTS: tl.constexpr, + ROW_REPEAT: tl.constexpr, + TILE: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + start = tl.load(starts + row // ROW_REPEAT * stride_start) if HAS_STARTS else 0 + end = tl.load(ends + row // ROW_REPEAT * stride_end) + edge = tl.load(flags + row * (nblocks + 1) + nblocks) + step = tl.num_programs(1) * TILE + tile_start = tl.program_id(1) * TILE + # The shared kernel also sanitizes columns at or past `end`. Nothing reads + # them: top_k_per_row_decode bounds its scan by the same row ends passed + # here, and persistent_topk/cooperative_topk clamp by the same lengths. + # Stopping at `end` is what makes the cost track the live context rather + # than max_model_len. + while tile_start < end: + cols = tile_start + tl.arange(0, TILE) + valid = (cols >= start) & (cols < end) & (cols < width) + block = (cols - start) // BLOCK_SIZE + keep = tl.load(flags + row * (nblocks + 1) + block, valid, other=0) + keep = (keep != 0) | ((cols == width - 1) & (edge != 0)) + tl.store( + logits + row * stride_row + cols * stride_col, + -float("inf"), + (cols < width) & ~(valid & keep), + ) + tile_start += step + + +def _apply_candidate_mask_strided( + logits: torch.Tensor, + row_ks: torch.Tensor | None, + row_ke: torch.Tensor, + candidate_blocks: torch.Tensor, + block_size: int, + row_repeat: int = 1, +) -> None: + """ROCm decode variant of ``apply_candidate_mask``. + + Same masking semantics over ``[0, end)``, but the grid is sized by a fixed + program count rather than by the logits width. Only worth using where the + width is the ``max_model_len`` workspace and the live context is far + shorter, i.e. the paged decode path below; the prefill chunks pass + chunk-sized logits and stay on the shared kernel. + """ + from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( + _candidate_flags_kernel, + ) + + rows, width = logits.shape + if not rows or not width: + return + nblocks = triton.cdiv(width, block_size) + flags = torch.empty((rows, nblocks + 1), device=logits.device, dtype=torch.uint8) + start_stride = row_ks.stride(0) if row_ks is not None else 0 + _candidate_flags_kernel[(rows,)]( + candidate_blocks, + row_ks, + flags, + *candidate_blocks.stride(), + start_stride, + width, + nblocks, + block_size, + candidate_blocks.shape[1], + row_ks is not None, + row_repeat, + ) + # Derived from width, which is a tensor shape, so the grid stays static and + # a FULL cudagraph capture remains valid across replays; only the loop trip + # count inside the kernel is data-dependent. The min keeps narrow widths + # from launching programs that would only fall through. + grid_cols = min(_MASK_GRID_COLS, triton.cdiv(width, _MASK_TILE)) + _mask_candidates_strided_kernel[(rows, grid_cols)]( + logits, + row_ks, + row_ke, + flags, + *logits.stride(), + start_stride, + row_ke.stride(0), + width, + nblocks, + block_size, + row_ks is not None, + row_repeat, + _MASK_TILE, + ) + + def _max_decode_logits_rows(num_batched_tokens: int) -> int: """Upper bound on decode rows the paged-MQA logits buffer can ever hold. @@ -1102,7 +1219,6 @@ def rocm_aiter_sparse_attn_indexer( if candidate_blocks is not None: from vllm.model_executor.layers.sparse_attn_indexer import ( - _apply_candidate_mask, _select_candidate_blocks, ) @@ -1123,7 +1239,7 @@ def rocm_aiter_sparse_attn_indexer( decode_candidates, ) else: - _apply_candidate_mask( + _apply_candidate_mask_strided( logits, row_starts, visible, From d392ac836e17e87f5fec6594ef5ee5b56db72a91 Mon Sep 17 00:00:00 2001 From: Michele Campi <215741962+MicheleCampi@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:15:28 +0200 Subject: [PATCH 048/139] [Doc][Metrics] Fix spec-decode PromQL examples to use the exposed names (#55016) Signed-off-by: Michele Campi --- vllm/v1/spec_decode/metrics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/v1/spec_decode/metrics.py b/vllm/v1/spec_decode/metrics.py index a3ccfb29e7..826dfbd16a 100644 --- a/vllm/v1/spec_decode/metrics.py +++ b/vllm/v1/spec_decode/metrics.py @@ -187,12 +187,12 @@ class SpecDecodingProm: 1 + ( rate(vllm:spec_decode_num_accepted_tokens_total[$interval]) / - rate(vllm:spec_decode_num_drafts[$interval])) + rate(vllm:spec_decode_num_drafts_total[$interval])) A per-position acceptance rate vector can be computed using - vllm:spec_decode_num_accepted_tokens_per_pos[$interval] / - vllm:spec_decode_num_drafts[$interval] + vllm:spec_decode_num_accepted_tokens_per_pos_total[$interval] / + vllm:spec_decode_num_drafts_total[$interval] """ _counter_cls = prometheus_client.Counter From d320a214316b52b08cb84a4745475370a1564e87 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 14:37:44 -0700 Subject: [PATCH 049/139] [Bugfix][Rust Frontend] Restore prost test dependency (#56866) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock --- rust/Cargo.lock | 1 + rust/src/server/Cargo.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 80edfe4ebc..1b7dae6568 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5972,6 +5972,7 @@ dependencies = [ "libc", "llm-multimodal", "openssl", + "prost", "prost-types", "rmp-serde", "rmpv", diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index b9abe19c63..caccbdefc5 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -59,6 +59,7 @@ async-openai = { workspace = true, features = ["full"] } bytes.workspace = true clap.workspace = true expect-test.workspace = true +prost.workspace = true serial_test.workspace = true tempfile.workspace = true tokio = { workspace = true, features = ["test-util"] } From bdad63c90aec1c97682b20f282ba2a959f8229c9 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 14 Sep 2026 16:40:35 -0500 Subject: [PATCH 050/139] [CI] Check target branch freshness before starting CI (#56169) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .github/workflows/run-ci-command.yml | 6 + .github/workflows/scripts/run_ci_command.py | 104 ++++++++-- .../workflows/scripts/test_run_ci_command.py | 196 ++++++++++++++++++ docs/contributing/README.md | 28 +++ 4 files changed, 322 insertions(+), 12 deletions(-) diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml index 747831c7d4..3435148ea4 100644 --- a/.github/workflows/run-ci-command.yml +++ b/.github/workflows/run-ci-command.yml @@ -21,11 +21,17 @@ jobs: (github.event.comment.body == '/ci run' || github.event.comment.body == '/ci run all' || github.event.comment.body == '/ci run nightly' || + github.event.comment.body == '/ci run --allow-stale' || + github.event.comment.body == '/ci run all --allow-stale' || + github.event.comment.body == '/ci run nightly --allow-stale' || github.event.comment.body == '/ci retry' || github.event.comment.body == '/ci cancel' || github.event.comment.body == '/amd-ci run' || github.event.comment.body == '/amd-ci run all' || github.event.comment.body == '/amd-ci run nightly' || + github.event.comment.body == '/amd-ci run --allow-stale' || + github.event.comment.body == '/amd-ci run all --allow-stale' || + github.event.comment.body == '/amd-ci run nightly --allow-stale' || github.event.comment.body == '/amd-ci retry' || github.event.comment.body == '/amd-ci cancel') runs-on: [self-hosted, linux, x64, vllm-runners] diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py index 6454f041ff..530bd697d0 100644 --- a/.github/workflows/scripts/run_ci_command.py +++ b/.github/workflows/scripts/run_ci_command.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import http.client import json import os import random @@ -22,6 +23,7 @@ COMMAND_RUN_AMD_CI_ALL = "/amd-ci run all" COMMAND_RUN_AMD_CI_NIGHTLY = "/amd-ci run nightly" COMMAND_RETRY_AMD_FAILED = "/amd-ci retry" COMMAND_CANCEL_AMD_CI = "/amd-ci cancel" +ALLOW_STALE_SUFFIX = " --allow-stale" RUN_CI_COMMAND_ENV = { COMMAND_RUN_CI: {}, COMMAND_RUN_CI_ALL: {"RUN_ALL": "1"}, @@ -30,20 +32,25 @@ RUN_CI_COMMAND_ENV = { COMMAND_RUN_AMD_CI_ALL: {"RUN_ALL": "1"}, COMMAND_RUN_AMD_CI_NIGHTLY: {"RUN_ALL": "1", "NIGHTLY": "1"}, } +ALLOW_STALE_COMMANDS = frozenset( + f"{command}{ALLOW_STALE_SUFFIX}" for command in RUN_CI_COMMAND_ENV +) +RUN_CI_COMMAND_ENV.update( + { + command: RUN_CI_COMMAND_ENV[command.removesuffix(ALLOW_STALE_SUFFIX)] + for command in ALLOW_STALE_COMMANDS + } +) UPSTREAM_CI_COMMANDS = frozenset( { - COMMAND_RUN_CI, - COMMAND_RUN_CI_ALL, - COMMAND_RUN_CI_NIGHTLY, + *(command for command in RUN_CI_COMMAND_ENV if command.startswith("/ci ")), COMMAND_RETRY_FAILED, COMMAND_CANCEL_CI, } ) AMD_CI_COMMANDS = frozenset( { - COMMAND_RUN_AMD_CI, - COMMAND_RUN_AMD_CI_ALL, - COMMAND_RUN_AMD_CI_NIGHTLY, + *(command for command in RUN_CI_COMMAND_ENV if command.startswith("/amd-ci ")), COMMAND_RETRY_AMD_FAILED, COMMAND_CANCEL_AMD_CI, } @@ -79,6 +86,10 @@ class ApiError(RuntimeError): self.status = status +class CiPreparationError(RuntimeError): + pass + + def rate_limit_jitter() -> float: return random.uniform(1, 5) @@ -137,6 +148,8 @@ class HttpTransport: None, f"API request failed: {error.reason}", ) from error + except (OSError, http.client.HTTPException) as error: + raise ApiError(None, f"API request failed: {error}") from error if not response_body: return None @@ -234,6 +247,17 @@ class GitHubClient: def get_pr(self, number: int) -> dict[str, Any]: return self._request(self._repo_path(f"/pulls/{number}")) + def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int: + base_ref = urllib.parse.quote(f"refs/heads/{base_ref}", safe="") + head_sha = urllib.parse.quote(head_sha, safe="") + response = self._request( + self._repo_path(f"/compare/{base_ref}...{head_sha}?per_page=1") + ) + behind = response.get("behind_by") if isinstance(response, dict) else None + if type(behind) is not int or behind < 0: + raise ApiError(None, "GitHub returned an invalid commits-behind count.") + return behind + def list_pulls_for_commit(self, commit: str) -> list[dict[str, Any]]: commit = urllib.parse.quote(commit, safe="") response = self._request(self._repo_path(f"/commits/{commit}/pulls")) @@ -747,6 +771,10 @@ def notify_authorized( ( f"✅ @{author}, CI is now available for this PR.\n\n" "- `/ci run` starts upstream CI; `/amd-ci run` starts AMD CI only.\n" + "- Your branch must contain every commit currently on its upstream " + "target branch. Merge or rebase onto the latest target branch, then " + "rerun the command. Append `--allow-stale` to a run command to test " + "an outdated branch at your own risk.\n" "- `/ci retry` retries failed jobs in the CI build for the current " "PR head. If the current head has no CI build, it starts a new CI " "build for the current head containing only jobs that failed in " @@ -794,6 +822,56 @@ def resolve_workflow_run_pr( return find_matching_pr(github.list_pulls_for_commit(head_sha)) +def prepare_pr_for_ci( + github: GitHubClient, + pr: Mapping[str, Any], + command: str, +) -> tuple[dict[str, Any], str]: + base_ref = pr["base"]["ref"] + no_action = "No new CI build was started." + try: + behind = github.get_commits_behind_base(base_ref, pr["head"]["sha"]) + current_pr = github.get_pr(pr["number"]) + except ApiError as error: + raise CiPreparationError( + f"Could not check the PR against upstream `{base_ref}`. " + f"{no_action} Comment `{command}` again. {error}" + ) from error + if ( + current_pr["state"] != "open" + or current_pr["head"]["sha"] != pr["head"]["sha"] + or current_pr["base"]["ref"] != pr["base"]["ref"] + ): + raise CiPreparationError( + f"The PR changed while checking its branch. {no_action} " + f"Comment `{command}` again." + ) + commits = "commit" if behind == 1 else "commits" + lag = f"This PR is {behind} {commits} behind upstream `{base_ref}`." + if command in ALLOW_STALE_COMMANDS: + warning = "" + if behind: + run_command = command.removesuffix(ALLOW_STALE_SUFFIX) + warning = ( + f"⚠️ {lag} Running CI at your own risk because " + "`--allow-stale` was requested; outdated CI configuration may " + "cause failures. Before merging, merge or rebase onto the latest " + f"`{base_ref}`, then rerun `{run_command}` on the latest PR commit." + ) + print(warning) + return current_pr, warning + + if behind: + raise CiPreparationError( + f"{lag} Your branch must contain every commit currently on upstream " + f"`{base_ref}`. {no_action} " + f"Merge or rebase onto the latest `{base_ref}`, then rerun `{command}`." + " To test this branch at your own risk, " + f"use `{command}{ALLOW_STALE_SUFFIX}`." + ) + return current_pr, "" + + def handle_run_ci( *, actor: str, @@ -828,12 +906,7 @@ def handle_run_ci( f"{ci_name} is already running for this commit: {active_build['web_url']}" ) - current_pr = github.get_pr(pr["number"]) - if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]: - return ( - "The PR head changed while processing the command. " - f"Comment `{command}` again." - ) + current_pr, warning = prepare_pr_for_ci(github, pr, command) build = buildkite.create_build( create_build_payload( @@ -846,6 +919,7 @@ def handle_run_ci( return ( f"Triggered [Buildkite {ci_name} #{build['number']}]({build['web_url']}) " f"for commit `{current_pr['head']['sha'][:12]}`." + + (f"\n\n{warning}" if warning else "") ) @@ -1099,6 +1173,12 @@ def run( raise ValueError(f"Unsupported CI command: {command}") add_reaction_safely(github, comment_id, "rocket") github.add_comment(issue_number, f"✅ {message}") + except CiPreparationError as error: + add_reaction_safely(github, comment_id, "-1") + github.add_comment( + issue_number, + f"❌ {error}\n\n{command_comment_marker(comment_id)}", + ) except Exception: add_reaction_safely(github, comment_id, "confused") raise diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py index ef93ea8559..9750cefaa3 100644 --- a/.github/workflows/scripts/test_run_ci_command.py +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -22,8 +22,10 @@ from run_ci_command import ( COMMAND_RUN_CI_ALL, COMMAND_RUN_CI_NIGHTLY, RETRY_STATES, + RUN_CI_COMMAND_ENV, ApiError, BuildkiteClient, + GitHubClient, HttpTransport, authorize, create_build_payload, @@ -85,6 +87,7 @@ class FakeGitHub: comments: list[str] | None = None, pulls_for_commit: list[dict[str, Any]] | None = None, prs: dict[int, dict[str, Any]] | None = None, + behind: int = 0, ) -> None: self.comments = comments or [] self.permission = permission @@ -95,6 +98,12 @@ class FakeGitHub: self.reviews = reviews or [] self.pulls_for_commit = pulls_for_commit or [] self.prs = prs or {self.pr["number"]: self.pr} + self.behind = behind + self.lag_queries: list[tuple[str, str]] = [] + + def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int: + self.lag_queries.append((base_ref, head_sha)) + return self.behind def get_pr(self, number: int) -> dict[str, Any]: try: @@ -322,6 +331,11 @@ class RunCiCommandTest(unittest.TestCase): for command in commands: with self.subTest(command=command): self.assertEqual(parse_command(command), command) + override = f"{command} --allow-stale" + if " run" in command: + self.assertEqual(parse_command(override), override) + else: + self.assertIsNone(parse_command(override)) self.assertIsNone(parse_command("/ci run please")) self.assertIsNone(parse_command("/ci run all please")) @@ -331,6 +345,8 @@ class RunCiCommandTest(unittest.TestCase): self.assertIsNone(parse_command("/amd-ci retry ")) self.assertIsNone(parse_command("/AMD-CI run")) self.assertIsNone(parse_command("/amdci run")) + self.assertIsNone(parse_command("/ci run --allow-stale all")) + self.assertIsNone(parse_command("/amd-ci run --allow-stale please")) def test_commands_select_only_their_configured_pipeline(self) -> None: cases = ( @@ -530,6 +546,177 @@ class RunCiCommandTest(unittest.TestCase): self.assertTrue(github.comments[0].startswith("✅ ")) self.assertIn("Buildkite CI #123", github.comments[0]) + def test_run_variants_require_current_base_without_changing_commit(self) -> None: + for command in RUN_CI_COMMAND_ENV: + if command.endswith(" --allow-stale"): + continue + for behind in (0, 1, 5): + with self.subTest(command=command, behind=behind): + github = FakeGitHub( + behind=behind, pr=make_pr(base={"ref": "release/1.0"}) + ) + buildkite = FakeBuildkite([[], []]) + run(make_event(command), github, buildkite) + + self.assertEqual( + github.lag_queries, + [("release/1.0", github.pr["head"]["sha"])], + ) + if behind == 0: + self.assertEqual( + buildkite.created_builds[0]["commit"], + github.pr["head"]["sha"], + ) + self.assertIn( + "](https://buildkite.example/builds/123)", + github.comments[0], + ) + else: + self.assertEqual(buildkite.created_builds, []) + self.assertTrue(github.comments[0].startswith("❌ ")) + self.assertIn(str(behind), github.comments[0]) + self.assertIn("upstream `release/1.0`", github.comments[0]) + self.assertIn("every commit", github.comments[0]) + self.assertIn( + "No new CI build was started.", github.comments[0] + ) + self.assertIn(command, github.comments[0]) + + def test_allow_stale_preserves_author_permissions_and_run_configuration( + self, + ) -> None: + cases = ( + (COMMAND_RUN_CI, "ci", {}), + (COMMAND_RUN_CI_ALL, "ci", {"RUN_ALL": "1"}), + (COMMAND_RUN_CI_NIGHTLY, "ci", {"RUN_ALL": "1", "NIGHTLY": "1"}), + (COMMAND_RUN_AMD_CI, "amd-ci", {}), + (COMMAND_RUN_AMD_CI_ALL, "amd-ci", {"RUN_ALL": "1"}), + ( + COMMAND_RUN_AMD_CI_NIGHTLY, + "amd-ci", + {"RUN_ALL": "1", "NIGHTLY": "1"}, + ), + ) + for original_command, pipeline, expected_env in cases: + command = f"{original_command} --allow-stale" + for ready, behind in ((False, 100), (True, 0), (True, 100)): + with self.subTest(command=command, ready=ready, behind=behind): + github = FakeGitHub( + permission="read", + behind=behind, + pr=make_pr( + base={"ref": "feature/stack"}, + labels=[{"name": "ready"}] if ready else [], + ), + ) + buildkite = FakeBuildkite([[], []]) + + run(make_event(command, "author"), github, buildkite) + + self.assertEqual(pipeline_for_command(command), pipeline) + if not ready: + self.assertEqual(buildkite.created_builds, []) + self.assertEqual(buildkite.list_calls, []) + self.assertEqual(github.lag_queries, []) + self.assertTrue(github.comments[0].startswith("❌ ")) + continue + self.assertEqual(len(buildkite.created_builds), 1) + payload = buildkite.created_builds[0] + self.assertEqual(payload["commit"], github.pr["head"]["sha"]) + self.assertEqual( + payload["pull_request_base_branch"], "feature/stack" + ) + self.assertEqual(payload["message"], f"PR #42 {command} by @author") + self.assertEqual( + payload["env"], + { + "VLLM_CI_GITHUB_COMMENT_ID": "99", + "VLLM_CI_TRIGGERED_BY": "author", + **expected_env, + }, + ) + self.assertEqual( + github.lag_queries, + [("feature/stack", github.pr["head"]["sha"])], + ) + self.assertIn( + "Buildkite AMD CI #123" + if pipeline == "amd-ci" + else "Buildkite CI #123", + github.comments[0], + ) + if behind: + self.assertIn("100 commits behind", github.comments[0]) + self.assertIn("upstream `feature/stack`", github.comments[0]) + self.assertIn("Before merging", github.comments[0]) + self.assertIn("latest PR commit", github.comments[0]) + self.assertIn(original_command, github.comments[0]) + + def test_base_lag_query_qualifies_branch_refs_and_validates_count(self) -> None: + for base_ref, encoded_ref in ( + ("main", "main"), + ("release/1.0", "release%2F1.0"), + ("feature/stack", "feature%2Fstack"), + ): + for count in (0, 5, None, -1, True, "5"): + with self.subTest(base_ref=base_ref, count=count): + transport = FakeTransport( + {} if count is None else {"behind_by": count} + ) + github = GitHubClient("token", "vllm-project/vllm", transport) + if type(count) is int and count >= 0: + self.assertEqual( + github.get_commits_behind_base(base_ref, "head-sha"), count + ) + else: + with self.assertRaises(ApiError): + github.get_commits_behind_base(base_ref, "head-sha") + self.assertEqual( + transport.calls[0]["url"], + "https://api.github.com/repos/vllm-project/vllm/compare/" + f"refs%2Fheads%2F{encoded_ref}...head-sha?per_page=1", + ) + self.assertEqual(transport.calls[0]["method"], "GET") + + def test_unknown_lag_or_changed_pr_prevents_runs(self) -> None: + for command in ( + COMMAND_RUN_CI, + COMMAND_RUN_AMD_CI, + f"{COMMAND_RUN_CI} --allow-stale", + f"{COMMAND_RUN_AMD_CI} --allow-stale", + ): + for fault in ("api", "head", "closed", "base"): + with self.subTest(command=command, fault=fault): + github = FakeGitHub() + buildkite = FakeBuildkite([[], []]) + + def check( + base_ref: str, + sha: str, + fault: str = fault, + github: FakeGitHub = github, + ) -> int: + if fault == "api": + raise ApiError(503, "Unavailable") + changed = make_pr() + if fault == "head": + changed["head"]["sha"] = "another-head" + elif fault == "closed": + changed["state"] = "closed" + else: + changed["base"]["ref"] = "another-base" + github.prs[42] = changed + return 0 + + with patch.object( + github, "get_commits_behind_base", side_effect=check + ): + run(make_event(command), github, buildkite) + self.assertEqual(buildkite.created_builds, []) + self.assertEqual(buildkite.retry_calls, []) + self.assertTrue(github.comments[0].startswith("❌ ")) + self.assertNotIn("Triggered", github.comments[0]) + def test_amd_run_ignores_blocked_builds(self) -> None: for metadata in ({}, {"github-comment-id": "98"}): with self.subTest(metadata=metadata): @@ -627,6 +814,7 @@ class RunCiCommandTest(unittest.TestCase): run(make_event(COMMAND_RUN_CI, "author"), github, buildkite) self.assertEqual(buildkite.list_calls, []) + self.assertEqual(github.lag_queries, []) self.assertEqual(github.reactions, ["eyes"]) self.assertTrue(github.comments[0].startswith("❌ ")) self.assertIn("approve the PR", github.comments[0]) @@ -848,6 +1036,7 @@ class RunCiCommandTest(unittest.TestCase): def test_ci_retry_retries_failed_jobs_while_build_is_running(self) -> None: github = FakeGitHub( + behind=100, permission="read", pr=make_pr(labels=[{"name": "ready"}]), ) @@ -867,10 +1056,12 @@ class RunCiCommandTest(unittest.TestCase): run(make_event(COMMAND_RETRY_FAILED, "author"), github, buildkite) self.assertEqual(buildkite.retry_calls, [(123, RETRY_STATES)]) + self.assertEqual(github.lag_queries, []) self.assertIn("Queued 3 failed job", github.comments[0]) def test_amd_ci_retry_retries_only_the_current_head_build(self) -> None: github = FakeGitHub( + behind=100, permission="read", pr=make_pr(labels=[{"name": "ready"}]), ) @@ -891,6 +1082,7 @@ class RunCiCommandTest(unittest.TestCase): run(make_event(COMMAND_RETRY_AMD_FAILED, "author"), github, buildkite) self.assertEqual(buildkite.retry_calls, [(321, RETRY_STATES)]) + self.assertEqual(github.lag_queries, []) self.assertIn("Buildkite AMD CI #321", github.comments[0]) def test_amd_ci_retry_requires_a_build_for_the_current_head(self) -> None: @@ -912,6 +1104,7 @@ class RunCiCommandTest(unittest.TestCase): def test_ci_retry_creates_filtered_build_for_new_head(self) -> None: github = FakeGitHub( + behind=100, permission="read", pr=make_pr(labels=[{"name": "ready"}]), ) @@ -953,6 +1146,7 @@ class RunCiCommandTest(unittest.TestCase): self.assertEqual(len(buildkite.created_builds), 1) payload = buildkite.created_builds[0] self.assertEqual(payload["commit"], "0123456789abcdef") + self.assertEqual(github.lag_queries, []) self.assertEqual(payload["message"], "PR #42 /ci retry by @author") self.assertEqual( json.loads(payload["env"]["VLLM_CI_ONLY_STEP_KEYS"]), @@ -1104,6 +1298,7 @@ class RunCiCommandTest(unittest.TestCase): run(make_event(COMMAND_CANCEL_CI), github, buildkite) self.assertEqual(buildkite.cancel_calls, [123, 124, 126]) + self.assertEqual(github.lag_queries, []) self.assertEqual( [request["branch"] for request in buildkite.list_requests], ["contributor:feature", "feature"], @@ -1152,6 +1347,7 @@ class RunCiCommandTest(unittest.TestCase): run(make_event(COMMAND_CANCEL_AMD_CI), github, buildkite) self.assertEqual(buildkite.cancel_calls, [322, 321]) + self.assertEqual(github.lag_queries, []) self.assertEqual( [request["branch"] for request in buildkite.list_requests], ["contributor:feature", "feature"], diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 55966f2854..7a1af9c1a7 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -315,6 +315,34 @@ review process: the `ready` label, the PR author can use `/ci run`, `/ci retry`, `/ci cancel`, or the corresponding `/amd-ci` variants. New commits do not start upstream CI automatically. +- Before creating a Buildkite build, `/ci run` and `/amd-ci run`, including + `all` and `nightly`, check the requested PR commit against its target branch + in the upstream repository. Your branch must contain every commit currently + on that branch: **zero commits behind**. This also applies to PRs targeting + release branches or another PR's branch. +- If the PR is behind, the bot reports the count without creating a new build. + Merge or rebase onto the latest target branch, then rerun the command. If the + target branch advances in the meantime, update your branch and try again. +- To run CI on an outdated branch at your own risk, append `--allow-stale` to + any run command, for example `/ci run --allow-stale` or + `/amd-ci run all --allow-stale`. The same authorization requirements apply. + The bot reports the lag and warns that outdated CI configuration may cause + failures. Before merging, merge or rebase onto the latest target branch and + rerun CI without `--allow-stale` on the latest PR commit. +- Retrying or cancelling builds does not check branch freshness. +- These commands do not modify your branch or enforce merge requirements. + Direct Buildkite launches are outside this workflow. No additional token is + needed. + +Maintainers must configure required Buildkite status checks on the latest PR +commit and enable GitHub's +[Require branches to be up to date before merging](https://docs.github.com/en/repositories/configuring-branches-and-merges-in-your-repository/managing-protected-branches/about-protected-branches#require-status-checks-before-merging) +protection on the relevant target branches to enforce the merge policy. Select +statuses that confirm the required tests completed, using Buildkite as their +expected source. An aggregate status that reports "passed and blocked" does not +by itself guarantee those tests ran; verify the pipeline's required and optional +blocked steps before choosing statuses or changing blocked-build reporting. +The CI command workflow does not configure these repository or pipeline settings. ### Pull Request Limits and Escalation From 00972dfd72988942138a7a6089eaee08580210b8 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:44:56 -0400 Subject: [PATCH 051/139] [ROCm][DSV4.1][Perf] Fold the mHC post step into the delayed pre projection (#56513) Signed-off-by: Fangzhou Ai Signed-off-by: Fangzhou Ai Co-authored-by: Cursor --- tests/kernels/test_mhc_kernels.py | 78 +++++++++- vllm/_aiter_ops.py | 175 ++++++++++++++++++----- vllm/model_executor/kernels/mhc/aiter.py | 21 ++- vllm/model_executor/layers/mhc.py | 82 ++++++++++- vllm/models/deepseek_v41/amd/model.py | 54 ++++--- 5 files changed, 346 insertions(+), 64 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 5049815380..f805fe39a5 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -983,8 +983,12 @@ def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried): ) expected = mhc_pre_delayed_torch(*args, pre_mix=pre_mix) - actual = object.__new__(MHCPreDelayedOp).forward_hip(*args, pre_mix=pre_mix) + residual_out, *actual = object.__new__(MHCPreDelayedOp).forward_hip( + *args, pre_mix=pre_mix + ) + # No post was requested, so the residual comes straight back. + assert residual_out is residual for i in (0, 1, 3): torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3) # The collapse is the same FP32 multiply-and-sum in both paths. @@ -1026,7 +1030,12 @@ def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch): broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1) broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous() expected = mhc_pre_delayed_torch(broadcast_residual, broadcast_fn, *args[2:], x=x) - actual = op.forward_hip(broadcast_residual, broadcast_fn, *args[2:], x=x) + # forward_hip leads with the residual; no post was requested here, so it is + # the caller's own tensor and the reference's four outputs follow it. + residual_out, *actual = op.forward_hip( + broadcast_residual, broadcast_fn, *args[2:], x=x + ) + assert residual_out is broadcast_residual assert not took_aiter, "the broadcast seam must not reach AITER" for i in range(4): torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3) @@ -1268,3 +1277,68 @@ def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch): assert layer.hc_attn_fn_broadcast is buffer expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1) assert torch.equal(layer.hc_attn_fn_broadcast, expected) + + +@pytest.mark.skipif( + not (current_platform.is_rocm() and HAS_AITER_MHC), + reason="AITER mHC required", +) +@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128]) +@pytest.mark.parametrize("carried", [False, True]) +def test_mhc_fused_post_pre_delayed_rocm_aiter( + num_tokens, carried, default_vllm_config +): + """Folding the post into the pre projection must match doing them apart. + + The residual is bfloat16, so the two accumulation orders are allowed to + land a rounding step apart; the tolerance here is one bfloat16 ULP. + """ + set_random_seed(0) + hc_mult, hidden_size = 4, 5120 + residual, fn, hc_scale, hc_base, _ = _rocm_mhc_inputs( + num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult + ) + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=DEVICE) + post_layer_mix = ( + torch.rand(num_tokens, hc_mult, 1, dtype=torch.float32, device=DEVICE) + 0.5 + ) + comb_res_mix = ( + torch.rand(num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=DEVICE) + + 0.2 + ) + pre_mix = ( + torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5 + if carried + else None + ) + rms_eps = hc_pre_eps = hc_sinkhorn_eps = 1e-6 + pre_args = (fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, 1.0, 20) + + next_residual = mhc_post_torch(x, residual, post_layer_mix, comb_res_mix) + expected = ( + next_residual, + *mhc_pre_delayed_torch(next_residual, *pre_args, pre_mix=pre_mix), + ) + actual = MHCPreDelayedOp().forward_hip( + residual, + *pre_args, + pre_mix=pre_mix, + sublayer_out=x, + post_layer_mix=post_layer_mix, + comb_res_mix=comb_res_mix, + ) + + for got, want in zip(actual, expected, strict=True): + torch.testing.assert_close(got, want.reshape(got.shape), atol=2e-2, rtol=8e-3) + + +@pytest.mark.skipif( + not (current_platform.is_rocm() and HAS_AITER_MHC), + reason="AITER mHC required", +) +def test_mhc_fused_post_pre_delayed_falls_back_for_large_batches(): + """AITER's heuristic hands large batches back to the unfused path.""" + from vllm._aiter_ops import rocm_aiter_ops + + assert not rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1) + assert rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(1 << 20) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 454aba175e..4cfb3b3884 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1676,6 +1676,75 @@ def _rocm_aiter_fp8_attn_fake( ) +def _mhc_delayed_pre_tail( + gemm_out: torch.Tensor, + sqrsum: torch.Tensor, + residual: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + pre_mix: torch.Tensor | None, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Finish a delayed mHC pre once AITER has produced the projection. + + Shared by the plain and the post-fused entry points, which differ only in + which kernel computes ``gemm_out`` / ``sqrsum`` and the residual they + project. ``mhc_pre_big_fuse`` also writes a collapse against its own + pre-mix, which the delayed formulation cannot use and so discards. + Returns ``(post_mix, comb_mix, layer_input, next_pre_mix)``. + """ + from aiter.ops.mhc import mhc_pre_big_fuse + + num_tokens, hc_mult, hidden_size = residual.shape + device = residual.device + with torch.device(device): + post_mix = torch.empty( + num_tokens, hc_mult, 1, dtype=torch.float32, device=device + ) + comb_mix = torch.empty( + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device + ) + unused_collapse = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=device + ) + mhc_pre_big_fuse( + post_mix, + comb_mix, + unused_collapse, + gemm_out, + sqrsum, + hc_scale, + hc_base, + residual, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) + + next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton( + gemm_out, + sqrsum, + hc_scale, + hc_base, + hc_mult, + hc_mult * hidden_size, + rms_eps, + hc_pre_eps, + ) + if pre_mix is None: + # Model entry selects residual stream zero. + layer_input = residual[:, 0] + else: + layer_input = torch.ops.vllm.hc_collapse_triton(residual, pre_mix) + return post_mix, comb_mix, layer_input, next_pre_mix + + # Global flag to ensure ops are registered only once _OPS_REGISTERED = False @@ -3611,6 +3680,10 @@ class rocm_aiter_ops: hc_post_mult_value: float, sinkhorn_repeat: int, pre_mix: torch.Tensor | None = None, + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + residual_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """mHC pre using the pre-mix carried from the previous sublayer. @@ -3625,6 +3698,15 @@ class rocm_aiter_ops: that redundant store is the price of not having a native delayed kernel, and is small next to the ~140 launches it replaces. + Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` + also applies the preceding post block. AITER then projects the new + residual in the same kernel that computes it + (``mhc_fused_post_pre_gemm_sqrsum``), which saves one launch; every + later stage is identical either way. That new residual is written to + ``residual_out``, which the caller must supply in that case; it is an + output buffer rather than a return value so the op never has to hand + back one of its own inputs on the unfused path. + Returns: post_mix: shape (..., hc_mult, 1), dtype torch.float32 comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32 @@ -3632,8 +3714,9 @@ class rocm_aiter_ops: next_pre_mix: shape (..., hc_mult), dtype torch.float32 """ from aiter.ops.mhc import ( + get_mhc_fused_post_pre_config, get_mhc_pre_splitk, - mhc_pre_big_fuse, + mhc_fused_post_pre_gemm_sqrsum, mhc_pre_gemm_sqrsum, ) @@ -3666,9 +3749,16 @@ class rocm_aiter_ops: pre_mix_flat = None if pre_mix is None else pre_mix.view(-1, hc_mult) + fuse_post = sublayer_out is not None # AITER's Python wrappers allocate without explicit device arguments. with torch.device(device): - splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size) + if fuse_post: + # Keyed on the per-stream K, i.e. hidden_size, not hc_mult * it. + splitk, tile_m, tile_n, tile_k = get_mhc_fused_post_pre_config( + num_tokens, hidden_size + ) + else: + splitk, tile_k = get_mhc_pre_splitk(num_tokens, hc_hidden_size) # AITER pads the GEMM output to a multiple of 32 columns. gemm_pad = torch.empty( splitk, @@ -3679,50 +3769,46 @@ class rocm_aiter_ops: ) gemm_out = gemm_pad[:, :, :hc_mult3] sqrsum = torch.empty(splitk, num_tokens, dtype=torch.float32, device=device) - mhc_pre_gemm_sqrsum(gemm_out, sqrsum, residual_flat, fn, tile_k, 0) + if fuse_post: + # Restated for the type checker, which cannot narrow the + # optional arguments through the fuse_post flag. + assert sublayer_out is not None + assert post_layer_mix is not None and comb_res_mix is not None + assert residual_out is not None + projected = residual_out.view(-1, hc_mult, hidden_size) + mhc_fused_post_pre_gemm_sqrsum( + gemm_out, + sqrsum, + projected, + sublayer_out.view(-1, hidden_size), + residual_flat, + # The kernel wants the post mix squeezed to (tokens, hc_mult). + post_layer_mix.view(-1, hc_mult, 1).squeeze(-1), + comb_res_mix.view(-1, hc_mult, hc_mult), + fn, + tile_m, + tile_n, + tile_k, + 0, + ) + else: + projected = residual_flat + mhc_pre_gemm_sqrsum(gemm_out, sqrsum, projected, fn, tile_k, 0) - post_mix = torch.empty( - num_tokens, hc_mult, 1, dtype=torch.float32, device=device - ) - comb_mix = torch.empty( - num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device - ) - unused_collapse = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=device - ) - mhc_pre_big_fuse( - post_mix, - comb_mix, - unused_collapse, - gemm_out, - sqrsum, - hc_scale, - hc_base, - residual_flat, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - ) - - next_pre_mix = torch.ops.vllm.mhc_pre_mix_triton( + post_mix, comb_mix, layer_input, next_pre_mix = _mhc_delayed_pre_tail( gemm_out, sqrsum, + projected, hc_scale, hc_base, - hc_mult, - hc_hidden_size, + pre_mix_flat, rms_eps, hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, ) - if pre_mix_flat is None: - # Model entry selects residual stream zero. - layer_input = residual_flat[:, 0] - else: - layer_input = torch.ops.vllm.hc_collapse_triton(residual_flat, pre_mix_flat) - return ( post_mix.view(*outer_shape, hc_mult, 1), comb_mix.view(*outer_shape, hc_mult, hc_mult), @@ -3730,6 +3816,21 @@ class rocm_aiter_ops: next_pre_mix.view(*outer_shape, hc_mult), ) + @staticmethod + def mhc_fused_post_pre_delayed_prefers_unfused(num_tokens: int) -> bool: + """True when AITER's own heuristic favours separate post and pre. + + Mirrors the ``fused_m_upper_bound`` table in ``aiter.ops.mhc``: the + fused GEMM wins for decode-sized batches and loses for large ones. + """ + from aiter.jit.utils.chip_info import get_gfx_runtime + + return num_tokens >= { + "gfx950": 1024, + "gfx942": 128, + "gfx1250": 1024, + }.get(get_gfx_runtime(), 1024) + @staticmethod def hc_head( hs_flat: torch.Tensor, diff --git a/vllm/model_executor/kernels/mhc/aiter.py b/vllm/model_executor/kernels/mhc/aiter.py index e04ce3322b..d34e29314d 100644 --- a/vllm/model_executor/kernels/mhc/aiter.py +++ b/vllm/model_executor/kernels/mhc/aiter.py @@ -115,6 +115,10 @@ def mhc_pre_delayed_aiter( hc_post_mult_value: float, sinkhorn_repeat: int, pre_mix: torch.Tensor | None = None, + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + residual_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """mHC pre with the pre-mix carried in from the previous sublayer. @@ -134,6 +138,13 @@ def mhc_pre_delayed_aiter( sinkhorn_repeat: number of sinkhorn iterations pre_mix: shape (..., hc_mult) from the previous sublayer, or None at model entry to select residual stream zero. + sublayer_out: attention or FFN output, shape (..., hidden_size). When + given with the two mixes below, the preceding post block is + applied here so AITER can fold it into the pre projection. + post_layer_mix: shape (..., hc_mult, 1), post gate for that block. + comb_res_mix: shape (..., hc_mult, hc_mult), residual comb for it. + residual_out: shape (..., hc_mult, hidden_size), written with the post + block's new residual. Required exactly when the post is folded in. Returns: post_mix: shape (..., hc_mult, 1), dtype torch.float32 @@ -156,6 +167,10 @@ def mhc_pre_delayed_aiter( hc_post_mult_value, sinkhorn_repeat, pre_mix, + sublayer_out, + post_layer_mix, + comb_res_mix, + residual_out, ) @@ -170,6 +185,10 @@ def _mhc_pre_delayed_aiter_fake( hc_post_mult_value: float, sinkhorn_repeat: int, pre_mix: torch.Tensor | None = None, + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + residual_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: hc_mult = residual.shape[-2] hidden_size = residual.shape[-1] @@ -313,7 +332,7 @@ direct_register_custom_op( direct_register_custom_op( op_name="mhc_pre_delayed_aiter", op_func=mhc_pre_delayed_aiter, - mutates_args=[], + mutates_args=["residual_out"], fake_impl=_mhc_pre_delayed_aiter_fake, ) direct_register_custom_op( diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index 14f1018597..73993fbc98 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -299,7 +299,13 @@ class MHCPreDelayedOp(CustomOp): Same gates as :class:`MHCPreOp`, but the stream collapse applies the caller's ``pre_mix`` and this sublayer's pre-mix is returned for the next - sublayer seam. Returns post_mix, comb_mix, layer_input, next_pre_mix. + sublayer seam. + + Passing ``sublayer_out`` / ``post_layer_mix`` / ``comb_res_mix`` also + applies the preceding post block, which lets AITER fold it into the pre + projection and saves a launch. Returns residual, post_mix, comb_mix, + layer_input, next_pre_mix; ``residual`` is the caller's own tensor when + no post is requested. """ # --8<-- [end:mhc_pre_delayed] @@ -307,6 +313,24 @@ class MHCPreDelayedOp(CustomOp): def enabled(cls) -> bool: return True + def __init__(self) -> None: + super().__init__() + # Built here, not lazily: CustomOp resolves dispatch against the + # current vLLM config, which is only set during model construction. + self._post = MHCPostOp() + + def _maybe_post( + self, + residual: torch.Tensor, + sublayer_out: torch.Tensor | None, + post_layer_mix: torch.Tensor | None, + comb_res_mix: torch.Tensor | None, + ) -> torch.Tensor: + if sublayer_out is None: + return residual + assert post_layer_mix is not None and comb_res_mix is not None + return self._post(sublayer_out, residual, post_layer_mix, comb_res_mix) + def forward_cuda( self, residual: torch.Tensor, @@ -322,8 +346,14 @@ class MHCPreDelayedOp(CustomOp): x: torch.Tensor | None = None, norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - return torch.ops.vllm.mhc_pre_delayed_tilelang( + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + residual = self._maybe_post( + residual, sublayer_out, post_layer_mix, comb_res_mix + ) + return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang( residual, fn, hc_scale, @@ -354,7 +384,10 @@ class MHCPreDelayedOp(CustomOp): x: torch.Tensor | None = None, norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual` # and folds no RMSNorm, so it cannot serve the model-entry broadcast # (which projects a narrower `x`) or a fused norm. Both are handled by @@ -364,7 +397,26 @@ class MHCPreDelayedOp(CustomOp): and norm_weight is None and _aiter_mhc_supported(residual, None, supports_norm=False) ): - return torch.ops.vllm.mhc_pre_delayed_aiter( + from vllm._aiter_ops import rocm_aiter_ops + + num_tokens = residual.numel() // (residual.shape[-1] * residual.shape[-2]) + # Folding the post in only pays while the residual still fits in + # cache; past that AITER's heuristic wants the separate kernels, + # which can ask for a non-temporal store and a large-m split-k. + if sublayer_out is not None and ( + rocm_aiter_ops.mhc_fused_post_pre_delayed_prefers_unfused(num_tokens) + ): + residual = self._maybe_post( + residual, sublayer_out, post_layer_mix, comb_res_mix + ) + sublayer_out = None + # The op writes the folded post's residual into this buffer rather + # than returning it, since on the unfused path it has none of its + # own and an op must not hand back one of its inputs. + residual_out = ( + torch.empty_like(residual) if sublayer_out is not None else None + ) + rest = torch.ops.vllm.mhc_pre_delayed_aiter( residual, fn, hc_scale, @@ -375,9 +427,17 @@ class MHCPreDelayedOp(CustomOp): hc_post_mult_value, sinkhorn_repeat, pre_mix, + sublayer_out, + post_layer_mix, + comb_res_mix, + residual_out, ) + return (residual if residual_out is None else residual_out), *rest + residual = self._maybe_post( + residual, sublayer_out, post_layer_mix, comb_res_mix + ) if HAS_TILELANG_MHC: - return torch.ops.vllm.mhc_pre_delayed_tilelang( + return residual, *torch.ops.vllm.mhc_pre_delayed_tilelang( residual, fn, hc_scale, @@ -392,6 +452,7 @@ class MHCPreDelayedOp(CustomOp): norm_weight, norm_eps, ) + # The post is already applied, so it must not be requested again. return self.forward_native( residual, fn, @@ -423,7 +484,13 @@ class MHCPreDelayedOp(CustomOp): x: torch.Tensor | None = None, norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + sublayer_out: torch.Tensor | None = None, + post_layer_mix: torch.Tensor | None = None, + comb_res_mix: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + residual = self._maybe_post( + residual, sublayer_out, post_layer_mix, comb_res_mix + ) post_mix, comb_mix, layer_input, next_pre_mix = ( mhc_kernels.mhc_pre_delayed_torch( residual, @@ -440,6 +507,7 @@ class MHCPreDelayedOp(CustomOp): ) ) return ( + residual, post_mix, comb_mix, _apply_mhc_norm(layer_input, norm_weight, norm_eps), diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index eaed15214f..b3eddc2024 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -274,7 +274,7 @@ class DeepseekV4DecoderLayer(nn.Module): # copies and the identity pre-mix selects copy 0. assert self.hc_attn_fn_broadcast is not None residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous() - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( residual, self.hc_attn_fn_broadcast, self.hc_attn_scale, @@ -288,7 +288,7 @@ class DeepseekV4DecoderLayer(nn.Module): ) else: residual = x - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( residual, self.hc_attn_fn, self.hc_attn_scale, @@ -301,18 +301,7 @@ class DeepseekV4DecoderLayer(nn.Module): pre_mix=pre_mix, ) else: - residual = self.mhc_post(x, residual, post_mix, res_mix) - if self.engram is not None and engram_hashes is not None: - # Engram injection happens between the previous sublayer's - # post and this block's pre, on the full hc stream, so the - # mix coefficients see the injected stream. - residual = self.engram( - residual, - engram_hashes[:, self.engram.layer_hash_index], - engram_mask, - ) - post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( - residual, + pre_args = ( self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base, @@ -321,8 +310,37 @@ class DeepseekV4DecoderLayer(nn.Module): self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, - pre_mix=pre_mix, ) + if self.engram is not None and engram_hashes is not None: + # Engram injection happens between the previous sublayer's + # post and this block's pre, on the full hc stream, so the + # mix coefficients see the injected stream. That read of the + # residual is what stops the post from fusing into the pre + # here, unlike the attention seam below. + residual = self.mhc_post(x, residual, post_mix, res_mix) + residual = self.engram( + residual, + engram_hashes[:, self.engram.layer_hash_index], + engram_mask, + ) + residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed( + residual, *pre_args, pre_mix=pre_mix + ) + else: + ( + residual, + post_mix, + res_mix, + x, + attn_pre, + ) = self.mhc_pre_delayed( + residual, + *pre_args, + pre_mix=pre_mix, + sublayer_out=x, + post_layer_mix=post_mix, + comb_res_mix=res_mix, + ) x = self.attn_norm(x) if self.use_sequence_parallel: @@ -332,8 +350,7 @@ class DeepseekV4DecoderLayer(nn.Module): if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual = self.mhc_post(x, residual, post_mix, res_mix) - post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( + residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed( residual, self.hc_ffn_fn, self.hc_ffn_scale, @@ -344,6 +361,9 @@ class DeepseekV4DecoderLayer(nn.Module): self.hc_post_alpha, self.hc_sinkhorn_iters, pre_mix=attn_pre, + sublayer_out=x, + post_layer_mix=post_mix, + comb_res_mix=res_mix, ) x = self.ffn_norm(x) x = self.ffn(x, input_ids) From 8e08cef46ebeef501a1b1593fd7da77909afe4ae Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Mon, 14 Sep 2026 17:58:33 -0500 Subject: [PATCH 052/139] [Bugfix] Redact credentials from benchmark logs (#56662) Signed-off-by: Taneem Ibrahim Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- benchmarks/multi_turn/benchmark_serving_multi_turn.py | 5 ++++- vllm/benchmarks/lib/utils.py | 11 +++++++++++ vllm/benchmarks/serve.py | 11 ++++++++--- 3 files changed, 23 insertions(+), 4 deletions(-) diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py index 5a60d9c668..f084a4d84c 100644 --- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py +++ b/benchmarks/multi_turn/benchmark_serving_multi_turn.py @@ -32,8 +32,11 @@ from bench_dataset import ( from bench_utils import TEXT_SEPARATOR, Color, logger from transformers import AutoTokenizer # type: ignore +from vllm.benchmarks.lib.utils import redact_sensitive_namespace + NUM_TOKENS_FROM_DATASET = 0 TERM_SIGNAL = None +_SENSITIVE_ARG_FIELDS = ("api_key", "header") class ConversationSampling(str, Enum): @@ -1548,7 +1551,7 @@ async def main() -> None: args = parser.parse_args() - logger.info(args) + logger.info(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS)) logger.info(f"{Color.GREEN}Input parameters:{Color.RESET}") logger.info(f"url={args.url}") diff --git a/vllm/benchmarks/lib/utils.py b/vllm/benchmarks/lib/utils.py index 7737621a4e..22409101e4 100644 --- a/vllm/benchmarks/lib/utils.py +++ b/vllm/benchmarks/lib/utils.py @@ -9,6 +9,17 @@ from contextlib import contextmanager from typing import Any +def redact_sensitive_namespace( + args: argparse.Namespace, fields: tuple[str, ...] +) -> argparse.Namespace: + """Return a copy of CLI arguments with selected values redacted.""" + redacted_args = argparse.Namespace(**vars(args)) + for field in fields: + if getattr(redacted_args, field, None) is not None: + setattr(redacted_args, field, "***") + return redacted_args + + def extract_field( args: argparse.Namespace, extra_info: dict[str, Any], field_name: str ) -> str: diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index aef7c7e7b4..0292ff939a 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -51,13 +51,18 @@ from vllm.benchmarks.lib.endpoint_request_func import ( RequestFuncOutput, ) from vllm.benchmarks.lib.ready_checker import wait_for_endpoint -from vllm.benchmarks.lib.utils import convert_to_pytorch_benchmark_format, write_to_json +from vllm.benchmarks.lib.utils import ( + convert_to_pytorch_benchmark_format, + redact_sensitive_namespace, + write_to_json, +) from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.gc_utils import freeze_gc_heap from vllm.utils.network_utils import join_host_port MILLISECONDS_TO_SECONDS_CONVERSION = 1000 +_SENSITIVE_ARG_FIELDS = ("header",) def _merge_overrides(base: dict | None, override: dict | None) -> dict | None: @@ -1536,7 +1541,7 @@ def save_to_pytorch_benchmark_format( # later if needed ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"] pt_records = convert_to_pytorch_benchmark_format( - args=args, + args=redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS), metrics={k: [results[k]] for k in metrics if k in results}, extra_info={ k: results[k] @@ -2013,7 +2018,7 @@ def main(args: argparse.Namespace) -> dict[str, Any]: async def main_async(args: argparse.Namespace) -> dict[str, Any]: - print(args) + print(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS)) if args.max_concurrency is not None and args.max_concurrency <= 0: raise ValueError("--max-concurrency must be greater than 0") From 3bb0a03f35269185b6358753e64932bb946cc531 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:58:48 -0700 Subject: [PATCH 053/139] [Model Runner V2] Acceptance estimation for adaptive verification (#52228) Signed-off-by: Giancarlo Delfin --- .../spec_decode/test_acceptance_estimator.py | 294 ++++++++++ .../test_gpu_autoregressive_speculator.py | 5 + ...st_gpu_extract_hidden_states_speculator.py | 2 + vllm/config/speculative.py | 12 +- vllm/v1/worker/gpu/model_runner.py | 4 + .../gpu/spec_decode/acceptance_estimator.py | 501 ++++++++++++++++++ .../gpu/spec_decode/dflash2/speculator.py | 6 + .../gpu/spec_decode/dspark/speculator.py | 34 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 70 ++- 9 files changed, 900 insertions(+), 28 deletions(-) create mode 100644 tests/v1/spec_decode/test_acceptance_estimator.py create mode 100644 vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py diff --git a/tests/v1/spec_decode/test_acceptance_estimator.py b/tests/v1/spec_decode/test_acceptance_estimator.py new file mode 100644 index 0000000000..a1f3da0940 --- /dev/null +++ b/tests/v1/spec_decode/test_acceptance_estimator.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""``predict`` writes one probability into two index spaces at once. + +Adaptive verification reads the estimate in this step's batch order, while the +next step's refit reads it in persistent request-state slots. A mix-up between +the two is silent -- every value stays a plausible probability, it just belongs +to another request -- so it is asserted here rather than left to acceptance +metrics to hint at. +""" + +import pytest +import torch + +from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import ( + OnlineAcceptanceEstimator, +) + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("CUDA required for acceptance estimator tests", allow_module_level=True) + +MAX_NUM_REQS = 16 +NUM_STEPS = 4 +VOCAB_SIZE = 2048 + + +def _make_estimator(device: torch.device) -> OnlineAcceptanceEstimator: + estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device) + # Distinct per-position intercepts, so reading the wrong one shows up. + estimator.slope.fill_(0.7) + estimator.intercepts[:] = torch.tensor([1.5, 1.0, 0.5, 0.0], device=device) + return estimator + + +def _expected_probs( + estimator: OnlineAcceptanceEstimator, + logits: torch.Tensor, + steps: torch.Tensor, + temperature: torch.Tensor, +) -> torch.Tensor: + """logit(q_max) by a different route than the kernel takes. + + The kernel shifts by the max and drops the unit term; here the argmax is + masked out and the complement is reduced with logsumexp, so agreement is + evidence about the value rather than about a shared formula. + """ + temp = torch.where(temperature > 0, temperature, torch.ones_like(temperature)) + scaled = logits.double() / temp[:, None].double() + top = scaled.max(dim=-1) + rest = scaled.scatter(1, top.indices[:, None], float("-inf")) + feature = (top.values - torch.logsumexp(rest, dim=-1)).clamp(-40.0, 40.0) + weight, bias = estimator.slope, estimator.intercepts + return torch.sigmoid(weight.double() * feature + bias[steps].double()).float() + + +@pytest.mark.parametrize("tokens_per_req", [1, NUM_STEPS]) +def test_predict_agrees_across_slot_and_batch_order(tokens_per_req: int): + """One row per request (per-step drafting) and a whole block at once. + + Which one it is is inferred from whether the draft step is per-row, so the + two cases differ only in how ``draft_step`` is shaped. + """ + device = torch.device("cuda") + torch.manual_seed(0) + estimator = _make_estimator(device) + num_reqs = 5 + + # Batch position i holds a shuffled persistent slot, as a served batch would. + slots = torch.randperm(MAX_NUM_REQS, device=device)[:num_reqs].to(torch.int32) + if tokens_per_req == 1: + idx_mapping = slots + draft_step = torch.tensor(2, device=device, dtype=torch.int32) + steps = draft_step.expand(num_reqs) + else: + idx_mapping = slots.repeat_interleave(NUM_STEPS) + draft_step = torch.arange(NUM_STEPS, device=device, dtype=torch.int32).repeat( + num_reqs + ) + steps = draft_step + batch_rows = torch.arange(num_reqs, device=device).repeat_interleave(tokens_per_req) + logits = torch.randn(idx_mapping.shape[0], VOCAB_SIZE, device=device) * 4.0 + + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device) + temperature = torch.rand(MAX_NUM_REQS, device=device) + 0.5 + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature) + + expected = _expected_probs( + estimator, logits, steps.long(), temperature[idx_mapping.long()] + ) + torch.testing.assert_close( + estimator.predictions[idx_mapping.long(), steps.long()], + expected, + atol=2e-3, + rtol=2e-3, + ) + torch.testing.assert_close( + confidence_probs[batch_rows, steps.long()], expected, atol=2e-3, rtol=2e-3 + ) + + +def test_predict_skips_cudagraph_padded_rows(): + """Padded rows carry slot -1 and must not land on a live request.""" + device = torch.device("cuda") + torch.manual_seed(0) + estimator = _make_estimator(device) + num_reqs, num_padded = 3, 6 + + slots = torch.tensor([7, 2, 11], device=device, dtype=torch.int32) + idx_mapping = torch.cat( + [slots, torch.full((num_padded,), -1, device=device, dtype=torch.int32)] + ) + draft_step = torch.tensor(1, device=device, dtype=torch.int32) + logits = torch.randn(num_reqs + num_padded, VOCAB_SIZE, device=device) * 4.0 + + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device) + temperature = torch.rand(MAX_NUM_REQS, device=device) + 0.5 + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature) + + live = (slots.long(), torch.full((num_reqs,), 1, device=device)) + expected = _expected_probs( + estimator, logits[:num_reqs], live[1], temperature[slots.long()] + ) + torch.testing.assert_close( + estimator.predictions[live], expected, atol=2e-3, rtol=2e-3 + ) + # Nothing outside the live (request, position) cells was touched. + untouched = torch.ones(MAX_NUM_REQS, NUM_STEPS, dtype=torch.bool, device=device) + untouched[live] = False + assert not estimator.predictions[untouched].any() + assert not confidence_probs[num_reqs:].any() + + +# A refit takes one undamped Newton step, so it needs a realistic number of +# observations per round to be well conditioned -- a handful of requests can +# send the step far past the optimum. This is the batch a served step sees. +REFIT_NUM_REQS = 512 + + +def _drive_round(estimator, features, num_accepted, slots): + """Feed one step's graded drafts through predict's buffers into `step`.""" + estimator.features[slots.long()] = features + estimator.predictions[slots.long()] = torch.sigmoid( + estimator.slope * features + estimator.intercepts + ) + estimator.step( + slots, + (num_accepted + 1).to(torch.int32), + (features.shape[1] - num_accepted).to(torch.int32), + ) + + +def test_refit_recovers_a_shared_slope_and_per_position_intercepts(): + """The fit must find the logistic that generated the observations. + + One slope is shared across draft positions and each keeps its own intercept, + so the arrowhead solve is only correct if it recovers both from data whose + positions differ in level but not in slope. + """ + device = torch.device("cuda") + torch.manual_seed(0) + true_slope = 0.6 + true_bias = torch.tensor([1.4, 0.9, 0.4, -0.1], device=device) + + estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device) + estimator.REFIT_INTERVAL = 5 + slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32) + for _ in range(60): + features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) * 8.0 + accepted = torch.rand_like(features) < torch.sigmoid( + true_slope * features + true_bias + ) + # Verification stops at the first rejection, so a position's label is + # conditional on every shallower one having been accepted. + num_accepted = (~accepted).float().argmax(dim=1) + num_accepted[accepted.all(dim=1)] = NUM_STEPS + _drive_round(estimator, features, num_accepted, slots) + + slope, bias = estimator.slope, estimator.intercepts + torch.testing.assert_close( + slope[0], torch.tensor(true_slope, device=device), atol=0.1, rtol=0 + ) + torch.testing.assert_close(bias, true_bias, atol=0.3, rtol=0) + + +def test_refit_survives_a_position_with_no_observations(): + """A draft position nobody reached must not stall the shared slope. + + Its statistics stay zero while the others accumulate, which is the degenerate + case the arrowhead denominator has to tolerate. + """ + device = torch.device("cuda") + torch.manual_seed(0) + estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device) + estimator.REFIT_INTERVAL = 5 + slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32) + before_slope = estimator.slope.clone() + before_intercepts = estimator.intercepts.clone() + + # Every request is rejected at position 1, so positions 2+ are never graded. + num_accepted = torch.ones(REFIT_NUM_REQS, device=device) + for _ in range(20): + features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) * 8.0 + _drive_round(estimator, features, num_accepted, slots) + + assert estimator.slope.isfinite().all() + assert estimator.intercepts.isfinite().all() + assert not torch.equal(estimator.slope, before_slope), "slope never moved" + # Positions 2 and 3 were never observed, so their intercepts must be untouched. + torch.testing.assert_close( + estimator.intercepts[2:], before_intercepts[2:], rtol=0, atol=0 + ) + + +def test_step_damping_scales_with_the_evidence(): + """A round moves a coefficient in proportion to the samples behind it. + + The Newton step itself barely changes with how much data a round carries -- + curvature and score both scale with it -- so scaling a round's statistics + isolates the damping factor n / (n + DAMPING_OBSERVATIONS). The slope/ + intercept coupling is zeroed so the two parameters can be checked apart. + """ + device = torch.device("cuda") + damping = OnlineAcceptanceEstimator.DAMPING_OBSERVATIONS + per_position_n = 20.0 + + def refit_with(scale): + estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device) + estimator.info[:, 1] = 0.0 # sum(w*margin): no slope/intercept coupling + estimator.info[:, 2] = 1.0 * scale # sum(w) + estimator.grad[:, 1] = 0.5 * scale # sum(y - p) + # The slope's information and score are accumulated as per-position + # partials that the refit sums, so spread each total over the positions. + estimator.info[:, 0] = 30.0 * scale / NUM_STEPS # sum(w*margin^2) + estimator.grad[:, 0] = 2.0 * scale / NUM_STEPS # sum((y - p)*margin) + estimator.counts[:] = per_position_n * scale + estimator._steps_since_refit = estimator.REFIT_INTERVAL - 1 + before_slope = estimator.slope.clone() + before_intercepts = estimator.intercepts.clone() + empty = torch.zeros(0, dtype=torch.int32, device=device) + estimator.step(empty, empty, empty) + return ( + (estimator.slope - before_slope)[0].item(), + (estimator.intercepts - before_intercepts)[0].item(), + ) + + for scale in (1.0, 10.0, 100.0): + total_n = per_position_n * scale * NUM_STEPS + slope_delta, bias_delta = refit_with(scale) + assert slope_delta == pytest.approx( + (2.0 / 30.0) * total_n / (total_n + damping), rel=1e-3 + ) + assert bias_delta == pytest.approx( + 0.5 * (per_position_n * scale) / (per_position_n * scale + damping), + rel=1e-3, + ) + + # A round with no observations at all must leave the coefficients alone. + assert refit_with(0.0) == (0.0, 0.0) + + +def test_feature_ignores_which_token_was_drawn(): + """The score must depend on the draft distribution, not on the draw. + + This is what keeps rejection sampling lossless: if the score read the drawn + token, trimming a draft would condition on the token being trimmed and tilt + the emitted distribution toward the drafter's confident modes. ``predict`` + takes no token argument, and this pins that -- two rows with identical + logits must score identically no matter what was sampled from them. + """ + device = torch.device("cuda") + torch.manual_seed(0) + estimator = _make_estimator(device) + + row = torch.randn(VOCAB_SIZE, device=device) * 4.0 + logits = row.expand(2, VOCAB_SIZE).contiguous() + idx_mapping = torch.tensor([0, 1], device=device, dtype=torch.int32) + draft_step = torch.tensor(0, device=device, dtype=torch.int32) + temperature = torch.ones(MAX_NUM_REQS, device=device) + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device) + + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature) + assert estimator.features[0, 0] == estimator.features[1, 0] + + # And the value is the top token's log-odds, not the top-2 margin. + scaled = row.double() + top = scaled.max(dim=-1) + rest = scaled.scatter(0, top.indices[None], float("-inf")) + torch.testing.assert_close( + estimator.features[0, 0].double(), + (top.values - torch.logsumexp(rest, dim=-1)).clamp(-40.0, 40.0), + atol=2e-3, + rtol=2e-3, + ) diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index 920c9a25f0..b9b3570b8a 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -123,6 +123,7 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected) draft_model_config=draft_model_config, use_local_argmax_reduction=False, draft_sample_method="greedy", + enable_adaptive_verification=False, ) vllm_config = SimpleNamespace( speculative_config=speculative_config, @@ -161,6 +162,7 @@ def test_mm_support_configured_after_model_load(monkeypatch): speculator.dtype = torch.float32 speculator.draft_model_config = draft_model_config speculator.supports_mm_inputs = False + speculator.use_acceptance_estimator = False checked_configs = [] @@ -195,6 +197,7 @@ def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch): speculator = object.__new__(_TestSpeculator) speculator.supports_mm_inputs = False speculator.inputs_embeds = None + speculator.use_acceptance_estimator = False speculator.vllm_config = SimpleNamespace(model_config=object()) speculator.max_num_tokens = 4 speculator.hidden_size = 3 @@ -219,6 +222,7 @@ def test_load_model_disables_mm_support_for_text_only_drafter(monkeypatch): speculator = object.__new__(_TestSpeculator) speculator.supports_mm_inputs = False speculator.inputs_embeds = None + speculator.use_acceptance_estimator = False speculator.vllm_config = SimpleNamespace(model_config=object()) draft_model = _TextOnlyDraftModel() speculator.test_draft_model = draft_model @@ -251,6 +255,7 @@ def test_multi_module_mm_support_configured_after_model_load(monkeypatch): speculator = object.__new__(MultiModuleMTPSpeculator) speculator.supports_mm_inputs = False speculator.inputs_embeds = None + speculator.use_acceptance_estimator = False speculator.cached_draft_input_embeds = None speculator.vllm_config = SimpleNamespace(model_config=object()) speculator.max_num_tokens = 4 diff --git a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py index d1aa0798fe..74464f542f 100644 --- a/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py +++ b/tests/v1/worker/test_gpu_extract_hidden_states_speculator.py @@ -60,6 +60,7 @@ def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch): layer_name = "cache_only_layers.2" speculator = object.__new__(ExtractHiddenStatesSpeculator) speculator.vllm_config = cast(Any, SimpleNamespace()) + speculator.acceptance_estimator = None speculator.num_hidden_states = 2 speculator.hidden_states = torch.zeros(4, 2, 3) speculator.draft_attn_layer_names = {layer_name} @@ -111,6 +112,7 @@ def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch): def test_propose_requires_aux_hidden_states(): speculator = object.__new__(ExtractHiddenStatesSpeculator) + speculator.acceptance_estimator = None speculator.num_hidden_states = 2 input_batch = cast( Any, SimpleNamespace(idx_mapping=torch.tensor([0], dtype=torch.int32)) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 1861266820..2396b26cb0 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1561,8 +1561,16 @@ class SpeculativeConfig: self.index_share_for_mtp_iteration ) - if self.method != "dspark" and self.enable_adaptive_verification: - raise ValueError("Adaptive verification only supported with DSpark") + if ( + self.method != "dspark" + and self.enable_adaptive_verification + and self.use_local_argmax_reduction + ): + raise ValueError( + "Adaptive verification estimates per-position acceptance from " + "the draft logits, which use_local_argmax_reduction never " + "materializes. Disable one of them." + ) return self diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4b1520f276..761893303a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -2120,6 +2120,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): if self.speculator is not None: assert self.sampler is not None + if isinstance(self.speculator, DraftModelSpeculator): + self.speculator.observe_verification( + input_batch.idx_mapping, num_sampled, num_rejected + ) # Let the target override the hidden state fed to the drafter # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The # target returns a persistent buffer sized at max_num_batched_tokens; diff --git a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py new file mode 100644 index 0000000000..46065d0546 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + _compute_max_and_sumexp, +) + +# The feature is a log-odds, log(q / (1 - q)), which a row carrying all of its +# mass on one token sends to infinity. Clamp it symmetrically. +_MAX_LOG_ODDS = 40.0 + +# Cold-start coefficients: the median slope and intercept fitted online across +# DeepSeek-V4-Flash, Muse-Glimmer-30B, MiMo-V2.5-Pro, Inkling and Kimi-K2.5. +# Log-loss is flat near the optimum, so one starting point serves all of them +# and a real fit replaces it within a refit or two. +_INIT_SLOPE = 0.5 +_INIT_BIAS = -0.2 + + +@triton.jit +def _accumulate_kernel( + info_ptr, + grad_ptr, + idx_mapping_ptr, + num_sampled_ptr, + num_rejected_ptr, + features_ptr, + features_stride, + pred_ptr, + pred_stride, + counts_ptr, + num_reqs, + BLOCK_R: tl.constexpr, +): + step = tl.program_id(0).to(tl.int64) + req_block = tl.arange(0, BLOCK_R) + req_mask = req_block < num_reqs + req_state_block = tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0).to( + tl.int64 + ) + + num_sampled = tl.load(num_sampled_ptr + req_block, mask=req_mask, other=0).to( + tl.int64 + ) + # Subtract bonus token. + num_accepted = tl.maximum(num_sampled - 1, 0) + num_rejected = tl.load(num_rejected_ptr + req_block, mask=req_mask, other=0).to( + tl.int64 + ) + num_admitted = num_accepted + num_rejected + + feature = tl.load( + features_ptr + req_state_block * features_stride + step, + mask=req_mask, + other=0.0, + ).to(tl.float32) + pred = tl.load( + pred_ptr + req_state_block * pred_stride + step, + mask=req_mask, + other=0.0, + ).to(tl.float32) + + # The accepted tokens and the first rejected draft token (if one exists) are + # "observed" and given labels. + observed = req_mask & (step <= num_accepted) & (step < num_admitted) + label = tl.where(step < num_accepted, 1.0, 0.0) + w = tl.where(observed, pred * (1.0 - pred), 0.0) + resid = tl.where(observed, label - pred, 0.0) + + # This block's contribution to the normal equations for design row + # x = [feature, eₖ]: the XᵀWX entries aₖ, bₖ, cₖ and the Xᵀr entries g₀ₖ, g₁ₖ. + # aₖ and g₀ₖ are partials of scalars shared by every position, which + # _refit_kernel sums. + a_k = tl.sum(w * feature * feature, axis=0) + b_k = tl.sum(w * feature, axis=0) + c_k = tl.sum(w, axis=0) + g0_k = tl.sum(resid * feature, axis=0) + g1_k = tl.sum(resid, axis=0) + count = tl.sum(tl.where(observed, 1.0, 0.0), axis=0) + tl.store(info_ptr + step * 3 + 0, tl.load(info_ptr + step * 3 + 0) + a_k) + tl.store(info_ptr + step * 3 + 1, tl.load(info_ptr + step * 3 + 1) + b_k) + tl.store(info_ptr + step * 3 + 2, tl.load(info_ptr + step * 3 + 2) + c_k) + tl.store(grad_ptr + step * 2 + 0, tl.load(grad_ptr + step * 2 + 0) + g0_k) + tl.store(grad_ptr + step * 2 + 1, tl.load(grad_ptr + step * 2 + 1) + g1_k) + tl.store(counts_ptr + step, tl.load(counts_ptr + step) + count) + + +@triton.jit +def _refit_kernel( + slope_ptr, + intercepts_ptr, + info_ptr, + info_row_stride, + grad_ptr, + grad_row_stride, + counts_ptr, + NUM_SPECULATIVE_STEPS: tl.constexpr, + L2: tl.constexpr, + DAMPING: tl.constexpr, + BLOCK: tl.constexpr, +): + k = tl.arange(0, BLOCK) + mask = k < NUM_SPECULATIVE_STEPS + + # Load the stats accumulated over this round to solve (XᵀWX + λI) Δθ = Xᵀr. + a_k = tl.load(info_ptr + k * info_row_stride + 0, mask=mask, other=0.0) + b_k = tl.load(info_ptr + k * info_row_stride + 1, mask=mask, other=0.0) + c_k = tl.load(info_ptr + k * info_row_stride + 2, mask=mask, other=0.0) + L2 + a = tl.sum(a_k, axis=0) + L2 + g0_k = tl.load(grad_ptr + k * grad_row_stride + 0, mask=mask, other=0.0) + g1_k = tl.load(grad_ptr + k * grad_row_stride + 1, mask=mask, other=0.0) + g0 = tl.sum(g0_k, axis=0) + n = tl.load(counts_ptr + k, mask=mask, other=0.0) + + # Load the current shared slope and per-position intercepts. + w = tl.load(slope_ptr) + bias = tl.load(intercepts_ptr + k, mask=mask, other=0.0) + + # Solve for the shared slope delta, Δω, in the following arrowhead system: + # ⎡ a+λ b₀ b₁ … bₙ₋₁ ⎤ ⎡ Δω ⎤ ⎡ g₀ ⎤ + # ⎢ b₀ c₀+λ ⎥ ⎢ Δβ₀ ⎥ ⎢ g₁₀ ⎥ + # ⎢ b₁ c₁+λ ⎥ ⎢ Δβ₁ ⎥ = ⎢ g₁₁ ⎥ + # ⎢ ⋮ ⋱ ⎥ ⎢ ⋮ ⎥ ⎢ ⋮ ⎥ + # ⎣ bₙ₋₁ cₙ₋₁+λ ⎦ ⎣ Δβₙ₋₁ ⎦ ⎣ g₁ₙ₋₁ ⎦ + # XᵀWX + λI Δθ Xᵀ(Y-P) + shrink = b_k * b_k / c_k + coupling = b_k * g1_k / c_k + step_w = (g0 - tl.sum(coupling, axis=0)) / (a - tl.sum(shrink, axis=0)) + + # The slope learns from every position, so it is damped by the round's + # total sample count. + total_n = tl.sum(n, axis=0) + step_w *= total_n / (total_n + DAMPING) + # Mask out NaNs and steps that would drive the slope negative. + step_w = tl.where((step_w == step_w) & (w + step_w >= 0.0), step_w, 0.0) + # Update the shared slope. + new_w = w + step_w + + # Solve for the per-position intercept deltas, Δβₖ, in the same system. + step_bias = (g1_k - b_k * step_w) / c_k + # Damp by the position's sample count. + step_bias *= n / (n + DAMPING) + # Mask out NaNs. + step_bias = tl.where(step_bias == step_bias, step_bias, 0.0) + # Update the per-position intercepts. + new_bias = bias + step_bias + + tl.store(slope_ptr, new_w) + tl.store(intercepts_ptr + k, new_bias, mask=mask) + + # Start the round clean by clearing the per-round statistics. + tl.store(info_ptr + k * info_row_stride + 0, 0.0, mask=mask) + tl.store(info_ptr + k * info_row_stride + 1, 0.0, mask=mask) + tl.store(info_ptr + k * info_row_stride + 2, 0.0, mask=mask) + tl.store(grad_ptr + k * grad_row_stride + 0, 0.0, mask=mask) + tl.store(grad_ptr + k * grad_row_stride + 1, 0.0, mask=mask) + tl.store(counts_ptr + k, 0.0, mask=mask) + + +@triton.jit +def _local_max_sumexp_kernel( + local_max_ptr, + local_max_stride, + local_sumexp_ptr, + local_sumexp_stride, + logits_ptr, + logits_stride, + idx_mapping_ptr, + idx_mapping_stride, + temperature_ptr, + num_tokens, + vocab_size, + BLOCK_SIZE: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + if token_idx >= num_tokens: + return + + block_idx = tl.program_id(1) + req_state_idx = tl.load(idx_mapping_ptr + token_idx * idx_mapping_stride).to( + tl.int64 + ) + if req_state_idx < 0: + # Cudagraph-padded requests carry -1. Write an empty partial so the + # buffer holds nothing uninitialized; _predict_kernel skips these rows. + tl.store( + local_max_ptr + token_idx * local_max_stride + block_idx, float("-inf") + ) + tl.store(local_sumexp_ptr + token_idx * local_sumexp_stride + block_idx, 0.0) + return + + # Draft logits are pre-temperature, and acceptance is decided on the scaled + # distribution the drafter actually sampled from, so scale first. A request + # with temp == 0 drafts greedily, so leave its logits alone rather than dividing + # by zero, matching what the sampler does. + temp = tl.load(temperature_ptr + req_state_idx).to(tl.float32) + temp = tl.where(temp > 0.0, temp, 1.0) + + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < vocab_size + logits = ( + tl.load( + logits_ptr + token_idx * logits_stride + offsets, + mask=mask, + other=float("-inf"), + ).to(tl.float32) + / temp + ) + block_max, block_sumexp = _compute_max_and_sumexp(logits) + tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, block_max) + tl.store( + local_sumexp_ptr + token_idx * local_sumexp_stride + block_idx, block_sumexp + ) + + +@triton.jit +def _predict_kernel( + features_ptr, + features_stride, + pred_ptr, + pred_stride, + conf_ptr, + conf_stride, + slope_ptr, + intercepts_ptr, + local_max_ptr, + local_max_stride, + local_sumexp_ptr, + local_sumexp_stride, + idx_mapping_ptr, + idx_mapping_stride, + step_ptr, + num_tokens, + vocab_num_blocks, + per_token_step: tl.constexpr, + NUM_SPECULATIVE_STEPS: tl.constexpr, + MAX_LOG_ODDS: tl.constexpr, + PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, +): + token_idx = tl.program_id(0).to(tl.int64) + if token_idx >= num_tokens: + return + + req_state_idx = tl.load(idx_mapping_ptr + token_idx * idx_mapping_stride).to( + tl.int64 + ) + if req_state_idx < 0: + # Cudagraph-padded requests carry -1. Skip them so that they don't + # scatter garbage over a live request's features. + return + + # A per-row step means the speculator sampled the whole block in one pass, so + # each request owns NUM_SPECULATIVE_STEPS consecutive rows; a step shared by + # every row means one row per request. Batch position, unlike req_state_idx, + # is only meaningful for the step that computed it, and is what adaptive + # verification indexes by. + if per_token_step: + step = tl.load(step_ptr + token_idx).to(tl.int64) + batch_idx = token_idx // NUM_SPECULATIVE_STEPS + else: + step = tl.load(step_ptr).to(tl.int64) + batch_idx = token_idx + + blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS) + blocks_mask = blocks < vocab_num_blocks + maxes = tl.load( + local_max_ptr + token_idx * local_max_stride + blocks, + mask=blocks_mask, + other=float("-inf"), + ) + sumexps = tl.load( + local_sumexp_ptr + token_idx * local_sumexp_stride + blocks, + mask=blocks_mask, + other=0.0, + ) + + # The feature, logit(q(m*)) for the row's max draft logit, m*, is: + # q(x) = e^(x - m*) / Σ_z e^(z - m*) + # logit(q(m*)) = log(q(m*) / (1 - q(m*))) + # = m* - log(Σ_{z≠m*} e^(z - m*)). + # It scores the draft distribution, not the token drawn from it, so it is + # a function of the prefix alone. Trimming a draft therefore never + # conditions on the token being trimmed, and rejection sampling stays + # lossless. Scoring the drawn token instead would select on the proposal + # and bias the emitted distribution toward the drafter's confident modes. + + # Rescale each block's partial to the global max and sum. The most likely + # token contributes exactly e^(m - m) = 1, so removing it leaves the sum + # over every other token. + global_max = tl.max(maxes, axis=0) + total_sumexp = tl.sum( + tl.where(blocks_mask, sumexps * tl.exp(maxes - global_max), 0.0), axis=0 + ) + complement_sumexp = tl.maximum(total_sumexp - 1.0, 0.0) + + # Shifted by m*, the numerator is e^(m* - m*) = 1, so the feature is just the + # negated complement log-sum-exponential. Clamp to the min/max log-odds. + feature = -tl.log(complement_sumexp) + feature = tl.minimum(tl.maximum(feature, -MAX_LOG_ODDS), MAX_LOG_ODDS) + tl.store(features_ptr + req_state_idx * features_stride + step, feature) + + # Predict the acceptance probability from the feature and current coefficients. + weight = tl.load(slope_ptr) + bias = tl.load(intercepts_ptr + step) + prob = tl.sigmoid(weight * feature + bias) + tl.store(pred_ptr + req_state_idx * pred_stride + step, prob) + tl.store(conf_ptr + batch_idx * conf_stride + step, prob) + + +class OnlineAcceptanceEstimator: + """Predicts per-position acceptance, and calibrates itself while serving. + + Lifecycle per step, driven by ``DraftModelSpeculator``: + + 1. ``step`` folds the previous step's drafts, now graded by the target, + into the IRLS accumulators, and periodically solves for new coefficients, + independently but identically on every rank. + 2. ``predict`` runs inside the captured draft graph, turning this step's + draft logits into acceptance probabilities for adaptive verification. + + Trimming starts immediately: survival is a running product of sigmoids, so it + decreases with draft position whatever the coefficients are, and an unfitted + estimator degrades to uniform-depth truncation rather than to anything harmful. + """ + + # Refit every this many steps, accumulating samples in between. + REFIT_INTERVAL = 100 + # Newton steps are damped by n / (n + DAMPING_OBSERVATIONS), so a round + # carrying this many observations moves a parameter half of a full step. It + # replaces a hard minimum-sample gate: a data-poor position keeps learning, + # just slowly, instead of freezing until it crosses a threshold. Lower values + # track a drifting workload faster, higher ones are steadier on thin data. + DAMPING_OBSERVATIONS = 50.0 + # Ridge on the Newton solve. + L2 = 1e-3 + + def __init__( + self, + max_num_reqs: int, + num_speculative_steps: int, + device: torch.device, + ): + self.num_speculative_steps = num_speculative_steps + self.device = device + self._steps_since_refit = 0 + self._refits = 0 + + # The learned coefficients used for predicting acceptance. + self.slope = torch.full((1,), _INIT_SLOPE, dtype=torch.float32, device=device) + self.intercepts = torch.full( + (num_speculative_steps,), _INIT_BIAS, dtype=torch.float32, device=device + ) + # Holds logit(max q), which is used as the feature for the logistic. + # Stored in stable slots keyed by persistent request-state index. + self.features = torch.zeros( + max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device + ) + # Holds the acceptance predictions made at draft time, computed as: + # p = σ(ω * f + βₖ) + # where f is feature (logit(max q)) at draft position k, ω is the learned + # shared slope, and βₖ is the learned per-position intercept. Stored in + # stable slots keyed by the persistent request-state index. + self.predictions = torch.zeros( + max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device + ) + + # Per-round Newton-IRLS statistics, cleared after each refit. For observation + # i the design row is xᵢ = [fᵢ, eₖ], the IRLS weight is wᵢ = pᵢ(1-pᵢ), and the + # residual is yᵢ - pᵢ, where yᵢ is the label (1 for accepted, 0 for rejected). + # XᵀWX is then the negative Hessian of the Bernoulli log-likelihood loss and + # Xᵀ(Y-P) is its gradient, so the Newton step for θ = [ω, β₀ … βₙ₋₁] is the + # the solution of (XᵀWX + λI) Δθ = Xᵀ(Y-P). + # + # Columns [aₖ, bₖ, cₖ] are each position's contribution to the XᵀWX arrowhead: + # ⎡ a b₀ b₁ … bₙ₋₁ ⎤ + # ⎢ b₀ c₀ ⎥ + # ⎢ b₁ c₁ ⎥ + # ⎢ ⋮ ⋱ ⎥ + # ⎣ bₙ₋₁ cₙ₋₁ ⎦ + # bₖ and cₖ are per-position, but a = Σₖ aₖ is one scalar shared by all of + # them. + self.info = torch.zeros( + num_speculative_steps, 3, dtype=torch.float32, device=device + ) + # Columns [g₀ₖ, g₁ₖ] are each position's contribution to the Xᵀ(Y-P) vector: + # ⎡ g₀ ⎤ + # ⎢ g₁₀ ⎥ + # ⎢ g₁₁ ⎥ + # ⎢ ⋮ ⎥ + # ⎣ g₁ₙ₋₁ ⎦ + # As above, g₁ₖ is per-position and g₀ = Σₖ g₀ₖ is shared. + self.grad = torch.zeros( + num_speculative_steps, 2, dtype=torch.float32, device=device + ) + # Number of graded drafts per-position in the current round, zeroed after + # each refit. _refit_kernel uses this to damp coefficient updates by: + # n / (n + DAMPING_OBSERVATIONS) + self.counts = torch.zeros( + num_speculative_steps, dtype=torch.float32, device=device + ) + + def step( + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + ) -> None: + # Accumulate the previous step's graded drafts into the IRLS statistics. + num_reqs = idx_mapping.shape[0] + _accumulate_kernel[(self.num_speculative_steps,)]( + self.info, + self.grad, + idx_mapping, + num_sampled, + num_rejected, + self.features, + self.features.stride(0), + self.predictions, + self.predictions.stride(0), + self.counts, + num_reqs, + BLOCK_R=triton.next_power_of_2(max(num_reqs, 1)), + ) + + self._steps_since_refit += 1 + if self._steps_since_refit < self.REFIT_INTERVAL: + return + self._steps_since_refit = 0 + + # Fit the coefficients to the accumulated statistics gathered over the + # course of the last REFIT_INTERVAL steps. + _refit_kernel[(1,)]( + self.slope, + self.intercepts, + self.info, + self.info.stride(0), + self.grad, + self.grad.stride(0), + self.counts, + NUM_SPECULATIVE_STEPS=self.num_speculative_steps, + L2=self.L2, + DAMPING=self.DAMPING_OBSERVATIONS, + BLOCK=triton.next_power_of_2(self.num_speculative_steps), + ) + self._refits += 1 + + def predict( + self, + logits: torch.Tensor, + idx_mapping: torch.Tensor, + draft_step: torch.Tensor, + confidence_probs: torch.Tensor, + temperature: torch.Tensor, + ) -> None: + num_tokens, vocab_size = logits.shape + VOCAB_BLOCK_SIZE = 4096 + num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE) + local_max = torch.empty( + num_tokens, num_blocks, dtype=torch.float32, device=self.device + ) + local_sumexp = torch.empty_like(local_max) + _local_max_sumexp_kernel[(num_tokens, num_blocks)]( + local_max, + local_max.stride(0), + local_sumexp, + local_sumexp.stride(0), + logits, + logits.stride(0), + idx_mapping, + idx_mapping.stride(0), + temperature, + num_tokens, + vocab_size, + BLOCK_SIZE=VOCAB_BLOCK_SIZE, + ) + _predict_kernel[(num_tokens,)]( + self.features, + self.features.stride(0), + self.predictions, + self.predictions.stride(0), + confidence_probs, + confidence_probs.stride(0), + self.slope, + self.intercepts, + local_max, + local_max.stride(0), + local_sumexp, + local_sumexp.stride(0), + idx_mapping, + idx_mapping.stride(0), + draft_step, + num_tokens, + num_blocks, + per_token_step=draft_step.dim() > 0, + NUM_SPECULATIVE_STEPS=self.num_speculative_steps, + MAX_LOG_ODDS=_MAX_LOG_ODDS, + PADDED_VOCAB_NUM_BLOCKS=triton.next_power_of_2(num_blocks), + ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index 1cd19acdab..7db5284497 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -213,5 +213,11 @@ class DFlash2Speculator(DFlashSpeculator): anchor_token_ids, ) self._sample_path(candidate_ids, scores, num_reqs) + if self.enable_adaptive_verification: + self._maybe_predict_acceptance( + self._selector_scores[:num_reqs].flatten(0, 1), + self.sample_idx_mapping[:num_sample], + self.sample_col[:num_sample], + ) if self.draft_logits is not None: self._cache_draft_logits(candidate_ids, num_sample) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 036c3014f3..411aa7f041 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -79,12 +79,7 @@ class DSparkSpeculator(DFlashSpeculator): self.draft_model_config.hf_config, "dspark_draft_topk", None ) - self.draft_token_confidence_probs = torch.empty_like( - self.draft_tokens, dtype=torch.float32 - ) - self.enable_adaptive_verification = ( - self.speculative_config.enable_adaptive_verification - ) + self.use_confidence_head: bool = False def load_draft_model( self, @@ -108,13 +103,14 @@ class DSparkSpeculator(DFlashSpeculator): dtype=self.draft_logits.dtype, device=self.device, ) - if self.enable_adaptive_verification and model.model.confidence_head is None: - raise ValueError( - "Adaptive verification needs a DSpark checkpoint with a confidence " - "head, and this one has none. Pass " - "enable_adaptive_verification=false in the speculative config to verify" - " a fixed number of drafts instead." - ) + self.use_confidence_head = ( + self.enable_adaptive_verification + and model.model.confidence_head is not None + ) + if self.use_confidence_head: + # The acceptance estimator is not needed when a trained confidence head + # is available. + self.use_acceptance_estimator = False return model def _sample_logits( @@ -124,8 +120,10 @@ class DSparkSpeculator(DFlashSpeculator): sample_pos: torch.Tensor, step: int, ) -> torch.Tensor: + self._maybe_predict_acceptance(logits, idx_map, self._step_cols[step]) if self.draft_logits is None: - return self.model.map_draft_to_target(logits.argmax(dim=-1)) + draft_ids = logits.argmax(dim=-1) + return self.model.map_draft_to_target(draft_ids) # Probabilistic sampling and rejection operate in target-vocabulary # space. A reduced draft vocabulary is scattered into its target rows. @@ -184,7 +182,7 @@ class DSparkSpeculator(DFlashSpeculator): for i in range(n_spec): # Sequential stage: Markov bias from the previously sampled token. markov_embed = self.model.markov_embed(prev) - if self.enable_adaptive_verification: + if self.use_confidence_head: confidence_markov_embeds.append(markov_embed) bias = self.model.markov_bias(markov_embed) logits_i = base_logits[:, i] + bias @@ -194,7 +192,7 @@ class DSparkSpeculator(DFlashSpeculator): self.draft_tokens[:num_reqs, i] = draft_sampled_i prev = draft_sampled_i - if self.enable_adaptive_verification: + if self.use_confidence_head: confidence = self.model.compute_confidence( sample_hidden, torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1), @@ -229,7 +227,7 @@ class DSparkSpeculator(DFlashSpeculator): for i in range(n_spec): markov_embed = self.model.markov_embed(prev) - if self.enable_adaptive_verification: + if self.use_confidence_head: confidence_markov_embeds.append(markov_embed) logits_i = self.model.apply_markov_bias_gathered( markov_embed, @@ -243,7 +241,7 @@ class DSparkSpeculator(DFlashSpeculator): self.draft_tokens[:num_reqs, i] = draft_sampled_i prev = draft_sampled_i - if self.enable_adaptive_verification: + if self.use_confidence_head: confidence = self.model.compute_confidence( sample_hidden, torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1), diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index f670441219..299cac6416 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -33,6 +33,9 @@ from vllm.v1.worker.gpu.dp_utils import DPSyncState from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import ( + OnlineAcceptanceEstimator, +) from vllm.v1.worker.utils import AttentionGroup if TYPE_CHECKING: @@ -158,6 +161,15 @@ class DraftModelSpeculator(BaseSpeculator): self.arange_np = np.arange(self.max_num_reqs + 1, dtype=np.int32) self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool) + self.draft_token_confidence_probs = torch.empty_like( + self.draft_tokens, dtype=torch.float32 + ) + self.enable_adaptive_verification = ( + self.speculative_config.enable_adaptive_verification + ) + self.use_acceptance_estimator = self.enable_adaptive_verification + self.acceptance_estimator: OnlineAcceptanceEstimator | None = None + self.draft_logits: torch.Tensor | None = None if self.speculative_config.draft_sample_method == "probabilistic": # Pre-temperature logits, cached from the previous decode step. @@ -227,6 +239,20 @@ class DraftModelSpeculator(BaseSpeculator): type(self.model).__name__, ) + if self.use_acceptance_estimator: + if self.use_local_argmax_reduction: + raise ValueError( + "Adaptive verification without a confidence head estimates " + "per-position acceptance from the draft logits, which " + "use_local_argmax_reduction never materializes. Disable one " + "of them." + ) + self.acceptance_estimator = OnlineAcceptanceEstimator( + self.max_num_reqs, + self.num_speculative_steps, + self.device, + ) + def set_eplb_state(self, eplb_state: EplbState) -> None: """Inject EPLB state after construction.""" self.eplb_state = eplb_state @@ -371,12 +397,6 @@ class DraftModelSpeculator(BaseSpeculator): "(communication: O(2*tp_size) vs O(vocab_size))." ) - def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.use_local_argmax_reduction: - return self.model.get_top_tokens(hidden_states) - logits = self.model.compute_logits(hidden_states) - return logits.argmax(dim=-1) - def sample_draft( self, hidden_states: torch.Tensor, @@ -408,8 +428,42 @@ class DraftModelSpeculator(BaseSpeculator): idx_mapping, temperature, ) - return sampled - return self._greedy_sample_draft(hidden_states) + elif self.use_local_argmax_reduction: + return self.model.get_top_tokens(hidden_states) + else: + logits = self.model.compute_logits(hidden_states) + sampled = logits.argmax(dim=-1) + self._maybe_predict_acceptance(logits, idx_mapping, draft_step) + return sampled + + def _maybe_predict_acceptance( + self, + logits: torch.Tensor, + idx_mapping: torch.Tensor, + draft_step: torch.Tensor, + ) -> None: + if self.acceptance_estimator is not None: + self.acceptance_estimator.predict( + logits, + idx_mapping, + draft_step, + self.draft_token_confidence_probs, + self.temperature, + ) + + def observe_verification( + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + ) -> None: + """Fold the target's verdict on the last drafts into the estimator. + + Must run before the next `propose`, which overwrites the per-slot + features the verdict grades. + """ + if self.acceptance_estimator is not None: + self.acceptance_estimator.step(idx_mapping, num_sampled, num_rejected) def prepare_watermarking( self, contexts: torch.Tensor, watermarking: torch.Tensor From ba2ae9f23961ac67bc5c055da8c26fbd660989c6 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Mon, 14 Sep 2026 18:11:11 -0500 Subject: [PATCH 054/139] [MRV2] Run pooling post processing on non-final PP ranks (#56666) Signed-off-by: Taneem Ibrahim Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/v1/worker/gpu_worker.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 3f457b9b7b..2a3531b61b 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1238,6 +1238,8 @@ class Worker(WorkerBase): ) self._pp_send_work = handles[1:] + if self.use_v2_model_runner and self.model_runner.is_pooling_model: + return self.model_runner.pool() # type: ignore return None def take_draft_token_ids(self) -> DraftTokenIds | None: From f0a61bd4386bf0d694cfe848e853f2450462e454 Mon Sep 17 00:00:00 2001 From: JinYan Su <751080330@qq.com> Date: Tue, 15 Sep 2026 07:48:23 +0800 Subject: [PATCH 055/139] [Rust][Benchmark] Support HF ShareGPT datasets in multi-turn mode (#51104) Signed-off-by: xiaguan <751080330@qq.com> Co-authored-by: OpenAI Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- rust/src/bench/README.md | 1 + rust/src/bench/src/config.rs | 38 +++++++++++++++++++++++++++----- rust/src/bench/src/multi_turn.rs | 26 +++++++++++++++++++--- 3 files changed, 56 insertions(+), 9 deletions(-) diff --git a/rust/src/bench/README.md b/rust/src/bench/README.md index d4e0c69b80..c8724d6cd6 100644 --- a/rust/src/bench/README.md +++ b/rust/src/bench/README.md @@ -667,6 +667,7 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n - `--dataset-name random` — synthetic conversations with controllable per-turn token lengths. Auto-sets `min_tokens` to enforce output length without `ignore_eos`. - `--dataset-name sharegpt` — loads all turns (not just the first two); filters for entries with ≥ 2 real turns. +- `--dataset-name hf` — downloads a ShareGPT-format HuggingFace config, then loads it like `sharegpt`. **Prefix sharing** (random dataset): when `--multi-turn-prefix-global-ratio` or `--multi-turn-prefix-conversation-ratio` is > 0, each turn sends a fixed-length message (no history accumulation) composed of a global prefix + per-conversation prefix + unique suffix. The two ratios must sum to < 1.0. diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 3a0936b5f4..b5a24d942f 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -549,13 +549,15 @@ impl BenchConfig { )); } - // Normalize and validate min/max turns. ShareGPT only consumes max_turns - // (the loader walks all available turns up to the cap), so the - // min/num/max coupling used for synthetic generation does not apply. - if args.dataset_name == DatasetName::ShareGpt { + // Normalize and validate min/max turns. Conversation datasets only + // consume max_turns (the loader walks all available turns up to the + // cap), so the min/num/max coupling used for synthetic generation + // does not apply. + if matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) { if args.multi_turn_max_turns == 1 { return Err(BenchError::Config( - "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(), + "--multi-turn-max-turns must be at least 2 for ShareGPT-format multi-turn" + .into(), )); } } else { @@ -722,7 +724,7 @@ impl BenchConfig { multi_turn_min_turns, multi_turn_max_turns, sharegpt_multi_turn_max_turns: if args.multi_turn - && args.dataset_name == DatasetName::ShareGpt + && matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) && args.multi_turn_max_turns != 0 { Some(args.multi_turn_max_turns) @@ -1067,6 +1069,30 @@ mod tests { assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20)); } + #[test] + fn test_hf_multi_turn_uses_conversation_turn_cap() { + let args = vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + "--dataset-name", + "hf", + "--dataset-path", + "org/sharegpt-dataset", + "--hf-subset", + "sharegpt", + "--multi-turn-max-turns", + "2", + ]; + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); + + assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2)); + } + #[test] fn test_sweep_summary_percentiles_default_empty() { let args = base_multi_turn_args(); diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index e0a3138170..8eb175a07f 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -186,9 +186,29 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result { - return Err(BenchError::Config( - "HF dataset multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), - )); + let tok = tokenizer + .as_ref() + .ok_or_else(|| BenchError::Config("HF dataset requires a tokenizer".into()))?; + let dataset_id = config.dataset_path.as_deref().ok_or_else(|| { + BenchError::Config("--dataset-path is required for --dataset-name hf".into()) + })?; + let (downloaded_path, _config, _split) = + crate::datasets::hf_dataset::download_hf_dataset( + dataset_id, + config.hf_subset.as_deref(), + config.hf_split.as_deref(), + config.num_prompts, + ) + .await?; + crate::datasets::multi_turn::load_sharegpt_multi_turn( + tok, + &downloaded_path, + config.num_prompts, + config.hf_output_len, + config.sharegpt_multi_turn_max_turns, + config.seed, + &config.request_id_prefix, + )? } }; From bbbd0a02c9d1bec965ee818af29a1319baefa69c Mon Sep 17 00:00:00 2001 From: zhaoguochun1995 <109069909+zhaoguochun1995@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:52:20 +0800 Subject: [PATCH 056/139] [Bugfix] Carry over queued work when materializing the dedicated stream (#56382) Signed-off-by: zhaoguochun1995 --- vllm/utils/torch_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index b453487a01..6f7dc14346 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -816,8 +816,11 @@ def current_stream() -> torch.cuda.Stream: # https://github.com/pytorch/pytorch/blob/42ad9edfb754743fdae3276ade43de000beb4f60/aten/src/ATen/cuda/CUDAGraph.cpp#L77 # for more details. Therefore, we create a dedicated stream per process. if current_platform.is_rocm() or current_platform.is_cuda(): + # Ensure new stream is ordered w.r.t. replaced stream's work. + new_stream = torch.cuda.Stream() + new_stream.wait_stream(torch.cuda.current_stream()) # torch.cuda.set_stream here is the alias of _pathed_set_stream - torch.cuda.set_stream(torch.cuda.Stream()) + torch.cuda.set_stream(new_stream) elif current_platform.is_cpu(): _current_stream_tls.value = _StreamPlaceholder() else: From a7576447b86454f5c5729958d921271392ced47f Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 14 Sep 2026 19:34:21 -0500 Subject: [PATCH 057/139] [ROCm][Bugfix] Initialize Ray NIXL agents for sharded RDT (#56687) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: Codex --- .../distributed/test_sharded_rdt_producer.py | 114 ++++++++++++++++++ vllm/distributed/nixl_utils.py | 17 +++ .../weight_transfer/sharded_rdt_common.py | 10 ++ .../weight_transfer/sharded_rdt_engine.py | 5 +- .../weight_transfer/sharded_rdt_trainer.py | 5 +- 5 files changed, 144 insertions(+), 7 deletions(-) diff --git a/tests/distributed/test_sharded_rdt_producer.py b/tests/distributed/test_sharded_rdt_producer.py index 62fc638762..5cfab1957d 100644 --- a/tests/distributed/test_sharded_rdt_producer.py +++ b/tests/distributed/test_sharded_rdt_producer.py @@ -20,17 +20,23 @@ covered by the GPU tests in `test_sharded_rdt_trainer.py`. import contextlib import gc +import os +import sys import threading import time from concurrent.futures import ThreadPoolExecutor from concurrent.futures import TimeoutError as FuturesTimeout +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock import pytest import torch import vllm.distributed.weight_transfer.sharded_rdt_trainer as trainer_mod +from vllm.distributed.nixl_utils import alias_nixl_for_ray from vllm.distributed.weight_transfer.sharded_rdt_common import ( buffer_alloc_bytes, + register_nixl_memory, ) from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( DEFAULT_GATHER_LOOKAHEAD, @@ -45,6 +51,114 @@ GROUP_A = ("model.layers.0.w", "model.layers.0.b") GROUP_B = ("model.layers.1.w",) +@pytest.fixture +def ray_nixl_transport(monkeypatch): + import ray + from ray.experimental.rdt import util + from ray.experimental.rdt.nixl_tensor_transport import NixlTensorTransport + + from vllm.distributed import nixl_utils + from vllm.platforms import current_platform + + transport = NixlTensorTransport() + package = ModuleType("nixl_rocm") + api = ModuleType("nixl_rocm._api") + api.__dict__.update(nixl_agent=Mock(), nixl_agent_config=Mock()) + package.__dict__["_api"] = api + modules = {package.__name__: package, api.__name__: api} + imported_rcache_limits = [] + + def import_nixl(name): + imported_rcache_limits.append(os.environ.get("UCX_RCACHE_MAX_UNRELEASED")) + monkeypatch.setitem(sys.modules, name, modules[name]) + return modules[name] + + nixl_names = ("nixl", "nixl._api", *modules) + for name in nixl_names: + monkeypatch.delitem(sys.modules, name, raising=False) + monkeypatch.delenv("UCX_RCACHE_MAX_UNRELEASED", raising=False) + importer = Mock(side_effect=import_nixl) + monkeypatch.setattr( + nixl_utils, "importlib", SimpleNamespace(import_module=importer) + ) + monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: True) + monkeypatch.setattr(util, "transport_managers", {"NIXL": transport}) + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + monkeypatch.setattr( + ray, + "get_runtime_context", + lambda: SimpleNamespace(get_actor_id=lambda: "rdt-actor"), + ) + try: + yield transport, api, importer, imported_rcache_limits + finally: + # The helper inserts aliases directly; remove them before monkeypatch + # restores any modules and environment values present before this test. + for name in nixl_names: + sys.modules.pop(name, None) + os.environ.pop("UCX_RCACHE_MAX_UNRELEASED", None) + + +def test_ray_registers_memory_through_rocm_alias(ray_nixl_transport): + """Ray's registration must reach the ROCm agent through the alias.""" + _, api, importer, _ = ray_nixl_transport + register_nixl_memory(torch.empty(4, dtype=torch.uint8)) + api.nixl_agent.return_value.register_memory.assert_called() + assert sys.modules["nixl"] is sys.modules["nixl_rocm"] + assert sys.modules["nixl._api"] is api + + importer.reset_mock() + alias_nixl_for_ray() + importer.assert_not_called() + + +@pytest.mark.parametrize("rcache_limit", [None, "4096"]) +def test_ray_nixl_sets_ucx_limit_before_import( + monkeypatch, ray_nixl_transport, rcache_limit +): + """UCX must see the default or the user's override during the first import.""" + _, _, _, imported_rcache_limits = ray_nixl_transport + if rcache_limit is not None: + monkeypatch.setenv("UCX_RCACHE_MAX_UNRELEASED", rcache_limit) + alias_nixl_for_ray() + assert imported_rcache_limits == [rcache_limit or "1024"] * 2 + + +@pytest.mark.parametrize("existing_api", [False, True]) +def test_ray_nixl_preserves_non_rocm_or_existing_imports( + monkeypatch, ray_nixl_transport, existing_api +): + from vllm.platforms import current_platform + + _, _, importer, _ = ray_nixl_transport + if existing_api: + api = ModuleType("nixl._api") + monkeypatch.setitem(sys.modules, "nixl._api", api) + else: + monkeypatch.setattr(current_platform, "is_rocm", lambda: False) + alias_nixl_for_ray() + importer.assert_not_called() + assert "UCX_RCACHE_MAX_UNRELEASED" not in os.environ + if existing_api: + assert sys.modules["nixl._api"] is api + else: + assert "nixl._api" not in sys.modules + + +def test_ray_nixl_alias_reports_missing_rocm_package(monkeypatch, ray_nixl_transport): + from vllm.distributed import nixl_utils + + transport, api, importer, _ = ray_nixl_transport + monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: False) + with pytest.raises(ImportError, match="nixl_rocm"): + register_nixl_memory(torch.empty(1)) + assert transport._nixl_agent is None + api.nixl_agent.assert_not_called() + importer.assert_not_called() + assert "nixl" not in sys.modules + assert "nixl._api" not in sys.modules + + @pytest.fixture def server_factory(monkeypatch): """Build real `_RDTProducerServer` instances without a GPU. diff --git a/vllm/distributed/nixl_utils.py b/vllm/distributed/nixl_utils.py index 3879bee0ce..c063b3ea74 100644 --- a/vllm/distributed/nixl_utils.py +++ b/vllm/distributed/nixl_utils.py @@ -91,9 +91,26 @@ def is_nixl_available() -> bool: return pkg in sys.modules or importlib.util.find_spec(pkg) is not None +def alias_nixl_for_ray() -> None: + """Let Ray's ``nixl._api`` import resolve to the ROCm implementation.""" + if not current_platform.is_rocm() or "nixl._api" in sys.modules: + return + + package_name = _get_nixl_package_name() + if not is_nixl_available(): + raise ImportError(f"Ray's NIXL transport on ROCm requires {package_name}") + + _maybe_set_ucx_rcache_limit() + package = importlib.import_module(package_name) + api = importlib.import_module(f"{package_name}._api") + sys.modules.setdefault("nixl", package) + sys.modules.setdefault("nixl._api", api) + + __all__ = [ "NixlWrapper", "nixl_agent_config", "nixlXferTelemetry", "is_nixl_available", + "alias_nixl_for_ray", ] diff --git a/vllm/distributed/weight_transfer/sharded_rdt_common.py b/vllm/distributed/weight_transfer/sharded_rdt_common.py index a5a7ed9233..70110517d3 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_common.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_common.py @@ -77,6 +77,16 @@ def check_ray_rdt_version() -> None: ) +def register_nixl_memory(tensor: torch.Tensor) -> None: + """Register a tensor with Ray using the platform's NIXL package.""" + from vllm.distributed.nixl_utils import alias_nixl_for_ray + + alias_nixl_for_ray() + from ray.experimental import register_nixl_memory as ray_register_nixl_memory + + ray_register_nixl_memory(tensor) + + def assign_producer_indices( num_producers: int, num_consumers: int, consumer_idx: int ) -> list[int]: diff --git a/vllm/distributed/weight_transfer/sharded_rdt_engine.py b/vllm/distributed/weight_transfer/sharded_rdt_engine.py index 2b8187007b..96dd53f0fd 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_engine.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_engine.py @@ -66,6 +66,7 @@ from vllm.distributed.weight_transfer.sharded_rdt_common import ( RdtRouter, buffer_alloc_bytes, check_ray_rdt_version, + register_nixl_memory, ) from vllm.distributed.weight_transfer.sharded_rdt_fake import ( BakeSink, @@ -588,8 +589,6 @@ class ShardedRDTWeightTransferEngine( plan = self._cached_plan if plan is None or not plan.chunks: return - from ray.experimental import register_nixl_memory - # (a) consumer receive buffers — one per ring slot, at the largest chunk. max_pack = max(c.pack_bytes for c in plan.chunks) alloc = buffer_alloc_bytes(max_pack, self._buffer_presize) @@ -740,7 +739,7 @@ class ShardedRDTWeightTransferEngine( any time after the metadata push. See the doc's "slot generation handshake". """ - from ray.experimental import register_nixl_memory, set_target_for_ref + from ray.experimental import set_target_for_ref if self._slot_read_done: with self._slot_cv: diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py index 60d0e60fb0..1be17c4fa0 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py @@ -47,6 +47,7 @@ from vllm.distributed.weight_transfer.sharded_rdt_common import ( ALLOWED_OPS, buffer_alloc_bytes, check_ray_rdt_version, + register_nixl_memory, ) from vllm.logger import init_logger @@ -334,8 +335,6 @@ class _RDTProducerServer: doc). The warmup buffer stays registered so the agent's CUDA-HMEM path stays initialized. """ - from ray.experimental import register_nixl_memory - self._nixl_warmup_buf = torch.zeros(1 << 20, dtype=torch.uint8, device="cuda") with self._reg_lock: register_nixl_memory(self._nixl_warmup_buf) @@ -510,8 +509,6 @@ class _RDTProducerServer: """Allocate + NIXL-register one serve slot: the single allocation seam, so registration cannot be skipped on either of the two paths that make buffers (the init-time reservation and the serve-path backstop).""" - from ray.experimental import register_nixl_memory - t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device) with self._reg_lock: register_nixl_memory(t) From 772e4f05fa929491f027e61a73876a95e3d86e64 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 18:04:46 -0700 Subject: [PATCH 058/139] [Bugfix][Rust] Fix HF multi-turn dataset integration (#56915) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- rust/src/bench/src/datasets/multi_turn.rs | 57 +++++++++++++++++++++++ rust/src/bench/src/multi_turn.rs | 21 +++++---- 2 files changed, 68 insertions(+), 10 deletions(-) diff --git a/rust/src/bench/src/datasets/multi_turn.rs b/rust/src/bench/src/datasets/multi_turn.rs index 97ec522e29..48ff11ff7b 100644 --- a/rust/src/bench/src/datasets/multi_turn.rs +++ b/rust/src/bench/src/datasets/multi_turn.rs @@ -350,6 +350,27 @@ pub fn load_sharegpt_multi_turn( .as_array() .ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?; + load_sharegpt_multi_turn_from_rows( + tokenizer, + entries, + num_conversations, + output_len_override, + max_turns, + seed, + request_id_prefix, + ) +} + +/// Load multi-turn conversations from deserialized ShareGPT rows. +pub fn load_sharegpt_multi_turn_from_rows( + tokenizer: &TokenizerKind, + entries: &[serde_json::Value], + num_conversations: usize, + output_len_override: Option, + max_turns: Option, + seed: u64, + request_id_prefix: &str, +) -> Result> { // Filter entries with at least 4 messages (2 turns: user+assistant+user+assistant) let mut filtered: Vec<&serde_json::Value> = entries .iter() @@ -507,6 +528,42 @@ fn gen_prompt_to_target_len( mod tests { use super::*; + #[test] + fn test_load_sharegpt_multi_turn_from_hf_rows() { + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); + let rows = vec![serde_json::json!({ + "conversations": [ + {"from": "human", "value": "hello"}, + {"from": "gpt", "value": "hi"}, + {"from": "human", "value": "how are you?"}, + {"from": "gpt", "value": "great"} + ] + })]; + + let conversations = + load_sharegpt_multi_turn_from_rows(&tokenizer, &rows, 1, Some(7), None, 42, "hf-") + .unwrap(); + let actual = conversations + .iter() + .map(|conversation| { + ( + conversation.conversation_id.as_str(), + conversation + .turns + .iter() + .map(|turn| (turn.user_message.as_ref(), turn.expected_output_len)) + .collect::>(), + ) + }) + .collect::>(); + + assert_eq!( + actual, + vec![("hf-conv-0", vec![("hello", 7), ("how are you?", 7)])] + ); + } + fn common_prefix_bytes(strings: &[&str]) -> usize { if strings.is_empty() { return 0; diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index 8eb175a07f..366fd4f6fb 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -192,17 +192,18 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result Date: Mon, 14 Sep 2026 18:07:26 -0700 Subject: [PATCH 059/139] [ROCm][DSv4.1][Perf] Dequantize the MXFP8 weight once when dot_scaled cannot be used (#56560) Signed-off-by: JohnQinAMD Co-authored-by: Claude Opus 5 (1M context) --- tests/kernels/test_minimax_m3_amd_ops.py | 63 +++++++++++++++++++ .../kernels/linear/mxfp8/rocm_native.py | 18 ++++-- 2 files changed, 76 insertions(+), 5 deletions(-) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 182672597a..ec8f1fd4ea 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -424,6 +424,69 @@ def test_mxfp8_linear_emulation_bf16_at_load( assert _relerr(out.float(), out_ref.float()) < 2e-2 +# --------------------------------------------------------------------------- # +# Native MXFP8 linear: the K % 128 != 0 weight is dequantized once at load +# --------------------------------------------------------------------------- # +@requires_gfx950 +# (5120, 576) and (5120, 288): DSv4.1 shared-expert down_proj at TP4 and TP8. +@pytest.mark.parametrize("shape", [(5120, 576), (5120, 288), (256, 2048)]) +@torch.inference_mode() +def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape): + """``dot_scaled`` tiles K by 128, so an unaligned weight must not reach it. + + ``process_weights_after_loading`` converts those to BF16 once; a weight that + never went through it has to fall back per step rather than take the + invalid path, which is what keying the dispatch on the weight dtype alone + would do. + """ + from vllm.model_executor.kernels.linear.mxfp8 import rocm_native + from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearLayerConfig, + ) + + N, K = shape + aligned = K % 128 == 0 + torch.manual_seed(0) + w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1 + w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False) + x = torch.randn(7, K, device=DEVICE, dtype=torch.bfloat16) * 0.5 + out_ref = torch.nn.functional.linear(x, dequant_mxfp8_to_bf16(w_fp8, w_scale)) + # dot_scaled re-quantizes the activation, the dequant paths do not. + tol = 5e-2 if aligned else 2e-3 + + def make_layer(): + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False) + layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False) + return layer + + kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + layer = make_layer() + kernel.process_weights_after_loading(layer) + assert (layer.weight.element_size() >= 2) is not aligned + assert _relerr(kernel.apply_weights(layer, x).float(), out_ref.float()) < tol + + # An unaligned weight that is still FP8 must not reach dot_scaled. + unprocessed = make_layer() + unprocessed.weight_scale = torch.nn.Parameter( + w_scale[:N, : K // 32].contiguous(), requires_grad=False + ) + calls = [] + real = rocm_native._mxfp8_dot_scaled_linear + + def _spy(*a, **k): + calls.append(1) + return real(*a, **k) + + rocm_native._mxfp8_dot_scaled_linear = _spy + try: + out = kernel.apply_weights(unprocessed, x) + finally: + rocm_native._mxfp8_dot_scaled_linear = real + assert bool(calls) is aligned + assert _relerr(out.float(), out_ref.float()) < tol + + # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ──────────────── # The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER # experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 9714f92dff..29739afae3 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -4,9 +4,9 @@ Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16); activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling -matrix cores. Falls back (via the kernel selector) to the BF16 -``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with -``K % 128 != 0``. +matrix cores. ``dot_scaled`` tiles K by 128; a weight whose K is not a multiple +of that is dequantized to BF16 once at load and served by a plain linear +instead, since ``can_implement`` does not filter on K. """ import torch @@ -194,6 +194,9 @@ def _select_cfg(M, N, K): return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3) +_DOT_SCALED_K_ALIGN = 128 + + class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``.""" @@ -219,6 +222,8 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): N, K = weight.shape scale_k = K // MXFP8_BLOCK_SIZE weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() + if K % _DOT_SCALED_K_ALIGN != 0: + weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale) layer.weight = Parameter(weight.contiguous(), requires_grad=False) layer.weight_scale = Parameter(weight_scale, requires_grad=False) @@ -235,10 +240,13 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): ) out_shape = (*x.shape[:-1], layer.weight.shape[0]) x2d = x.reshape(-1, x.shape[-1]) - if x2d.shape[-1] % 128 == 0: + if layer.weight.element_size() >= 2: + out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype)) + elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0: out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale) else: - # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise. + # process_weights_after_loading dequantizes these weights once, so + # this runs only if it did not; dot_scaled would be invalid here. w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale) out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype) out = out.reshape(out_shape) From 84e5f3edac29d4b8cf7266d7882d7cb5f8dd0245 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 15 Sep 2026 09:17:41 +0800 Subject: [PATCH 060/139] [Bugfix] Validate routed-expert prompt offsets before engine submission (#56844) Signed-off-by: aoshen02 Co-authored-by: Bhushan Asati Co-authored-by: Megha Agarwal <19240983+meghaagr13@users.noreply.github.com> --- tests/test_request_input_bounds.py | 30 +++++++++++++++++++++++ tests/v1/core/test_scheduler.py | 39 ++++++++++++++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 2 +- vllm/v1/engine/input_processor.py | 22 ++++++++++------- 4 files changed, 83 insertions(+), 10 deletions(-) diff --git a/tests/test_request_input_bounds.py b/tests/test_request_input_bounds.py index 44f536e7fb..990f10cfeb 100644 --- a/tests/test_request_input_bounds.py +++ b/tests/test_request_input_bounds.py @@ -12,6 +12,7 @@ import subprocess import sys from collections.abc import Callable from typing import Protocol +from unittest.mock import Mock import pytest from pydantic import ValidationError @@ -27,10 +28,39 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.exceptions import VLLMValidationError from vllm.sampling_params import BeamSearchParams from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding +from vllm.v1.engine.input_processor import InputProcessor pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] +@pytest.mark.parametrize("offset", [-1, 0, 2, 3, 4, 2**127]) +def test_routed_experts_offset_validated_before_engine_submission(offset): + """Reject invalid offsets even when routed-expert output is disabled.""" + processor = Mock(spec=InputProcessor) + processor.tokenizer = None + processor.generation_config_fields = {} + processor.renderer = Mock() + processor.renderer.get_eos_token_id.return_value = None + processor.vllm_config = Mock() + params = SamplingParams(routed_experts_prompt_start=offset) + prompt = {"type": "token", "prompt_token_ids": [1, 2, 3]} + + if not 0 <= offset <= 3: + with pytest.raises(VLLMValidationError, match="routed_experts_prompt_start"): + InputProcessor.process_inputs( + processor, "invalid", prompt, params, ("generate",) + ) + params.routed_experts_prompt_start = 0 + + request = InputProcessor.process_inputs( + processor, "valid", prompt, params, ("generate",) + ) + assert ( + request.sampling_params.routed_experts_prompt_start + == params.routed_experts_prompt_start + ) + + # --- Stop strings: public requests cap the number of stop strings --------- diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 3a972e682e..6ca6a8cf87 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -51,6 +51,7 @@ from vllm.v1.outputs import ( ECConnectorOutput, KVConnectorOutput, ModelRunnerOutput, + RoutedExpertsLists, SamplingMaskLists, make_empty_encoder_model_runner_output, ) @@ -131,6 +132,44 @@ def test_add_requests(): assert len(scheduler.waiting) == i + 1 +def test_routed_experts_prompt_start_at_prompt_end(): + """A prompt-end offset should return empty routing data, not crash.""" + scheduler = create_scheduler() + (request,) = create_requests(num_requests=1) + request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens + scheduler.add_request(request) + scheduler_output = scheduler.schedule() + num_tokens = scheduler_output.num_scheduled_tokens[request.request_id] + + manager = Mock() + manager.routed_experts_by_slot = np.empty((0, 1, 1), dtype=np.uint8) + manager.get.return_value = np.empty((0, 1, 1), dtype=np.uint8) + scheduler.enable_return_routed_experts = True + scheduler.routed_experts_mgr = manager + scheduler._re_block_ids = {request.request_id: []} + + outputs = scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[0]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + routed_experts=RoutedExpertsLists( + routing_data=np.zeros((num_tokens, 1, 1), dtype=np.uint8), + slot_mapping=np.arange(num_tokens), + ), + ), + ) + + manager.get.assert_called_once_with( + [], request.num_prompt_tokens, token_start=request.num_prompt_tokens + ) + assert outputs[request.client_index].outputs[0].routed_experts.shape == (0, 1, 1) + + def test_finish_request(): scheduler = create_scheduler() requests = create_requests(num_requests=10) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 1b17eb3b15..3af527124d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2101,7 +2101,7 @@ class Scheduler(SchedulerInterface): prompt_start = ( request.sampling_params.routed_experts_prompt_start ) - assert prompt_start < request.num_prompt_tokens + assert prompt_start <= request.num_prompt_tokens else: prompt_start = 0 routed_experts = self.routed_experts_mgr.get( diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 4cc433cbfe..3f86f3c8b2 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -354,12 +354,21 @@ class InputProcessor: if isinstance(params, SamplingParams): # TODO: can we avoid cloning here in multiproc case? sampling_params = params.clone() + prompt_len = length_from_prompt_token_ids_or_embeds( + prompt_token_ids, prompt_embeds + ) + if not 0 <= sampling_params.routed_experts_prompt_start <= prompt_len: + raise VLLMValidationError( + f"routed_experts_prompt_start must be between 0 and " + f"the prompt length ({prompt_len}), inclusive.", + parameter="routed_experts_prompt_start", + value=sampling_params.routed_experts_prompt_start, + ) # If unset max tokens, then generate up to the max_model_len. if sampling_params.max_tokens is None: - seq_len = length_from_prompt_token_ids_or_embeds( - prompt_token_ids, prompt_embeds + sampling_params.max_tokens = ( + self.model_config.max_model_len - prompt_len ) - sampling_params.max_tokens = self.model_config.max_model_len - seq_len sampling_params.update_from_generation_config( self.generation_config_fields, @@ -368,12 +377,7 @@ class InputProcessor: if self.tokenizer is not None: sampling_params.update_from_tokenizer(self.tokenizer) if sampling_params.trace_decode_token_ids: - self._normalize_trace_replay_params( - sampling_params, - length_from_prompt_token_ids_or_embeds( - prompt_token_ids, prompt_embeds - ), - ) + self._normalize_trace_replay_params(sampling_params, prompt_len) else: pooling_params = params.clone() From ceaa9f95ddbae0203ae444ec864bbbfb68b9b046 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 14 Sep 2026 18:56:04 -0700 Subject: [PATCH 061/139] [Minor] Invert MRV1 specdec method fallback logic (#56899) Signed-off-by: Nick Hill --- vllm/config/vllm.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7d2baea26c..ad8f077148 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2829,16 +2829,16 @@ class VllmConfig: unsupported.append("pipeline parallelism with external_launcher") if speculative_config is not None: - # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet - if speculative_config.method in ("ngram", "ngram_gpu"): - unsupported.append("ngram/ngram_gpu speculative decoding") - elif speculative_config.method not in ( - "eagle", - "eagle3", - "mtp", - "dflash", - "dspark", - "extract_hidden_states", + if speculative_config.method in ( + # https://github.com/vllm-project/vllm/pull/40704 + "ngram", + "ngram_gpu", + # https://github.com/vllm-project/vllm/pull/43091 + "draft_model", + "suffix", + "medusa", + "mlp_speculator", + "custom_class", ): unsupported.append(f"speculative method '{speculative_config.method}'") From 31928987546ff6454f9a7d945f3ff132c5f9c531 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:57:59 -0700 Subject: [PATCH 062/139] [ROCm][CI] Prepare TheRock image for CI (#56921) Signed-off-by: Matthew Wong Co-authored-by: Andreas Karatzas --- docker/Dockerfile.rock | 9 +++++++++ docker/Dockerfile.rock_base | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index c3ce785c67..14ba523b07 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rock @@ -875,6 +875,15 @@ RUN mkdir src && mv vllm src/vllm # Catch GPU<->CPU syncs in execute_model/sample_tokens during tests. ENV VLLM_GPU_SYNC_CHECK=error +FROM test AS test_smoke +ARG ROCM_SMOKE_ID +RUN --network=none \ + bash /vllm-workspace/.buildkite/scripts/rocm/smoke-test-image.sh --inside \ + && printf '%s\n' "${ROCM_SMOKE_ID:-local}" > /vllm-smoke-ok + +FROM scratch AS export_test_smoke +COPY --from=test_smoke /vllm-smoke-ok / + # ----------------------- # Final vLLM image FROM mori_base AS final_common diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base index 1b1d089881..80b6aaca29 100644 --- a/docker/Dockerfile.rock_base +++ b/docker/Dockerfile.rock_base @@ -146,6 +146,15 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \ fi; \ done \ + && for f in librocprofiler-sdk.so.1 librocprofiler-register.so.0 \ + librocprofiler-sdk-roctx.so.1; do \ + if [ -e "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}" ] \ + && [ -e "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" ]; then \ + rm -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ + ln -sfn "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" \ + "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ + fi; \ + done \ && ldconfig # Install sccache if USE_SCCACHE is enabled (for release builds) From a08dffe21534916f218dcc93012dbaa01f942691 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Tue, 15 Sep 2026 10:23:14 +0800 Subject: [PATCH 063/139] [XPU][CI] update gpt-oss package version (#56783) Signed-off-by: zhenwei-intel --- .buildkite/intel_jobs/lm_eval_intel.yaml | 3 ++- requirements/test/xpu.txt | 19 ++----------------- 2 files changed, 4 insertions(+), 18 deletions(-) diff --git a/.buildkite/intel_jobs/lm_eval_intel.yaml b/.buildkite/intel_jobs/lm_eval_intel.yaml index 11d3edc90e..1ef00a4794 100644 --- a/.buildkite/intel_jobs/lm_eval_intel.yaml +++ b/.buildkite/intel_jobs/lm_eval_intel.yaml @@ -23,5 +23,6 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'cd tests && + 'uv pip install "gpt-oss[eval]==0.0.5" && + cd tests && pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt' diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index b7cf49649b..2e2fbd88ed 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -257,7 +257,7 @@ googleapis-common-protos==1.75.1 # google-api-core # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -gpt-oss==0.0.9 +gpt-oss==0.0.8 # via -r requirements/test/xpu.in graphql-core==3.2.11 # via hypothesis-graphql @@ -406,10 +406,6 @@ jsonschema-rs==0.49.9 # via schemathesis jsonschema-specifications==2025.9.1 # via jsonschema -jupyter-client==8.10.0 - # via gpt-oss -jupyter-core==5.9.1 - # via jupyter-client kiwisolver==1.5.1 # via matplotlib lark==1.2.2 @@ -672,9 +668,7 @@ pillow==12.3.0 # mistral-common # torchvision platformdirs==4.11.3 - # via - # jupyter-core - # pooch + # via pooch pluggy==1.6.0 # via # pytest @@ -813,7 +807,6 @@ pytest-timeout==2.4.0 python-dateutil==2.9.0.post0 # via # botocore - # jupyter-client # matplotlib # pandas # typepy @@ -849,7 +842,6 @@ pyzmq==27.1.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # jupyter-client rapidfuzz==3.14.5 # via # -r requirements/test/xpu.in @@ -1074,8 +1066,6 @@ torchvision==0.28.0+xpu # via # open-clip-torch # timm -tornado==6.5.8 - # via jupyter-client tqdm==4.70.0 # via # -r requirements/test/../common.txt @@ -1091,10 +1081,6 @@ tqdm==4.70.0 # pqdm # sentence-transformers # transformers -traitlets==5.16.1 - # via - # jupyter-client - # jupyter-core transformers==5.16.1 # via # -c requirements/common.txt @@ -1146,7 +1132,6 @@ typing-extensions==4.16.0 # grpcio # httpx2 # huggingface-hub - # jupyter-client # lm-eval # mcp # mcp-types From 17bc3e106acbaeb552892893373659bb68d5020e Mon Sep 17 00:00:00 2001 From: liranschour Date: Tue, 15 Sep 2026 05:38:18 +0300 Subject: [PATCH 064/139] [KVConnector][P2P] Configurable unbound-store timeout and one-RTT rejection of a late fetch (#53453) Signed-off-by: Liran Schour --- docs/features/kv_offloading_usage.md | 22 ++++ .../v1/kv_offload/tiering/p2p/test_manager.py | 79 ++++++++++++- vllm/v1/kv_offload/tiering/p2p/manager.py | 106 ++++++++++++++++-- 3 files changed, 193 insertions(+), 14 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 54a54d8988..a5c8684b42 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -203,9 +203,31 @@ Block content hashes must match across instances for peers to exchange blocks (s | `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. | | `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. | | `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. | +| `unbound_store_timeout_s` | no | `60` | Seconds a producer holds stored blocks for a consumer that has not fetched them yet. Raise it for deployments whose prefills outlast the default; the blocks keep primary-tier CPU slots pinned for the whole window. Once it expires a late fetch is rejected in a single round trip, so the consumer falls back to local prefill immediately. | The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process. +A producer parks a request's blocks until the consumer's `FetchMsg` arrives; if none arrives +within `unbound_store_timeout_s` the blocks are released so they stop pinning primary-tier +slots. Raise it when prefills legitimately take longer than the default: + +```bash +vllm serve \ + --kv-transfer-config '{ + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "spec_name": "TieringOffloadingSpec", + "secondary_tiers": [ + { + "type": "p2p", + "unbound_store_timeout_s": 180 + } + ] + } + }' +``` + #### Environment Variables Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence. diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index dc0e384d7c..aceec1cd8f 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -117,6 +117,8 @@ def _make_manager() -> P2PSecondaryTierManager: mgr._sessions = {} mgr._kv_to_session = {} mgr._unbound_stores = {} + mgr._reaped_stores = {} + mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S mgr._failed_serve_ctxs = [] return mgr @@ -759,9 +761,9 @@ class TestGetFinished: assert "req-fresh" in mgr._unbound_stores def test_unbound_store_reaped_after_timeout(self): - """Unbound stores past _UNBOUND_STORE_TIMEOUT_S surface as failed - and their kv_request_id lands in _failed_req_ids so a late - FetchMsg/lookup doesn't try to satisfy them.""" + """Unbound stores past the reap deadline surface as failed and their + kv_request_id lands in _reaped_stores so a late FetchMsg is rejected + instead of parking demand nothing can satisfy.""" from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = self._make() @@ -779,7 +781,7 @@ class TestGetFinished: # 2 baseline + 2 buffered stores assert JobResult(job_id=10, success=False) in results assert JobResult(job_id=11, success=False) in results - assert "req-stale" in mgr._failed_req_ids + assert "req-stale" in mgr._reaped_stores def test_submit_store_parks_unbound_batch(self): """submit_store on an unbound id appends a batch with a fresh @@ -1204,6 +1206,75 @@ class TestBidirectionalManager: assert 201 in b_ok, f"B loads succeeded: {b_ok}" assert 200 in b_ok, f"B stores succeeded: {b_ok}" + def test_late_fetch_after_reap_fails_immediately(self): + """A fetch for a reaped kv_request_id fails in one round trip. + + The producer drops parked blocks once unbound_store_timeout_s + expires. A consumer whose fetch arrives after that must learn on the + next tick via TransferDoneMsg(success=False) — not stall for + _LOAD_TIMEOUT_S and then take the abort path. + """ + from vllm.v1.kv_offload.tiering.p2p.session.client import _LOAD_TIMEOUT_S + + mgr_a, mgr_b = _build_paired_managers() + kv_id = "req-late" + + # Record every message type leaving the consumer, so the assertion + # below can prove no abort was ever needed. + sent_from_a: list[str] = [] + drain = mgr_a._control._drain_outbound_to + + def recording_drain(peer_local_id: str): + out = drain(peer_local_id) + sent_from_a.extend(msg.get("type") for _, msg in out) + return out + + mgr_a._control._drain_outbound_to = recording_drain # type: ignore[method-assign] + + # Producer parks the blocks, then reaps them before any fetch lands. + mgr_b.submit_store( + _job_metadata( + job_id=200, + keys=[b"b-block"], + chunk_ids=[0], + kv_params={"remote_decoder": {"kv_request_id": kv_id}}, + ) + ) + mgr_b._unbound_store_timeout_s = 0.0 + reap_results = list(mgr_b.get_finished_jobs()) + assert JobResult(job_id=200, success=False) in reap_results + assert kv_id in mgr_b._reaped_stores + + # Consumer now asks for the blocks that no longer exist. + consumer_params = { + "remote_prefiller": { + "kv_request_id": kv_id, + "remote_host": "B", + "remote_port": 2, + }, + } + mgr_a.on_new_request(_req_context(consumer_params)) + mgr_a.submit_load( + _job_metadata( + job_id=101, + keys=[b"b-block"], + chunk_ids=[0], + kv_params=consumer_params, + ) + ) + + started = time.monotonic() + all_a: list[JobResult] = [] + for _ in range(8): + all_a.extend(list(mgr_a.get_finished_jobs())) + list(mgr_b.get_finished_jobs()) + elapsed = time.monotonic() - started + + assert JobResult(job_id=101, success=False) in all_a, all_a + assert kv_id in mgr_a._failed_req_ids + assert "abort_fetch" not in sent_from_a, sent_from_a + assert elapsed < _LOAD_TIMEOUT_S + # --------------------------------------------------------------------------- # _accept_new_peers — duplicate connection rejection diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 2cd4901710..ba1973b767 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -47,10 +47,18 @@ logger = init_logger(__name__) # Reap unbound store batches that have been parked without a FetchMsg # binding them to a session for longer than this. Protects against the # prefiller buffering blocks for a decoder that never asks (decoder died, -# network partition, lost kv_request_id). Must be longer than the per-store -# deadline so the store-timeout path fires first for individual jobs. +# network partition, lost kv_request_id). Applies only while a batch is +# unbound: binding hands it to the session, which times the job out under +# _STORE_TIMEOUT_S instead. +# Overridable per tier via the ``unbound_store_timeout_s`` config key. _UNBOUND_STORE_TIMEOUT_S = 60.0 +# How long a reaped kv_request_id is remembered so a fetch that arrives +# after the reap can be rejected immediately. Must exceed the consumer's +# load timeout (session.client._LOAD_TIMEOUT_S): past that the consumer has +# already given up on its own, so there is nobody left to reject. +_REAPED_ID_RETENTION_S = 60.0 + # Time we wait during shutdown for inflight transfers to drain via # cancel(mode="wait") before falling back to mode="immediate". Bounded # so a wedged peer can't hang shutdown. @@ -208,6 +216,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): port: int | None = None, backends: list[str] | None = None, num_threads: int = 4, + unbound_store_timeout_s: float = _UNBOUND_STORE_TIMEOUT_S, **kwargs: Any, ) -> None: """Initialize the P2P secondary tier manager. @@ -245,9 +254,29 @@ class P2PSecondaryTierManager(SecondaryTierManager): num_threads: NIXL agent worker threads for the UCX-only branch. Ignored when ``backends`` contains a non-UCX entry. + unbound_store_timeout_s: Seconds a producer holds stored blocks + for a consumer that has not fetched them yet, before + ``_reap_unbound_stores`` drops them. Raise it for + deployments whose prefills outlast the default; the blocks + keep primary-tier slots pinned for the whole window. **kwargs: Reserved for future tier-specific options. + + Raises: + ValueError: If ``unbound_store_timeout_s`` is not a positive + number, or anything convertible to one. """ super().__init__(offloading_spec, primary_kv_view, tier_type) + try: + timeout_s = float(unbound_store_timeout_s) + except (TypeError, ValueError): + timeout_s = float("nan") + if not timeout_s > 0: + raise ValueError( + f"unbound_store_timeout_s must be a positive number, got " + f"{unbound_store_timeout_s!r}" + ) + self._unbound_store_timeout_s = timeout_s + # Block hashes chain from NONE_HASH (see v1/core/kv_cache_utils.py). # Peers whose seeds differ compute different hashes for identical # content, so lookups silently miss and no KV crosses the wire. The @@ -303,8 +332,16 @@ class P2PSecondaryTierManager(SecondaryTierManager): # kv_request_id → list of batches submit_store'd before any peer # asked for that id. Drained into a session by _on_session_fetch # when the corresponding FetchMsg arrives, or surfaced as failures - # by _reap_unbound_stores after _UNBOUND_STORE_TIMEOUT_S. + # by _reap_unbound_stores after _unbound_store_timeout_s. self._unbound_stores: dict[str, list[_UnboundStoreBatch]] = {} + # kv_request_id → time its parked batches were reaped. Nothing under + # one of these ids can be served any more, so a FetchMsg for it is + # rejected on the spot by _poll_once instead of parking demand until + # the consumer's own load timeout, and submit_store fails rather than + # re-parking under it — a prefill that outran the timeout would + # otherwise keep pinning fresh primary-tier slots, one timeout at a + # time. Entries are pruned after _REAPED_ID_RETENTION_S. + self._reaped_stores: dict[str, float] = {} self._finished_jobs: list[JobResult] = [] # kv_request_ids that hit a transport/session failure; On load lookup() @@ -390,7 +427,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): kv_transfer_params; if a session has bound the id, finish it. If no session has bound the id yet, this is a no-op: parked batches in `_unbound_stores` are left in place and cleaned up only by - `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`. + `_reap_unbound_stores` after `_unbound_store_timeout_s`. """ source = req_context.get_state(P2PSourceInfo) dest = req_context.get_state(P2PDestInfo) @@ -447,6 +484,22 @@ class P2PSecondaryTierManager(SecondaryTierManager): self._finished_jobs.append(JobResult(job_id=job_id, success=False)) return + # This id's earlier batches were already reaped, so the consumer has + # either given up or will be rejected by the fetch path. Parking more + # would pin primary-tier slots for another full timeout on blocks + # nobody can fetch, so fail the job now — a prefill that outran the + # timeout falls back to local recompute on the decoder either way. + if kv_request_id in self._reaped_stores: + logger.warning( + "P2P %s: submit_store for reaped kv_request_id=%s job_id=%d " + "— failing without parking", + self._local_id, + kv_request_id, + job_id, + ) + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + return + # Fast path: a session has already received FetchMsg for this id, # so we can route the batch straight into its ServerRole. session = self._kv_to_session.get(kv_request_id) @@ -707,13 +760,27 @@ class P2PSecondaryTierManager(SecondaryTierManager): """Time out submit_store batches that no peer has ever fetched. Walks `_unbound_stores` for entries whose oldest batch is older - than `_UNBOUND_STORE_TIMEOUT_S`. Drops the kv_request_id, surfaces - every batched job as failed, and adds the id to `_failed_req_ids` - so a late inbound FetchMsg short-circuits to a clean rejection. + than `_unbound_store_timeout_s`. Drops the kv_request_id, surfaces + every batched job as failed, and records the id in `_reaped_stores` + so a late inbound FetchMsg is rejected in one round trip and a late + `submit_store` fails instead of re-parking under the reaped id. + + Also expires `_reaped_stores` entries whose consumer can no longer + be waiting, which is why this runs before the empty check below. """ + now = time.monotonic() + if self._reaped_stores: + retention_deadline = now - _REAPED_ID_RETENTION_S + for kid in [ + kid + for kid, reaped_at in self._reaped_stores.items() + if reaped_at <= retention_deadline + ]: + del self._reaped_stores[kid] + if not self._unbound_stores: return - deadline = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S + deadline = now - self._unbound_store_timeout_s expired: list[str] | None = None for kid, batches in self._unbound_stores.items(): # Batches are appended in arrival order, so the head is oldest. @@ -725,7 +792,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): return for kid in expired: batches = self._unbound_stores.pop(kid) - self._failed_req_ids.add(kid) + self._reaped_stores[kid] = now for batch in batches: self._finished_jobs.append( JobResult(job_id=batch.job_id, success=False) @@ -735,7 +802,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): "without a fetch — failing %d job(s)", self._local_id, kid, - _UNBOUND_STORE_TIMEOUT_S, + self._unbound_store_timeout_s, len(batches), ) @@ -779,6 +846,24 @@ class P2PSecondaryTierManager(SecondaryTierManager): # inline in dispatch, so the replayed add_stored_blocks calls # match that demand and submit transfers immediately. for kv_request_id in result.new_fetch_ids: + if kv_request_id in self._reaped_stores: + # This id's blocks were already reaped, so no + # submit_store will ever satisfy the demand on_fetch + # just recorded. Finalize the round now: the peer gets + # TransferDoneMsg(success=False) on this tick and falls + # back to local prefill, instead of parking until its + # own load timeout expires. The id is deliberately left + # unbound so any later submit_store keeps taking the + # reject path. Batches re-parked between the reap and the + # retention prune are failed here rather than left to a + # second reap, which would pin their slots for another + # full timeout. + session.finish_request(kv_request_id) + for batch in self._unbound_stores.pop(kv_request_id, ()): + self._finished_jobs.append( + JobResult(job_id=batch.job_id, success=False) + ) + continue self._kv_to_session[kv_request_id] = session for batch in self._unbound_stores.pop(kv_request_id, ()): session.add_stored_blocks( @@ -810,6 +895,7 @@ class P2PSecondaryTierManager(SecondaryTierManager): JobResult(job_id=batch.job_id, success=False) ) self._unbound_stores.clear() + self._reaped_stores.clear() self._control.close() self._data.close() From 567f745c1b14ffaf09bfa71b7c4d0177a7994c04 Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Mon, 14 Sep 2026 19:43:42 -0700 Subject: [PATCH 065/139] [Bugfix] Export weight-cache IPC tensors separately for each client (#56472) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Signed-off-by: Isotr0py Co-authored-by: Codex Co-authored-by: Isotr0py --- .../model_loader/weight_cache/daemon.py | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index b19fb96272..f0af5a284a 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -93,11 +93,14 @@ def export_entries( ``named_parameters``/``named_buffers`` are iterated with ``remove_duplicate=False`` so tied weights (e.g. ``lm_head.weight`` sharing storage with ``embed_tokens.weight``) are not silently dropped. Each unique - tensor is exported once; every additional name that refers to the same + tensor is exported once per call; every additional name that refers to the same tensor object is recorded in the returned alias map so the client can re-establish the shared identity instead of allocating uninitialized memory for it. + CUDA reduction arguments must be exported separately for each consumer so + that PyTorch registers a reference for each IPC mapping's lifetime. + Returns: A ``(entries, aliases)`` pair where ``entries`` maps a canonical name to its ``TensorEntry`` and ``aliases`` maps each duplicate name to its @@ -167,8 +170,6 @@ class WeightCacheDaemon: self.distributed_init_method = distributed_init_method self.socket_dir = socket_dir self.model: torch.nn.Module | None = None - self.entries: dict[str, TensorEntry] = {} - self.aliases: dict[str, str] = {} # Fingerprint before loading: process_weights_after_loading may # mutate hf_config.quantization_config. self.cache_config = WeightCacheKey.from_model_config( @@ -190,17 +191,11 @@ class WeightCacheDaemon: with set_current_vllm_config(self.vllm_config): ensure_model_parallel_initialized(tp_size, 1) self.model = get_daemon_model(self.vllm_config) - self._export_entries() logger.info( - "Weight cache daemon rank %d cached %d tensors", + "Weight cache daemon rank %d loaded model", self.tp_rank, - len(self.entries), ) - def _export_entries(self) -> None: - assert self.model is not None - self.entries, self.aliases = export_entries(self.model) - def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> None: """Serve requests until terminated. @@ -295,28 +290,28 @@ class WeightCacheDaemon: logger.warning("WeightCacheKey mismatch on fields: %s", mismatched) send_msg(conn, {"status": "mismatch", "fields": mismatched}) return - if not self.entries: + if self.model is None: send_msg(conn, {"status": "error", "message": "Weights were released"}) return + gpu_uuid = get_current_device_uuid() + entries, aliases = export_entries(self.model) send_msg( conn, { "status": "ok", - "entries": self.entries, - "aliases": self.aliases, - "gpu_uuid": get_current_device_uuid(), + "entries": entries, + "aliases": aliases, + "gpu_uuid": gpu_uuid, }, ) logger.info_once( "Weight cache daemon rank %d sent %d tensors (+%d aliases) to engine", self.tp_rank, - len(self.entries), - len(self.aliases), + len(entries), + len(aliases), ) def _handle_release(self, conn: socket.socket) -> None: - self.entries.clear() - self.aliases.clear() self.model = None torch.accelerator.empty_cache() logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank) From e45bb43984b60fbbddd6d2ba4acd6e4a5525f49c Mon Sep 17 00:00:00 2001 From: Liqian Lin Date: Tue, 15 Sep 2026 10:48:38 +0800 Subject: [PATCH 066/139] [Bugfix][KV Offload] Ignore pending chunks in invalid sliding windows (#56486) Signed-off-by: Liqian Lin --- .../offloading_connector/test_scheduler.py | 78 ++++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 12 ++- 2 files changed, 86 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 26664e3ac5..e7a83cb7b5 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -74,7 +74,7 @@ from vllm.v1.kv_offload.base import ( make_offload_key, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager -from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus @@ -1708,6 +1708,15 @@ class TestMaximalPrefixLookup: class TestSlidingWindowLookup: + def test_pending_chunk_behind_gap_does_not_defer(self): + """A pending suffix too short for the window cannot improve the hit.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT, 4: LookupResult.HIT_PENDING} + ) + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 2 + ) + def test_all_hit_exact_window(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2 @@ -1826,6 +1835,73 @@ class TestSlidingWindowLookup: # --------------------------------------------------------------------------- +@pytest.mark.parametrize("async_scheduling", [False, True]) +def test_swa_load_does_not_wait_for_unusable_pending_suffix( + request_runner, async_scheduling +): + """A CPU store beyond a missing SWA chunk must not block a ready prefix.""" + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["swa"], + SlidingWindowSpec( + block_size=4, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=8, + ), + ), + ] + runner = request_runner( + block_size=4, + num_gpu_blocks=32, + async_scheduling=async_scheduling, + kv_cache_groups=groups, + ) + sched = runner.connector_scheduler + manager = CPUOffloadingManager(num_chunks=16) + sched.manager = manager + runner.new_request(token_ids=list(range(17))) + state = sched._req_status["0"] + state.update_offload_keys() + full, swa = (g.offload_keys for g in state.group_states) + ready = manager.prepare_store(full[:4] + swa[:2], state.req_context) + assert ready is not None + manager.complete_store(ready.keys_to_store, state.req_context) + pending = manager.prepare_store(swa[3:4], state.req_context) + assert pending is not None + + # The last SWA chunk is still being written; the preceding one is missing. + output = runner.scheduler.schedule() + metadata = output.kv_connector_metadata + assert isinstance(metadata, OffloadingConnectorMetadata) + [(job_id, load)] = metadata.load_jobs.items() + assert isinstance(load.dst_spec, GPULoadStoreSpec) + assert load.dst_spec.group_sizes == [2, 2] + assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING + + runner.scheduler.update_from_output( + output, + ModelRunnerOutput.with_kv_conn_output_only( + KVConnectorOutput( + finished_recving={"0"}, + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 1} + ), + ) + ), + ) + resumed = runner.scheduler.schedule() + assert resumed.num_scheduled_tokens == {"0": 9} + assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING + + @pytest.mark.parametrize("alignment_chunk_count", [4, 8, 64]) @pytest.mark.parametrize("sliding_window_chunks", [1, 2, 3]) @pytest.mark.parametrize("is_eagle_group", [False, True]) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 86e3510939..2fb3c5bb90 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -685,6 +685,7 @@ class OffloadingConnectorScheduler: The first run may need a larger window for a partial rightmost chunk. Returns 0 on miss, None if the backend deferred a lookup.""" defer_lookup = False + pending_in_window = False consecutive_hits = 0 required_window = initial_window_size or sliding_window_size for idx in range(len(keys) - 1, -1, -1): @@ -695,7 +696,7 @@ class OffloadingConnectorScheduler: # Block is in cache, just not readable yet — counts # as hit for the consecutive streak. Don't break: # keep scanning to let manager kick off async lookups. - defer_lookup = True + pending_in_window = True consecutive_hits += 1 case LookupResult.RETRY: # Block location uncertain — does not count as hit. @@ -703,13 +704,18 @@ class OffloadingConnectorScheduler: # async lookups. defer_lookup = True consecutive_hits = 0 + pending_in_window = False required_window = sliding_window_size case LookupResult.MISS: consecutive_hits = 0 + # This gap rules out the incomplete window to its right. + pending_in_window = False required_window = sliding_window_size if consecutive_hits == required_window: - return idx + required_window if not defer_lookup else None - return consecutive_hits if not defer_lookup else None + return ( + None if defer_lookup or pending_in_window else idx + required_window + ) + return None if defer_lookup or pending_in_window else consecutive_hits def _touch(self, req_status: RequestOffloadState): for group_config, group_state in zip( From a529c1a748d3e3e21b3f97883b578e126892abf1 Mon Sep 17 00:00:00 2001 From: wuhangxian <1391938827@qq.com> Date: Tue, 15 Sep 2026 10:52:43 +0800 Subject: [PATCH 067/139] [Bugfix] Handle bare and malformed tool call openers in Gemma4 parser (#53444) Signed-off-by: wuhangxian <1391938827@qq.com> Co-authored-by: Chauncey --- tests/tool_parsers/test_gemma4_tool_parser.py | 108 ++++++++++++++++++ vllm/parser/gemma4.py | 8 ++ 2 files changed, 116 insertions(+) diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 8d74f04319..97ed0cd69d 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -897,3 +897,111 @@ class TestStreamingExtraction: } assert args_text.count("replace_all") == 1 + + +class TestBareAndMalformedOpeners: + """Tests for bare and malformed tool call openers. + + See https://github.com/vllm-project/vllm/issues/53431. + """ + + def test_bare_opener_non_streaming(self, parser, mock_request): + model_output = '<|tool_call>:get_weather{location:<|"|>London<|"|>}' + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"location": "London"} + + def test_documented_form_still_works(self, parser, mock_request): + model_output = ( + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + + def test_bare_opener_multiple_args(self, parser, mock_request): + model_output = ( + "<|tool_call>:get_weather{" + 'location:<|"|>Tokyo<|"|>,unit:<|"|>celsius<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "get_weather" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"location": "Tokyo", "unit": "celsius"} + + def test_malformed_no_brace_bounded(self, parser, mock_request): + model_output = "<|tool_call>call:bad_func no brace" + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + name = result.tool_calls[0].function.name + assert "" not in name + + def test_malformed_then_well_formed(self, parser, mock_request): + model_output = ( + "<|tool_call>call:bad_func no brace" + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 2 + assert result.tool_calls[0].function.name == "bad_func no brace" + assert result.tool_calls[1].function.name == "get_weather" + assert "London" in result.tool_calls[1].function.arguments + + def test_bare_malformed_then_well_formed(self, parser, mock_request): + model_output = ( + "<|tool_call>:bad_func no brace" + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 2 + assert result.tool_calls[0].function.name == "bad_func no brace" + assert result.tool_calls[1].function.name == "get_weather" + + def test_three_mixed_calls(self, parser, mock_request): + model_output = ( + "<|tool_call>:bad no brace" + '<|tool_call>call:search{input:<|"|>hello<|"|>}' + "<|tool_call>:set{flag:true,count:5}" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 3 + assert result.tool_calls[0].function.name == "bad no brace" + assert result.tool_calls[1].function.name == "search" + assert result.tool_calls[2].function.name == "set" + + def test_colon_in_content_not_broken(self, parser, mock_request): + model_output = ( + "Hello: world" + '<|tool_call>:get_weather{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.content is not None + assert "Hello" in result.content + + def test_bare_opener_empty_args(self, parser, mock_request): + model_output = "<|tool_call>:set_status{}" + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "set_status" diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index da1718332b..51c4696636 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -303,6 +303,7 @@ def gemma4_config() -> ParserEngineConfig: "TOOL_START": TOOL_CALL_START, "TOOL_END": TOOL_CALL_END, "CALL_PREFIX": "call:", + "COLON": ":", "OPEN_BRACE": "{", "OPEN_PAREN": "(", }, @@ -348,6 +349,13 @@ def gemma4_config() -> ParserEngineConfig: ParserState.TOOL_NAME, (), ), + # Bare opener: some checkpoints emit "<|tool_call>:name{...}" + # without the "call" prefix. Match the colon directly so the + # tool call is not silently dropped. + (ParserState.TOOL_PREAMBLE, "COLON"): Transition( + ParserState.TOOL_NAME, + (), + ), (ParserState.TOOL_NAME, "OPEN_BRACE"): Transition( ParserState.TOOL_ARGS, (), From e6eb0d120c480aa5116aafc2d39f19c8ebcf3181 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 14 Sep 2026 20:12:55 -0700 Subject: [PATCH 068/139] [Model][DSv4.1] Store the whole KV in MXFP8 (FlashMLA V4.1 record) (#56893) Signed-off-by: Yongye Zhu Signed-off-by: Yongye Zhu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) --- cmake/external_projects/flashmla.cmake | 85 +++++--- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 94 ++++++--- csrc/libtorch_stable/ops.h | 2 +- csrc/libtorch_stable/torch_bindings.cpp | 2 +- tests/kernels/test_compressor_kv_cache.py | 130 ++++++++++++ ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 98 +++++++++ .../nvidia/ops/dequant_gather_k_cutedsl.py | 62 +++--- vllm/models/deepseek_v41/amd/dspark.py | 2 + vllm/models/deepseek_v41/attention.py | 43 +++- .../deepseek_v41/common/ops/cache_utils.py | 194 +++++++++++++++++- .../common/ops/fused_compress_quant_cache.py | 76 ++++++- vllm/models/deepseek_v41/nvidia/dspark.py | 2 + .../attention/backends/mla/flashmla_sparse.py | 11 + vllm/v1/attention/backends/mla/sparse_swa.py | 18 +- 14 files changed, 709 insertions(+), 110 deletions(-) diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 382fd0289a..cfa8243f72 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -19,7 +19,7 @@ else() FetchContent_Declare( flashmla GIT_REPOSITORY https://github.com/vllm-project/FlashMLA - GIT_TAG 6bc49418c5ead572ff0339191ddf3b155749e183 + GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -76,43 +76,77 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_SOURCES ${flashmla_SOURCE_DIR}/csrc/api/api.cpp + ${flashmla_SOURCE_DIR}/csrc/api/sparse_prefill.cpp + ${flashmla_SOURCE_DIR}/csrc/api/sparse_decode.cpp + ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp + ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp # Misc kernels for decoding - ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu - ${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu # sm90 dense decode - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu # sm90 sparse decode - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu # sm90 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu - # sm100 dense prefill & backward - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu + # sm100 dense prefill (inference-only: dense backward stays out) + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu # sm100 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu + + # sm100 fused norm + rope + attn + rope + cast + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu # sm100 sparse decode - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu ) set(FlashMLA_Extension_SOURCES @@ -125,7 +159,6 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_INCLUDES ${flashmla_SOURCE_DIR}/csrc ${flashmla_SOURCE_DIR}/csrc/kerutils/include - ${flashmla_SOURCE_DIR}/csrc/sm90 ${flashmla_SOURCE_DIR}/csrc/cutlass/include ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include ) diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 06fe7108bd..254b9783bc 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -23,9 +23,15 @@ * cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin * second along last dim; each half is rope_dim/2 = 32 values) * - * Cache layout per paged-cache block (block_size tokens): - * [0, bs*576): token data, 448 fp8 + 128 bf16 each - * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + * Cache layout per paged-cache block (block_size tokens), selected by the + * `kv_mxfp8` argument: + * V4 (584 B/token, kv_mxfp8=false) + * [0, bs*576): token data, 448 fp8 + 128 bf16 each + * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + * V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the + * whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims) + * [0, bs*512): token data, 512 fp8 each + * [bs*512, bs*512 + bs*16): UE8M0 scales, 16 per token */ #include "torch_utils.h" @@ -91,6 +97,10 @@ constexpr int kQuantBlock = 64; constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7 constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad) constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576 +// V4.1 MXFP8 record: every dim quantized, one UE8M0 scale per 32 dims. +constexpr int kMxfp8QuantBlock = 32; +constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock; // 16 +constexpr int kMxfp8TokenDataBytes = kHeadDim; // 512 // FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw // 240.0) to match the rest of vLLM's FNUZ pipeline. #if defined(USE_ROCM) && defined(__gfx942__) @@ -141,6 +151,11 @@ __device__ __forceinline__ uint4 packFp8E4M3x16(float const* values, // ──────────────────────────────────────────────────────────────────────────── // Small inline helpers // ──────────────────────────────────────────────────────────────────────────── +__device__ __forceinline__ float warp2MaxAbs(float val) { + // Reduce absolute max across 2 consecutive lanes (32 dims = one MXFP8 block). + return fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 1)); +} + __device__ __forceinline__ float warp4MaxAbs(float val) { // Reduce absolute max across 4 consecutive lanes (lane id & 3 group). float peer = __shfl_xor_sync(FINAL_MASK, val, 1); @@ -176,6 +191,10 @@ __device__ __forceinline__ float warpSum(float val) { // slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant // + paged-cache // insert) +// +// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a +// runtime argument rather than doubling every (head count, q-norm) +// instantiation of a kernel whose hot path is the Q branch. template __device__ __forceinline__ void processDeepseekV4Slot( uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx, @@ -184,7 +203,7 @@ __device__ __forceinline__ void processDeepseekV4Slot( uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping, int64_t const* __restrict__ position_ids, float const* __restrict__ cos_sin_cache, int const cache_block_size, - int const kv_block_stride) { + int const kv_block_stride, bool const kv_mxfp8) { using Converter = vllm::_typeConvert; bool const isKV = (slotIdx == kNumHeadsQPadded); bool const isPadQ = !isKV && (slotIdx >= num_heads_q); @@ -302,19 +321,24 @@ __device__ __forceinline__ void processDeepseekV4Slot( *reinterpret_cast(dst) = out0; *reinterpret_cast(dst + 8) = out1; } else { - // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert. + // ── KV: FP8 quant + cache insert. V4 keeps the RoPE dims in bf16; + // V4.1 (kv_mxfp8) quantizes them too. int64_t const slot_id = slot_mapping[tokenIdx]; if (slot_id >= 0) { + int const token_data_bytes = + kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes; + int const scale_bytes = + kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken; int64_t const block_idx = slot_id / cache_block_size; int64_t const pos_in_block = slot_id % cache_block_size; uint8_t* block_base = k_cache + block_idx * static_cast(kv_block_stride); - uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes; + uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes; uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim; uint8_t* token_scale_ptr = block_base + - static_cast(cache_block_size) * kTokenDataBytes + - pos_in_block * kScaleBytesPerToken; + static_cast(cache_block_size) * token_data_bytes + + pos_in_block * scale_bytes; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -326,11 +350,13 @@ __device__ __forceinline__ void processDeepseekV4Slot( for (int i = 0; i < kElemsPerLane; i++) { local_absmax = fmaxf(local_absmax, fabsf(elements[i])); } - float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f); + float const absmax = fmaxf( + kv_mxfp8 ? warp2MaxAbs(local_absmax) : warp4MaxAbs(local_absmax), + 1e-4f); float const exponent = ceilf(log2f(absmax / kFp8Max)); float const inv_scale = exp2f(-exponent); - if (!is_rope_lane) { + if (kv_mxfp8 || !is_rope_lane) { uint8_t out_bytes[kElemsPerLane]; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -347,13 +373,19 @@ __device__ __forceinline__ void processDeepseekV4Slot( *reinterpret_cast(token_fp8_ptr + dim_base) = *reinterpret_cast(out_bytes); - if ((laneId & 3) == 0) { - int const q_block_idx = laneId >> 2; - float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); - token_scale_ptr[q_block_idx] = static_cast(encoded); - } - if (laneId == 0) { - token_scale_ptr[kNumQuantBlocks] = 0; + float const encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); + if (kv_mxfp8) { + // One scale per 32 dims = one per lane pair. + if ((laneId & 1) == 0) { + token_scale_ptr[laneId >> 1] = static_cast(encoded); + } + } else { + if ((laneId & 3) == 0) { + token_scale_ptr[laneId >> 2] = static_cast(encoded); + } + if (laneId == 0) { + token_scale_ptr[kNumQuantBlocks] = 0; + } } } else { uint4 out0, out1; @@ -421,7 +453,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full int const num_heads_q, // live Q heads (input layout) int const cache_block_size, // tokens per paged-cache block - int const kv_block_stride) { // bytes per paged-cache block + int const kv_block_stride, // bytes per paged-cache block + bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) // BF16 _typeConvert specialization is unavailable on pre-Ampere. The // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a @@ -477,7 +510,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( processDeepseekV4Slot( v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size, - kv_block_stride); + kv_block_stride, kv_mxfp8); #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaTriggerProgrammaticLaunchCompletion(); @@ -506,7 +539,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( float const* __restrict__ cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, - int const kv_block_stride) { + int const kv_block_stride, bool const kv_mxfp8) { #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) if constexpr (std::is_same_v) { return; @@ -564,7 +597,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( processDeepseekV4Slot( v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, - cos_sin_cache, cache_block_size, kv_block_stride); + cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8); // ── Buffer rotation: hand the prefetched LDGs to the next iter. v0_curr = v0_next; @@ -590,7 +623,7 @@ static void launchFusedDeepseekV4Templated( uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids, float const* cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, - int const cache_block_size, int const kv_block_stride, + int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8, cudaStream_t stream) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; @@ -630,7 +663,8 @@ static void launchFusedDeepseekV4Templated( scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, - num_heads_q, cache_block_size, kv_block_stride); + num_heads_q, cache_block_size, kv_block_stride, + kv_mxfp8); } else { config.gridDim = dim3(num_tokens_full); cudaLaunchKernelEx( @@ -639,7 +673,7 @@ static void launchFusedDeepseekV4Templated( scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, - kv_block_stride); + kv_block_stride, kv_mxfp8); } #else // ROCm: use standard kernel launch syntax (no PDL/stream serialization) @@ -649,7 +683,7 @@ static void launchFusedDeepseekV4Templated( <<>>( q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, - cache_block_size, kv_block_stride); + cache_block_size, kv_block_stride, kv_mxfp8); #endif } @@ -663,19 +697,19 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const num_heads_q_padded, int const cache_block_size, int const kv_block_stride, - bool const apply_q_norm, cudaStream_t stream) { + bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) { #define DISPATCH(N) \ case N: \ if (apply_q_norm) { \ launchFusedDeepseekV4Templated( \ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, stream); \ + num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ } else { \ launchFusedDeepseekV4Templated( \ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, stream); \ + num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ } \ return; @@ -967,7 +1001,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& position_ids, // [N] int64 torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 int64_t q_head_padded, // padded Q head count for output - double eps, int64_t cache_block_size, bool apply_q_norm) { + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) { STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(), "q_in must be contiguous CUDA"); STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(), @@ -1036,7 +1070,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( cos_sin_cache.const_data_ptr(), static_cast(eps), num_tokens_full, num_tokens_insert, num_heads_q, num_heads_q_padded, cache_block_size_i, kv_block_stride, - apply_q_norm, stream); + apply_q_norm, kv_mxfp8, stream); }); return q_out; } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 68a24ca434..e19f49d603 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -264,7 +264,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, torch::stable::Tensor const& position_ids, torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm); + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8); void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor& q, torch::stable::Tensor const& kv, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 110f8595ee..54548dfd05 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -430,7 +430,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True) -> Tensor"); + "bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor"); // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 0f1976cf68..7e092229fe 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -215,6 +215,136 @@ def test_v41_fused_save_compress_and_insert( torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0) +def _mxfp8_record_reference( + latent_row: torch.Tensor, cos_sin_row: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference V4.1 record for one token: 512 fp8 bytes + 16 UE8M0 scales. + + Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the + RoPE tail first, then scale every 32-dim tile, RoPE tiles included. + """ + row = latent_row.float() + c, s = cos_sin_row.chunk(2) + rotated = row.clone() + even, odd = row[448::2], row[449::2] + rotated[448::2] = even * c - odd * s + rotated[449::2] = odd * c + even * s + quantized, scales = _ue8m0_reference(rotated, 32, 448.0) + return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("compress_ratio", [1, 2]) +def test_v41_rope_insert_mxfp8_record(compress_ratio: int): + """The 528-byte V4.1 record quantizes the RoPE dims instead of keeping them + in bf16, so a page is [block_size x 512 fp8][block_size x 16 UE8M0]. + + Byte-exact against the reference quantizer, including the rows and page + padding the kernel must leave alone. + """ + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(7) + device = "cuda" + cache_block = 64 + cache_stride = math.ceil(cache_block * 528 / 512) * 512 + num_tokens = 12 + + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(64, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + + cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device) + cache = cache_backing.as_strided((2, cache_block, 528), (cache_stride, 528, 1)) + # Leave one token unmapped so the kernel's negative-slot skip is covered. + slots = torch.arange(num_tokens, dtype=torch.int64, device=device) + slots[3] = -1 + + rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio) + + expected = torch.full_like(cache_backing, 165) + for t in range(num_tokens): + slot = slots[t].item() + pos = positions[t].item() + if slot < 0 or (pos + 1) % compress_ratio: + continue + page, row = divmod(slot, cache_block) + values, scales = _mxfp8_record_reference( + latent[t], cos_sin[pos // compress_ratio * compress_ratio] + ) + expected[page, row * 512 : (row + 1) * 512] = values + scale_offset = cache_block * 512 + row * 16 + expected[page, scale_offset : scale_offset + 16] = scales + + torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_v41_mxfp8_cache_round_trip(): + """quantize_and_insert -> dequantize_and_gather recovers the V4.1 record. + + The RoPE dims now go through fp8, so they carry quantization error too -- + the check is that every dim survives within one MXFP8 tile step, and that + the bytes on the way match the reference quantizer. + """ + from vllm.models.deepseek_v41.common.ops import ( + dequantize_and_gather_k_cache, + quantize_and_insert_k_cache, + ) + + torch.manual_seed(11) + device = "cuda" + block_size = 64 + num_tokens = 70 + num_blocks = 4 + page_bytes = math.ceil(block_size * 528 / 512) * 512 + + k = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + cache = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + quantize_and_insert_k_cache( + k, cache, slot_mapping, block_size=block_size, bytes_per_token=528 + ) + + for t in (0, 1, block_size, num_tokens - 1): + values, scales = _ue8m0_reference(k[t].float(), 32, 448.0) + page, row = divmod(t, block_size) + torch.testing.assert_close( + cache[page, row * 512 : (row + 1) * 512], + values.view(torch.uint8), + rtol=0, + atol=0, + ) + scale_offset = block_size * 512 + row * 16 + torch.testing.assert_close( + cache[page, scale_offset : scale_offset + 16], + (scales.log2() + 127).to(torch.uint8), + rtol=0, + atol=0, + ) + + out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device) + dequantize_and_gather_k_cache( + out, + cache.view(num_blocks, block_size, 528), + seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device), + gather_lens=None, + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, -1 + ), + block_size=block_size, + offset=0, + ) + # Half an e4m3 ULP at the top of the range (16 units) bounds the error. + tile_amax = k.float().abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) + tolerance = 16.0 * (tile_amax / 448.0).log2().ceil().exp2() + error = (out[0].float() - k.float()).abs() + assert (error <= tolerance.repeat_interleave(32, dim=-1)).all() + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) @pytest.mark.parametrize("store_fp8", [False, True]) diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 0e25884368..68eabfb94e 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -424,6 +424,104 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int): ) +# ── Test 2c: DeepSeek V4.1 MXFP8 record (kv_mxfp8=True) ────────────────────── + +V41_HEAD_BYTES = HEAD_DIM + HEAD_DIM // 32 # 512 fp8 + 16 UE8M0 scales = 528 + + +@pytest.mark.parametrize("num_tokens", [1, 17, 64]) +@pytest.mark.parametrize("block_size", [32, 64]) +def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int): + """``kv_mxfp8=True`` writes DeepSeek's V4.1 record: all 512 dims as fp8 e4m3 + with one UE8M0 scale per 32 dims, no bf16 RoPE tail. + + Both sides quantize the same bf16-rounded rotated row, so the bytes must + match the Triton encoder exactly apart from the documented 1-ULP bf16 + rounding split on the rotation itself. + """ + from vllm.models.deepseek_v41.common.ops import ( + dequantize_and_gather_k_cache as v41_dequantize_and_gather_k_cache, + ) + from vllm.models.deepseek_v41.common.ops import ( + quantize_and_insert_k_cache as v41_quantize_and_insert_k_cache, + ) + + torch.manual_seed(3) + device = "cuda" + dtype = torch.bfloat16 + max_pos = 4096 + + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + + k_cache_ref = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + v41_quantize_and_insert_k_cache( + apply_rope_gptj_last_k(kv, positions, cos_sin_cache), + k_cache_ref, + slot_mapping, + block_size=block_size, + use_fnuz=USE_FNUZ, + bytes_per_token=V41_HEAD_BYTES, + ) + + k_cache_fused = torch.zeros_like(k_cache_ref) + q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_dummy, + kv, + k_cache_fused, + slot_mapping, + positions, + cos_sin_cache, + 64, + 1e-6, + block_size, + True, + True, # kv_mxfp8 + ) + + # A pre-quantization bf16 tie can land the two encoders on adjacent e4m3 + # codes, so compare after dequantizing and allow one step at the tile scale. + def _dequant(cache_2d): + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=dtype, device=device) + v41_dequantize_and_gather_k_cache( + out, + cache_2d.view(num_blocks, block_size, V41_HEAD_BYTES), + torch.tensor([num_tokens], dtype=torch.int32, device=device), + None, + torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(0), + block_size, + offset=0, + use_fnuz=USE_FNUZ, + ) + return out[0] + + rec_ref = _dequant(k_cache_ref).float() + rec_fused = _dequant(k_cache_fused).float() + tile_scale = ( + (rec_ref.abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) / FP8_MAX) + .log2() + .ceil() + .exp2() + ) + # One e4m3 ULP at the top of the range is 32 quantized units. + step = 32.0 * tile_scale.repeat_interleave(32, dim=-1) + assert ((rec_fused - rec_ref).abs() <= step).all() + # The NoPE dims see no rotation at all, so they must be byte-identical. + nope_bytes = slice(0, block_size * HEAD_DIM) + ref_rows = k_cache_ref[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + fused_rows = k_cache_fused[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + torch.testing.assert_close( + fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0 + ) + + # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ───────────────────── diff --git a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py index fcb12740f2..5920ed290c 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py @@ -23,10 +23,11 @@ from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( class DequantGatherKCacheKernel( VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"] ): - # Hard-coded for DSv4. + # DSv4 geometry. The record format (how much of the row is quantized and + # how wide a scale group is) comes from the compile key: V4 quantizes the + # 448 NoPE dims in groups of 64 and keeps a bf16 RoPE tail, V4.1 quantizes + # all 512 dims in groups of 32 (MXFP8) and has no tail. head_dim = 512 - fp8_dim = 448 - group_size = 64 # 1 scale per 64 elems num_warps = 4 tb_size = num_warps * 32 num_stages = 4 @@ -35,17 +36,20 @@ class DequantGatherKCacheKernel( class CompileKey: block_size: int has_gather_lens: bool + fp8_dim: int = 448 + group_size: int = 64 # dims per scale @staticmethod def kernel(compile_key: CompileKey) -> Any: head_dim = DequantGatherKCacheKernel.head_dim - fp8_dim = DequantGatherKCacheKernel.fp8_dim - group_size = DequantGatherKCacheKernel.group_size + fp8_dim = compile_key.fp8_dim + group_size = compile_key.group_size num_warps = DequantGatherKCacheKernel.num_warps tb_size = DequantGatherKCacheKernel.tb_size num_stages = DequantGatherKCacheKernel.num_stages bf16_dim = head_dim - fp8_dim data_dim = fp8_dim + bf16_dim * 2 + scale_bytes = head_dim // group_size block_size = compile_key.block_size @cute.jit @@ -62,8 +66,8 @@ class DequantGatherKCacheKernel( ): op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) - cp8_atom = cute.make_copy_atom( - cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64 + cp_scale_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=scale_bytes * 8 ) page_id = block_table[req_id, pos // block_size] block_offset = pos % block_size @@ -88,7 +92,7 @@ class DequantGatherKCacheKernel( ) elif idx == cutlass.const_expr(data_dim // 16): cute.copy( - cp8_atom, + cp_scale_atom, k_scale[page_id, block_offset, None], s_kscale[None, stage_id], ) @@ -119,8 +123,8 @@ class DequantGatherKCacheKernel( )[None, warp_id, None] s_kscale = smem.allocate_tensor( Uint8, - cute.make_layout((8, num_warps, num_stages)), - byte_alignment=8, + cute.make_layout((scale_bytes, num_warps, num_stages)), + byte_alignment=scale_bytes, )[None, warp_id, None] # Prepare for 16B cp.async, also for BF16 smem loads later. @@ -239,14 +243,16 @@ class DequantGatherKCacheKernel( # Last 64 elems are BF16 tail, corresponds to dequant1 of last # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + - # 8x 16B for BF16. - if lane_id + 32 >= fp8_dim // 8: - idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 - cute.copy( - cp16_atom, - s_kdata_16B_slice[(None, idx), compute_stage], - dequant1, - ) + # 8x 16B for BF16. The V4.1 record is fully quantized and has + # no tail. + if cutlass.const_expr(bf16_dim > 0): # noqa: SIM102 + if lane_id + 32 >= fp8_dim // 8: + idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 + cute.copy( + cp16_atom, + s_kdata_16B_slice[(None, idx), compute_stage], + dequant1, + ) # Store two 16B BF16 chunks per lane: first half, then second half. dst = out_slice[req_id, offset + i, (None, lane_id)] @@ -269,7 +275,7 @@ class DequantGatherKCacheKernel( ): # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] # block is actually a concat of - # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. + # [block_size, fp8_dim + bf16_dim * 2] and [block_size, scale_bytes]. k_data = cute.make_tensor( k_cache.iterator, layout=cute.make_layout( @@ -282,8 +288,8 @@ class DequantGatherKCacheKernel( k_scale = cute.make_tensor( k_cache.iterator + (block_size * data_dim), layout=cute.make_layout( - (k_cache.shape[0], block_size, 8), - stride=(k_cache.stride[0], 8, 1), + (k_cache.shape[0], block_size, scale_bytes), + stride=(k_cache.stride[0], scale_bytes, 1), ), ) @@ -311,10 +317,14 @@ class DequantGatherKCacheKernel( *, block_size: int, has_gather_lens: bool, + fp8_dim: int = 448, + group_size: int = 64, ) -> CompileKey: return self.CompileKey( block_size=block_size, has_gather_lens=has_gather_lens, + fp8_dim=fp8_dim, + group_size=group_size, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -345,9 +355,11 @@ class DequantGatherKCacheKernel( def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_reqs = cute.sym_int() head_dim = self.head_dim - fp8_dim = self.fp8_dim + fp8_dim = compile_key.fp8_dim block_size = compile_key.block_size - head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8 + head_bytes = ( + fp8_dim + (head_dim - fp8_dim) * 2 + head_dim // compile_key.group_size + ) out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) k_cache = cute.runtime.make_fake_tensor( @@ -384,10 +396,14 @@ class DequantGatherKCacheKernel( block_table: torch.Tensor, block_size: int, offset: int, + fp8_dim: int = 448, + group_size: int = 64, ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]: compile_key = self.dispatch( block_size=block_size, has_gather_lens=gather_lens is not None, + fp8_dim=fp8_dim, + group_size=group_size, ) launch_args = ( out, diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index 3a6cef8aa7..78f07b6b11 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -259,6 +259,8 @@ def _insert_context_kv( attn.padded_heads, attn.eps, block_size, + True, # apply_q_norm (unused: the query is a discarded dummy) + attn.kv_mxfp8, ) elif cache_dtype == torch.bfloat16: swa_3d = swa_cache.view(-1, block_size, attn.head_dim) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 131cf4d97e..fa26351681 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -51,6 +51,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.models.utils import extract_layer_index from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope from vllm.models.deepseek_v41.compressor import DeepseekCompressor +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, @@ -109,6 +110,16 @@ def _fill_short_context_topk_indices( ) +# Which packed fp8_ds_mla record a V4.1 layer writes. FlashMLA decodes +# DeepSeek's V4.1 record -- all 512 dims (RoPE included) as fp8 e4m3 with one +# UE8M0 scale per 32 dims, 512 data bytes + 16 scale bytes per token, pages +# rounded to the kernel's 512 B TMA stride -- only in its SM100 sparse-decode +# kernels. Every other arch keeps the V4 record: 448 fp8 NoPE + 64 bf16 RoPE +# plus 7 UE8M0 scales of 64 dims and a pad byte (584 B, 576 B pages). +def _use_v41_mxfp8_kv_record() -> bool: + return current_platform.is_device_capability_family(100) + + def _resolve_dsv4_kv_cache_dtype( use_fp8_ds_mla_layout: bool, kv_cache_dtype: str, @@ -445,6 +456,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config ) + self.kv_mxfp8 = _use_v41_mxfp8_kv_record() + self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584 + self.kv_page_alignment = 512 if self.kv_mxfp8 else 576 self.swa_cache_layer = DeepseekV4SWACache( head_dim=self.head_dim, @@ -454,6 +468,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): cache_config=cache_config, backend_cls=self.swa_backend_cls, block_size=32, + packed_bytes_per_token=self.kv_bytes_per_token, + packed_page_alignment=self.kv_page_alignment, ) # The attention layer itself was already registered with the @@ -879,6 +895,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.eps, swa_metadata.block_size, False, + self.kv_mxfp8, ) # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV @@ -934,9 +951,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # DeepseekV4SWACache. if not self.is_kv_source: return None - # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B - # alignment; plain bf16 / per-tensor fp8 rows use natural element-size - # pages. + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up + # to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows + # use natural element-size pages. uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, @@ -945,12 +962,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, tokens_per_state=self.compress_ratio, cache_dtype_str=self.kv_cache_dtype, - alignment=576 if uses_fp8_ds_mla_layout else 512, + alignment=self.kv_page_alignment if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), - # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token; - # head_size stays semantic (512). - state_content_bytes=584 if uses_fp8_ds_mla_layout else None, + # Packed record width; head_size stays semantic (512). + state_content_bytes=( + self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None + ), ) def _compressed_kv_cache(self) -> torch.Tensor: @@ -998,8 +1016,15 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): head_size=self.head_dim, dtype=self.dtype, tokens_per_state=self.compress_ratio, - # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). - alignment=576 if uses_fp8_ds_mla_layout else 512, + # The indexer page shares a physical block with the main KV cache + # (the block stride is the sum of every page in the group), so it + # has to carry the same padding: the packed record's TMA stride, + # else 512B for FlashInfer sparse (#44577). + alignment=( + 576 + if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() + else 512 + ), ) def forward(self): ... diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 64462eb73b..51743016bf 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -36,6 +36,18 @@ from vllm.triton_utils import tl, triton from vllm.utils.import_utils import has_cutedsl from vllm.utils.math_utils import next_power_of_2 +# Per-token byte width of the two paged fp8 records. Both put a page's whole +# data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the +# width the KV-cache spec publishes -- identifies the record: +# V4 (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64 +# dims each plus a pad byte. +# V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0 +# scales of 32 dims each. FlashMLA's ``ModelType::V41``. +V4_BYTES_PER_TOKEN = 584 +V41_BYTES_PER_TOKEN = 528 +V41_QUANT_BLOCK = 32 +V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16 + @triton.jit def quantize_and_insert_k_kernel( @@ -163,6 +175,67 @@ def quantize_and_insert_k_kernel( tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) +@triton.jit +def _quantize_and_insert_k_mxfp8_kernel( + k_ptr, # [num_tokens, 512] bf16 + slot_mapping_ptr, # [num_tokens] int64 + k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) + num_tokens, + head_dim: tl.constexpr, # 512, RoPE dims included + scale_dim: tl.constexpr, # 16 + quant_block: tl.constexpr, # 32 dims per UE8M0 scale + cache_block_size: tl.constexpr, # paged cache block size, in tokens + block_stride: tl.constexpr, # total bytes per block (padded) + fp8_max: tl.constexpr, + use_fnuz: tl.constexpr = False, +): + """Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache. + + One program per token. Unlike the V4 record there is no bf16 tail, so every + dim is scaled. + """ + pid = tl.program_id(0) + if pid >= num_tokens: + return + + slot_idx = tl.load(slot_mapping_ptr + pid) + if slot_idx == -1: + return + + block_idx = slot_idx // cache_block_size + pos_in_block = slot_idx % cache_block_size + + # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks. + cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride + token_data_ptr = cache_block_ptr + pos_in_block * head_dim + token_scale_ptr = ( + cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim + ) + + d = tl.arange(0, head_dim) + x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32) + + tiles = tl.reshape(x, (scale_dim, quant_block)) + # Floor the absmax like the CUDA encoder so an all-zero tile still lands on + # a finite scale. + amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4) + exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max))) + scaled = tl.clamp( + tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max + ) + if use_fnuz: # noqa: SIM108 + fp8 = scaled.to(tl.float8e4b8) + else: + fp8 = scaled.to(tl.float8e4nv) + tl.store( + token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,)) + ) + + # UE8M0 encoding: stored_value = exponent + 127 (bias). + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0) + tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8)) + + def quantize_and_insert_k_cache( k: torch.Tensor, # [num_tokens, 512] bf16 k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 @@ -170,16 +243,13 @@ def quantize_and_insert_k_cache( block_size: int = 64, is_ue8m0: bool = True, use_fnuz: bool = False, + bytes_per_token: int = V4_BYTES_PER_TOKEN, ): """ Quantize K tensor and insert into paged K cache. - K Cache block layout (block_size=64 tokens): - - First 64 * 576 = 36864 bytes: Token data - - Each token: 448 bytes (fp8) + 128 bytes (bf16) - - Next 64 * 8 = 512 bytes: Scales - - Each token: 8 bytes (uint8 scales, 7 real + 1 padding) - - Padded to multiple of 576 + ``bytes_per_token`` picks the record (see the module header): the V4 one, + or V4.1's all-dims MXFP8 one. ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3, @@ -196,6 +266,26 @@ def quantize_and_insert_k_cache( num_tokens = slot_mapping.shape[0] block_stride = k_cache.stride(0) # bytes per block + if bytes_per_token == V41_BYTES_PER_TOKEN: + if use_fnuz and not current_platform.is_fp8_fnuz(): + raise ValueError("use_fnuz=True requires a platform using FNUZ FP8") + _quantize_and_insert_k_mxfp8_kernel[(num_tokens,)]( + k, + slot_mapping, + k_cache, + num_tokens, + head_dim=512, + scale_dim=V41_NUM_SCALES, + quant_block=V41_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=block_stride, + fp8_max=get_fp8_min_max()[1] + if use_fnuz + else torch.finfo(torch.float8_e4m3fn).max, + use_fnuz=use_fnuz, + ) + return + TOKEN_FP8_DIM = 448 TOKEN_BF16_DIM = 64 TOKEN_SCALE_DIM = 8 @@ -343,6 +433,69 @@ def _dequantize_and_gather_k_kernel( tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) +@triton.jit +def _dequantize_and_gather_k_mxfp8_kernel( + out_ptr, + out_stride0, + out_stride1, + k_cache_ptr, + seq_lens_ptr, + block_table_ptr, + offset, + gather_lens_ptr, + max_blocks_per_seq: tl.constexpr, + head_dim: tl.constexpr, # 512 + scale_dim: tl.constexpr, # 16 + quant_block: tl.constexpr, # 32 + cache_block_size: tl.constexpr, + block_stride: tl.constexpr, + use_fnuz: tl.constexpr = False, +): + """Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place.""" + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + if gather_lens_ptr is not None: # noqa: SIM108 + gather_len = tl.load(gather_lens_ptr + batch_idx) + else: + gather_len = seq_len + start_pos = seq_len - gather_len + + d = tl.arange(0, head_dim) + s = tl.arange(0, scale_dim) + for i in range(worker_id, gather_len, num_workers): + pos = start_pos + i + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq + physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 + + # int64: physical_block_idx * block_stride can exceed 2^31. + cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + token_data_ptr = cache_block_ptr + pos_in_block * head_dim + token_scale_ptr = ( + cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim + ) + + x_uint8 = tl.load(token_data_ptr + d) + if use_fnuz: + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block)) + + # UE8M0: scale = 2^(stored_value - 127). + encoded = tl.load(token_scale_ptr + s) + scale = tl.exp2(encoded.to(tl.float32) - 127.0) + dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) + + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) + + def dequantize_and_gather_k_cache_triton( # [num_reqs, max_num_tokens, head_size] out: torch.Tensor, @@ -358,6 +511,28 @@ def dequantize_and_gather_k_cache_triton( offset: int, use_fnuz: bool = False, ) -> None: + num_reqs = seq_lens.shape[0] + NUM_WORKERS = 128 + if k_cache.shape[-1] == V41_BYTES_PER_TOKEN: + _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + head_dim=512, + scale_dim=V41_NUM_SCALES, + quant_block=V41_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=k_cache.stride(0), + use_fnuz=use_fnuz, + ) + return + TOKEN_FP8_DIM = 448 TOKEN_BF16_DIM = 64 TOKEN_SCALE_DIM = 8 @@ -365,8 +540,6 @@ def dequantize_and_gather_k_cache_triton( FP8_MAX = 448.0 TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 - num_reqs = seq_lens.shape[0] - NUM_WORKERS = 128 _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( out, out.stride(0), @@ -408,6 +581,8 @@ def dequantize_and_gather_k_cache( ) -> None: """Dequantize and gather a paged DSv4 K cache. + The record is read off ``k_cache.shape[-1]``; see the module header. + ``use_fnuz`` MUST match the encoder of the specific cache being read: ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere), ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder @@ -419,6 +594,7 @@ def dequantize_and_gather_k_cache( _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, ) + mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( out=out, k_cache=k_cache, @@ -427,6 +603,8 @@ def dequantize_and_gather_k_cache( block_table=block_table, block_size=block_size, offset=offset, + fp8_dim=512 if mxfp8 else 448, + group_size=V41_QUANT_BLOCK if mxfp8 else 64, ) return diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 5950c4b6ba..96f15849dd 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -232,10 +232,13 @@ def rope_quant_insert( The BF16 latent supplies both NoPE quantization and RoPE input. It is read only for valid slots at group boundaries. The cache dtype selects the - layout: ``uint8`` is the fp8_ds_mla paged layout (576 value bytes and eight - segregated UE8M0 scale bytes per token, including one zero padding scale); - ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows - read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``. + layout: ``uint8`` is the fp8_ds_mla paged layout, whose record the + per-token byte width names -- 584 B for V4 (576 value bytes and eight + segregated UE8M0 scale bytes, including one zero padding scale) or 528 B + for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 + scales of 32 dims each). ``bfloat16`` and ``float8_e4m3fn`` are the plain + [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the + per-tensor ``fp8_scale``. """ assert compress_ratio in (1, 2) assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16 @@ -246,8 +249,15 @@ def rope_quant_insert( return launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {} if kv_cache.dtype == torch.uint8: - assert kv_cache.shape[-1] == 584 - _rope_quant_insert_kernel[(num_tokens,)]( + assert kv_cache.shape[-1] in (584, 528), ( + f"unsupported paged KV record width {kv_cache.shape[-1]}" + ) + kernel = ( + _rope_quant_insert_mxfp8_kernel + if kv_cache.shape[-1] == 528 + else _rope_quant_insert_kernel + ) + kernel[(num_tokens,)]( latent, positions, cos_sin_cache, @@ -338,6 +348,60 @@ def _rope_quant_insert_kernel( tl.store(rope_dst + d - 448, rotated.to(tl.bfloat16), d >= 448) +@triton.jit +def _rope_quant_insert_mxfp8_kernel( + latent, + positions, + cos_sin, + cache, + cache_slots, + COS_STRIDE: tl.constexpr, + CACHE_STRIDE: tl.constexpr, + CACHE_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, +): + """V4.1 record: RoPE first, then MXFP8-quantize all 512 dims. + + The RoPE dims are quantized here too, so unlike the V4 kernel the rotation + has to happen before the scales are picked. + """ + t = tl.program_id(0) + slot = tl.load(cache_slots + t) + if slot < 0: + return + position = tl.load(positions + t) + if (position + 1) % COMPRESS_RATIO != 0: + return + d = tl.arange(0, 512) + normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32) + + # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there. + even, odd = tl.split(tl.reshape(normed, (256, 2))) + pair = tl.arange(0, 256) - 224 + cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE + c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32) + s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32) + rotated = tl.interleave(even * c - odd * s, odd * c + even * s) + if SANITIZE_CACHE_NANS: + rotated = tl.where(rotated == rotated, rotated, 0.0) + + page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE + values = page + (slot % CACHE_BLOCK) * 512 + scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16 + + quant = tl.reshape(rotated, (16, 32)) + amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4) + exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0))) + scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1)) + fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv) + tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,))) + + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded) + tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8)) + + @triton.jit def _rope_plain_insert_kernel( latent, diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 3c711681aa..88646404f4 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -264,6 +264,8 @@ def _insert_context_kv( attn.padded_heads, attn.eps, block_size, + True, # apply_q_norm (unused: the query is a discarded dummy) + attn.kv_mxfp8, ) elif cache_dtype == torch.bfloat16: swa_3d = swa_cache.view(-1, block_size, attn.head_dim) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index d8c4368401..9c440776ae 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -91,6 +91,17 @@ Bytes, structured as: The first `ue8m0` is the scale for the first 64 `float8_e4m3` values, the second for the next 64, and so on. +For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as: +- **First 512 bytes:** all 512 dims as `float8_e4m3`. Unlike V4 the 64 RoPE + dims are quantized too, so there is no `bfloat16` part. +- **Last 16 bytes:** Scale factors, containing 16 `ue8m0` values, one per 32 + consecutive `float8_e4m3` values (i.e. MXFP8). + +The V4 and V4.1 records are not laid out token-by-token within a page: a page +holds all its data rows first and all its scale rows after, so a page is +`block_size * bytes_per_token` bytes rounded up to the decode kernel's TMA +stride (576 B for V4, 512 B for V4.1). + In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each token's KV cache is 352 Bytes, structured as: - **First 256 bytes:** 512 `e2m1` NoPE values packed 2/byte (low nibble = diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 62c3921b5e..c4f3cd8b4b 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -79,6 +79,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): cache_config: CacheConfig, backend_cls: "type[AttentionBackend] | None" = None, block_size: int = 64, + packed_bytes_per_token: int = 584, + packed_page_alignment: int = 576, ): super().__init__() self.backend_cls = backend_cls or DeepseekSparseSWABackend @@ -88,6 +90,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): self.prefix = prefix self.cache_config = cache_config self.dtype = dtype + self.packed_bytes_per_token = packed_bytes_per_token + self.packed_page_alignment = packed_page_alignment compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -105,8 +109,9 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): self.kv_cache = kv_cache.squeeze(1) def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous - # bf16/fp8 cache uses the natural element-size page. + # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode + # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural + # element-size page. uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" return SlidingWindowMLASpec( block_size=self.block_size, @@ -115,10 +120,11 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): dtype=self.dtype, sliding_window=self.window_size, cache_dtype_str=self.cache_config.cache_dtype, - # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales) - state_content_bytes=584 if uses_fp8_ds_mla_layout else None, - # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). - alignment=576 if uses_fp8_ds_mla_layout else 512, + state_content_bytes=( + self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None + ), + # FlashMLA packing stride; 512B for FlashInfer sparse (#44577). + alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), ) From 381c61e008682945d43d6c8808ed633477364b4b Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Mon, 14 Sep 2026 20:15:19 -0700 Subject: [PATCH 069/139] [Bugfix][NIXL] Report a full prefix cache hit as finished (#56640) Signed-off-by: Yizheng Jiao Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- .../kv_connector/unit/test_nixl_connector.py | 72 +++++++++++++++++++ .../unit/test_nixl_connector_hma.py | 43 ++++++++++- .../kv_connector/v1/nixl/base_scheduler.py | 10 ++- .../kv_connector/v1/nixl/metadata.py | 7 ++ .../kv_connector/v1/nixl/pull_scheduler.py | 5 +- .../kv_connector/v1/nixl/pull_worker.py | 7 ++ .../kv_connector/v1/nixl/push_scheduler.py | 4 +- 7 files changed, 142 insertions(+), 6 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 9b54f626f2..298886d160 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -341,6 +341,10 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv(): assert req_meta.remote.request_id == f"prefill-{42}" # do_remote_prefill is consumed by request_finished to prevent re-issuing. assert request.kv_transfer_params["do_remote_prefill"] is False + # The scheduler is not waiting on this recv -- the request is already gone + # from self.requests, so reporting it would trip `assert req_id in + # self.requests` in _update_from_kv_xfer_finished. + assert req_meta.awaiting_kvs is False @patch( @@ -2695,6 +2699,74 @@ class FailingNixlWrapper(FakeNixlWrapper): return super().check_xfer_state(handle) +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FailingNixlWrapper, +) +@pytest.mark.parametrize("awaiting_kvs", [True, False]) +@pytest.mark.parametrize("notification_fails", [False, True]) +def test_empty_recv_is_reported_only_when_awaited( + default_vllm_config, dist_init, awaiting_kvs, notification_fails +): + """An empty recv is reported iff the scheduler parked the request on it. + + _read_blocks returns early when the local block list is empty -- D already + holds the KV and only the producer needs telling. Two kinds of caller reach + that path, and they need opposite handling: + + Parked (awaiting_kvs=True, a full prefix cache hit). Without a report the + request is named in neither _recving_transfers nor _failed_recv_reqs and + never reaches finished_recving. The scheduler has no other way to release a + WAITING_FOR_REMOTE_KVS request, and there is no timeout, so it sits in + skipped_waiting holding its blocks for the life of the process. + + Notify-only (awaiting_kvs=False): request_finished seeding an empty recv to + free P's blocks for a request aborted before it was scheduled, or a readback + on a request that stays RUNNING. Reporting either one crashes the engine + core -- _update_from_kv_xfer_finished asserts the id is still in + self.requests and that the request is parked or finished. + + A failed notification changes neither: the KV is local either way, and the + producer frees its own blocks on a timeout. + """ + vllm_config = create_vllm_config() + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, + connector.engine_id, + hand_shake_latency=0.0, + kv_cache_config=connector._kv_cache_config, + ) + connector.connector_worker.nixl_wrapper.fail_send_notif = notification_fails + + request_id = f"test_empty_recv_awaited_{awaiting_kvs}" + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id=request_id, + local_block_ids=(), # empty: the whole prompt is already cached locally + kv_transfer_params={ + "remote_block_ids": [[20, 21, 22]], + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + "remote_request_id": f"prefill-{request_id}", + "remote_host": "localhost", + "remote_port": 1234, + "remote_tp_size": 1, + }, + awaiting_kvs=awaiting_kvs, + ) + connector.bind_connector_metadata(metadata) + dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) + connector.start_load_kv(dummy_ctx) + connector.bind_connector_metadata(NixlConnectorMetadata()) + time.sleep(0.1) + connector.start_load_kv(dummy_ctx) + + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert (request_id in done_recving) is awaiting_kvs + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 259e9536c1..60eb2d6614 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -149,8 +149,49 @@ def test_update_state_after_alloc_tracks_cached_blocks_per_group(): scheduler.update_state_after_alloc(request, blocks, num_external_tokens=2) - _, _, local_num_computed_blocks = scheduler._reqs_need_recv[request.request_id] + _, _, local_num_computed_blocks, awaiting_kvs = scheduler._reqs_need_recv[ + request.request_id + ] assert local_num_computed_blocks == (2, 1) + # There are external tokens to pull, so the scheduler parks the request and + # is waiting to see it in finished_recving. + assert awaiting_kvs is True + + +@pytest.mark.cpu_test +def test_full_local_hit_is_not_awaited_by_the_scheduler(): + """A recv with num_external_tokens == 0 must not be flagged as awaited. + + The request keeps running -- there is nothing to pull, only a notification + owed to the producer -- so the worker must not report it in + finished_recving. Doing so trips `assert RequestStatus.is_finished` in + _update_from_kv_xfer_finished, which only tolerates a finished recv for a + request that is parked or already done.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, + ) + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.kv_cache_utils import KVCacheBlock + + scheduler = object.__new__(NixlPullConnectorScheduler) + scheduler._reqs_in_batch = set() + scheduler._reqs_need_save = {} + scheduler._reqs_need_recv = {} + scheduler.use_host_buffer = False + scheduler.is_bidirectional_kv_xfer_enabled = False + scheduler._is_hma_required = False + scheduler.kv_cache_config = MagicMock( + select_transfer_block_ids=lambda block_ids: block_ids + ) + + blocks = KVCacheBlocks(blocks=([KVCacheBlock(block_id=0, _block_hash=object())],)) + request = create_request(do_remote_prefill=True) + + scheduler.update_state_after_alloc(request, blocks, num_external_tokens=0) + + _, local_block_ids, _, awaiting_kvs = scheduler._reqs_need_recv[request.request_id] + assert not local_block_ids # full local hit: nothing to pull + assert awaiting_kvs is False @pytest.mark.cpu_test diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 5c01d99696..9cfd100730 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -113,7 +113,7 @@ class NixlBaseConnectorScheduler: # New requests are added by update_state_after_alloc in # the scheduler. Used to make metadata passed to Worker. self._reqs_need_recv: dict[ - ReqId, tuple[Request, BlockIds, tuple[int, ...]] + ReqId, tuple[Request, BlockIds, tuple[int, ...], bool] ] = {} self._reqs_need_save: dict[ReqId, Request] = {} # Reqs to send and their expiration time @@ -459,13 +459,19 @@ class NixlBaseConnectorScheduler: meta = NixlConnectorMetadata() # Loop through scheduled reqs and convert to ReqMeta. - for req_id, (req, block_ids, cached) in self._reqs_need_recv.items(): + for req_id, ( + req, + block_ids, + cached, + awaiting_kvs, + ) in self._reqs_need_recv.items(): assert req.kv_transfer_params is not None meta.add_new_req_to_recv( request_id=req_id, local_block_ids=block_ids, kv_transfer_params=req.kv_transfer_params, local_num_computed_blocks=cached, + awaiting_kvs=awaiting_kvs, ) if self.use_host_buffer: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index f61fdfa835..be32a6969d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -246,6 +246,9 @@ class ReqMeta: remote_block_size: int | None = None # Remote producer pipeline-parallel size (push mode, D side). pp_size: int = 1 + # True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS + # and expects it in finished_recving; notify-only recvs must not be reported. + awaiting_kvs: bool = False class NixlConnectorMetadata(KVConnectorMetadata): @@ -275,6 +278,7 @@ class NixlConnectorMetadata(KVConnectorMetadata): local_block_ids: BlockIds, kv_transfer_params: dict[str, Any], local_num_computed_blocks: tuple[int, ...] = (), + awaiting_kvs: bool = False, ) -> ReqMeta: return ReqMeta( local_block_ids=local_block_ids, @@ -285,6 +289,7 @@ class NixlConnectorMetadata(KVConnectorMetadata): remote_block_size=kv_transfer_params.get("remote_block_size"), pp_size=kv_transfer_params.get("pp_size", 1), local_num_computed_blocks=local_num_computed_blocks, + awaiting_kvs=awaiting_kvs, ) def add_new_req_to_save( @@ -303,11 +308,13 @@ class NixlConnectorMetadata(KVConnectorMetadata): local_block_ids: BlockIds, kv_transfer_params: dict[str, Any], local_num_computed_blocks: tuple[int, ...] = (), + awaiting_kvs: bool = False, ): req = self._add_new_req( local_block_ids, kv_transfer_params, local_num_computed_blocks, + awaiting_kvs, ) req.remote = RemoteMeta( block_ids=kv_transfer_params["remote_block_ids"], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 446126c328..9b9012a532 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -172,6 +172,9 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler): request, local_block_ids, local_num_computed_blocks, + # Parked in WAITING_FOR_REMOTE_KVS only when there is + # something to pull; a full local hit stays RUNNING. + num_external_tokens > 0, ) else: @@ -224,7 +227,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler): # To avoid stranding the prefill blocks in the prefill instance, # we must add empty block_ids to _reqs_need_recv so that our # worker side will notify and free blocks in the prefill instance. - self._reqs_need_recv[request.request_id] = (request, [], ()) + self._reqs_need_recv[request.request_id] = (request, [], (), False) params["do_remote_prefill"] = False return False, None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 59205b943a..f787639cfe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -303,6 +303,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): local_dram_handle=local_dram_handle, remote_xfer_side_handle=remote_xfer_side_handle, expected_consumers=plan.local_consumers, + awaiting_kvs=meta.awaiting_kvs, ): return @@ -327,6 +328,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): local_dram_handle: int | None, remote_xfer_side_handle: int, expected_consumers: int, + awaiting_kvs: bool, ) -> bool: """ Post a READ point-to-point xfer request from a single local worker to @@ -387,6 +389,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): remote_agent_name=agent_name, ) self.xfer_stats.record_failed_notification() + # Report even on notification failure: the KV is already local, and + # an unreported parked request would hold its blocks forever. + # Notify-only recvs must stay unreported (scheduler asserts). + if awaiting_kvs: + self._recving_transfers.setdefault(request_id, []) return True if read_spec.block_ids_by_region: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index c11a63bb34..a007e0c5a3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py @@ -198,7 +198,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler): # ReqMeta without a KeyError — the actual remote block IDs are # learned by P over the NIXL handshake at WRITE time. params["remote_block_ids"] = () - self._reqs_need_recv[request.request_id] = (request, local_block_ids, ()) + self._reqs_need_recv[request.request_id] = (request, local_block_ids, (), False) # Mark as processed so a re-entry (e.g. preemption + reschedule) # doesn't re-stage the registration. @@ -240,7 +240,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler): # recv so the worker emits a notif that lets P free them. # Seed remote_block_ids so add_new_req_to_recv won't KeyError. params["remote_block_ids"] = () - self._reqs_need_recv[request.request_id] = (request, [], ()) + self._reqs_need_recv[request.request_id] = (request, [], (), False) params["do_remote_prefill"] = False return False, None From a72c046463cecc35737f353b2b6cc0f44bad3f29 Mon Sep 17 00:00:00 2001 From: Miguel <60073809+migarci2@users.noreply.github.com> Date: Tue, 15 Sep 2026 05:19:08 +0200 Subject: [PATCH 070/139] [Bugfix] Avoid nested rerank cancellation handlers (#55602) Signed-off-by: Miguel Garcia --- vllm/entrypoints/pooling/scoring/api_router.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/entrypoints/pooling/scoring/api_router.py b/vllm/entrypoints/pooling/scoring/api_router.py index 90be32ba24..83ed600fdd 100644 --- a/vllm/entrypoints/pooling/scoring/api_router.py +++ b/vllm/entrypoints/pooling/scoring/api_router.py @@ -91,7 +91,6 @@ async def do_rerank(request: RerankRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation async def do_rerank_v1(request: RerankRequest, raw_request: Request): logger.warning_once( "To indicate that the rerank API is not part of the standard OpenAI" @@ -110,6 +109,5 @@ async def do_rerank_v1(request: RerankRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation async def do_rerank_v2(request: RerankRequest, raw_request: Request): return await do_rerank(request, raw_request) From 9d14fd3095eb169044e2da0a2a98dd1f2c77ce22 Mon Sep 17 00:00:00 2001 From: 404mario <18063257856@163.com> Date: Tue, 15 Sep 2026 12:04:22 +0800 Subject: [PATCH 071/139] [Kernel][MoE] DeepEP v2: async finalize to overlap shared experts with combine (#52781) Signed-off-by: Mario <404mario@users.noreply.github.com> Co-authored-by: Mario <404mario@users.noreply.github.com> Co-authored-by: Claude Fable 5 --- .../moe/test_deepep_v2_async_finalize.py | 122 ++++++++++++++++++ .../fused_moe/prepare_finalize/deepep_v2.py | 36 ++++-- 2 files changed, 150 insertions(+), 8 deletions(-) create mode 100644 tests/kernels/moe/test_deepep_v2_async_finalize.py diff --git a/tests/kernels/moe/test_deepep_v2_async_finalize.py b/tests/kernels/moe/test_deepep_v2_async_finalize.py new file mode 100644 index 0000000000..193d8906f3 --- /dev/null +++ b/tests/kernels/moe/test_deepep_v2_async_finalize.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Contract tests for DeepEPV2PrepareAndFinalize.finalize_async. + +Hermetic: fake buffer/event, single process, no GPU collectives — only +requires deep_ep to be importable. + + T1 finalize_async -> combine is issued with async_with_compute_stream=True + and the join is deferred to the receiver, which waits + on the combine event (device-side) before the copy. + T2 DBO active -> combine falls back to async_with_compute_stream=False; + the receiver still produces the correct output. + T3 finalize (sync) -> combine is issued synchronously and the output is + copied immediately. +""" + +import pytest +import torch + +from vllm.utils.import_utils import has_deep_ep_v2 + +requires_deep_ep_v2 = pytest.mark.skipif( + not has_deep_ep_v2(), + reason="Requires DeepEP v2 (ElasticBuffer)", +) + +if has_deep_ep_v2(): + import vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 as _dv2 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import ( + DeepEPV2PrepareAndFinalize, + ) + from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceContiguous, + ) + + +class _FakeEvent: + def __init__(self, has_event: bool): + self.event = object() if has_event else None + self.waited = False + + def current_stream_wait(self): + assert self.event is not None + self.waited = True + + +class _FakeBuffer: + def __init__(self, out: torch.Tensor): + self.out = out + self.calls: list[dict] = [] + self.last_event: _FakeEvent | None = None + + def combine( + self, + x, + handle, + topk_weights, + async_with_compute_stream, + allocate_on_comm_stream, + ): + self.calls.append({"async_with_compute_stream": async_with_compute_stream}) + # DeepEP returns a completion event only for an async combine. + self.last_event = _FakeEvent(has_event=async_with_compute_stream) + return self.out, None, self.last_event + + +def _make_pf(out: torch.Tensor): + pf = DeepEPV2PrepareAndFinalize( + buffer=_FakeBuffer(out), + num_dispatchers=1, + dp_size=1, + rank_expert_offset=0, + num_experts=8, + num_topk=2, + ) + pf.handles[0] = object() + return pf + + +def _run(pf, output: torch.Tensor, do_async: bool): + empty = torch.empty(0, dtype=torch.bfloat16) + weights = torch.empty(0, 2) + ids = torch.empty(0, 2, dtype=torch.int64) + fn = pf.finalize_async if do_async else pf.finalize + return fn(output, empty, weights, ids, False, TopKWeightAndReduceContiguous()) + + +@requires_deep_ep_v2 +def test_finalize_async_defers_join_to_receiver(): + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + recv = _run(pf, dst, do_async=True) + assert pf.buffer.calls[0]["async_with_compute_stream"] is True + assert not torch.equal(dst, out_src), "join must not happen before receiver" + recv() + assert pf.buffer.last_event.waited, "receiver must join via event wait" + assert torch.equal(dst, out_src) + + +@requires_deep_ep_v2 +def test_finalize_async_under_dbo_uses_sync_combine(monkeypatch): + monkeypatch.setattr(_dv2, "dbo_enabled", lambda: True) + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + recv = _run(pf, dst, do_async=True) + assert pf.buffer.calls[0]["async_with_compute_stream"] is False + recv() + assert not pf.buffer.last_event.waited, "no event to wait on for sync combine" + assert torch.equal(dst, out_src) + + +@requires_deep_ep_v2 +def test_finalize_sync(): + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + ret = _run(pf, dst, do_async=False) + assert ret is None + assert pf.buffer.calls[0]["async_with_compute_stream"] is False + assert torch.equal(dst, out_src) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index d834c6eff7..b3a146274d 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -21,6 +21,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, + dbo_enabled, ) @@ -89,8 +90,10 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): use cudagraphs anyway - Provides expert_tokens_meta for efficient batched expert kernels - Both modes use async_with_compute_stream=False (synchronous from - caller's perspective). The ElasticBuffer handles comm internally. + Dispatch always uses async_with_compute_stream=False. finalize_async + issues the combine with async_with_compute_stream=True (except under + DBO) so the modular kernel can overlap the shared-expert FFN with the + combine a2a; the returned receiver joins via a device-side event wait. """ @staticmethod @@ -465,15 +468,31 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): f"got {fused_expert_output.dtype}" ) + # DBO drives its own hook/receiver schedule; keep the combine + # synchronous there (the receiver then only performs the copy). + combine_async = do_async and not dbo_enabled() combined_x, _, event = self.buffer.combine( x=fused_expert_output, handle=handle, topk_weights=None, - async_with_compute_stream=False, + async_with_compute_stream=combine_async, + allocate_on_comm_stream=combine_async, ) - output.copy_(combined_x, non_blocking=True) - return None + if do_async: + # The combine ran on DeepEP's comm stream; the modular kernel + # issues the shared-expert FFN before calling the receiver, which + # joins via a device-side cudaStreamWaitEvent (no host sync, so + # this is safe inside a captured region). + def _receiver(): + if event.event is not None: + event.current_stream_wait() + output.copy_(combined_x, non_blocking=True) + + return _receiver + else: + output.copy_(combined_x, non_blocking=True) + return None def finalize_async( self, @@ -484,16 +503,17 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): apply_router_weight_on_input: bool, weight_and_reduce_impl: mk.TopKWeightAndReduce, ) -> Callable: - self._finalize( + receiver = self._finalize( output, fused_expert_output, topk_weights, topk_ids, apply_router_weight_on_input, weight_and_reduce_impl, - False, + True, ) - return lambda: None + assert receiver is not None + return receiver def finalize( self, From 876595275548dff6056ff0942592dbe93b070193 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 21:18:10 -0700 Subject: [PATCH 072/139] [CI/Build] Retry CPU image stream interruptions (#56597) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex Co-authored-by: Sherlock Co-authored-by: OpenAI Codex --- .buildkite/scripts/hardware_ci/run-cpu-test.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index 74c7749d50..e8d0a24bc7 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -59,7 +59,7 @@ docker pull ubuntu:25.04 || true # building the docker image echo "--- :docker: Building Docker image" -BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:' +BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:|The read operation timed out|BrokenPipeError:.*Broken pipe|HTTP/[0-9.]+ stream [0-9]+ was not closed cleanly|fetch-pack: unexpected disconnect|fatal: early EOF|invalid index-pack output' BUILD_MAX_ATTEMPTS=4 # 1 initial + 3 retries BUILD_RETRY_WAITS=(10 20 40) # seconds to wait before retry 1/2/3 build_log="$(mktemp)" From a491b6039d8d23314efc30a2113000ea987e409e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 21:20:14 -0700 Subject: [PATCH 073/139] [CI] Drop root privileges for the OTel GPU sampler (#56941) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/scripts/ci-otel/ci_otel.sh | 53 ++++++++++++- .../scripts/ci-otel/tests/test_ci_otel.py | 79 +++++++++++++++++++ 2 files changed, 129 insertions(+), 3 deletions(-) diff --git a/.buildkite/scripts/ci-otel/ci_otel.sh b/.buildkite/scripts/ci-otel/ci_otel.sh index c9897e732e..7f1b700fe6 100755 --- a/.buildkite/scripts/ci-otel/ci_otel.sh +++ b/.buildkite/scripts/ci-otel/ci_otel.sh @@ -65,6 +65,39 @@ _ci_otel_python() { fi } +# The Docker test commands run as root, while /workdir is owned by the +# Buildkite checkout user. Run the GPU sampler as that unprivileged identity so +# any future project-code import cannot leave root-owned files in the checkout. +_ci_otel_prepare_gpu_user() { + _CI_INFRA_GPU_UID="" + _CI_INFRA_GPU_GID="" + [ "$(id -u 2>/dev/null || :)" = "0" ] || return 0 + command -v setpriv >/dev/null 2>&1 || return 1 + + _CI_INFRA_OTEL_WORKSPACE_DIR="${CI_INFRA_OTEL_WORKSPACE_DIR:-/workdir}" + [ -d "${_CI_INFRA_OTEL_WORKSPACE_DIR}" ] || return 1 + # The owner fields are numeric and whitespace-free on GNU stat. + # shellcheck disable=SC2046 + set -- $(stat -c '%u %g' -- "${_CI_INFRA_OTEL_WORKSPACE_DIR}" 2>/dev/null || :) + [ "$#" -eq 2 ] || return 1 + case "$1:$2" in + *[!0-9:]* | 0:* | *:0) return 1 ;; + esac + + _CI_INFRA_GPU_UID="$1" + _CI_INFRA_GPU_GID="$2" + chown "${_CI_INFRA_GPU_UID}:${_CI_INFRA_GPU_GID}" \ + "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1 + chmod 0700 "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1 + if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then + # Permit traversal to the private child spool, but not directory listing. + chmod 0711 "${CI_INFRA_OTEL_RUNTIME_DIR}" || return 1 + fi + setpriv --reuid="${_CI_INFRA_GPU_UID}" \ + --regid="${_CI_INFRA_GPU_GID}" --clear-groups --no-new-privs \ + -- test -w "${CI_INFRA_OTEL_SPOOL_DIR}" +} + ci_otel_start() { _CI_INFRA_OTEL_COMMAND_INDEX="$1" _CI_INFRA_OTEL_COMMAND_LABEL="${2:-command ${_CI_INFRA_OTEL_COMMAND_INDEX}}" @@ -87,9 +120,23 @@ ci_otel_start() { if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] && [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] && command -v nvidia-smi >/dev/null 2>&1; then - PYTHONDONTWRITEBYTECODE=1 \ - "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" &2 || : + fi fi } diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index ae51010f9e..c66530989f 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -987,6 +987,85 @@ def test_gpu_periodic_batches_are_bounded_even_when_upload_fails(monkeypatch): assert [len(batch[0].events) for batch in batches] == [1, 1, 1] +def _write_executable(path: Path, body: str) -> None: + path.write_text(f"#!/bin/sh\n{body}") + path.chmod(0o755) + + +@pytest.mark.parametrize( + "workspace_owner,expect_sampler", + [("997 996", True), ("0 0", False)], +) +def test_root_shell_runs_gpu_sampler_as_checkout_owner( + tmp_path, workspace_owner, expect_sampler +): + bin_dir = tmp_path / "bin" + workspace = tmp_path / "workdir" + runtime = tmp_path / "runtime" + log = tmp_path / "calls.log" + bin_dir.mkdir() + workspace.mkdir() + + _write_executable( + bin_dir / "id", + '[ "$1" = "-u" ] && { echo 0; exit 0; }; exec /usr/bin/id "$@"\n', + ) + _write_executable( + bin_dir / "stat", + f"printf '%s\\n' '{workspace_owner}'\n", + ) + for command in ("chown", "chmod"): + _write_executable(bin_dir / command, "exit 0\n") + _write_executable( + bin_dir / "setpriv", + 'printf "setpriv %s\\n" "$*" >> "$CI_OTEL_TEST_LOG"\n' + 'while [ "$#" -gt 0 ] && [ "$1" != "--" ]; do shift; done\n' + '[ "$#" -gt 0 ] && shift\n' + 'exec "$@"\n', + ) + _write_executable( + bin_dir / "python3", + '[ "$1" = "-c" ] && exit 0\n' + '[ "$2" = "new-context" ] && {\n' + ' echo "01010101010101010101010101010101 0202020202020202 - 1"\n' + " exit 0\n" + "}\n" + 'case "$1" in *ci_gpu.py) echo gpu >> "$CI_OTEL_TEST_LOG";; esac\n' + "exit 0\n", + ) + _write_executable(bin_dir / "nvidia-smi", "exit 0\n") + shell = ( + f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' + f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0" + ) + + result = subprocess.run( + ["/bin/sh", "-c", shell], + check=False, + capture_output=True, + text=True, + env={ + **os.environ, + "PATH": f"{bin_dir}:{os.environ['PATH']}", + "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), + "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime), + "CI_INFRA_OTEL_WORKSPACE_DIR": str(workspace), + "CI_OTEL_TEST_LOG": str(log), + }, + ) + + assert result.returncode == 0, result.stderr + calls = log.read_text() if log.exists() else "" + if expect_sampler: + expected = "--reuid=997 --regid=996 --clear-groups --no-new-privs --" + assert calls.count("setpriv") == 2 + assert calls.count(expected) == 2 + assert "gpu" in calls + else: + assert calls == "" + assert "non-root GPU sampler unavailable" in result.stderr + + @pytest.mark.parametrize("sampling", ["1", "0"]) def test_gpu_sampler_stops_with_command_and_preserves_failure_status( tmp_path, sampling From caf56ce9dac8614b11d682d1cfd288cfbe473a43 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 21:25:30 -0700 Subject: [PATCH 074/139] [CI] Pass the scale-out endpoint flag in EC E2E (#56819) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: Tyler Michael Smith Co-authored-by: OpenAI Codex Co-authored-by: Tyler Michael Smith --- docs/usage/security.md | 2 +- .../scale_out/ec_integration/run_scale_out_ec_e2e_test.sh | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index 3c7f1e3f9b..5d4d429c6f 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -170,7 +170,7 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/messages/count_tokens` - Count tokens for Anthropic messages - `/v1/responses` - Create a response -- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render` unless explicitly disabled) - `/v1/responses/{response_id}` - Retrieve a response - `/v1/responses/{response_id}/cancel` - Cancel a response - `/v1/score` - Scoring API diff --git a/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh b/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh index b99b1c3ca5..462ada22d7 100755 --- a/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh +++ b/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh @@ -181,8 +181,9 @@ run_scale_out_ec() { # Encode-only EC producer instance. It runs the vision encoder and # publishes embeddings through the ECExampleConnector shared storage. echo "Starting encode instance on GPU $GPU_E, port $ENCODE_PORT" - env CUDA_VISIBLE_DEVICES="$GPU_E" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + env CUDA_VISIBLE_DEVICES="$GPU_E" \ vllm serve "$MODEL" \ + --enable-scale-out \ --port "$ENCODE_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ @@ -204,8 +205,9 @@ run_scale_out_ec() { # features plus ec_transfer_params and loads the embeddings that the # encode instance published. echo "Starting prefill instance on GPU $GPU_PD, port $PREFILL_PORT" - env CUDA_VISIBLE_DEVICES="$GPU_PD" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + env CUDA_VISIBLE_DEVICES="$GPU_PD" \ vllm serve "$MODEL" \ + --enable-scale-out \ --port "$PREFILL_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ From 45779ef99f8c7f4bd59f86a48acccc860fc3c97c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 15 Sep 2026 06:27:10 +0200 Subject: [PATCH 075/139] [Bugfix][NIXL] Fix multiple handles xfer race (#56104) Signed-off-by: NickLucche Co-authored-by: Nick Hill --- .../kv_connector/unit/test_nixl_connector.py | 161 +++++++++--------- .../unit/test_nixl_push_connector.py | 73 +++++++- .../kv_connector/v1/nixl/base_worker.py | 113 ++++++------ .../kv_connector/v1/nixl/pull_worker.py | 17 +- .../kv_connector/v1/nixl/push_worker.py | 20 ++- 5 files changed, 233 insertions(+), 151 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 298886d160..39fb0bb54f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -4,6 +4,7 @@ import contextlib import inspect import os +import queue import tempfile import textwrap import time @@ -1952,17 +1953,43 @@ def test_mixed_memory_local_descriptors_split_by_memory_type(): assert memory_types == ["DRAM", "VRAM"] -def test_mixed_memory_read_notifies_after_both_transfers_finish(): +@pytest.fixture +def recv_worker(): + """Receive lifecycle state without distributed or device initialization.""" worker = object.__new__(NixlConnectorWorker) + worker.transfer_topo = MagicMock() + worker.transfer_topo.block_size_ratio.return_value = 1 + worker.transfer_topo.get_engine_info.return_value = SimpleNamespace( + remote_block_size=16, remote_physical_blocks_per_logical=1 + ) + worker._physical_blocks_per_logical_kv_block = 1 + worker._recving_metadata = {"request": MagicMock(local_block_ids=([1, 2, 3],))} + worker._recving_transfers = defaultdict(list) + worker._failed_recv_reqs = queue.Queue() + worker._recv_failures = set() + worker._replicated_pcp_done_sending = set() + worker._invalid_block_ids = queue.Queue() + worker._pending_recv_notifs = {} + worker._reqs_to_send = {} + worker._replicated_pcp_done_sending = set() + worker._is_hma_required = False + worker._has_mamba = False + worker.use_host_buffer = False + worker.enable_permute_local_kv = False + worker.enable_heterogeneous_attn_post_process = False + worker.tp_rank = 0 + worker._log_failure = MagicMock() # type: ignore[method-assign] + worker.xfer_stats = NixlKVConnectorStats() + worker.nixl_wrapper = MagicMock() + worker.nixl_wrapper.get_xfer_telemetry.return_value = get_default_xfer_telemetry() + return worker + + +def test_mixed_memory_read_notifies_after_both_transfers_finish(recv_worker): + worker = recv_worker worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} worker._dram_src_handles_by_block_size = {16: 10} - worker._recving_metadata = {"request": MagicMock()} - worker._recving_transfers = defaultdict(list) - worker._pending_recv_notifs = {} - worker._failed_inflight_recvs = set() - worker.xfer_stats = MagicMock() - worker.nixl_wrapper = MagicMock() worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102] worker.nixl_wrapper.check_xfer_state.return_value = "DONE" @@ -1989,27 +2016,20 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish(): np.testing.assert_array_equal(device_read.args[4], [7]) worker.nixl_wrapper.send_notif.assert_not_called() - assert worker._pop_done_transfers(worker._recving_transfers) == {"request"} + assert worker.get_finished() == (set(), {"request"}) worker.nixl_wrapper.send_notif.assert_called_once_with( "prefill", notif_msg=b"request:1" ) -def test_mixed_memory_read_failure_does_not_notify_producer(): +def test_mixed_memory_read_failure_does_not_notify_producer(recv_worker): """A failed half of a split READ must suppress its success notification.""" - worker = object.__new__(NixlConnectorWorker) - worker._recving_metadata = {"request": MagicMock()} + worker = recv_worker worker._recving_transfers = {"request": [101, 102]} worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]} - worker._failed_inflight_recvs = set() - worker._is_hma_required = True - worker._failed_recv_reqs = MagicMock() - worker._log_failure = MagicMock() # type: ignore[method-assign] - worker.xfer_stats = MagicMock() - worker.nixl_wrapper = MagicMock() worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"] - worker._pop_done_transfers(worker._recving_transfers) + assert worker.get_finished() == (set(), {"request"}) assert "request" not in worker._pending_recv_notifs worker.nixl_wrapper.send_notif.assert_not_called() @@ -2969,13 +2989,10 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init): # Wait for handshake to fail time.sleep(0.3) - # Check that blocks were marked invalid - invalid_blocks = connector.get_block_ids_with_load_errors() - assert invalid_blocks == {1, 2, 3} - - # Check that request appears in get_finished + assert connector.get_block_ids_with_load_errors() == set() _, done_recving = connector.get_finished(finished_req_ids=set()) assert request_id in done_recving + assert connector.get_block_ids_with_load_errors() == {1, 2, 3} @patch( @@ -3026,6 +3043,7 @@ def test_transfer_setup_failure_returns_finished( time.sleep(0.1) connector.start_load_kv(dummy_ctx) + assert connector.get_block_ids_with_load_errors() == set() results = connector.get_transfer_results(finished_req_ids=set()) assert request_id in results.finished_recving assert results.failed_recving == {request_id} @@ -3248,74 +3266,49 @@ def test_failed_request_skips_kv_postprocessing( assert invalid_blocks == {1, 2, 3} -@patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", - FailingNixlWrapper, -) -def test_handles_failing_in_separate_polls_do_not_kill_the_engine( - default_vllm_config, dist_init -): - """Failures across polls report once, after the last writer is terminal.""" - vllm_config = create_vllm_config() - connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) - ) - connector.connector_worker = FakeNixlConnectorWorker( - vllm_config, connector.engine_id, hand_shake_latency=0 - ) - worker = connector.connector_worker +@pytest.mark.parametrize("sibling_state", ["ERR", "DONE"]) +def test_recv_failure_waits_for_sibling_transfer(recv_worker, sibling_state): + """Failure must not release blocks while a sibling can still write to them.""" + worker = recv_worker + worker._recving_transfers["request"] = [101, 102] + worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]} + worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "PROC", sibling_state] - request_id = "test_two_handles_failing_in_separate_polls" - metadata = NixlConnectorMetadata() - metadata.add_new_req_to_recv( - request_id=request_id, - local_block_ids=([1, 2, 3],), - kv_transfer_params={ - "remote_block_ids": ([4, 5, 6],), - "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, - "remote_request_id": f"prefill-{request_id}", - "remote_host": "localhost", - "remote_port": 1234, - "remote_tp_size": 1, - }, - ) - connector.bind_connector_metadata(metadata) - dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) - connector.start_load_kv(dummy_ctx) - connector.bind_connector_metadata(NixlConnectorMetadata()) - time.sleep(0.1) - connector.start_load_kv(dummy_ctx) + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() + assert "request" in worker._recving_metadata + worker.nixl_wrapper.release_xfer_handle.assert_called_once_with(101) - # Give the request a second handle, as a second remote rank would. - handles = worker._recving_transfers[request_id] - assert len(handles) == 1 - first, second = handles[0], object() - worker._recving_transfers[request_id] = [first, second] + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == {"request"} + assert worker.get_block_ids_with_load_errors() == {1, 2, 3} + worker.nixl_wrapper.send_notif.assert_not_called() - # Poll 1: the first handle has failed, the second is still running. - def first_failed(handle): - return "ERR" if handle is first else "PROC" + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() - with patch.object( - worker.nixl_wrapper, "check_xfer_state", side_effect=first_failed - ): - _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id not in done_recving - assert request_id in worker._recving_metadata - assert worker._recving_transfers[request_id] == [second] - assert connector.get_block_ids_with_load_errors() == set() +def test_recv_failure_waits_for_unpollable_handle_to_be_released(recv_worker): + """A polling exception and failed release do not prove that DMA has stopped.""" + worker = recv_worker + worker._recving_transfers["request"] = [101] + worker.nixl_wrapper.check_xfer_state.side_effect = [RuntimeError("poll"), "DONE"] + worker.nixl_wrapper.release_xfer_handle.side_effect = [ + RuntimeError("release"), + None, + ] - # Only the last terminal handle permits the scheduler to reuse the blocks. - with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"): - _, done_recving = connector.get_finished(finished_req_ids=set()) + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() + assert "request" in worker._recving_metadata - assert done_recving == {request_id} - assert request_id not in worker._recving_transfers - assert request_id not in worker._recving_metadata - assert connector.get_block_ids_with_load_errors() == {1, 2, 3} - _, done_recving = connector.get_finished(finished_req_ids=set()) - assert done_recving == set() + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == {"request"} + assert worker.get_block_ids_with_load_errors() == {1, 2, 3} + assert worker.nixl_wrapper.release_xfer_handle.call_count == 2 def _set_test_speculative_config( diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 2c7f0cb866..423b0e9dec 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -351,7 +351,7 @@ class _StubWriterWorker(NixlPushConnectorWorker): # Base worker fields touched by start_load_kv / _get_new_notifs. w._recving_metadata = {} w._pending_recv_notifs = {} - w._failed_inflight_recvs = set() + w._recv_failures = set() w._recving_transfers = defaultdict(list) w._reqs_to_process = set() w._reqs_to_send = {} @@ -877,6 +877,68 @@ class TestPushWriterNotifs: assert evicted == ["req-done"] assert w._push_writer_wake.is_set() + @staticmethod + def _pollable_worker() -> _StubWriterWorker: + """Stub worker with the extra base-worker state the real + ``get_transfer_results`` needs to poll ``_sending_transfers``.""" + w = _StubWriterWorker.fresh() + w.transfer_topo = MagicMock() + w.nixl_wrapper = MagicMock() + w.xfer_stats = MagicMock() + w._log_failure = MagicMock() # type: ignore[method-assign] + w._failed_recv_reqs = queue.Queue() + w._invalid_block_ids = queue.Queue() + w._pending_recv_notifs = {} + w._replicated_pcp_done_sending = set() + w.use_host_buffer = False + return w + + def _make_sending_req(self, w: _StubWriterWorker) -> str: + request_id = "req-send-1" + w._sending_transfers[request_id] = [101, 102] + w._reqs_to_send[request_id] = time.perf_counter() + 60 + w._reqs_to_process.add(request_id) + return request_id + + def test_failed_push_send_is_not_reported_done_sending(self): + """A failed P-side WRITE must not surface as finished sending. + + Reporting it would free P-side blocks as if the KV had been + delivered while D still waits for the data; the request must be + left to the lease / watchdog instead. + """ + w = self._pollable_worker() + request_id = self._make_sending_req(w) + w.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"] + + results = w.get_transfer_results() + + assert request_id not in results.finished_sending + # Lease tracking is preserved so the watchdog can reschedule. + assert request_id in w._reqs_to_send + assert request_id in w._reqs_to_process + # Both handles were still cleaned up. + assert request_id not in w._sending_transfers + + def test_completed_push_send_waits_for_consumer_notif(self): + """P-side sends are completed by consumer notifs / lease expiry, + not by local WRITE-handle completion.""" + w = self._pollable_worker() + request_id = self._make_sending_req(w) + w.nixl_wrapper.check_xfer_state.return_value = "DONE" + + results = w.get_transfer_results() + + assert request_id not in results.finished_sending + assert request_id in w._reqs_to_send + assert request_id not in w._sending_transfers + + # The consumer notif is what completes the send. + w._pending_completion_notifs.put(f"{request_id}:1".encode()) + results = w.get_transfer_results() + assert request_id in results.finished_sending + assert request_id not in w._reqs_to_send + # ----------------------------------------------------------------- # # Negative / error-path tests # @@ -1079,7 +1141,7 @@ class TestPushWriterNegative: # Wake set so the writer drains NIXL notifs even when idle. assert w._push_writer_wake.is_set() - def test_failed_push_transfer_finishes_sending(self): + def test_failed_push_transfer_does_not_evict_writer_state(self): w = _StubWriterWorker.fresh() w.nixl_wrapper = MagicMock() w.nixl_wrapper.check_xfer_state.return_value = "ERR" @@ -1096,11 +1158,12 @@ class TestPushWriterNegative: ): results = w.get_transfer_results() - assert results.finished_sending == {"req-failed"} + assert results.finished_sending == set() assert results.finished_recving == set() assert "req-failed" not in w._sending_transfers - assert "req-failed" not in w._reqs_to_send - assert "req-failed" not in w._reqs_to_process + assert "req-failed" in w._reqs_to_send + assert "req-failed" in w._reqs_to_process + assert w._evict_finished_inbox.empty() w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17) def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 58ab5850ea..6f8f1e55dc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -692,7 +692,7 @@ class NixlBaseConnectorWorker: # Uses Queue for thread-safe cross-thread coordination with the # background handshake thread, matching the _ready_requests pattern. self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue() - self._failed_inflight_recvs: set[ReqId] = set() + self._recv_failures: set[ReqId] = set() self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {} # Handshake metadata of this worker for NIXL transfers. @@ -1212,7 +1212,7 @@ class NixlBaseConnectorWorker: error=e, meta=meta, ) - self._handle_failed_transfer(req_id, None) + self._failed_recv_reqs.put(req_id) fut.add_done_callback(request_ready) @@ -2618,22 +2618,32 @@ class NixlBaseConnectorWorker: """ assert self.transfer_topo is not None done_sending = self._get_new_notifs() - done_recving = self._pop_done_transfers(self._recving_transfers) + done_recving, newly_failed = self._pop_done_transfers(self._recving_transfers) + if newly_failed: + self._recv_failures.update(newly_failed) done_sending.update(self._replicated_pcp_done_sending) self._replicated_pcp_done_sending.clear() - # Drain queue of requests where handshake or transfer setup failed. - failed_recv_reqs = set[ReqId]() + # Process receive failures reported by background threads. while not self._failed_recv_reqs.empty(): try: - failed_recv_reqs.add(self._failed_recv_reqs.get_nowait()) + req_id = self._failed_recv_reqs.get_nowait() except queue.Empty: break + self._handle_failed_transfer(req_id, None, self._recv_failures) + + failed_recv_reqs: set[ReqId] = set() + if self._recv_failures: + # Keep failure bookkeeping independent of healthy request count. + for req_id in tuple(self._recv_failures): + if req_id not in self._recving_metadata: + self._recv_failures.remove(req_id) + elif req_id not in self._recving_transfers: + self._recv_failures.remove(req_id) + failed_recv_reqs.add(req_id) + done_recving.update(failed_recv_reqs) - # Add failed requests to done_recving for scheduler tracking - # (blocks are already marked invalid, scheduler will handle recompute) - done_recving.update(failed_recv_reqs) if len(done_sending) > 0 or len(done_recving) > 0: logger.debug( "Rank %s, get_finished: %s requests done sending " @@ -2654,12 +2664,17 @@ class NixlBaseConnectorWorker: # Skip KV sync and post-processing for failed requests if req_id in failed_recv_reqs: + self._pending_recv_notifs.pop(req_id, None) + # TODO (NickLucche) handle failed transfer for HMA. + if not self._is_hma_required: + self._invalid_block_ids.put(set(meta.local_block_ids[0])) logger.warning( "Skipping KV post-processing for failed request %s", req_id, ) continue + self._send_pending_recv_notifs(req_id) assert meta.remote is not None if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) @@ -2783,22 +2798,26 @@ class NixlBaseConnectorWorker: new_expiry, ) - def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: - """ - Pop completed xfers by checking for DONE state. + def _pop_done_transfers( + self, transfers: dict[str, list[int]] + ) -> tuple[set[str], set[str]]: + """Poll transfers, retaining handles until they can be released. + Args: - transfers: dict of req_id -> list[running_xfer] + transfers: Outstanding handles by request, updated in place. + Returns: - set of req_ids that have all done xfers + Requests with no outstanding handles and requests with observed + failures. Failed requests may still have outstanding handles. """ done_req_ids: set[str] = set() + failed_req_ids: set[str] = set() for req_id, handles in list(transfers.items()): in_progress = [] for handle in handles: try: xfer_state = self.nixl_wrapper.check_xfer_state(handle) if xfer_state == "DONE": - # Get telemetry from NIXL res = self.nixl_wrapper.get_xfer_telemetry(handle) self.xfer_stats.record_transfer(res) self.nixl_wrapper.release_xfer_handle(handle) @@ -2807,58 +2826,52 @@ class NixlBaseConnectorWorker: else: self._log_failure( failure_type="transfer_failed", - msg="Deferring request completion until its last " - "xfer is terminal", req_id=req_id, xfer_state=xfer_state, ) - self._handle_failed_transfer(req_id, handle) + if not self._handle_failed_transfer( + req_id, handle, failed_req_ids + ): + in_progress.append(handle) except Exception as e: self._log_failure( failure_type="transfer_exception", - msg="Handle is unpollable; treating it as terminal", req_id=req_id, error=e, ) - self._handle_failed_transfer(req_id, handle) + if not self._handle_failed_transfer(req_id, handle, failed_req_ids): + in_progress.append(handle) if not in_progress: - del transfers[req_id] done_req_ids.add(req_id) - if req_id in self._failed_inflight_recvs: - self._failed_inflight_recvs.remove(req_id) - self._report_failed_recv(req_id) - else: - self._send_pending_recv_notifs(req_id) + del transfers[req_id] else: transfers[req_id] = in_progress - return done_req_ids + return done_req_ids, failed_req_ids - def _handle_failed_transfer(self, req_id: str, handle: int | None): - """ - Handle a failed transfer by marking all (logical) blocks as invalid and - recording the failure. - - Args: - req_id: The request ID. - handle: The transfer handle. - """ - # A sibling READ may still be writing these blocks. Retain metadata - # and defer invalidation until every handle is terminal. - self._pending_recv_notifs.pop(req_id, None) - if req_id in self._recving_transfers: - self._failed_inflight_recvs.add(req_id) - else: - self._report_failed_recv(req_id) - if handle is not None: + def _try_release_xfer_handle(self, req_id: str, handle: int) -> bool: + """Release a handle, returning False if the caller must retain it.""" + try: self.nixl_wrapper.release_xfer_handle(handle) - self.xfer_stats.record_failed_transfer() + except Exception as e: + # A status error does not guarantee that the backend stopped DMA. + self._log_failure( + failure_type="transfer_release_failed", + msg="Retaining handle and blocks until release succeeds", + req_id=req_id, + error=e, + ) + return False + return True - def _report_failed_recv(self, req_id: str) -> None: - if (meta := self._recving_metadata.get(req_id)) is not None: - if not self._is_hma_required: - self._invalid_block_ids.put(set(meta.local_block_ids[0])) - self._failed_recv_reqs.put(req_id) + def _handle_failed_transfer( + self, req_id: str, handle: int | None, failed_req_ids: set[str] | None = None + ) -> bool: + """Record a failure and release its handle, returning False to retain it.""" + self.xfer_stats.record_failed_transfer() + if failed_req_ids is not None: + failed_req_ids.add(req_id) + return handle is None or self._try_release_xfer_handle(req_id, handle) def _send_pending_recv_notifs(self, req_id: str) -> None: """Send notifications deferred by split DRAM/VRAM reads.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index f787639cfe..04caebbd1f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -149,7 +149,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): engine_id, ) self.xfer_stats.record_kv_expired_req() - self._handle_failed_transfer(req_id, None) + self._handle_failed_transfer(req_id, None, self._recv_failures) return if any(len(group) > 0 for group in meta.local_block_ids): @@ -493,16 +493,19 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): self._recving_transfers[request_id].append(handle) return True except Exception as e: - # mark all (logical) blocks for this request as invalid self._log_failure( failure_type="transfer_setup_failed", req_id=request_id, - msg="Marking blocks as invalid", + msg="Deferring failure reporting until outstanding transfers finish", error=e, dst_engine_id=dst_engine_id, remote_rank=remote_rank, ) - self._handle_failed_transfer(request_id, handle) + if not self._handle_failed_transfer( + request_id, handle, self._recv_failures + ): + assert handle is not None + self._recving_transfers[request_id].append(handle) return False def _read_blocks_mixed( @@ -544,7 +547,8 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): ) except Exception: for handle in handles: - self.nixl_wrapper.release_xfer_handle(handle) + if not self._try_release_xfer_handle(request_id, handle): + self._recving_transfers[request_id].append(handle) raise self._pending_recv_notifs.setdefault(request_id, []).append( @@ -555,7 +559,8 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): self.nixl_wrapper.transfer(handle) except Exception: for unstarted in handles[i:]: - self.nixl_wrapper.release_xfer_handle(unstarted) + if not self._try_release_xfer_handle(request_id, unstarted): + self._recving_transfers[request_id].append(unstarted) raise self._recving_transfers[request_id].append(handle) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 603321260d..8b34e7ffbb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -341,7 +341,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): self._log_failure( failure_type="push_reg_handshake_failed", req_id=rid, error=e ) - self._handle_failed_transfer(rid, None) + self._failed_recv_reqs.put(rid) return # Re-queue for the writer to send now that the handshake is done. self._reg_send_inbox.put((rid, rd)) @@ -363,7 +363,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): engine_id, req_id, ) - self._handle_failed_transfer(req_id, None) + self._failed_recv_reqs.put(req_id) return for rank, agent_name in agents.items(): try: @@ -722,9 +722,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): # don't have a ``_recving_metadata`` entry to invalidate, so # we just release the handle and let the engine reschedule # via the lease / watchdog. - if handle is not None: - self.nixl_wrapper.release_xfer_handle(handle) - self.xfer_stats.record_failed_transfer() + if not self._handle_failed_transfer(request_id, handle): + return handle return None # --- Notification handling on engine main thread ------------------ # @@ -803,7 +802,16 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): # ``_pop_done_transfers`` mutates ``_sending_transfers``; the # writer thread also appends to it, so guard the pop. with self._sending_transfers_lock: - done_pushing = self._pop_done_transfers(self._sending_transfers) + done_pushing, failed_pushing = self._pop_done_transfers( + self._sending_transfers + ) + # A failed send must never be reported as done: its blocks + # are freed via the lease / watchdog instead. + done_pushing = { + req_id + for req_id in done_pushing - failed_pushing + if req_id in self._recving_metadata + } for req_id in done_pushing: self._reqs_to_send.pop(req_id, None) self._reqs_to_process.discard(req_id) From f84b0c4bceee120936030dfb0c90af0a818263ed Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 14 Sep 2026 21:28:32 -0700 Subject: [PATCH 076/139] [Model Runner V2] Tolerate lack of pinned memory support (#56908) Signed-off-by: Nick Hill --- vllm/v1/worker/cpu/buffer_utils.py | 8 +++--- vllm/v1/worker/gpu/buffer_utils.py | 43 +++++++++++++++++++++++++----- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/vllm/v1/worker/cpu/buffer_utils.py b/vllm/v1/worker/cpu/buffer_utils.py index 8ee257a28f..cdc76ab8b4 100644 --- a/vllm/v1/worker/cpu/buffer_utils.py +++ b/vllm/v1/worker/cpu/buffer_utils.py @@ -4,13 +4,11 @@ from collections.abc import Sequence import torch -from vllm.utils.platform_utils import is_uva_available - class UvaBuffer: def __init__(self, size: int | Sequence[int], dtype: torch.dtype): - if not is_uva_available(): - raise RuntimeError("UVA is not available") self.cpu = torch.zeros(size, dtype=dtype, device="cpu") self.np = self.cpu.numpy() - self.uva = self.cpu + + def uva(self, n: int | None = None): + return self.cpu[:n] if n is not None else self.cpu diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 2c1dac2d07..0066f80907 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -6,13 +6,17 @@ from functools import partial import numpy as np import torch +from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import ( + PIN_MEMORY, async_tensor_h2d, get_accelerator_view_from_cpu_tensor, ) +logger = init_logger(__name__) + # Default round-robin depth for the UVA buffer pools. Must be >= the number of # concurrent in-flight steps (engine batch_queue_size). _DEFAULT_MAX_CONCURRENCY = 2 @@ -37,7 +41,7 @@ def async_copy_to_gpu( out = torch.empty_like(x, device=device) # pin_memory() is no-op if the memory is already pinned. - pinned = x.pin_memory() + pinned = x.pin_memory() if PIN_MEMORY else x return out.copy_(pinned, non_blocking=True) @@ -47,7 +51,30 @@ class UvaBuffer: raise RuntimeError("UVA is not available") self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True) self.np = self.cpu.numpy() - self.uva = get_accelerator_view_from_cpu_tensor(self.cpu) + self._uva = get_accelerator_view_from_cpu_tensor(self.cpu) + + def uva(self, n: int | None = None) -> torch.Tensor: + return self._uva[:n] if n is not None else self._uva + + +class NonUvaBuffer: + """Explicit-copy fallback for platforms without pinned memory.""" + + def __init__(self, size: int | Sequence[int], dtype: torch.dtype): + from vllm.platforms import current_platform + + logger.warning_once( + "Pinned memory is not available on this platform; falling back " + "to device memory for UVA buffers." + ) + self.cpu = torch.zeros(size, dtype=dtype, device="cpu") + self.np = self.cpu.numpy() + self._uva = torch.zeros(size, dtype=dtype, device=current_platform.device_type) + + def uva(self, n: int | None = None) -> torch.Tensor: + if n is None: + return self._uva.copy_(self.cpu) + return self._uva[:n].copy_(self.cpu[:n]) class UvaBufferPool: @@ -69,7 +96,8 @@ class UvaBufferPool: self.max_concurrency = max_concurrency # UVA buffers for concurrency - self._uva_bufs = [UvaBuffer(size, dtype) for _ in range(max_concurrency)] + self._buffer_cls = UvaBuffer if is_uva_available() else NonUvaBuffer + self._uva_bufs = [self._buffer_cls(size, dtype) for _ in range(max_concurrency)] # Current buffer index self._curr = 0 @@ -80,12 +108,12 @@ class UvaBufferPool: n = len(x) if n > buf.cpu.shape[0]: capacity = 1 << (n - 1).bit_length() - buf = UvaBuffer((capacity, *buf.cpu.shape[1:]), self.dtype) + buf = self._buffer_cls((capacity, *buf.cpu.shape[1:]), self.dtype) self._uva_bufs[self._curr] = buf # CPU-to-CPU copy dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np dst[:n] = x - return buf.uva[:n] + return buf.uva(n) def copy_to_gpu( self, @@ -141,6 +169,9 @@ class StagedWriteTensor: self.device = device self.max_concurrency = max_concurrency + # Without UVA, large-but-cold tensors live on the GPU directly. + uva_instead_of_gpu = uva_instead_of_gpu and is_uva_available() + if not uva_instead_of_gpu: # Create a GPU tensor (default) self.gpu = torch.zeros(size, dtype=dtype, device=device) @@ -148,7 +179,7 @@ class StagedWriteTensor: # For a large but not-frequently-accessed tensor, we can use UVA instead of # GPU to save GPU memory self._uva_buf = UvaBuffer(size, dtype) - self.gpu = self._uva_buf.uva + self.gpu = self._uva_buf.uva() self._staged_write_indices: list[int] = [] self._staged_write_starts: list[int] = [] From 7dbc3d62c388df67d2764f28c90c51cc11ff0329 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 14 Sep 2026 21:29:15 -0700 Subject: [PATCH 077/139] [Cleanup] Remove vestigial `tpu_input_batch.py` (#56898) Signed-off-by: Nick Hill --- .../v1/e2e/general/test_mamba_prefix_cache.py | 7 +- vllm/v1/worker/lora_model_runner_mixin.py | 6 +- vllm/v1/worker/mamba_utils.py | 11 +- vllm/v1/worker/tpu_input_batch.py | 559 ------------------ 4 files changed, 9 insertions(+), 574 deletions(-) delete mode 100644 vllm/v1/worker/tpu_input_batch.py diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index aa1f730afb..e388b73468 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -31,9 +31,8 @@ from vllm.v1.request import Request from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.worker import mamba_utils -from vllm.v1.worker.gpu_input_batch import CachedRequestState +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner -from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch from vllm.v1.worker.mamba_utils import get_mamba_groups @@ -315,7 +314,7 @@ def get_fake_process_mamba_fn( action: tuple[int, int], kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], - input_batch: GPUInputBatch, + input_batch: InputBatch, ): assert copy_info is not None if action == (-1, -1): @@ -371,7 +370,7 @@ def get_fake_process_mamba_fn( kv_cache_config: KVCacheConfig, cache_config: CacheConfig, mamba_state_idx: dict[str, int], - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], diff --git a/vllm/v1/worker/lora_model_runner_mixin.py b/vllm/v1/worker/lora_model_runner_mixin.py index 27e850b66a..5e7f1f70c5 100644 --- a/vllm/v1/worker/lora_model_runner_mixin.py +++ b/vllm/v1/worker/lora_model_runner_mixin.py @@ -6,7 +6,6 @@ Define LoRA functionality mixin for model runners. from collections.abc import Callable from contextlib import contextmanager -from typing import TypeAlias import numpy as np import torch @@ -19,10 +18,7 @@ from vllm.lora.layers import LoRAMapping, LoRAMappingType from vllm.lora.request import LoRARequest from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager from vllm.model_executor.models import supports_lora -from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch -from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch - -InputBatch: TypeAlias = TPUInputBatch | GPUInputBatch +from vllm.v1.worker.gpu_input_batch import InputBatch logger = init_logger(__name__) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 724c714057..1deb89bda6 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -26,8 +26,7 @@ from vllm.v1.kv_cache_interface import ( UniformTypeKVCacheSpecs, ) from vllm.v1.utils import CpuGpuBuffer -from vllm.v1.worker.gpu_input_batch import CachedRequestState -from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch logger = init_logger(__name__) @@ -1436,7 +1435,7 @@ def preprocess_mamba( kv_cache_config: KVCacheConfig, cache_config: CacheConfig, mamba_state_idx: dict[str, int], - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, @@ -1541,7 +1540,7 @@ def preprocess_mamba( def postprocess_mamba_all( scheduler_output: SchedulerOutput, kv_cache_config: KVCacheConfig, - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], num_spec_tokens: int, @@ -1572,7 +1571,7 @@ def postprocess_mamba_all( def preprocess_mamba_all_specdec( scheduler_output: SchedulerOutput, - input_batch: GPUInputBatch, + input_batch: InputBatch, mamba_state_idx: dict[str, int], num_reqs: int, prev_last_scheduled_idx_buf: CpuGpuBuffer, @@ -1591,7 +1590,7 @@ def postprocess_mamba_align_gpu( num_reqs: int, num_accepted_tokens_gpu: torch.Tensor, num_accepted_tokens_cpu_tensor: torch.Tensor, - input_batch: GPUInputBatch, + input_batch: InputBatch, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, diff --git a/vllm/v1/worker/tpu_input_batch.py b/vllm/v1/worker/tpu_input_batch.py deleted file mode 100644 index c16156a5e3..0000000000 --- a/vllm/v1/worker/tpu_input_batch.py +++ /dev/null @@ -1,559 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Datastructures defining a TPU input batch - -from typing import cast - -import numpy as np -import torch - -from vllm.lora.request import LoRARequest -from vllm.sampling_params import SamplingType -from vllm.utils import length_from_prompt_token_ids_or_embeds -from vllm.utils.collection_utils import swap_dict_values -from vllm.v1.outputs import LogprobsTensors -from vllm.v1.worker.block_table import MultiGroupBlockTable -from vllm.v1.worker.gpu_input_batch import CachedRequestState - -_SAMPLING_EPS = 1e-5 - - -class InputBatch: - def __init__( - self, - max_num_reqs: int, - max_model_len: int, - max_num_batched_tokens: int, - device: torch.device, - pin_memory: bool, - vocab_size: int, - block_sizes: list[int], # The block_size of each kv cache group - kernel_block_sizes: list[int], - max_num_blocks_per_req: list[int], - ): - self.max_num_reqs = max_num_reqs - self.max_model_len = max_model_len - self.max_num_batched_tokens = max_num_batched_tokens - self.device = device - self.pin_memory = pin_memory - self.vocab_size = vocab_size - - self._req_ids: list[str | None] = [] - self.req_id_to_index: dict[str, int] = {} - - # TODO(woosuk): This buffer could be too large if max_model_len is big. - # Find a way to reduce the CPU memory usage. - # This buffer is not directly transferred to the GPU, so it does not - # need to be pinned. - self.token_ids_cpu_tensor = torch.zeros( - (max_num_reqs, max_model_len), - device="cpu", - dtype=torch.int32, - pin_memory=False, - ) - self.token_ids_cpu = self.token_ids_cpu_tensor.numpy() - self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32) - self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32) - self.num_computed_tokens_cpu_tensor = torch.zeros( - (max_num_reqs,), - device="cpu", - dtype=torch.int32, - pin_memory=pin_memory, - ) - self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy() - - # Block table. - self.block_table = MultiGroupBlockTable( - max_num_reqs=max_num_reqs, - max_num_batched_tokens=max_num_batched_tokens, - pin_memory=pin_memory, - device=device, - block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, - max_num_blocks=max_num_blocks_per_req, - ) - - # Sampling-related. - self.temperature = torch.empty( - (max_num_reqs,), dtype=torch.float32, device=device - ) - self.temperature_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.temperature_cpu = self.temperature_cpu_tensor.numpy() - self.greedy_reqs: set[str] = set() - self.random_reqs: set[str] = set() - - self.top_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) - self.top_p_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.top_p_cpu = self.top_p_cpu_tensor.numpy() - self.top_p_reqs: set[str] = set() - - self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device) - self.top_k_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory - ) - self.top_k_cpu = self.top_k_cpu_tensor.numpy() - self.top_k_reqs: set[str] = set() - - self.min_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) - self.min_p_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.min_p_cpu = self.min_p_cpu_tensor.numpy() - self.min_p_reqs: set[str] = set() - - # Frequency penalty related data structures - self.frequency_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.frequency_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy() - self.frequency_penalties_reqs: set[str] = set() - - # Presence penalty related data structures - self.presence_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.presence_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy() - self.presence_penalties_reqs: set[str] = set() - - # Repetition penalty related data structures - self.repetition_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.repetition_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy() - self.repetition_penalties_reqs: set[str] = set() - - # req_index -> (min_tokens, stop_token_ids) - self.min_tokens: dict[int, tuple[int, set[int]]] = {} - - # lora related - self.request_lora_mapping = np.zeros((self.max_num_reqs,), dtype=np.int64) - self.lora_id_to_request_ids: dict[int, set[str]] = {} - self.lora_id_to_lora_request: dict[int, LoRARequest] = {} - - # req_index -> generator - # NOTE(woosuk): The indices of the requests that do not have their own - # generator should not be included in the dictionary. - self.generators: dict[int, torch.Generator] = {} - - self.num_logprobs: dict[str, int] = {} - - # To accumulate prompt logprobs tensor chunks across prefill steps. - self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {} - - self.logit_bias: list[dict[int, float] | None] = [None] * max_num_reqs - self.has_allowed_token_ids: set[str] = set() - # NOTE(lufang): In the mask tensor, if the corresponding token allowed, - # the value is False. Since we use masked_fill_ to set -inf. - self.allowed_token_ids_mask: torch.Tensor | None = None - self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None - - # req_index -> bad_words_token_ids - self.bad_words_token_ids: dict[int, list[list[int]]] = {} - - self.req_output_token_ids: list[list[int] | None] = [] - - @property - def req_ids(self) -> list[str]: - # None elements should only be present transiently - # while performing state updates to the batch. - return cast(list[str], self._req_ids) - - def add_request( - self, - request: "CachedRequestState", - req_index: int | None = None, - ) -> None: - if req_index is None: - req_index = self.num_reqs - assert req_index < self.max_num_reqs - - req_id = request.req_id - if req_index == len(self._req_ids): - self._req_ids.append(req_id) - self.req_output_token_ids.append(request.output_token_ids) - else: - self._req_ids[req_index] = req_id - self.req_output_token_ids[req_index] = request.output_token_ids - - self.req_id_to_index[req_id] = req_index - - # Copy the prompt token ids and output token ids. - num_prompt_tokens = length_from_prompt_token_ids_or_embeds( - request.prompt_token_ids, request.prompt_embeds - ) - # TODO: copy prompt_embeds - self.num_prompt_tokens[req_index] = num_prompt_tokens - self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids - start_idx = num_prompt_tokens - end_idx = start_idx + len(request.output_token_ids) - self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids - # Number of tokens without spec decode tokens. - self.num_tokens_no_spec[req_index] = request.num_tokens - - self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens - self.block_table.add_row(request.block_ids, req_index) - - sampling_params = request.sampling_params - assert sampling_params is not None, "pooling requests not supported yet" - if sampling_params.sampling_type == SamplingType.GREEDY: - # Should avoid division by zero later when apply_temperature. - self.temperature_cpu[req_index] = 0.0 - self.greedy_reqs.add(req_id) - else: - self.temperature_cpu[req_index] = sampling_params.temperature - self.random_reqs.add(req_id) - - self.top_p_cpu[req_index] = sampling_params.top_p - if sampling_params.top_p < 1: - self.top_p_reqs.add(req_id) - top_k = sampling_params.top_k - if 0 < top_k < self.vocab_size: - self.top_k_reqs.add(req_id) - else: - top_k = self.vocab_size - self.top_k_cpu[req_index] = top_k - self.min_p_cpu[req_index] = sampling_params.min_p - self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty - if sampling_params.min_p > _SAMPLING_EPS: - self.min_p_reqs.add(req_id) - if sampling_params.frequency_penalty != 0.0: - self.frequency_penalties_reqs.add(req_id) - self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty - if sampling_params.presence_penalty != 0.0: - self.presence_penalties_reqs.add(req_id) - self.repetition_penalties_cpu[req_index] = sampling_params.repetition_penalty - if sampling_params.repetition_penalty != 1.0: - self.repetition_penalties_reqs.add(req_id) - if sampling_params.min_tokens: - self.min_tokens[req_index] = ( - sampling_params.min_tokens, - sampling_params.all_stop_token_ids, - ) - - # NOTE(woosuk): self.generators should not include the requests that - # do not have their own generator. - if request.generator is not None: - self.generators[req_index] = request.generator - - if sampling_params.logprobs is not None: - self.num_logprobs[req_id] = sampling_params.logprobs - if sampling_params.logit_bias is not None: - self.logit_bias[req_index] = sampling_params.logit_bias - - if sampling_params.allowed_token_ids: - self.has_allowed_token_ids.add(req_id) - if self.allowed_token_ids_mask_cpu_tensor is None: - # Lazy allocation for this tensor, which can be large. - # False means we don't fill with -inf. - self.allowed_token_ids_mask = torch.zeros( - self.max_num_reqs, - self.vocab_size, - dtype=torch.bool, - device=self.device, - ) - self.allowed_token_ids_mask_cpu_tensor = torch.zeros( - self.max_num_reqs, self.vocab_size, dtype=torch.bool, device="cpu" - ) - self.allowed_token_ids_mask_cpu_tensor[req_index] = True - # False means we don't fill with -inf. - self.allowed_token_ids_mask_cpu_tensor[req_index][ - sampling_params.allowed_token_ids - ] = False - - if sampling_params.bad_words_token_ids: - self.bad_words_token_ids[req_index] = sampling_params.bad_words_token_ids - - # Add request lora ID - if request.lora_request: - lora_id = request.lora_request.lora_int_id - if lora_id not in self.lora_id_to_request_ids: - self.lora_id_to_request_ids[lora_id] = set() - - self.request_lora_mapping[req_index] = lora_id - self.lora_id_to_request_ids[lora_id].add(request.req_id) - self.lora_id_to_lora_request[lora_id] = request.lora_request - else: - # No LoRA - self.request_lora_mapping[req_index] = 0 - - def remove_request(self, req_id: str) -> int | None: - """This method must always be followed by a call to condense().""" - - req_index = self.req_id_to_index.pop(req_id, None) - if req_index is None: - return None - self._req_ids[req_index] = None - self.req_output_token_ids[req_index] = None - - self.greedy_reqs.discard(req_id) - self.random_reqs.discard(req_id) - self.top_p_reqs.discard(req_id) - self.top_k_reqs.discard(req_id) - self.min_p_reqs.discard(req_id) - self.min_tokens.pop(req_index, None) - self.frequency_penalties_reqs.discard(req_id) - self.presence_penalties_reqs.discard(req_id) - self.repetition_penalties_reqs.discard(req_id) - self.generators.pop(req_index, None) - self.num_logprobs.pop(req_id, None) - self.in_progress_prompt_logprobs_cpu.pop(req_id, None) - - # LoRA - lora_id = self.request_lora_mapping[req_index] - if lora_id != 0: - self.lora_id_to_request_ids[lora_id].discard(req_id) - if len(self.lora_id_to_request_ids[lora_id]) == 0: - self.lora_id_to_request_ids.pop(lora_id) - self.lora_id_to_lora_request.pop(lora_id) - self.request_lora_mapping[req_index] = 0 - - self.logit_bias[req_index] = None - self.has_allowed_token_ids.discard(req_id) - if self.allowed_token_ids_mask_cpu_tensor is not None: - # False means we don't fill with -inf. - self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False) - self.bad_words_token_ids.pop(req_index, None) - return req_index - - def swap_states(self, i1: int, i2: int) -> None: - old_id_i1 = self._req_ids[i1] - old_id_i2 = self._req_ids[i2] - self._req_ids[i1], self._req_ids[i2] = self._req_ids[i2], self._req_ids[i1] # noqa - self.req_output_token_ids[i1], self.req_output_token_ids[i2] = ( - self.req_output_token_ids[i2], - self.req_output_token_ids[i1], - ) - assert old_id_i1 is not None and old_id_i2 is not None - self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] = ( - self.req_id_to_index[old_id_i2], - self.req_id_to_index[old_id_i1], - ) - self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = ( - self.num_tokens_no_spec[i2], - self.num_tokens_no_spec[i1], - ) - self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] = ( - self.num_prompt_tokens[i2], - self.num_prompt_tokens[i1], - ) - self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = ( - self.num_computed_tokens_cpu[i2], - self.num_computed_tokens_cpu[i1], - ) - self.temperature_cpu[i1], self.temperature_cpu[i2] = ( - self.temperature_cpu[i2], - self.temperature_cpu[i1], - ) - self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1] - self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1] - self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = ( - self.frequency_penalties_cpu[i2], - self.frequency_penalties_cpu[i1], - ) - self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] = ( - self.presence_penalties_cpu[i2], - self.presence_penalties_cpu[i1], - ) - self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = ( - self.repetition_penalties_cpu[i2], - self.repetition_penalties_cpu[i1], - ) - self.min_p_cpu[i1], self.min_p_cpu[i2] = self.min_p_cpu[i2], self.min_p_cpu[i1] - - # NOTE: the following is unsafe - # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\ - # self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...] - # instead, we need to temporarily copy the data for one of the indices - n = max(self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2]) - tmp = self.token_ids_cpu[i1, :n].copy() - self.token_ids_cpu[i1, :n] = self.token_ids_cpu[i2, :n] - self.token_ids_cpu[i2, :n] = tmp - - swap_dict_values(self.generators, i1, i2) - swap_dict_values(self.min_tokens, i1, i2) - swap_dict_values(self.bad_words_token_ids, i1, i2) - - self.request_lora_mapping[i1], self.request_lora_mapping[i2] = ( - self.request_lora_mapping[i2], - self.request_lora_mapping[i1], - ) - self.logit_bias[i1], self.logit_bias[i2] = ( - self.logit_bias[i2], - self.logit_bias[i1], - ) - - if self.allowed_token_ids_mask_cpu_tensor is not None: - ( - self.allowed_token_ids_mask_cpu_tensor[i1], - self.allowed_token_ids_mask_cpu_tensor[i2], - ) = ( - self.allowed_token_ids_mask_cpu_tensor[i2], - self.allowed_token_ids_mask_cpu_tensor[i1], - ) - self.block_table.swap_row(i1, i2) - - def condense(self, empty_req_indices: list[int]) -> None: - """Move non-empty requests down into lower, empty indices. - - Args: - empty_req_indices: empty batch indices, sorted descending. - """ - num_reqs = self.num_reqs - if num_reqs == 0: - # The batched states are empty. - self._req_ids.clear() - self.req_output_token_ids.clear() - return - - # NOTE(woosuk): This function assumes that the empty_req_indices - # is sorted in descending order. - last_req_index = num_reqs + len(empty_req_indices) - 1 - while empty_req_indices: - # Find the largest non-empty index. - while last_req_index in empty_req_indices: - last_req_index -= 1 - - # Find the smallest empty index. - empty_index = empty_req_indices.pop() - if empty_index >= last_req_index: - break - - # Swap the states. - req_id = self._req_ids[last_req_index] - output_token_ids = self.req_output_token_ids[last_req_index] - assert req_id is not None - self._req_ids[empty_index] = req_id - self._req_ids[last_req_index] = None - self.req_output_token_ids[empty_index] = output_token_ids - self.req_output_token_ids[last_req_index] = None - self.req_id_to_index[req_id] = empty_index - - num_tokens = self.num_tokens_no_spec[last_req_index] - self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[ - last_req_index, :num_tokens - ] - self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[ - last_req_index - ] - self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index] - self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[ - last_req_index - ] - self.block_table.move_row(last_req_index, empty_index) - self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index] - self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index] - self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index] - self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[ - last_req_index - ] - self.presence_penalties_cpu[empty_index] = self.presence_penalties_cpu[ - last_req_index - ] - self.repetition_penalties_cpu[empty_index] = self.repetition_penalties_cpu[ - last_req_index - ] - self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index] - generator = self.generators.pop(last_req_index, None) - if generator is not None: - self.generators[empty_index] = generator - - min_token = self.min_tokens.pop(last_req_index, None) - if min_token is not None: - self.min_tokens[empty_index] = min_token - - self.request_lora_mapping[empty_index] = self.request_lora_mapping[ - last_req_index - ] - - self.logit_bias[empty_index] = self.logit_bias[last_req_index] - - if self.allowed_token_ids_mask_cpu_tensor is not None: - self.allowed_token_ids_mask_cpu_tensor[empty_index] = ( - self.allowed_token_ids_mask_cpu_tensor[last_req_index] - ) - - bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None) - if bad_words_token_ids is not None: - self.bad_words_token_ids[empty_index] = bad_words_token_ids - # Decrement last_req_index since it is now empty. - last_req_index -= 1 - - # Trim lists to the batch size. - del self._req_ids[self.num_reqs :] - del self.req_output_token_ids[self.num_reqs :] - - def make_lora_inputs( - self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray - ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: - """ - Given the num_scheduled_tokens for each request in the batch, return - datastructures used to activate the current LoRAs. - Returns: - 1. prompt_lora_mapping: A tuple of size self.num_reqs where, - prompt_lora_mapping[i] is the LoRA id to use for the ith prompt. - 2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens) - where, token_lora_mapping[i] is the LoRA id to use for ith token. - 3. lora_requests: Set of relevant LoRA requests. - """ - - req_lora_mapping = self.request_lora_mapping[: self.num_reqs] - prompt_lora_mapping = tuple(req_lora_mapping) - token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens)) - active_lora_requests: set[LoRARequest] = set( - self.lora_id_to_lora_request.values() - ) - - return prompt_lora_mapping, token_lora_mapping, active_lora_requests - - @property - def num_reqs(self) -> int: - return len(self.req_id_to_index) - - @property - def all_greedy(self) -> bool: - return len(self.random_reqs) == 0 - - @property - def all_random(self) -> bool: - return len(self.greedy_reqs) == 0 - - @property - def no_top_p(self) -> bool: - return len(self.top_p_reqs) == 0 - - @property - def no_top_k(self) -> bool: - return len(self.top_k_reqs) == 0 - - @property - def no_min_p(self) -> bool: - return len(self.min_p_reqs) == 0 - - @property - def no_penalties(self) -> bool: - return ( - len(self.presence_penalties_reqs) == 0 - and len(self.frequency_penalties_reqs) == 0 - and len(self.repetition_penalties_reqs) == 0 - ) - - @property - def max_num_logprobs(self) -> int | None: - return max(self.num_logprobs.values()) if self.num_logprobs else None - - @property - def no_allowed_token_ids(self) -> bool: - return len(self.has_allowed_token_ids) == 0 From eb61707851f35281aa82c3ba7714f1ce9c1860e0 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 15 Sep 2026 12:34:41 +0800 Subject: [PATCH 078/139] [Bugfix] Fix stale HPC QK-norm weights after weight refit (#56706) Signed-off-by: aoshen02 --- .../model_loader/test_reload.py | 107 ++++++++++++++++++ vllm/model_executor/layers/hpc/hpc_ihc.py | 51 +++------ vllm/model_executor/layers/hpc/hpc_module.py | 3 - vllm/model_executor/layers/hpc/rope_norm.py | 63 +++-------- vllm/model_executor/model_loader/utils.py | 10 -- vllm/model_executor/models/hy_v3.py | 7 +- 6 files changed, 143 insertions(+), 98 deletions(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index f19ec35560..b4a544d2fc 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc +import importlib.machinery import inspect +import sys +import types from unittest.mock import Mock from weakref import WeakKeyDictionary, ref @@ -817,6 +820,110 @@ def test_layerwise_reload_updates_loaded_non_persistent_buffers(monkeypatch): assert "0.scale" not in model.state_dict() +@pytest.fixture +def hpc_rope_norm(monkeypatch, default_vllm_config): + """Import HpcRopeNorm with the external ``hpc`` package stubbed out.""" + if "hpc" not in sys.modules: + stub = types.ModuleType("hpc") + stub.__spec__ = importlib.machinery.ModuleSpec("hpc", loader=None) + stub.QuantType = types.SimpleNamespace( # type: ignore[attr-defined] + QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR=types.SimpleNamespace(value=0) + ) + monkeypatch.setitem(sys.modules, "hpc", stub) + from vllm.model_executor.layers.hpc import rope_norm + + monkeypatch.setattr(rope_norm, "_hpc_rope_norm_instances", {}) + return rope_norm + + +def test_hpc_rope_norm_kernel_sees_refit_norm_weights(monkeypatch, hpc_rope_norm): + """The fused HPC kernel is handed the live QK-norm weights after a refit. + + Drives the production ``_forward_impl`` with a recording ``hpc`` stub. The + Q/K norm weights it receives must be the model's own float32 parameters, + so a layerwise reload that rewrites them in place (same storage) is what + the kernel sees. Previously the kernel read separate mirrors that no + reload path refreshed. + """ + from vllm.model_executor.layers.layernorm import RMSNorm + + head_dim, num_heads, num_kv_heads, block_size = 128, 8, 1, 4 + layer = torch.nn.Module() + layer.q_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32) + layer.k_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32) + layer.hpc_rope_norm = hpc_rope_norm.HpcRopeNorm( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + cos_sin_cache=torch.ones(16, head_dim), + use_qk_norm=True, + fallback_qnorm=layer.q_norm, + fallback_knorm=layer.k_norm, + kv_cache_dtype="auto", + layer_name="hpc_test_layer", + ) + model = torch.nn.Sequential(layer) + rnorm = layer.hpc_rope_norm + + calls: list[dict] = [] + + def record(*args, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(sys.modules["hpc"], "rope_norm_store_kv", record, raising=False) + + def kernel_norm_weights(): + q_size, kv_size = num_heads * head_dim, num_kv_heads * head_dim + qkv = torch.zeros(1, q_size + 2 * kv_size, dtype=torch.bfloat16) + kv_cache = torch.zeros( + 2, num_kv_heads, block_size, 2 * head_dim, dtype=torch.bfloat16 + ) + attn_layer = types.SimpleNamespace( + _k_scale=torch.ones(1), _v_scale=torch.ones(1) + ) + attn_metadata = types.SimpleNamespace( + num_actual_tokens=1, + num_decodes=1, + num_decode_tokens=1, + num_prefills=0, + num_prefill_tokens=0, + max_query_len=1, + decode_query_len=1, + qo_indptr=None, + qo_indptr_decode=None, + slot_mapping=torch.tensor([4]), + seq_lens=torch.tensor([1]), + block_table_tensor=torch.tensor([[1]]), + hpc_kv_written=False, + ) + output = torch.zeros(1, q_size, dtype=torch.bfloat16) + rnorm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output) + return calls[-1]["q_norm_weight"], calls[-1]["k_norm_weight"] + + def loaded(value): + return torch.full((head_dim,), value, dtype=torch.bfloat16) + + default_weight_loader(layer.q_norm.weight, loaded(0.5)) + default_weight_loader(layer.k_norm.weight, loaded(0.25)) + q, k = kernel_norm_weights() + assert torch.equal(q, loaded(0.5).float()) + assert torch.equal(k, loaded(0.25).float()) + q_ptr, k_ptr = q.data_ptr(), k.data_ptr() + + record_metadata_for_reloading(model) + initialize_layerwise_reload(model) + layer.q_norm.weight.weight_loader(layer.q_norm.weight, loaded(2.0)) + layer.k_norm.weight.weight_loader(layer.k_norm.weight, loaded(3.0)) + finalize_layerwise_reload(model, model_config=None) + + q, k = kernel_norm_weights() + assert q.dtype == k.dtype == torch.float32 + assert torch.equal(q, loaded(2.0).float()) + assert torch.equal(k, loaded(3.0).float()) + assert (q.data_ptr(), k.data_ptr()) == (q_ptr, k_ptr) + assert not hasattr(rnorm, "qnorm_weight") + + @pytest.mark.parametrize( "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] ) diff --git a/vllm/model_executor/layers/hpc/hpc_ihc.py b/vllm/model_executor/layers/hpc/hpc_ihc.py index 1a8681aa8e..9a8065a052 100644 --- a/vllm/model_executor/layers/hpc/hpc_ihc.py +++ b/vllm/model_executor/layers/hpc/hpc_ihc.py @@ -20,9 +20,8 @@ Constraints: - x / residual must be bfloat16; hc_fn weights stay float32 (the checkpoint keeps them out of fp8 quantization via modules_to_not_convert) -The fused modules own no parameters: they hold references to the weights of the -eager layer they replace, bound in ``process_weights_after_loading`` (invoked -generically for every ``HpcModule`` by the model loader). +The fused modules own no parameters: their forwards read the weights of the +eager layer they replace directly via the captured owner reference. NOTE: The reference implementation also offers a cross-layer post+pre fusion (``HpcIHCPostPre``) that folds one segment's post into the next segment's pre. @@ -147,41 +146,26 @@ class HpcIHCPre(HpcModule): object.__setattr__(self, "_fallback_op", fallback_op) # Same anti-cycle idiom for the fused RMSNorm owner (a sibling module). object.__setattr__(self, "_norm_owner", norm_owner) - self.weight = None - self.hc_scale = None - self.hc_base = None - self.rms_weight = None - self.rms_eps = 0.0 @classmethod def support(cls, hc_mult: int, hidden_size: int) -> bool: return _ihc_supported(hc_mult, hidden_size) - def process_weights_after_loading(self, model=None): - # hc_fn is a ReplicatedLinear whose weight is [2 * hc_mult, hc_mult * d]. - # It stays float32 -- the HPC kernel consumes fp32 weights directly, so - # there is no conversion (and no precision loss) here. - owner = self._fallback_op - self.weight = owner.hc_fn.weight - self.hc_scale = owner.hc_scale - self.hc_base = owner.hc_base - if self._norm_owner is not None: - self.rms_weight = self._norm_owner.weight - self.rms_eps = self._norm_owner.variance_epsilon - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: import hpc + owner = self._fallback_op + norm_owner = self._norm_owner return hpc.fuse_ihc_pre( x, - self.weight, - self.hc_scale, - self.hc_base, + owner.hc_fn.weight, + owner.hc_scale, + owner.hc_base, self.norm_eps, self.hc_eps, self.magnitude, - self.rms_weight, - self.rms_eps, + norm_owner.weight if norm_owner is not None else None, + norm_owner.variance_epsilon if norm_owner is not None else 0.0, ) @@ -233,29 +217,20 @@ class HpcIHCHead(HpcModule): self.norm_eps = norm_eps # See HpcIHCPre: keep the owner out of the module tree. object.__setattr__(self, "_fallback_op", fallback_op) - self.weight = None - self.hc_scale = None - self.hc_base = None @classmethod def support(cls, hc_mult: int, hidden_size: int) -> bool: return _ihc_supported(hc_mult, hidden_size) - def process_weights_after_loading(self, model=None): - # hc_head_fn weight is [hc_mult, hc_mult * d], float32. - owner = self._fallback_op - self.weight = owner.hc_head_fn.weight - self.hc_scale = owner.hc_head_scale - self.hc_base = owner.hc_head_base - def forward(self, x: torch.Tensor) -> torch.Tensor: import hpc + owner = self._fallback_op return hpc.fuse_ihc_head( x, - self.weight, - self.hc_scale, - self.hc_base, + owner.hc_head_fn.weight, + owner.hc_head_scale, + owner.hc_head_base, self.norm_eps, self.hc_eps, ) diff --git a/vllm/model_executor/layers/hpc/hpc_module.py b/vllm/model_executor/layers/hpc/hpc_module.py index e1efad832e..c3a0dab676 100644 --- a/vllm/model_executor/layers/hpc/hpc_module.py +++ b/vllm/model_executor/layers/hpc/hpc_module.py @@ -11,8 +11,5 @@ class HpcModule(nn.Module): def support(cls, *args, **kwargs): return True - def process_weights_after_loading(self, model): - pass - def forward(self, *args, **kwargs): pass diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py index c1503c754c..02f06ed89c 100644 --- a/vllm/model_executor/layers/hpc/rope_norm.py +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -90,13 +90,13 @@ class HpcRopeNorm(CustomOp, HpcModule): """HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant). Registered as a sub-module in model layers (e.g. HunYuanAttention). - Norm weights are extracted from fallback norm modules via - process_weights_after_loading() after all weights are loaded. + The QK-Norm weights are read directly from the fallback norm modules + (built in float32 when HPC is active), eliminating the derived copy. forward() is dispatched by CustomOp framework: - In compiled mode: forward_cuda() calls torch.ops.vllm.hpc_rope_norm_forward as a splitting point — internal Python control flow is opaque - to torch.compile and not captured by CUDA Graph. + to torch.compile (its kernel launches can still be CUDA-graph captured). - In eager/native mode: forward_native() falls back to forward_cuda(). """ @@ -124,8 +124,6 @@ class HpcRopeNorm(CustomOp, HpcModule): self.num_kv_heads = num_kv_heads self.head_dim = head_dim - self.use_qk_norm = use_qk_norm - self.q_size = num_heads * head_dim self.kv_size = num_kv_heads * head_dim @@ -139,31 +137,18 @@ class HpcRopeNorm(CustomOp, HpcModule): self.head_per_group = num_heads // num_kv_heads - # Pre-allocate norm weight tensors as Parameters so they are tracked by - # CuMemAllocator (for sleep/wake_up) and have stable addresses for CUDA - # Graph replay. process_weights_after_loading() updates them inplace via - # copy_() so refit does not invalidate captured graph tensor pointers. - # Shape is [head_dim] to match the HPC kernel's q/k_norm_weight layout. - if use_qk_norm and fallback_qnorm is not None: - self.qnorm_weight: torch.nn.Parameter | None = torch.nn.Parameter( - torch.empty(head_dim, dtype=torch.float32), - requires_grad=False, - ) - else: - self.qnorm_weight = None - if use_qk_norm and fallback_knorm is not None: - self.knorm_weight: torch.nn.Parameter | None = torch.nn.Parameter( - torch.empty(head_dim, dtype=torch.float32), - requires_grad=False, - ) - else: - self.knorm_weight = None - self.use_fp8 = "fp8" in kv_cache_dtype # The RMSNorm/RoPE ordering is model dependent (e.g. HunYuan V3 applies # QK-Norm before RoPE -> NORM_THEN_ROPE), so it is supplied by the # caller. When QK-Norm is disabled the policy is forced to NONE. self.qk_norm_policy = qk_norm_policy if use_qk_norm else QkNormPolicy.NONE + if self.qk_norm_policy != QkNormPolicy.NONE and ( + fallback_qnorm is None or fallback_knorm is None + ): + raise ValueError( + "QK-Norm is enabled but fallback_qnorm/fallback_knorm were not " + "supplied; the fused kernel reads their weights directly." + ) # Register layer_name + add self to the global instance registry so the # module-level custom op (hpc_rope_norm_forward) can route back here. @@ -215,20 +200,6 @@ class HpcRopeNorm(CustomOp, HpcModule): logger.info_once("enable hpc rope_norm") return True - def process_weights_after_loading(self, model: torch.nn.Module = None) -> None: - """Copy norm weights (float32) from fallback norm modules inplace. - - Uses copy_() to preserve tensor addresses for CUDA Graph / refit - compatibility. Called by the model's load_weights() after all weights - are loaded (and generically from the model loader for DummyModelLoader - / sleep-wake_up reload paths). - """ - if self.use_qk_norm: - if self.fallback_qnorm is not None and self.qnorm_weight is not None: - self.qnorm_weight.data.copy_(self.fallback_qnorm.weight.data.float()) - if self.fallback_knorm is not None and self.knorm_weight is not None: - self.knorm_weight.data.copy_(self.fallback_knorm.weight.data.float()) - def register_layer_name(self, layer_name: str) -> None: """Register layer_name and add self to the global registry. @@ -430,12 +401,14 @@ class HpcRopeNorm(CustomOp, HpcModule): k_scale = attn_layer._k_scale.reshape(1) v_scale = attn_layer._v_scale.reshape(1) - q_norm_weight = ( - self.qnorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None - ) - k_norm_weight = ( - self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None - ) + if self.qk_norm_policy != QkNormPolicy.NONE: + assert self.fallback_qnorm is not None + assert self.fallback_knorm is not None + q_norm_weight = self.fallback_qnorm.weight + k_norm_weight = self.fallback_knorm.weight + else: + q_norm_weight = None + k_norm_weight = None # --- Prefill --- if num_prefill_reqs > 0: diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index e12910865e..11e6be5cf8 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -15,7 +15,6 @@ import vllm.envs as envs from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.attention import is_deferred_attention_layer -from vllm.model_executor.layers.hpc import HpcModule from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -154,15 +153,6 @@ def process_weights_after_loading( with device_loading_context(module, target_device): module.process_weights_after_loading(model_config.dtype) - # Process HPC modules (HpcRopeNorm, etc.) that rely on - # process_weights_after_loading being called from the model's - # load_weights(). When using DummyModelLoader (e.g. profiling or - # sleep/wake_up reload), the model's load_weights() is not called, so we - # must handle HPC modules here generically. - for _, module in model.named_modules(): - if isinstance(module, HpcModule): - module.process_weights_after_loading(model) - # Model-level post-load hook, after the per-layer quant finalize. if hasattr(model, "process_weights_after_loading"): model.process_weights_after_loading() diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 0b54c26f47..248ae04747 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -300,8 +300,11 @@ class HYV3Attention(nn.Module): prefix=f"{prefix}.attn", ) if self.use_qk_norm: - self.q_norm = RMSNorm(self.head_dim, rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, rms_norm_eps) + # The HPC fused kernel reads norm weights directly in float32; + # the fallback path keeps the model default dtype. + norm_dtype = torch.float32 if rope_support else None + self.q_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype) + self.k_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype) # HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant). # HunYuan V3 applies QK-Norm *before* RoPE, so NORM_THEN_ROPE. From 7b7e117c291c023ede8fb504b00913566271e509 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 21:54:57 -0700 Subject: [PATCH 079/139] [CI] Extend DSv4 engine readiness timeout on main (#56735) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .buildkite/test_areas/disaggregated.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 5ab0e836eb..226d7ac94a 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -287,6 +287,7 @@ steps: PREFILL_BLOCK_SIZE: "256" DECODE_BLOCK_SIZE: "256" MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" + VLLM_ENGINE_READY_TIMEOUT_S: "900" VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8" # The V4 indexer only supports block-outermost layouts; let it pick. PREFILLER_KV_LAYOUT: "" From cd10ed6f9f6b37a8ace9cf380007e66fe12ec0c3 Mon Sep 17 00:00:00 2001 From: cjackal <44624812+cjackal@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:58:52 +0900 Subject: [PATCH 080/139] [feature] [xgrammar] support `patternProperties`/`propertyNames`/`unevaluatedProperties` kw for object types (#42904) Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> Co-authored-by: Artem Perevedentsev Co-authored-by: Claude Sonnet 5 --- tests/v1/structured_output/test_utils.py | 138 +++++++++++++++++- tests/v1/structured_output/test_validation.py | 9 +- vllm/v1/structured_output/backend_xgrammar.py | 41 +++++- 3 files changed, 172 insertions(+), 16 deletions(-) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index 11b3e5c01b..f5f8020220 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from xgrammar import Grammar +from xgrammar.testing import _is_grammar_accept_string from vllm.v1.structured_output.backend_xgrammar import ( has_xgrammar_unsupported_json_features, @@ -10,13 +12,19 @@ from vllm.v1.structured_output.backend_xgrammar import ( pytestmark = pytest.mark.cpu_test +def grammar_accepts(schema: dict, text: str) -> bool: + return _is_grammar_accept_string(Grammar.from_json_schema(schema), text) + + @pytest.fixture def unsupported_string_schemas(): return [ {"type": "string", "format": "non_existing_format"}, + # TODO(arpera): # pattern/format is compiled but length bounds are silently dropped, # so the combination must be rejected instead of producing quietly # wrong output + # https://github.com/mlc-ai/xgrammar/issues/749 {"type": "string", "pattern": "^a+$", "maxLength": 2}, {"type": "string", "pattern": "^a+$", "minLength": 3}, {"type": "string", "format": "email", "maxLength": 10}, @@ -48,15 +56,93 @@ def unsupported_array_schemas(): @pytest.fixture -def unsupported_object_schemas(): +def unsupported_property_names_combinations(): return [ - {"type": "object", "propertyNames": {"pattern": "^[a-z]+$"}}, - {"type": "object", "patternProperties": {"^S": {"type": "string"}}}, + # TODO(arpera): + # this case is not covered by xgrammar, so we should report this bug to xgrammar + # xgrammar drops propertyNames whenever patternProperties is present, + # so "grade_12" matches the pattern and escapes the name constraint + { + "type": "object", + "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, + "propertyNames": {"pattern": "^grade_[0-9]$"}, # does NOT match "grade_12" + }, + # propertyNames makes xgrammar discard the sibling additionalProperties + # value schema: https://github.com/mlc-ai/xgrammar/issues/826 + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "additionalProperties": {"type": "integer"}, + }, + # propertyNames is a string schema that conventionally omits "type", so + # it escapes the string check while xgrammar drops its length bound all + # the same: https://github.com/mlc-ai/xgrammar/issues/749 + { + "type": "object", + "propertyNames": {"pattern": "^a+$", "maxLength": 2}, + }, + # xgrammar emits named `properties` separately from `propertyNames` and + # only applies `propertyNames` in its additional-properties branch, so a + # key declared in `properties` escapes the name constraint regardless of + # what `propertyNames` contains. + { + "type": "object", + "properties": {"Bad": {"type": "integer"}}, + "required": ["Bad"], + "propertyNames": {"pattern": "^[a-z]+$"}, + }, + { + "type": "object", + "properties": {"Bad": {"type": "integer"}}, + "propertyNames": {"enum": ["good"]}, + }, + # xgrammar's propertyNames-only branch falls back to an unconstrained + # value type, so it drops unevaluatedProperties whether it restricts + # values (a schema) or forbids them (false). + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "unevaluatedProperties": {"type": "integer"}, + }, + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "unevaluatedProperties": False, + }, ] @pytest.fixture -def supported_schema(): +def unsupported_pattern_properties_combinations(): + return [ + # JSON Schema requires a property matched by both `properties` and + # `patternProperties` to satisfy both (conjunction), but xgrammar + # compiles them as alternative branches, so satisfying either one is + # enough. + { + "type": "object", + "properties": {"x": {"type": "string"}}, + "patternProperties": {"^x$": {"type": "integer"}}, + }, + # The same alternative-branches problem applies to overlapping + # patternProperties patterns: a key matching both patterns only has to + # satisfy one of them. + { + "type": "object", + "patternProperties": { + "^a[a-z]*$": {"type": "string"}, + "^[a-z]*z$": {"type": "integer"}, + }, + }, + ] + + +@pytest.fixture +def supported_frankenstein_schema(): + # IMPORTANT(arpera): + # Do NOT add more keywords here! This schema is overcrowded enough that a new + # keyword can end up having no effect on the compiled grammar while the + # test still passes. Give a new keyword its own fixture and named test. return { "type": "object", "properties": { @@ -88,6 +174,19 @@ def supported_schema(): } +@pytest.fixture +def property_names_schema(): + return {"type": "object", "propertyNames": {"pattern": "^[a-z_]+$"}} + + +@pytest.fixture +def pattern_properties_schema(): + return { + "type": "object", + "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, + } + + @pytest.mark.parametrize( "schema_type", [ @@ -95,7 +194,8 @@ def supported_schema(): "unsupported_integer_schemas", "unsupported_number_schemas", "unsupported_array_schemas", - "unsupported_object_schemas", + "unsupported_property_names_combinations", + "unsupported_pattern_properties_combinations", ], ) def test_unsupported_json_features_by_type(schema_type, request): @@ -106,7 +206,29 @@ def test_unsupported_json_features_by_type(schema_type, request): ) -def test_supported_json_features(supported_schema): - assert not has_xgrammar_unsupported_json_features(supported_schema), ( - "Schema should be supported" +@pytest.mark.parametrize( + "schema_type", + [ + "supported_frankenstein_schema", + "pattern_properties_schema", + "property_names_schema", + ], +) +def test_supported_json_features(schema_type, request): + schema = request.getfixturevalue(schema_type) + assert not has_xgrammar_unsupported_json_features(schema), ( + f"Schema should be supported: {schema}" ) + + +def test_property_names_constrains_keys(property_names_schema): + assert grammar_accepts(property_names_schema, '{"score": 5}') + assert grammar_accepts(property_names_schema, '{"score": "seven"}') + assert not grammar_accepts(property_names_schema, '{"Score": 5}') + assert not grammar_accepts(property_names_schema, '{"score_1": 5}') + + +def test_pattern_properties_constrains_keys_and_values(pattern_properties_schema): + assert grammar_accepts(pattern_properties_schema, '{"grade_1": 5}') + assert not grammar_accepts(pattern_properties_schema, '{"other": 5}') + assert not grammar_accepts(pattern_properties_schema, '{"grade_1": "five"}') diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index af417bc03e..a231894384 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -137,7 +137,7 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): @pytest.mark.parametrize( "schema, expected_backend", [ - # multipleOf is unsupported by xgrammar, patternProperties also by guidance. + # multipleOf is unsupported by xgrammar. ( { "type": "object", @@ -145,8 +145,13 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): }, "guidance", ), + # patternProperties + properties is also unsupported by guidance. ( - {"type": "object", "patternProperties": {"^a": {"type": "string"}}}, + { + "type": "object", + "properties": {"a": {"type": "string"}}, + "patternProperties": {"^a$": {"type": "string"}}, + }, "outlines", ), ], diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 3809ba2d91..1ba7e2734d 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -236,6 +236,12 @@ STRING_SUPPORTED_FORMATS = { } +def _has_pattern_and_length_bounds(schema: dict[str, Any]) -> bool: + return ("pattern" in schema or "format" in schema) and ( + "minLength" in schema or "maxLength" in schema + ) + + def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool: """Check if JSON schema contains features unsupported by xgrammar.""" @@ -269,16 +275,39 @@ def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool: # the compiled EBNF: pattern/format grammars come out byte-identical # with and without the length keywords, while maxLength alone lowers # to {0, N} correctly. + if obj.get("type") == "string" and _has_pattern_and_length_bounds(obj): + return True + + # propertyNames validates names, so it is a string schema even when it + # omits "type", which is the form that escapes the check above. if ( - obj.get("type") == "string" - and ("pattern" in obj or "format" in obj) - and ("minLength" in obj or "maxLength" in obj) + obj.get("type") == "object" + and isinstance(obj.get("propertyNames"), dict) + and _has_pattern_and_length_bounds(obj["propertyNames"]) ): return True - # Unsupported keywords for objects - if obj.get("type") == "object" and any( - key in obj for key in ("patternProperties", "propertyNames") + # FIXME: propertyNames conflicts with properties/patternProperties/ + # additionalProperties/unevaluatedProperties under xgrammar. + # https://github.com/mlc-ai/xgrammar/issues/826 + if ( + obj.get("type") == "object" + and "propertyNames" in obj + and ( + "properties" in obj + or "patternProperties" in obj + or isinstance(obj.get("additionalProperties"), dict) + or obj.get("unevaluatedProperties", True) is not True + ) + ): + return True + + # FIXME: multiple patternProperties, or patternProperties alongside + # properties, conflict under xgrammar. + if ( + obj.get("type") == "object" + and isinstance(obj.get("patternProperties"), dict) + and ("properties" in obj or len(obj["patternProperties"]) > 1) ): return True From ec6b14c0448cef6173ec1c331dc2f251ff41591a Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 14 Sep 2026 23:05:59 -0700 Subject: [PATCH 081/139] [CI] Select the supported DCP backend for PCP eval (#56677) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml index f57467ac74..b9d2675431 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml @@ -11,6 +11,7 @@ server_args: >- --moe-backend flashinfer_cutlass --prefill-context-parallel-size 4 --decode-context-parallel-size 4 + --dcp-comm-backend ag_rs --enable-expert-parallel --kv-cache-dtype fp8 env: From f547c23ec971fd1c78f0d89d15ac1957ae407471 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Mon, 14 Sep 2026 23:18:24 -0700 Subject: [PATCH 082/139] [Bugfix][Kimi-K3] Keep transient checkpoints out of prefix-cache eviction (#56794) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: zjy0516 Co-authored-by: zjy0516 Co-authored-by: Kimi Code --- .../test_mamba_eagle_resume_checkpoint.py | 9 ++- .../test_partial_prefix_cache_hits.py | 71 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 16 ++++- 3 files changed, 94 insertions(+), 2 deletions(-) diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index fd003fdce6..059b4005f0 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -299,7 +299,14 @@ def test_a_junction_past_the_prompt_falls_back_and_registers_nothing(): running = request.num_computed_tokens assert not mamba.req_to_blocks["r"][running // block_size].is_null request.shared_prefix_boundary = running - assert mamba._cache_partial_tail_block(request, running) is None + assert ( + mamba._cache_partial_tail_block( + request, + running, + retention_interval=manager.coordinator.retention_interval, + ) + is None + ) @pytest.mark.parametrize("eagle_group", [None, 0]) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 8012c82f78..393ec48465 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -615,6 +615,77 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle(): assert num_computed == 112 +@pytest.mark.parametrize( + "retention_interval,transient_published", + [(None, True), (0, False)], +) +def test_internal_checkpoint_publication_respects_retention( + retention_interval, transient_published +): + """With retention_interval=0, a mid-prompt internal checkpoint is + request-local and must not enter the prefix cache; the prompt-end + checkpoint keeps the existing publication behavior. Dense retention + (None) is unchanged. + """ + hash_block_size = 16 + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=hash_block_size, + full_block_size=hash_block_size, + mamba_block_size=32, + num_prefill_checkpoint_blocks=1, + ) + manager.coordinator.retention_interval = retention_interval + + request = make_request("producer", list(range(240)), hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(request) + + # Mid-prompt chunk: exports a transient checkpoint at 112. + assert manager.allocate_slots(request, 128, num_computed, computed_blocks) + transient_hash = request.block_hashes[112 // hash_block_size - 1] + transient_hit = manager.block_pool.get_cached_block(transient_hash, [1]) + assert (transient_hit is not None) == transient_published + + request.num_computed_tokens = 128 + manager.new_step_starts() + + # Final chunk: the prompt-end checkpoint at 224 stays published. + assert manager.allocate_slots(request, 112) is not None + end_hash = request.block_hashes[224 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(end_hash, [1]) is not None + + +def test_transient_checkpoint_evicts_retained_boundary_hash(): + """With retention_interval=0, the checkpoint slot can coincide with a + reachable boundary block that this step's full-block pass just hashed. + The checkpoint state overwrites the slot, so the stale boundary hash + must be evicted rather than left pointing at the wrong state. + """ + hash_block_size = 16 + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=hash_block_size, + full_block_size=hash_block_size, + mamba_block_size=32, + num_prefill_checkpoint_blocks=1, + ) + manager.coordinator.retention_interval = 0 + + request = make_request("producer", list(range(240)), hash_block_size, sha256) + # Block 2 ends at the shared-prefix junction 96, so the full-block pass + # hashes it as a retained boundary. + request.shared_prefix_boundary = 96 + computed_blocks, num_computed, _ = manager.get_computed_blocks(request) + assert manager.allocate_slots(request, 128, num_computed, computed_blocks) + + # The same slot is the transient checkpoint block (state@112): both the + # boundary hash and the checkpoint itself must stay unpublished. + boundary_hash = request.block_hashes[96 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(boundary_hash, [1]) is None + checkpoint_hash = request.block_hashes[112 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None + + def test_eagle_block_aligned_checkpoint_replaces_newer_hash(): hash_block_size = mamba_block_size = 32 manager = make_full_mamba_manager( diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index d4e68d8adc..fc5dc41048 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -2003,7 +2003,9 @@ class MambaManager(SingleTypeKVCacheManager): ) num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) if self.mamba_cache_mode == "align": - partial_hash = self._cache_partial_tail_block(request, num_tokens) + partial_hash = self._cache_partial_tail_block( + request, num_tokens, retention_interval=retention_interval + ) if partial_hash is not None: self.cached_blocks_this_step.add(partial_hash) if num_cached_blocks_after > num_cached_blocks_before: @@ -2038,6 +2040,7 @@ class MambaManager(SingleTypeKVCacheManager): self, request: Request, num_tokens: int, + retention_interval: int | None, ) -> BlockHashWithGroupId | None: hash_block_size = self.block_pool.hash_block_size # Re-key the reserved block at its exported checkpoint boundary. @@ -2047,6 +2050,17 @@ class MambaManager(SingleTypeKVCacheManager): blocks = self.req_to_blocks[request.request_id] assert 0 <= checkpoint_idx < len(blocks) checkpoint_block = blocks[checkpoint_idx] + if ( + retention_interval == 0 + and num_tokens < request.num_prompt_tokens + and checkpoint_position != request.shared_prefix_boundary + ): + # retention_interval == 0 keeps this transient checkpoint + # request-local. The slot may carry a hash from this step's + # full-block pass; that must go too, since the checkpoint + # state is about to overwrite the block. + self.block_pool._maybe_evict_cached_block(checkpoint_block) + return None if checkpoint_block.block_hash_num_tokens == checkpoint_position: return None return self.block_pool.cache_partial_block( From 79e205e8cdf7864b8b9f2d23c908b887f2f5badb Mon Sep 17 00:00:00 2001 From: Jyotirmoy Roy Date: Tue, 15 Sep 2026 12:05:13 +0530 Subject: [PATCH 083/139] [Doc] Clarify ITL vs TPOT Prometheus metrics (#55283) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Jyotirmoy Roy Signed-off-by: Cenab Batu Bora <40250311+cenab@users.noreply.github.com> Signed-off-by: Itay Alroy Signed-off-by: Rita Brugarolas Brufau Signed-off-by: Josiah Davis Signed-off-by: simondanielsson Signed-off-by: Nick Hill Signed-off-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Signed-off-by: Julien Denize <40604584+juliendenize@users.noreply.github.com> Signed-off-by: Wei Zhao Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Signed-off-by: wzhao18 Signed-off-by: Matthew Wong Signed-off-by: Micah Williamson Signed-off-by: JiataiWang Signed-off-by: 子华 Signed-off-by: Colin McNamara Signed-off-by: Tyler Michael Smith Signed-off-by: Marceli Fylcek Signed-off-by: LearningMachine621 <147347705+LearningMachine621@users.noreply.github.com> Signed-off-by: jamesETsmith Signed-off-by: Cursor Signed-off-by: Yichao Zhu Signed-off-by: wangyicong Signed-off-by: Kushal <72650064+kushaldabbe@users.noreply.github.com> Signed-off-by: Tianyu Guo Signed-off-by: Summer Yang Signed-off-by: Lucas Wilkinson Signed-off-by: hungh Signed-off-by: wang.yuqi Signed-off-by: zhicheng Liu Signed-off-by: Andrii Skliar Signed-off-by: Alexander Salikov Signed-off-by: Luca Motz Signed-off-by: pmanczak Signed-off-by: Fadi Arafeh Signed-off-by: Julien Darve Signed-off-by: Bugen Zhao Signed-off-by: NickLucche Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Signed-off-by: khluu Signed-off-by: Kevin H. Luu Signed-off-by: Tony Lin Signed-off-by: Yongye Zhu Signed-off-by: Jee Jee Li Signed-off-by: andyluo7 Signed-off-by: zjy0516 Signed-off-by: JinYan Su <751080330@qq.com> Signed-off-by: nicole-lihui Signed-off-by: waizuichougou <2082431897@qq.com> Signed-off-by: waizuichougou <105572299+waizuichougou@users.noreply.github.com> Signed-off-by: zRzRzRzRzRzRzR Signed-off-by: 김재억 Signed-off-by: xiaoming <1259730330@qq.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Guy Stone Signed-off-by: Djordje Ramic Signed-off-by: Isotr0py Signed-off-by: yewentao256 Signed-off-by: Lucain Pouget Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> Signed-off-by: Sören Dréano Signed-off-by: Chris Eastwood Signed-off-by: Andy Lo Signed-off-by: mgoin Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com> Signed-off-by: Lyle Lai Signed-off-by: Vincent Signed-off-by: Giancarlo Delfin Signed-off-by: xawang Signed-off-by: Alec Flowers Signed-off-by: Fangzhou Ai Signed-off-by: Yizheng Jiao Signed-off-by: Tahsin Tunan Signed-off-by: Divakar Verma Signed-off-by: zhenwei-intel Signed-off-by: Biswa Panda Signed-off-by: LopezCastroRoberto Signed-off-by: Matthew Bonanni Signed-off-by: Lucas Wilkinson Signed-off-by: Matthew Bonanni Signed-off-by: Nicolò Lucchesi Signed-off-by: Canlin Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com> Signed-off-by: walterbm Signed-off-by: arcusbuilds Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: Kunshang Ji Signed-off-by: Taneem Ibrahim Signed-off-by: JartX Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Signed-off-by: Artem Perevedentsev Signed-off-by: Joe Cotant Signed-off-by: yangzeyu <532183776@qq.com> Signed-off-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Signed-off-by: Raphael Rialland Signed-off-by: Simon Veitner Signed-off-by: Tomas Ruiz Signed-off-by: patryk.saffer Signed-off-by: patryk.saffer Signed-off-by: PatrykSaffer Signed-off-by: Chao Lei Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Signed-off-by: Ayushman Singh Signed-off-by: Markus Hartikainen Signed-off-by: Woosuk Kwon Signed-off-by: surajm20061998 Signed-off-by: Roger Wang Signed-off-by: Ylean <61880605+Woolgathererer@users.noreply.github.com> Signed-off-by: Juntian Liu Signed-off-by: Kebe Signed-off-by: linitra24 Signed-off-by: Shenglei Fu Signed-off-by: pavelzak Signed-off-by: fai Signed-off-by: Stu Cao Signed-off-by: <> Signed-off-by: Zhewen Li Signed-off-by: CedricHuang <38417461+CedricHwong@users.noreply.github.com> Signed-off-by: siyu Signed-off-by: Wang Zupeng Signed-off-by: Chaojun Zhang Signed-off-by: hotragn Signed-off-by: Isotr0py <2037008807@qq.com> Signed-off-by: specture724 Signed-off-by: lxy-alexander Signed-off-by: BWAAEEEK Signed-off-by: Jared Wen Signed-off-by: Chinmay Kulkarni Signed-off-by: jiang1.li Signed-off-by: Jared Wen Signed-off-by: Matthias Gehre Signed-off-by: liuqihao Signed-off-by: Martin Hickey Signed-off-by: jiangkuaixue123 Signed-off-by: Juan Pérez de Algaba Signed-off-by: reidliu41 Signed-off-by: Xiaozhen Bi Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: R Signed-off-by: Alex Signed-off-by: Ronald1995 Signed-off-by: arkohut <39525455+arkohut@users.noreply.github.com> Signed-off-by: Russell Bryant Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Signed-off-by: Orestis Zambounis Signed-off-by: Claude Signed-off-by: Orestis Zambounis Signed-off-by: Speculators Bot Signed-off-by: Speculators Bot Signed-off-by: zqlcode Signed-off-by: zql <37731799+zqlcode@users.noreply.github.com> Signed-off-by: Thillai Chithambaram Signed-off-by: ElizaWszola Signed-off-by: snordmann Signed-off-by: Ianniu <985683179yzy@gmail.com> Signed-off-by: Luossu Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Signed-off-by: Shiyang Chen Signed-off-by: Fangzhou Ai Signed-off-by: Michele Campi Signed-off-by: xiaguan <751080330@qq.com> Signed-off-by: zhaoguochun1995 Signed-off-by: JohnQinAMD Signed-off-by: aoshen02 Signed-off-by: Liran Schour Signed-off-by: Liqian Lin Signed-off-by: wuhangxian <1391938827@qq.com> Signed-off-by: Yongye Zhu Signed-off-by: Miguel Garcia Signed-off-by: Mario <404mario@users.noreply.github.com> Signed-off-by: aoshen02 Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> Co-authored-by: Cenab Batu Bora <40250311+cenab@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Itay Alroy <75032521+itayalroy@users.noreply.github.com> Co-authored-by: Tyler Michael Smith Co-authored-by: Rita Brugarolas Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Josiah Davis Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: Nick Hill Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com> Co-authored-by: Julien Denize <40604584+juliendenize@users.noreply.github.com> Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: Wei Zhao Co-authored-by: Matt <156021403+mawong-amd@users.noreply.github.com> Co-authored-by: Micah Williamson Co-authored-by: Jiatai Wang Co-authored-by: OpenAI Codex Co-authored-by: shaohuaxi Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Colin McNamara Co-authored-by: Marceli Fylcek Co-authored-by: 焦恺文 <15529091557@163.com> Co-authored-by: James E T Smith Co-authored-by: Cursor Co-authored-by: YukioZzz Co-authored-by: Yicong Wang Co-authored-by: Kushal <72650064+kushaldabbe@users.noreply.github.com> Co-authored-by: Tianyu Guo Co-authored-by: Lucas Wilkinson Co-authored-by: Summer Yang Co-authored-by: Hung Hoang (黄兴） Co-authored-by: wang.yuqi Co-authored-by: Zhicheng Liu <2084597944@qq.com> Co-authored-by: Andrii Skliar Co-authored-by: Andrii Skliar Co-authored-by: Alexander Co-authored-by: lucamotz Co-authored-by: Cyrus Leung Co-authored-by: pmanczak Co-authored-by: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Co-authored-by: Li, Jiang Co-authored-by: JulienDarve <86800349+JulienDarve@users.noreply.github.com> Co-authored-by: Bugen Zhao Co-authored-by: Nicolò Lucchesi Co-authored-by: Andreas Karatzas Co-authored-by: Lucas Wilkinson Co-authored-by: Kevin H. Luu Co-authored-by: Tony Lin Co-authored-by: Yongye Zhu Co-authored-by: Jee Jee Li Co-authored-by: andyluo7 Co-authored-by: zjy0516 Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: JinYan Su <751080330@qq.com> Co-authored-by: Nicole LiHui 🥜 Co-authored-by: Chauncey Co-authored-by: waizuichougou <105572299+waizuichougou@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py Co-authored-by: Yuxuan Zhang Co-authored-by: 김재억 Co-authored-by: labAxiaoming <34019940+labAxiaoming@users.noreply.github.com> Co-authored-by: Guy Stone Co-authored-by: djramic Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Lucain Co-authored-by: Sören Dréano <71752785+SorenDreano@users.noreply.github.com> Co-authored-by: Chris Eastwood <106503529+eastwood-c@users.noreply.github.com> Co-authored-by: Andy Lo Co-authored-by: Misha Goin Co-authored-by: chengchengpei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: Lyle Lai Co-authored-by: mevince <26759093+mevince@users.noreply.github.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: amd-xavierwang Co-authored-by: Alec <35311602+alec-flowers@users.noreply.github.com> Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Co-authored-by: Fangzhou Ai Co-authored-by: Yizheng Jiao Co-authored-by: Tahsin Tunan Co-authored-by: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Co-authored-by: liuzhenwei Co-authored-by: Biswa Panda Co-authored-by: Woosuk Kwon Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Matthew Bonanni Co-authored-by: Lucas Wilkinson Co-authored-by: Canlin Guo Co-authored-by: junuxyz <216036880+junuxyz@users.noreply.github.com> Co-authored-by: Walter Beller-Morales Co-authored-by: Flora Feng <4florafeng@gmail.com> Co-authored-by: Srijan Keshri Co-authored-by: Shengqi Chen Co-authored-by: Kunshang Ji Co-authored-by: Taneem Ibrahim Co-authored-by: JartX Co-authored-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com> Co-authored-by: Artem Perevedentsev Co-authored-by: jcotant-inferact Co-authored-by: Zeyu Yang <532183776@qq.com> Co-authored-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Co-authored-by: Raphaël Rialland <36076211+TQCB@users.noreply.github.com> Co-authored-by: Simon Veitner Co-authored-by: Tomas Ruiz Co-authored-by: PatrykSaffer Co-authored-by: patryk.saffer Co-authored-by: patryk.saffer Co-authored-by: Patryk Saffer Co-authored-by: Chao Lei Co-authored-by: Ayushman Singh Co-authored-by: Markus Hartikainen Co-authored-by: Jee Jee Li Co-authored-by: Suraj Co-authored-by: Kimi Code Co-authored-by: Roger Wang Co-authored-by: Ylean <61880605+Woolgathererer@users.noreply.github.com> Co-authored-by: Juntian Liu Co-authored-by: Kebe Co-authored-by: linitra24 Co-authored-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Co-authored-by: Pavel Zakharov Co-authored-by: stu-cao Co-authored-by: Stu Cao Co-authored-by: Zhewen Li Co-authored-by: Zhewen Li Co-authored-by: CedricHuang <38417461+CedricHwong@users.noreply.github.com> Co-authored-by: siyu Co-authored-by: Zupeng Wang Co-authored-by: Chaojun Zhang Co-authored-by: Hotragn Pettugani <103170876+Hotragn@users.noreply.github.com> Co-authored-by: Anjie Hou <149605198+specture724@users.noreply.github.com> Co-authored-by: Alexander Lee Co-authored-by: JooHo Lee <96564470+BWAAEEEK@users.noreply.github.com> Co-authored-by: Jared Wen Co-authored-by: Chinmay-Kulkarni-AMD Co-authored-by: Matthias Gehre Co-authored-by: BruceLoveDecimal <39156883+BruceLoveDecimal@users.noreply.github.com> Co-authored-by: liuqihao Co-authored-by: Martin Hickey Co-authored-by: jiangkuaixue123 Co-authored-by: Juan Pérez de Algaba <124347725+jperezdealgaba@users.noreply.github.com> Co-authored-by: Reid <61492567+reidliu41@users.noreply.github.com> Co-authored-by: Xiaozhen Bi Co-authored-by: Ganesh R Co-authored-by: AlexHuang Co-authored-by: Or Ozeri Co-authored-by: Kimi Code Co-authored-by: Ronald Co-authored-by: arkohut <39525455+arkohut@users.noreply.github.com> Co-authored-by: Russell Bryant Co-authored-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Co-authored-by: Benjamin Chislett Co-authored-by: Orestis Zambounis Co-authored-by: Speculators Bot Co-authored-by: zql <37731799+zqlcode@users.noreply.github.com> Co-authored-by: zqlcode Co-authored-by: Thillai Chithambaram <79466435+thillai-c@users.noreply.github.com> Co-authored-by: Mark McLoughlin Co-authored-by: Samuel Nordmann Co-authored-by: ElizaWszola Co-authored-by: Zhiyao(Ian) Yu <105995086+Ianniu123@users.noreply.github.com> Co-authored-by: 鐘天楽 Co-authored-by: Jorge Gallego Feliciano Co-authored-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: sychen52 <41452870+sychen52@users.noreply.github.com> Co-authored-by: Michele Campi <215741962+MicheleCampi@users.noreply.github.com> Co-authored-by: Sherlock Co-authored-by: zhaoguochun1995 <109069909+zhaoguochun1995@users.noreply.github.com> Co-authored-by: JohnQinAMD Co-authored-by: aoshen02 Co-authored-by: Bhushan Asati Co-authored-by: Megha Agarwal <19240983+meghaagr13@users.noreply.github.com> Co-authored-by: liranschour Co-authored-by: Liqian Lin Co-authored-by: wuhangxian <1391938827@qq.com> Co-authored-by: Yongye Zhu Co-authored-by: Miguel <60073809+migarci2@users.noreply.github.com> Co-authored-by: 404mario <18063257856@163.com> Co-authored-by: Mario <404mario@users.noreply.github.com> Co-authored-by: Sherlock Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com> --- docs/design/metrics.md | 23 ++- .../prometheus_grafana/grafana.json | 168 +++++++++++++++++- 2 files changed, 186 insertions(+), 5 deletions(-) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index c7bc85756e..3d108890d7 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -33,7 +33,8 @@ In v1, an extensive set of metrics are exposed via a Prometheus-compatible `/met - `vllm:request_prompt_tokens` (Histogram) - Histogram of input prompt token counts. - `vllm:request_generation_tokens` (Histogram) - Histogram of generation token counts. - `vllm:time_to_first_token_seconds` (Histogram) - Time to first token (TTFT). -- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency. +- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency (time between consecutive streamed outputs). +- `vllm:request_time_per_output_token_seconds` (Histogram) - Per-request Time Per Output Token (TPOT). - `vllm:e2e_request_latency_seconds` (Histogram) - End-to-end request latency. - `vllm:request_prefill_time_seconds` (Histogram) - Request prefill time. - `vllm:request_decode_time_seconds` (Histogram) - Request decode time. @@ -49,7 +50,8 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of - `vllm:e2e_request_latency_seconds_bucket` - End to end request latency measured in seconds. - `vllm:prompt_tokens` - Prompt tokens. - `vllm:generation_tokens` - Generation tokens. -- `vllm:inter_token_latency_seconds` - Inter-token latency (Time Per Output Token, TPOT) in seconds. +- `vllm:inter_token_latency_seconds` - Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains exactly one token, but differs from request-level TPOT when an output contains multiple tokens or when aggregation weights differ. +- `vllm:request_time_per_output_token_seconds` - Per-request Time Per Output Token (TPOT) in seconds, computed as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. It is recorded as zero for requests that generate no more than one token. - `vllm:time_to_first_token_seconds` - Time to First Token (TTFT) latency in seconds. - `vllm:num_requests_running` (also, `_swapped` and `_waiting`) - Number of requests in the RUNNING, WAITING, and SWAPPED states. - `vllm:kv_cache_usage_perc` - Percentage of used cache blocks by vLLM. @@ -63,6 +65,20 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of See [the PR which added this Dashboard](https://github.com/vllm-project/vllm/pull/2316) for interesting and useful background on the choices made here. +`vllm:inter_token_latency_seconds` records one sample per streamed output event +(the wall-clock gap between successive outputs), whereas +`vllm:request_time_per_output_token_seconds` is recorded once per finished request +as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. These two +metrics differ when an output bundles multiple tokens (e.g. speculative +decoding) or when mean ITL and mean per-request TPOT are aggregated over +differing weights. See [Understanding the latency metrics](../benchmarking/cli.md#understanding-the-latency-metrics) +for the benchmark definitions. Requests that generate no more than one token +are recorded with a TPOT of zero, while `vllm bench serve` excludes them from +its TPOT statistics. Therefore, the Prometheus histogram and benchmark TPOT +statistics can differ. Use `vllm:request_time_per_output_token_seconds` for +request-level TPOT and `vllm:inter_token_latency_seconds` when you specifically +want the inter-output latency distribution. + ### Prometheus Client Library Prometheus support was initially added [using the aioprometheus library](https://github.com/vllm-project/vllm/pull/1890), but a switch was made quickly to [prometheus_client](https://github.com/vllm-project/vllm/pull/2730). The rationale is discussed in both linked PRs. @@ -238,8 +254,7 @@ statistics relating to that iteration: iteration. - The prefill intervals for any requests that completed prefill in this iteration. -- The inter-token intervals (Time Per Output Token, TPOT), for all - requests included in this iteration. +- The inter-token intervals, for all requests included in this iteration. - The Time-To-First-Token (TTFT) for any requests that completed prefill in this iteration. However, we calculate this interval relative to when the request was first received by the frontend diff --git a/examples/observability/prometheus_grafana/grafana.json b/examples/observability/prometheus_grafana/grafana.json index 1c89d45938..6144f2bfd3 100644 --- a/examples/observability/prometheus_grafana/grafana.json +++ b/examples/observability/prometheus_grafana/grafana.json @@ -318,7 +318,7 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "description": "Inter token latency in seconds.", + "description": "Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains one token, but differs from request-level TPOT under speculative decoding or varying aggregation weights.", "fieldConfig": { "defaults": { "color": { @@ -479,6 +479,172 @@ "title": "Inter Token Latency", "type": "timeseries" }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "Per-request Time Per Output Token (TPOT) in seconds, computed as (end-to-end latency - TTFT) divided by the number of output tokens minus one. Requests with no more than one output token are recorded as zero; vllm bench serve excludes these requests from its TPOT statistics.", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + }, + "unit": "s" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 48 + }, + "id": 17, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P99", + "range": true, + "refId": "A", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.95, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P95", + "range": true, + "refId": "B", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.9, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P90", + "range": true, + "refId": "C", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.5, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P50", + "range": true, + "refId": "D", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "rate(vllm:request_time_per_output_token_seconds_sum{model_name=\"$model_name\"}[$__rate_interval])\n/\nrate(vllm:request_time_per_output_token_seconds_count{model_name=\"$model_name\"}[$__rate_interval])", + "hide": false, + "instant": false, + "legendFormat": "Mean", + "range": true, + "refId": "E" + } + ], + "title": "Time Per Output Token (TPOT)", + "type": "timeseries" + }, { "datasource": { "type": "prometheus", From 51918e252af4dfff381001f42b7ebe6a0454501b Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Tue, 15 Sep 2026 14:51:35 +0800 Subject: [PATCH 084/139] [Bugfix][SparseMLA] Fix piecewise cudagraph capture crash in index group (#56825) Signed-off-by: Jared Wen Co-authored-by: Kimi Code --- .../v1/attention/test_sparse_mla_backends.py | 78 ++++++++++++++++++- vllm/forward_context.py | 9 +++ vllm/v1/attention/backends/mla/index_group.py | 15 +++- vllm/v1/hisparse/runtime.py | 4 +- 4 files changed, 100 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 2ca839d59c..6d2c4fbd58 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -22,10 +22,17 @@ from tests.v1.attention.utils import ( create_vllm_config, ) from vllm import _custom_ops as ops -from vllm.config import HiSparseConfig, SpeculativeConfig, set_current_vllm_config +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture +from vllm.config import ( + CUDAGraphMode, + HiSparseConfig, + SpeculativeConfig, + set_current_vllm_config, +) from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( HiSparseConnectorWorker, ) +from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention import mla_attention from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( @@ -897,6 +904,75 @@ def test_triton_convert_req_index_to_global_index_decode_only( torch.testing.assert_close(result, reference_result, rtol=0, atol=0) +def test_index_group_convert_during_piecewise_capture(): + """In piecewise cudagraph mode the indexer is captured while the convert + runs in the following eager break. The convert's side stream must not + wait on ``logical_topk_ready`` there: an event recorded inside a captured + segment is graph-local, and an eager wait on it raises + cudaErrorInvalidValue.""" + device = torch.device(DEVICE_TYPE) + num_tokens, num_topk, num_requests, blocks_per_req, block_size = 8, 128, 4, 4, 16 + + logical_topk_indices = torch.randint( + 0, + block_size * blocks_per_req, + (num_tokens, num_topk), + dtype=torch.int32, + device=device, + ) + builder = SparseMLAIndexGroupBuilder(logical_topk_indices) + group, layer_index = builder.register_layer(is_index_producing_layer=True) + assert layer_index == 0 and group.has_indexer + + req_id = torch.arange(num_tokens, dtype=torch.int32, device=device) % num_requests + block_table = torch.arange( + num_requests * blocks_per_req, dtype=torch.int32, device=device + ).view(num_requests, blocks_per_req) + attn_metadata = SimpleNamespace( + req_id_per_token=req_id, block_table=block_table, block_size=block_size + ) + expected = triton_convert_req_index_to_global_index( + req_id, + block_table, + logical_topk_indices, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=num_topk, + ) + + vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B") + marker = torch.zeros(1, device=device) + capture_stream = torch.cuda.Stream(device=device) + # torch.cuda.stream routes through vLLM's patched torch.cuda.set_stream, + # so current_stream() tracks the capture stream as in production. + with ( + torch.cuda.stream(capture_stream), + set_forward_context( + None, vllm_config, cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE + ), + ): + capture = BreakableCUDAGraphCapture() + with capture: + marker.add_(1) + # Record inside a captured segment, convert in the eager + # break — mirrors mla.py -> unified_mla_attention_with_output. + group.set_logical_topk_ready(layer_index) + capture.add_eager( + lambda: group.convert_logical_to_physical_topk( + layer_index, + logical_topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + ) + marker.add_(1) + capture.replay() + capture_stream.synchronize() + + result = group.physical_topk_indices[:num_tokens] + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + @pytest.mark.parametrize("block_size", [16]) @pytest.mark.skipif( torch.cuda.get_device_capability() < (9, 0), diff --git a/vllm/forward_context.py b/vllm/forward_context.py index f57fc2c950..8d9eb5d7db 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -209,6 +209,15 @@ def is_forward_context_available() -> bool: return _forward_context is not None +def in_piecewise_cudagraph() -> bool: + """Whether the current forward runs in piecewise cudagraph mode (graph + segments separated by eager breaks), at capture or replay time.""" + return ( + is_forward_context_available() + and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE + ) + + def create_forward_context( attn_metadata: Any, vllm_config: VllmConfig, diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 8ba0171835..1b8455a232 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -10,6 +10,7 @@ import torch from vllm import _custom_ops as ops from vllm.config import VllmConfig +from vllm.forward_context import in_piecewise_cudagraph from vllm.utils.torch_utils import current_stream from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, @@ -56,8 +57,16 @@ class SparseMLAIndexGroup: return layer_index def set_logical_topk_ready(self, layer_index: int) -> None: - if layer_index == 0 and self.has_indexer: - self.logical_topk_ready.record(current_stream()) + if layer_index != 0 or not self.has_indexer: + return + # In piecewise cudagraph mode the consumer converts on the side stream + # from an eager break; an event recorded inside a captured segment is + # graph-local and any eager wait on it raises cudaErrorInvalidValue. + # The convert side falls back to wait_stream there instead (and never + # touches this event), so skip the record as well. + if in_piecewise_cudagraph(): + return + self.logical_topk_ready.record(current_stream()) def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None: pass @@ -129,7 +138,7 @@ class SparseMLAIndexGroup: valid_topk_counts = self.valid_topk_counts[:num_tokens] compute_stream = current_stream() if layer_index == 0: - if self.has_indexer: + if self.has_indexer and not in_piecewise_cudagraph(): self.side_stream.wait_event(self.logical_topk_ready) else: self.side_stream.wait_stream(compute_stream) diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 314d3ba86b..5100379083 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -16,7 +16,7 @@ import torch from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed import get_tp_group -from vllm.forward_context import get_forward_context +from vllm.forward_context import get_forward_context, in_piecewise_cudagraph from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.math_utils import round_up @@ -975,7 +975,7 @@ class HiSparseRuntime: ) -> None: group = self.index_group compute_stream = current_stream() - if group.logical_topk_ready is not None: + if group.logical_topk_ready is not None and not in_piecewise_cudagraph(): group.copy_stream.wait_event(group.logical_topk_ready) else: group.copy_stream.wait_stream(compute_stream) From 000c7df9ffd3e470980fd4cd6b8ec1b0585500ff Mon Sep 17 00:00:00 2001 From: Moein Khazraee <33970824+mkhazraee@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:55:04 -0700 Subject: [PATCH 085/139] [KV Offload] Add KVCR secondary-tier adapter (#53624) Signed-off-by: Moein Khazraee Signed-off-by: Moein Co-authored-by: Olga Andreeva <124622579+oandreeva-nv@users.noreply.github.com> Co-authored-by: Olga Andreeva --- .buildkite/scripts/install-kv-offload.sh | 18 + .buildkite/test_areas/misc.yaml | 1 + tests/v1/kv_offload/tiering/test_factory.py | 6 +- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 544 +++++++++++++++ vllm/v1/kv_offload/tiering/factory.py | 6 + vllm/v1/kv_offload/tiering/kvcr/__init__.py | 2 + vllm/v1/kv_offload/tiering/kvcr/manager.py | 635 ++++++++++++++++++ 7 files changed, 1211 insertions(+), 1 deletion(-) create mode 100644 .buildkite/scripts/install-kv-offload.sh create mode 100644 tests/v1/kv_offload/tiering/test_kvcr_tier.py create mode 100644 vllm/v1/kv_offload/tiering/kvcr/__init__.py create mode 100644 vllm/v1/kv_offload/tiering/kvcr/manager.py diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh new file mode 100644 index 0000000000..46d756b621 --- /dev/null +++ b/.buildkite/scripts/install-kv-offload.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euo pipefail + +if [ -z "${CUDA_VERSION:-}" ]; then + exit 0 +fi + +# TODO: Use the PyPI wheel once it is published. +uv pip install --system \ + "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0" + +# Keep only the NIXL wheel matching the CI image's CUDA runtime. +NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p') +uv pip uninstall --system nixl-cu12 nixl-cu13 +uv pip install --system --no-deps "nixl-cu${CUDA_VERSION%%.*}==${NIXL_VERSION}" diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index b45b26b2f2..ed82c7d71d 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -130,6 +130,7 @@ steps: source_file_dependencies: *v1-core-kv-metrics-deps commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh - pytest -v -s v1/kv_offload - pytest -v -s v1/simple_kv_offload mirror: diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py index f5a576d254..c391ac2016 100644 --- a/tests/v1/kv_offload/tiering/test_factory.py +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -10,6 +10,7 @@ These tests verify: 3. Error paths — missing tier_type, unknown tier_type, duplicate registration. """ +from importlib.util import find_spec from unittest.mock import MagicMock import pytest @@ -42,8 +43,11 @@ def _make_mock_args(): def test_pre_registered_tiers_can_be_imported(): - """CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager.""" + """CI sentinel: registered tiers import and yield SecondaryTierManager.""" for tier_type in SecondaryTierFactory._registry: + # KVCR is an optional external dependency and may not be installed. + if tier_type == "kvcr" and find_spec("kvcr") is None: + continue cls = SecondaryTierFactory._registry[tier_type]() assert issubclass(cls, SecondaryTierManager) diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py new file mode 100644 index 0000000000..38a2ae044c --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -0,0 +1,544 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Collection, Iterable, Mapping +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +pytest.importorskip("kvcr") + +from kvcr import ROUTER_HINT_KEY, KVCRBindings +from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig +from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy +from kvcr.types import ( + BlockKey, + CacheTier, + InventoryEvent, + MemDescriptor, + OpEntryResult, + OpEntryStatus, + OpHandle, + QueryStatus, +) + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadKey, + ReqContext, + make_offload_key, +) +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob +from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager +from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager + + +def _op_entries( + entries: Mapping[BlockKey, bool], +) -> dict[BlockKey, OpEntryResult]: + return { + key: OpEntryResult(OpEntryStatus.SUCCESS if success else OpEntryStatus.FAILED) + for key, success in entries.items() + } + + +class _StubControlChannel: + def __init__(self, endpoint: str) -> None: + self.endpoint = endpoint + + def send(self, endpoint: str, message: bytes) -> bool: + return True + + def recv(self) -> list[bytes]: + return [] + + def close(self) -> None: + pass + + +class RecordingKVCR: + def __init__(self) -> None: + self.config: KVCRConfig | None = None + self.guard_config: KVCRGuardConfig | None = None + self.backend_configs = KVCRBackendConfigs() + self.constructor_bindings: KVCRBindings | None = None + self.nixl_agent_name = "recording" + self.framework_control: _StubControlChannel | None = None + self.inventory_sink = None + self.query_status = QueryStatus.MISS + self.stats: OffloadingConnectorStats | None = None + self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + self.discard_hint_calls: list[str] = [] + self.deliver_calls: list[ + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None] + ] = [] + self.deposit_calls: list[ + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]] + ] = [] + self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = [] + self._next_op_handle = 1 + + def submit_hint(self, *args: Any, **kwargs: Any) -> None: + self.submit_hint_calls.append((args, kwargs)) + + def discard_hint(self, request_id: str) -> None: + self.discard_hint_calls.append(request_id) + + def query( + self, + keys: Collection[BlockKey], + request_id: str | None = None, + ) -> list[tuple[QueryStatus, CacheTier | None]]: + tier = { + QueryStatus.HIT: CacheTier.LOCAL_G2, + QueryStatus.FETCHING: CacheTier.LOCAL_G2, + QueryStatus.FETCHABLE: CacheTier.REMOTE_G2, + QueryStatus.MISS: None, + }[self.query_status] + return [(self.query_status, tier) for _ in keys] + + def deliver( + self, + blocks: Mapping[BlockKey, list[MemDescriptor]], + request_id: str | None = None, + ) -> OpHandle: + op_handle = self._next_op_handle + self._next_op_handle += 1 + self.deliver_calls.append((op_handle, dict(blocks), request_id)) + self.complete(op_handle, {key: True for key in blocks}) + return op_handle + + def deposit( + self, + blocks: Mapping[BlockKey, list[MemDescriptor]], + ) -> OpHandle: + op_handle = self._next_op_handle + self._next_op_handle += 1 + self.deposit_calls.append((op_handle, dict(blocks))) + self.complete(op_handle, {key: True for key in blocks}) + return op_handle + + def complete( + self, + op_handle: OpHandle, + entries: Mapping[BlockKey, bool], + ) -> None: + self.completed.append((op_handle, _op_entries(entries))) + + def poll_completed( + self, + ) -> Iterable[tuple[OpHandle, dict[BlockKey, OpEntryResult]]]: + completed = self.completed + self.completed = [] + return completed + + def get_stats(self) -> OffloadingConnectorStats | None: + stats = self.stats + self.stats = None + return stats + + def close(self) -> None: + pass + + +class _ExternalPolicy(FIFOPolicy): + pass + + +def _job( + job_id: int, + req_context: ReqContext, + key: OffloadKey | None = None, + chunk_id: int = 0, +) -> TransferJob: + return TransferJob( + job_id=job_id, + keys=[key if key is not None else OffloadKey(b"k0")], + chunk_ids=np.array([chunk_id], dtype=np.int64), + is_promotion=True, + req_context=req_context, + ) + + +def _make_tier( + monkeypatch, + kvcr: RecordingKVCR, + *, + enable_telemetry: bool = False, + secondary_g2_slots: int = 0, + kvcr_service_socket_path: str | None = None, + compatibility_digest: str | None = None, + enable_kv_cache_events: bool = False, + self_describing_kv_events: bool = False, + policy: str | None = None, + g3: dict[str, object] | None = None, + control_ports: list[int] | None = None, + data_parallel_rank_local: int | None = None, +) -> KVCRSecondaryTierManager: + def make_control(_bind_host, bind_port, advertise_host): + return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}") + + def make_kvcr(config, bindings, backend_configs, guard_config): + kvcr.config = config + kvcr.guard_config = guard_config + kvcr.nixl_agent_name = config.nixl_agent_name + kvcr.backend_configs = backend_configs + kvcr.constructor_bindings = bindings + kvcr.framework_control = bindings.framework_control + kvcr.inventory_sink = bindings.inventory_sink + return kvcr + + monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr) + monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control) + return KVCRSecondaryTierManager( + offloading_spec=SimpleNamespace( + config=SimpleNamespace( + parallel=SimpleNamespace( + data_parallel_rank_local=data_parallel_rank_local, + ) + ), + kv_events_config=SimpleNamespace( + enable_kv_cache_events=enable_kv_cache_events, + self_describing_kv_events=self_describing_kv_events, + ), + ), + primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)), + tier_type="kvcr", + router_capabilities=["router_hint"], + control_host="127.0.0.1", + control_ports=control_ports if control_ports is not None else [7777], + control_advertise_host="127.0.0.1", + enable_telemetry=enable_telemetry, + secondary_g2_slots=secondary_g2_slots, + kvcr_service_socket_path=kvcr_service_socket_path, + compatibility_digest=compatibility_digest, + policy=policy, + g3=g3, + local_dram_backend="UCX", + remote_fw_dram_backend="UCX", + ) + + +def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): + """Keep the control endpoint and guard pool aligned to the local DP rank.""" + kvcr = RecordingKVCR() + tier = _make_tier( + monkeypatch, + kvcr, + kvcr_service_socket_path="/tmp/kvcr.sock", + compatibility_digest="Opaque-Digest", + secondary_g2_slots=1, + control_ports=[7001, 7002], + data_parallel_rank_local=1, + ) + + assert kvcr.framework_control is not None + assert kvcr.framework_control.endpoint == "tcp://127.0.0.1:7002" + assert kvcr.guard_config == KVCRGuardConfig( + kvcr_service_socket_path="/tmp/kvcr.sock", + guard_index=1, + compatibility_digest="Opaque-Digest", + ) + assert kvcr.config is not None + assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)] + assert kvcr.backend_configs.local_dram is None + + +def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path): + """Convert user-provided G3 paths to KVCR's typed configuration.""" + kvcr = RecordingKVCR() + path = tmp_path / "g3.data" + + _make_tier( + monkeypatch, + kvcr, + g3={"paths": [str(path)], "capacity_bytes_per_file": 64}, + ) + + assert kvcr.backend_configs.g3 == G3Options( + paths=(path,), + capacity_bytes_per_file=64, + ) + + +def test_kvcr_tier_adapts_request_and_load(monkeypatch): + """Check hint forwarding, key conversion, load descriptors, and cleanup.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + router_hint = {"opaque": True} + ctx = ReqContext( + req_id="req", + kv_transfer_params={ROUTER_HINT_KEY: router_hint, "unrelated": object()}, + ) + key = make_offload_key((123).to_bytes(8, "big"), 0) + same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7) + other_key = make_offload_key((124).to_bytes(8, "big"), 0) + + tier.on_new_request(ctx) + + assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": router_hint})] + + bindings = kvcr.constructor_bindings + assert bindings is not None + assert bindings.key_adapter is not None + monkeypatch.setenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "0") + decode = bindings.key_adapter.decode + hash_123 = (123).to_bytes(8, "big") + assert decode(BlockKey(bytes(key))) == hash_123 + assert decode(BlockKey(bytes(same_hash_other_group))) == hash_123 + assert decode(BlockKey(bytes(other_key))) == (124).to_bytes(8, "big") + + tier.submit_load(_job(7, ctx, key=key, chunk_id=2)) + + # Check load destinations and job completion independently of hint parsing. + assert len(kvcr.submit_hint_calls) == 1 + _, blocks, request_id = kvcr.deliver_calls[0] + assert request_id == "req" + assert list(blocks) == [key] + (descriptor,) = blocks[key] + assert descriptor.end_point_name == kvcr.nixl_agent_name + assert descriptor.addr == tier._primary_base_addr + 2 * 16 + assert descriptor.size == 16 + assert list(tier.get_finished_jobs()) == [JobResult(7, True)] + + # Here we verify that request cleanup discards the request-scoped hint in + # KVCR. + tier.on_request_finished(ctx) + assert kvcr.discard_hint_calls == ["req"] + + +def test_kvcr_tier_allows_request_without_router_hint(monkeypatch): + """Keep router hints optional for requests from non-hint-aware routers.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + + tier.on_new_request(ReqContext(req_id="req")) + + assert kvcr.submit_hint_calls == [] + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + (QueryStatus.MISS, LookupResult.MISS), + (QueryStatus.HIT, LookupResult.HIT), + (QueryStatus.FETCHABLE, LookupResult.HIT), + (QueryStatus.FETCHING, LookupResult.RETRY), + ], +) +def test_kvcr_tier_maps_query_status(monkeypatch, status, expected): + """Map KVCR cache states to the scheduler's lookup contract.""" + kvcr = RecordingKVCR() + kvcr.query_status = status + tier = _make_tier(monkeypatch, kvcr) + + assert tier.lookup(OffloadKey(b"k0"), ReqContext(req_id="req")) is expected + + +def test_kvcr_tier_serves_primary_pin_request(monkeypatch): + """Hold primary-tier hits until KVCR releases the corresponding pin.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2")) + hit_keys = (keys[0], keys[2]) + chunk_ids = {keys[0]: 1, keys[2]: 5} + lifecycle: list[str] = [] + + class Parent: + def on_new_request(self, req_context): + lifecycle.append("new") + return SimpleNamespace() + + def lookup(self, key, req_context): + return LookupResult.HIT if key in hit_keys else LookupResult.MISS + + def create_store_job(self, requested_keys, req_context): + return TransferJob( + job_id=11, + keys=requested_keys, + chunk_ids=np.array([chunk_ids[key] for key in requested_keys]), + is_promotion=True, + req_context=req_context, + ) + + def on_request_finished(self, req_context): + lifecycle.append("finished") + + bindings = kvcr.constructor_bindings + assert bindings is not None + request = bindings.request_pin(keys) + + tier.serve_external_requests(Parent()) + + [(queued_request, result)] = bindings.poll_pin_results() + assert queued_request == request + assert result is not None + pin_handle, descriptors = result + assert descriptors[keys[1]] is None + block_descriptors = descriptors[keys[2]] + assert block_descriptors is not None + (descriptor,) = block_descriptors + assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride) + assert lifecycle == ["new", "finished"] + + polls = 0 + + def release_pin(): + nonlocal polls + polls += 1 + bindings.release_pin(pin_handle) + return [] + + monkeypatch.setattr(kvcr, "poll_completed", release_pin) + tier.drain_jobs() + + assert polls == 1 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)] + + +def test_kvcr_telemetry_is_opt_in_and_namespaced_at_vllm_boundary(monkeypatch): + """Keep telemetry opt-in and namespace metrics only at the vLLM boundary.""" + assert KVCRSecondaryTierManager.build_metric_definitions({}) == {} + definitions = KVCRSecondaryTierManager.build_metric_definitions( + {"enable_telemetry": True} + ) + assert "vllm:kvcr_duration_seconds" in definitions + assert "vllm:kvcr_transfer_blocks" in definitions + + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr, enable_telemetry=True) + bindings = kvcr.constructor_bindings + assert bindings is not None + assert bindings.stats_factory is not None + stats = bindings.stats_factory() + stats.increase_counter( + "kvcr_transfer_blocks", + 2, + ("remote_deliver",), + ) + kvcr.stats = stats + + returned = tier.get_stats() + + assert returned is stats + assert returned.reduce() == {"vllm:kvcr_transfer_blocks:('remote_deliver',)": 2} + + +@pytest.mark.parametrize( + ("policy", "expected_type"), + [ + ("fifo", FIFOPolicy), + ("lru", LRUPolicy), + ("g3_fifo", G3FIFOPolicy), + ("g3_lru", G3LRUPolicy), + (f"{__name__}._ExternalPolicy", _ExternalPolicy), + ], +) +def test_kvcr_tier_passes_selected_policy(monkeypatch, policy, expected_type): + """Resolve every built-in and fully qualified external policy.""" + kvcr = RecordingKVCR() + _make_tier(monkeypatch, kvcr, policy=policy) + + bindings = kvcr.constructor_bindings + assert bindings is not None + assert type(bindings.policy) is expected_type + + +@pytest.mark.parametrize( + ("socket_path", "digest"), + [("/tmp/kvcr.sock", None), (None, "Opaque-Digest")], +) +def test_kvcr_tier_requires_complete_service_config(monkeypatch, socket_path, digest): + """Reject partial guard configuration before connecting to the service.""" + with pytest.raises(ValueError, match="configured together"): + _make_tier( + monkeypatch, + RecordingKVCR(), + kvcr_service_socket_path=socket_path, + compatibility_digest=digest, + ) + + +def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): + """Cover store descriptors and inventory translation at the KVCR boundary.""" + kvcr = RecordingKVCR() + tier = _make_tier( + monkeypatch, + kvcr, + secondary_g2_slots=2, + enable_kv_cache_events=True, + self_describing_kv_events=True, + ) + local_dram = kvcr.backend_configs.local_dram + assert local_dram is not None + assert [(name, size) for name, _, size in local_dram.pools] == [ + ("", 2 * tier._primary_row_stride) + ] + + key = OffloadKey(b"k0") + tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2)) + + _, blocks = kvcr.deposit_calls[0] + assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)] + + assert kvcr.inventory_sink is not None + kvcr.inventory_sink( + InventoryEvent((BlockKey(bytes(key)),), CacheTier.LOCAL_G2, False) + ) + kvcr.inventory_sink(InventoryEvent((BlockKey(bytes(key)),), CacheTier.G3, False)) + events = list(tier.take_events()) + assert [ + (event.keys, event.medium, event.ownership, event.removed) for event in events + ] == [ + ([key], Medium.CPU, "kvcr", False), + ([key], Medium.STORAGE, "kvcr", False), + ] + assert all(event.removal_expected for event in events) + tier.shutdown() + + +def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch): + """Require tier-aware events whenever KVCR owns local cache inventory.""" + with pytest.raises(ValueError, match="self_describing_kv_events"): + _make_tier( + monkeypatch, + RecordingKVCR(), + secondary_g2_slots=1, + enable_kv_cache_events=True, + ) + + +def test_kvcr_tier_accumulates_block_results(monkeypatch): + """Wait for every block result and preserve partial success while draining.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + keys = [OffloadKey(b"k0"), OffloadKey(b"k1")] + tier.submit_load( + TransferJob( + job_id=13, + keys=keys, + chunk_ids=np.array([0, 1], dtype=np.int64), + is_promotion=True, + req_context=ReqContext(req_id="req"), + ) + ) + + op_handle, _, _ = kvcr.deliver_calls[0] + kvcr.completed.clear() + kvcr.complete(op_handle, {keys[0]: True}) + assert list(tier.get_finished_jobs()) == [] + + kvcr.complete(op_handle, {keys[1]: False}) + tier.drain_jobs() + + [result] = tier.get_finished_jobs() + assert result.job_id == 13 + assert not result.success + assert result.successful_keys == {keys[0]} diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 17f8430204..3310167b87 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -123,3 +123,9 @@ SecondaryTierFactory.register_tier( "vllm.v1.kv_offload.tiering.obj.manager", "ObjectStoreSecondaryTierManager", ) + +SecondaryTierFactory.register_tier( + "kvcr", + "vllm.v1.kv_offload.tiering.kvcr.manager", + "KVCRSecondaryTierManager", +) diff --git a/vllm/v1/kv_offload/tiering/kvcr/__init__.py b/vllm/v1/kv_offload/tiering/kvcr/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/v1/kv_offload/tiering/kvcr/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py new file mode 100644 index 0000000000..00186279a4 --- /dev/null +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -0,0 +1,635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM secondary-tier adapter for KVCR.""" + +import ctypes +import mmap +import socket +import time +import uuid +from collections.abc import Collection, Iterable, Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +from kvcr import ( + DURATION_METRIC, + KVCR, + ROUTER_HINT_CAPABILITIES, + ROUTER_HINT_KEY, + STATE_METRIC, + TRANSFER_BLOCKS_METRIC, + TRANSFER_BYTES_METRIC, + KVCRBindings, +) +from kvcr.config import ( + FrameworkDramInput, + G3Options, + KeyAdapter, + KVCRBackendConfigs, + KVCRConfig, + KVCRGuardConfig, + LocalDramOptions, + RemoteFWDramOptions, +) +from kvcr.control_channels import ZmqPeerControlChannel +from kvcr.policy import ( + FIFOPolicy, + G3FIFOPolicy, + G3LRUPolicy, + KVCachePolicy, + LRUPolicy, +) +from kvcr.types import ( + BlockKey, + CacheTier, + InventoryEvent, + MemDescriptor, + OpHandle, + PinRequestId, + PinResult, + QueryStatus, +) +from typing_extensions import override + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.v1.core.kv_cache_utils import ( + BlockHash, + maybe_convert_block_hash, +) +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadingCounterMetadata, + OffloadingEvent, + OffloadingGaugeMetadata, + OffloadingHistogramMetadata, + OffloadingMetricMetadata, + OffloadKey, + ReqContext, + RequestOffloadingContext, + get_offload_block_hash, +) +from vllm.v1.kv_offload.tiering.base import ( + JobResult, + ParentManager, + SecondaryTierManager, + TransferJob, +) + +if TYPE_CHECKING: + from vllm.v1.kv_offload.base import OffloadingSpec + + +_REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES + +logger = init_logger(__name__) + +_BUILTIN_POLICIES: dict[str, type[KVCachePolicy]] = { + "fifo": FIFOPolicy, + "lru": LRUPolicy, + "g3_fifo": G3FIFOPolicy, + "g3_lru": G3LRUPolicy, +} + + +def _resolve_policy(name: str | None) -> KVCachePolicy | None: + if name is None: + return None + policy_type = _BUILTIN_POLICIES.get(name) + if policy_type is None: + if "." not in name: + raise ValueError( + f"Unknown KVCR policy: {name!r}. " + f"Supported: {list(_BUILTIN_POLICIES)}; external policies " + "must use a fully qualified module.Class path." + ) + policy_type = resolve_obj_by_qualname(name) + if not isinstance(policy_type, type) or not issubclass( + policy_type, KVCachePolicy + ): + raise TypeError(f"{name} is not a KVCachePolicy") + return policy_type() + + +def _vllm_metric_name(name: str) -> str: + return f"vllm:{name}" + + +def _kvcr_metric_definitions() -> dict[str, OffloadingMetricMetadata]: + return { + _vllm_metric_name(DURATION_METRIC): OffloadingHistogramMetadata( + documentation="KVCR operation and stage duration in seconds.", + labelnames=("scope", "result"), + ), + _vllm_metric_name(TRANSFER_BYTES_METRIC): OffloadingCounterMetadata( + documentation="Bytes transferred by KVCR operations.", + labelnames=("operation",), + ), + _vllm_metric_name(TRANSFER_BLOCKS_METRIC): OffloadingCounterMetadata( + documentation="KV blocks transferred by successful KVCR operations.", + labelnames=("operation",), + ), + _vllm_metric_name(STATE_METRIC): OffloadingGaugeMetadata( + documentation="Current KVCR metadata and operation counts.", + labelnames=("resource",), + ), + } + + +class _FrameworkPinAdapter: + """Adapt KVCR framework pin callbacks to primary-tier store jobs.""" + + def __init__(self, tier: "KVCRSecondaryTierManager") -> None: + self._tier = tier + self._next_request_id = 0 + self._next_pin_request_id = 0 + self._pending_pins: dict[PinRequestId, tuple[BlockKey, ...]] = {} + self._completed_pins: list[tuple[PinRequestId, PinResult]] = [] + self._pin_jobs: dict[str, int] = {} + self._pin_job_results: list[JobResult] = [] + + def request_pin(self, keys: Collection[BlockKey]) -> PinRequestId: + # The parent callback is valid only from serve_external_requests(). + request = PinRequestId(self._next_pin_request_id) + self._next_pin_request_id += 1 + self._pending_pins[request] = tuple(keys) + return request + + def process_pending(self, parent: ParentManager) -> None: + pending = self._pending_pins + self._pending_pins = {} + self._completed_pins.extend( + (request, self._resolve_pin(parent, keys)) + for request, keys in pending.items() + ) + + def poll_pin_results(self) -> list[tuple[PinRequestId, PinResult]]: + completed = self._completed_pins + self._completed_pins = [] + return completed + + def cancel_pin_request(self, request: PinRequestId) -> None: + self._pending_pins.pop(request, None) + + def _resolve_pin( + self, parent: ParentManager, keys: Collection[BlockKey] + ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None: + offload_keys = tuple(OffloadKey(bytes(key)) for key in keys) + request_id = f"kvcr-source:{self._next_request_id}" + self._next_request_id += 1 + req_context = ReqContext(req_id=request_id) + started = False + job: TransferJob | None = None + pin_handle: str | None = None + try: + parent.on_new_request(req_context) + started = True + hit_keys = tuple( + key + for key in offload_keys + if parent.lookup(key, req_context) is LookupResult.HIT + ) + if not hit_keys: + return None + + job = parent.create_store_job(hit_keys, req_context) + job_keys = tuple(job.keys) + chunk_ids = tuple(int(chunk_id) for chunk_id in job.chunk_ids) + if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys): + return None + + descriptors: dict[BlockKey, list[MemDescriptor] | None] = { + BlockKey(bytes(key)): None for key in offload_keys + } + descriptors.update( + ( + BlockKey(bytes(key)), + [self._tier._make_descriptor(chunk_id)], + ) + for key, chunk_id in zip(job_keys, chunk_ids) + ) + pin_handle = f"kvcr-job-{job.job_id}" + self._pin_jobs[pin_handle] = job.job_id + return pin_handle, descriptors + except Exception: + logger.warning("KVCR source acquisition failed", exc_info=True) + return None + finally: + if started: + try: + parent.on_request_finished(req_context) + except Exception: + logger.warning("KVCR source request cleanup failed", exc_info=True) + if job is not None and pin_handle is None: + self._pin_job_results.append( + JobResult(job_id=job.job_id, success=False) + ) + + def release_pin(self, pin_handle: str) -> bool: + job_id = self._pin_jobs.pop(pin_handle, None) + if job_id is not None: + self._pin_job_results.append(JobResult(job_id=job_id, success=True)) + return True + + def has_active_pins(self) -> bool: + return bool(self._pin_jobs) + + def take_pin_job_results(self) -> list[JobResult]: + results = self._pin_job_results + self._pin_job_results = [] + return results + + +class _VllmKeyAdapter: + """Adapt vLLM's local key format to KVCR.""" + + def encode(self, framework_key: object) -> BlockKey: + if not isinstance(framework_key, bytes): + raise TypeError("vLLM offload keys must be bytes") + return BlockKey(framework_key) + + def decode(self, key: BlockKey) -> int | bytes: + block_hash = BlockHash(get_offload_block_hash(OffloadKey(key))) + return maybe_convert_block_hash(block_hash) + + +# Job ID, remaining blocks, aggregate success, and successful load keys. +_JobState = tuple[int, int, bool, set[OffloadKey] | None] + + +class KVCRSecondaryTierManager(SecondaryTierManager): + """Secondary tier wrapper around the KVCR KV P2P API.""" + + @classmethod + @override + def build_metric_definitions( + cls, extra_config: dict[str, Any] + ) -> dict[str, OffloadingMetricMetadata]: + if not extra_config.get("enable_telemetry", False): + return {} + return _kvcr_metric_definitions() + + def __init__( + self, + offloading_spec: "OffloadingSpec", + primary_kv_view: memoryview, + tier_type: str, + router_capabilities: Iterable[str] | None = None, + control_host: str = "0.0.0.0", + control_ports: list[int] | None = None, + control_advertise_host: str | None = None, + eager_ctrl_connect: bool = True, + opportunistic_query: bool = True, + enable_telemetry: bool = False, + operation_timeout_ms: int = 1000, + abandon_timeout_ms: int = 5000, + metadata_retry_interval_ms: int = 100, + secondary_g2_slots: int = 0, + kvcr_service_socket_path: str | None = None, + compatibility_digest: str | None = None, + policy: str | None = None, + g3: dict[str, Any] | None = None, + local_dram_backend: str = "UCX", + remote_fw_dram_backend: str = "UCX", + ) -> None: + super().__init__(offloading_spec, primary_kv_view, tier_type) + selected_policy = _resolve_policy(policy) + if (kvcr_service_socket_path is None) != (compatibility_digest is None): + raise ValueError( + "kvcr_service_socket_path and compatibility_digest must be " + "configured together" + ) + missing = _REQUIRED_ROUTER_CAPABILITIES - set(router_capabilities or ()) + if missing: + raise ValueError( + f"KVCR secondary tier requires router capabilities: {sorted(missing)}" + ) + events_enabled = offloading_spec.kv_events_config.enable_kv_cache_events + self_describing_events = ( + offloading_spec.kv_events_config.self_describing_kv_events + ) + if ( + events_enabled + and ( + secondary_g2_slots > 0 + or kvcr_service_socket_path is not None + or g3 is not None + ) + and not self_describing_events + ): + raise ValueError( + "KVCR local G2/G3 inventory requires self_describing_kv_events" + ) + self._key_adapter: KeyAdapter = _VllmKeyAdapter() + advertise_host = control_advertise_host or socket.gethostname() + dp_local_rank = offloading_spec.config.parallel.data_parallel_rank_local + if dp_local_rank is None: + if control_ports is None or len(control_ports) != 1: + raise ValueError( + "control_ports must contain exactly one port when " + "data_parallel_rank_local is unset" + ) + dp_local_rank = 0 + if control_ports is None or not 0 <= dp_local_rank < len(control_ports): + raise ValueError( + "control_ports must contain a port for " + f"local data-parallel rank {dp_local_rank}" + ) + if any(not 1 <= port <= 65535 for port in control_ports): + raise ValueError("control_ports must be between 1 and 65535") + control = ZmqPeerControlChannel( + control_host, + control_ports[dp_local_rank], + advertise_host, + ) + # Give colocated workers distinct transport listen ports. + with socket.socket() as _s: + _s.bind(("", 0)) + _nixl_listen_port = _s.getsockname()[1] + self._primary_base_addr = ctypes.addressof( + ctypes.c_char.from_buffer(primary_kv_view) + ) + if primary_kv_view.strides is None: + raise ValueError("primary KV memoryview must expose strides") + self._primary_row_stride = int(primary_kv_view.strides[0]) + if secondary_g2_slots < 0: + raise ValueError("secondary_g2_slots must be non-negative") + if ( + secondary_g2_slots or kvcr_service_socket_path is not None + ) and not events_enabled: + logger.warning( + "KVCR local DRAM is enabled, but KV cache events are disabled; " + "local DRAM inventory will not be published" + ) + local_mapping: mmap.mmap | None = None + local_dram: LocalDramOptions | None = None + if secondary_g2_slots: + if kvcr_service_socket_path is not None: + logger.warning( + "secondary_g2_slots is ignored when " + "kvcr_service_socket_path is configured" + ) + else: + local_mapping = mmap.mmap( + -1, secondary_g2_slots * self._primary_row_stride + ) + local_dram = LocalDramOptions( + pools=[ + ( + "", + ctypes.addressof(ctypes.c_char.from_buffer(local_mapping)), + len(local_mapping), + ) + ], + backend=local_dram_backend, + ) + guard_config = ( + KVCRGuardConfig( + kvcr_service_socket_path=kvcr_service_socket_path, + guard_index=dp_local_rank, + compatibility_digest=compatibility_digest, + ) + if kvcr_service_socket_path is not None and compatibility_digest is not None + else None + ) + g3_config = None + if g3 is not None: + g3_values = dict(g3) + g3_values["paths"] = tuple(Path(path) for path in g3_values["paths"]) + g3_config = G3Options(**g3_values) + nixl_agent_name = f"KVCR-{uuid.uuid4()}" + self._framework_pin_adapter = _FrameworkPinAdapter(self) + self._inventory_events: list[OffloadingEvent] = [] + + try: + self._kvcr = KVCR( + KVCRConfig( + pool_layouts=[("", self._primary_row_stride)], + enable_telemetry=enable_telemetry, + operation_timeout_ms=operation_timeout_ms, + abandon_timeout_ms=abandon_timeout_ms, + nixl_agent_name=nixl_agent_name, + nixl_listen_port=_nixl_listen_port, + ), + KVCRBindings( + request_pin=self._framework_pin_adapter.request_pin, + poll_pin_results=self._framework_pin_adapter.poll_pin_results, + release_pin=self._framework_pin_adapter.release_pin, + cancel_pin_request=(self._framework_pin_adapter.cancel_pin_request), + framework_control=control, + key_adapter=self._key_adapter, + inventory_sink=(self._record_inventory if events_enabled else None), + stats_factory=( + OffloadingConnectorStats if enable_telemetry else None + ), + policy=selected_policy, + ), + KVCRBackendConfigs( + framework_dram=FrameworkDramInput( + self._primary_base_addr, primary_kv_view.nbytes + ), + local_dram=local_dram, + g3=g3_config, + remote_fw_dram=RemoteFWDramOptions( + eager_ctrl_connect=eager_ctrl_connect, + opportunistic_query=opportunistic_query, + metadata_retry_interval_ms=metadata_retry_interval_ms, + backend=remote_fw_dram_backend, + ), + ), + guard_config, + ) + except BaseException: + control.close() + if local_mapping is not None: + local_mapping.close() + raise + self._local_dram_mmap = local_mapping + self._finished_jobs: list[JobResult] = [] + self._jobs_by_op: dict[OpHandle, _JobState] = {} + + @override + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + block_key = self._key_adapter.encode(key) + status, _ = self._kvcr.query((block_key,), req_context.req_id)[0] + if status is QueryStatus.FETCHING: + return LookupResult.RETRY + if status in (QueryStatus.HIT, QueryStatus.FETCHABLE): + return LookupResult.HIT + return LookupResult.MISS + + @override + def submit_load(self, job_metadata: TransferJob) -> None: + blocks = { + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + for key, chunk_id in zip( + job_metadata.keys, job_metadata.chunk_ids, strict=True + ) + } + if not blocks: + self._finished_jobs.append( + JobResult(job_id=job_metadata.job_id, success=True) + ) + return + + op_handle = self._kvcr.deliver( + blocks, request_id=job_metadata.req_context.req_id + ) + self._jobs_by_op[op_handle] = ( + job_metadata.job_id, + len(blocks), + True, + set(), + ) + + @override + def submit_store(self, job_metadata: TransferJob) -> None: + blocks = { + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + for key, chunk_id in zip( + job_metadata.keys, job_metadata.chunk_ids, strict=True + ) + } + if not blocks: + self._finished_jobs.append( + JobResult(job_id=job_metadata.job_id, success=True) + ) + return + op_handle = self._kvcr.deposit(blocks) + self._jobs_by_op[op_handle] = ( + job_metadata.job_id, + len(blocks), + True, + None, + ) + + @override + def get_finished_jobs(self) -> Iterable[JobResult]: + results = self._finished_jobs + self._finished_jobs = [] + results.extend(self._poll_finished_jobs()) + return results + + @override + def get_stats(self) -> OffloadingConnectorStats | None: + stats = cast(OffloadingConnectorStats | None, self._kvcr.get_stats()) + if stats is not None: + # Add the framework namespace only at the vLLM boundary. + for metrics in stats.data.values(): + for name in tuple(metrics): + metrics[_vllm_metric_name(name)] = metrics.pop(name) + return stats + + @override + def serve_external_requests(self, parent: ParentManager) -> None: + self._framework_pin_adapter.process_pending(parent) + + def _poll_finished_jobs(self) -> list[JobResult]: + results: list[JobResult] = [] + for op_handle, entries in self._kvcr.poll_completed(): + job_state = self._jobs_by_op.get(op_handle) + if job_state is None: + continue + job_id, remaining, success, successful_keys = job_state + remaining -= len(entries) + success = success and all(entry.success for entry in entries.values()) + if successful_keys is not None: + successful_keys.update( + OffloadKey(bytes(key)) + for key, entry in entries.items() + if entry.success + ) + if remaining > 0: + self._jobs_by_op[op_handle] = ( + job_id, + remaining, + success, + successful_keys, + ) + continue + self._jobs_by_op.pop(op_handle, None) + results.append( + JobResult( + job_id=job_id, + success=success, + successful_keys=successful_keys if not success else None, + ) + ) + results.extend(self._framework_pin_adapter.take_pin_job_results()) + return results + + @override + def has_pending_work(self) -> bool: + # Keep scheduler polling: control messages may arrive without a load job. + return True + + @override + def take_events(self) -> Iterable[OffloadingEvent]: + events = self._inventory_events + self._inventory_events = [] + return events + + @override + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + params = getattr(req_context, "kv_transfer_params", None) + if isinstance(params, Mapping): + hint = params.get(ROUTER_HINT_KEY) + if hint is not None: + self._kvcr.submit_hint( + request_id=req_context.req_id, + hints=hint, + ) + return RequestOffloadingContext() + + @override + def on_request_finished(self, req_context: ReqContext) -> None: + self._kvcr.discard_hint(req_context.req_id) + + @override + def drain_jobs(self) -> None: + while self._jobs_by_op or self._framework_pin_adapter.has_active_pins(): + drained = self._poll_finished_jobs() + if drained: + self._finished_jobs.extend(drained) + else: + time.sleep(0.001) + + @override + def shutdown(self) -> None: + self.drain_jobs() + # TODO: Keep registered buffers alive if KVCR.close() fails. + self._kvcr.close() + if self._local_dram_mmap is not None: + self._local_dram_mmap.close() + self._local_dram_mmap = None + + def _record_inventory(self, event: InventoryEvent) -> None: + if event.tier is CacheTier.LOCAL_G2: + medium = Medium.CPU + elif event.tier is CacheTier.G3: + medium = Medium.STORAGE + else: + raise ValueError(f"unsupported KVCR inventory tier: {event.tier}") + self._inventory_events.append( + OffloadingEvent( + keys=[OffloadKey(bytes(key)) for key in event.keys], + medium=medium, + removed=event.removed, + ownership="kvcr", + removal_expected=True, + ) + ) + + def _make_descriptor(self, chunk_id: int) -> MemDescriptor: + return MemDescriptor( + end_point_name=self._kvcr.config.nixl_agent_name, + mem_type="DRAM", + addr=self._primary_base_addr + chunk_id * self._primary_row_stride, + size=self._primary_row_stride, + device_Id=0, + info="", + ) From 8263ea12bd8fa7584277f5200d521193259707b7 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Tue, 15 Sep 2026 00:57:18 -0700 Subject: [PATCH 086/139] [Build] Move DeepGEMM pin to the vLLM fork and update the Mega MoE call convention (#56876) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) --- cmake/external_projects/deepgemm.cmake | 16 ++++++++++------ docker/Dockerfile | 1 + tools/build_deepgemm_C.py | 2 +- tools/install_deepgemm.sh | 8 +++++--- vllm/models/kimi_k3/nvidia/model.py | 6 ++++-- vllm/utils/deep_gemm.py | 6 +++++- 6 files changed, 26 insertions(+), 13 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 72f7f5eb66..2ff69bd8b6 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -28,9 +28,11 @@ if(DEEPGEMM_SRC_DIR) message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() # Keep in sync with tools/install_deepgemm.sh - set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") - # Pinned to the tip of the nv_dev branch (SM120 support). - set(_DEEPGEMM_UPSTREAM_TAG "8b1392b978f5a03c828dd1711090d7fb50958b8a") + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") + # Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120 + # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU + # alpha/beta, and SiTU for FP8/FP4 Mega MoE. + set(_DEEPGEMM_UPSTREAM_TAG "ad1f1726aa540a76c1d26d6a120effb8de21eaa4") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -51,7 +53,7 @@ else() BINARY_DIR "${_deepgemm_bin}" GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}" GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}" - GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" + GIT_SUBMODULES "third-party/cutlass" "third-party/deep_jit" GIT_PROGRESS TRUE ) endif() @@ -119,13 +121,15 @@ if(DEEPGEMM_ARCHS) message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}") # add_custom_command does no implicit header scanning; glob explicitly so - # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild. + # header-only edits in DeepGEMM/cutlass/deep_jit re-trigger the rebuild. file(GLOB_RECURSE _dg_headers "${deepgemm_SOURCE_DIR}/csrc/*.h" "${deepgemm_SOURCE_DIR}/csrc/*.hpp" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp" - "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh") + "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh" + "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.h" + "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.hpp") set(_dg_markers) set(_dg_seen_soabis) diff --git a/docker/Dockerfile b/docker/Dockerfile index 9b344552c5..b3e32777e7 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -131,6 +131,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ git \ curl \ sudo \ + elfutils-devel \ rdma-core-devel \ && dnf clean all \ && rm -rf /var/cache/dnf diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 67a527405e..f4438c6755 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -50,7 +50,7 @@ includes = [ str(src / "deep_gemm/include"), str(src / "third-party/cutlass/include"), str(src / "third-party/cutlass/tools/util/include"), - str(src / "third-party/fmt/include"), + str(src / "third-party/deep_jit/include"), *cpp_extension.include_paths(device_type="cuda"), ] diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index f6a64d8796..db81bf4be3 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -6,9 +6,11 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake -DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" -# NOTE: This is currently targeting the nv_dev branch tip due to sm120 support -DEEPGEMM_GIT_REF="8b1392b978f5a03c828dd1711090d7fb50958b8a" +DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" +# NOTE: This targets the vLLM fork's dev branch tip, which carries the sm120 +# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE +# activations +DEEPGEMM_GIT_REF="ad1f1726aa540a76c1d26d6a120effb8de21eaa4" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index c4f0092f8d..6a26ce78b4 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -532,8 +532,10 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): symm_buffer, activation_clamp=activation_clamp, activation=self.activation, - situ_beta=self.activation_beta, - situ_linear_beta=self.activation_linear_beta, + # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the + # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. + activation_alpha=self.activation_beta or 1.0, + activation_beta=self.activation_linear_beta or 0.0, fast_math=fast_math, ) return y diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 3fc70c0f5d..d515d5131e 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -280,7 +280,11 @@ def _lazy_init() -> None: _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None) _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None) _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None) - _grouped_masked_impl = getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) + # DeepGEMM 2.8.0 dropped the legacy `fp8_*` aliases; keep both spellings so + # an externally pinned older DeepGEMM still resolves. + _grouped_masked_impl = getattr( + _dg, "m_grouped_fp8_gemm_nt_masked", None + ) or getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None) # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`. From 2e2fdaa99a512b58e6fa3b9fd24a07f7699d694c Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 01:00:40 -0700 Subject: [PATCH 087/139] [CI] Extend Arm CPU kernel shard timeout (#56601) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock --- .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 052063f2f6..77e0f3e07b 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -112,7 +112,9 @@ print_packages export CONTAINER_NAME export -f kernel_tests model_tests serving_tests case "$SHARD_ID" in - 0) timeout 30m bash -c kernel_tests ;; + # The kernel shard contains several large parametrized suites and can exceed + # 30 minutes on otherwise healthy Arm runners. + 0) timeout 60m bash -c kernel_tests ;; 1) timeout 30m bash -c model_tests ;; 2) timeout 30m bash -c serving_tests ;; esac From ca67438c088d5ada7da901b77f9aeaebc1b9542c Mon Sep 17 00:00:00 2001 From: shaohuaxi Date: Tue, 15 Sep 2026 01:17:30 -0700 Subject: [PATCH 088/139] [Bugfix][Frontend] Handle aborted requests in beam search (#56249) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 子华 Co-authored-by: Codex --- tests/samplers/test_beam_search.py | 85 +++++++++++++++++++ tests/samplers/test_beam_search_online.py | 54 +++++++++++- .../generate/beam_search/offline.py | 14 ++- .../generate/beam_search/online.py | 15 ++-- .../entrypoints/generate/beam_search/utils.py | 3 +- 5 files changed, 162 insertions(+), 9 deletions(-) diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py index 91dee5634b..8bcf4c48ab 100644 --- a/tests/samplers/test_beam_search.py +++ b/tests/samplers/test_beam_search.py @@ -6,13 +6,17 @@ Run `pytest tests/samplers/test_beam_search.py`. """ import json +from types import SimpleNamespace +from unittest.mock import Mock import jsonschema import pytest from transformers import AutoModelForSeq2SeqLM +from vllm import CompletionOutput, RequestOutput from vllm.assets.audio import AudioAsset from vllm.entrypoints.llm import LLM +from vllm.logprobs import Logprob, SampleLogprobs from vllm.platforms import current_platform from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams @@ -45,6 +49,87 @@ MM_BEAM_WIDTHS = [2] MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"] +@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)]) +@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]]) +def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts( + monkeypatch, + abort_after: int, + prompt_token: int, + terminal_logprobs: SampleLogprobs | None, +) -> None: + """Return tokens and scores from the last completed step for aborted prompts. + + Other prompts in the batch continue generating. + """ + + def run_requests(prompts, **kwargs): + results = [] + for prompt in prompts: + tokens = prompt["prompt_token_ids"] + if tokens[0] == prompt_token: + assert len(tokens) <= abort_after + 1, "Continued after abort" + aborted = ( + tokens[0] == prompt_token + and len(tokens) == abort_after + 1 + and tokens[-1] != 12 + ) + results.append( + RequestOutput( + request_id="inner", + prompt=None, + prompt_token_ids=tokens, + prompt_logprobs=None, + finished=True, + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[] if aborted else [11], + cumulative_logprob=None, + logprobs=terminal_logprobs + if aborted + else [{11: Logprob(-0.1), 12: Logprob(-0.2)}], + finish_reason="abort" if aborted else "length", + ) + ], + ) + ) + return results + + llm = LLM.__new__(LLM) + llm.llm_engine = Mock() + tokenizer = SimpleNamespace( + eos_token_id=0, + decode=lambda tokens, skip_special_tokens=False: " ".join( + str(token) for token in tokens if not skip_special_tokens or token != 0 + ), + ) + llm.renderer = Mock(get_tokenizer=Mock(return_value=tokenizer)) + monkeypatch.setattr(llm, "_preprocess_cmpl", lambda prompts: prompts) + monkeypatch.setattr(llm, "_render_and_run_requests", run_requests) + prompts = [ + {"type": "token", "prompt_token_ids": [token]} for token in [prompt_token, 2] + ] + params = BeamSearchParams(beam_width=2, max_tokens=3) + outputs = llm.beam_search(prompts, params) + + aborted, normal = outputs + expected_tokens = ( + [[prompt_token, 11], [prompt_token, 12]] if abort_after else [[prompt_token]] + ) + assert [beam.tokens for beam in aborted.sequences] == expected_tokens + for beam in aborted.sequences: + assert beam.finish_reason == "abort" + assert beam.text == tokenizer.decode( + beam.tokens, skip_special_tokens=params.skip_special_tokens + ) + assert len(beam.logprobs) == abort_after + assert aborted.sequences[0].cum_logprob == pytest.approx(-0.1 * abort_after) + assert len(normal.sequences) == 2 + assert all(len(beam.tokens) == 4 for beam in normal.sequences) + assert all(beam.finish_reason != "abort" for beam in normal.sequences) + + @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) @pytest.mark.parametrize("max_tokens", MAX_TOKENS) diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index 102e66dcbe..ccb40cc659 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -5,7 +5,7 @@ import pytest from vllm import CompletionOutput, RequestOutput from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin -from vllm.logprobs import Logprob +from vllm.logprobs import Logprob, SampleLogprobs from vllm.sampling_params import BeamSearchParams @@ -108,3 +108,55 @@ async def test_beam_search_respects_skip_special_tokens( assert outputs[0].outputs[0].text == expected_text assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)]) +@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]]) +async def test_beam_search_abort_returns_partial_outputs( + monkeypatch, + abort_after: int, + prompt_token: int, + terminal_logprobs: SampleLogprobs | None, +) -> None: + """Abort returns each beam's tokens and scores from the last completed step.""" + calls = 0 + + async def generate(prompt, *args, **kwargs): + nonlocal calls + calls += 1 + depth = len(prompt["prompt_token_ids"]) - 1 + assert depth <= abort_after, "Beam search continued after abort" + result = await anext(_EngineClient().generate(prompt, *args, **kwargs)) + output = result.outputs[0] + output.token_ids = [11] + output.logprobs = [{11: Logprob(-0.1), 12: Logprob(-0.2)}] + output.finish_reason = "length" + if depth == abort_after and prompt["prompt_token_ids"][-1] != 12: + output.token_ids = [] + output.logprobs = terminal_logprobs + output.finish_reason = "abort" + yield result + + serving = _Serving() + monkeypatch.setattr(serving.engine_client, "generate", generate) + prompt = {"type": "token", "prompt": "prompt", "prompt_token_ids": [prompt_token]} + outputs = [ + output + async for output in serving.beam_search( + prompt, "request", BeamSearchParams(beam_width=2, max_tokens=4) + ) + ] + + assert calls == 1 + 2 * abort_after + assert len(outputs) == 1 + assert outputs[0].finished + assert outputs[0].request_id == "request" + expected_tokens = [[11], [12]] if abort_after else [[]] + assert [output.token_ids for output in outputs[0].outputs] == expected_tokens + for output in outputs[0].outputs: + assert output.finish_reason == "abort" + assert output.text == _Tokenizer().decode(list(output.token_ids)) + assert output.logprobs is not None + assert len(output.logprobs) == abort_after + assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1 * abort_after) diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index b919435f65..73b9ebe06b 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -293,6 +293,18 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): allowed_sets[i] = set(entry[1]) for (start, end), instance in zip(instance_start_and_end, instances_batch): + instance_output = output[start:end] + if any( + result is not None and result.outputs[0].finish_reason == "abort" + for result in instance_output + ): + for beam, result in zip(all_beams[start:end], instance_output): + if result is not None: + beam.finish_reason = "abort" + instance.completed.append(beam) + instance.beams = [] + continue + instance_new_beams = [] for i in range(start, end): current_beam = all_beams[i] @@ -302,8 +314,6 @@ class BeamSearchOfflineMixin(OfflineInferenceMixin): continue if result.outputs[0].logprobs is not None: - # if logprobs is None, the sequence completed - # due to max-model-len or abortion. logprobs = result.outputs[0].logprobs[0] allowed = allowed_sets[i] for token_id, logprob_obj in logprobs.items(): diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index ac1a585e8f..da02b974c0 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -99,11 +99,7 @@ class BeamSearchOnlineMixin(ABC): output = [x[0] for x in await asyncio.gather(*tasks)] - candidates = [] - # Iterate through all beam inference results - for i, result in enumerate(output): - current_beam = all_beams[i] - + for result in output: # check for error finish reason and abort beam search if result.outputs[0].finish_reason == "error": # yield error output and terminate beam search @@ -126,6 +122,15 @@ class BeamSearchOnlineMixin(ABC): ) return + if any(result.outputs[0].finish_reason == "abort" for result in output): + for beam in all_beams: + beam.finish_reason = "abort" + break + + candidates = [] + for i, result in enumerate(output): + current_beam = all_beams[i] + if result.outputs[0].logprobs is not None: logprobs = result.outputs[0].logprobs[0] for token_id, logprob_obj in logprobs.items(): diff --git a/vllm/entrypoints/generate/beam_search/utils.py b/vllm/entrypoints/generate/beam_search/utils.py index 47f69cbf7a..61f97963b6 100644 --- a/vllm/entrypoints/generate/beam_search/utils.py +++ b/vllm/entrypoints/generate/beam_search/utils.py @@ -150,7 +150,8 @@ def get_beam_search_score( if tokens[-1] == eos_token_id: seq_len -= 1 - return cumulative_logprob / (seq_len**length_penalty) + # An aborted beam may contain only an EOS prompt token. + return cumulative_logprob / (max(seq_len, 1) ** length_penalty) def create_sort_beams_key_function(eos_token_id: int, length_penalty: float): From 676650397ff71cd802e7e0e30c2e46a215ab86ee Mon Sep 17 00:00:00 2001 From: Qiu Chunshuo Date: Tue, 15 Sep 2026 16:27:46 +0800 Subject: [PATCH 089/139] [Feature][PCP] Support decode-only FULL CUDA graphs (#53867) Signed-off-by: QiuChunshuo Signed-off-by: QiuChunshuo Co-authored-by: Codex Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_gpu_pcp_manager.py | 159 +++++++++++++++++++++++- vllm/v1/worker/gpu/cudagraph_utils.py | 9 +- vllm/v1/worker/gpu/input_batch.py | 8 +- vllm/v1/worker/gpu/model_runner.py | 35 +++--- vllm/v1/worker/gpu/pcp_manager.py | 114 ++++++++++++++--- 5 files changed, 283 insertions(+), 42 deletions(-) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index dc935929d4..d3ec562aa7 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -1,15 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock import numpy as np import pytest import torch +from vllm.config import CUDAGraphMode from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_context from vllm.v1.worker.gpu.pcp_manager import PCPManager @@ -20,6 +24,58 @@ def _copy_to_cpu(value, out=None, device=None): return tensor +def _make_config(cudagraph_mode: CUDAGraphMode): + return SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=2, + decode_context_parallel_size=1, + pipeline_parallel_size=1, + dcp_comm_backend="ag_rs", + ), + model_config=SimpleNamespace( + use_mla=True, + is_encoder_decoder=False, + hf_text_config=SimpleNamespace(), + ), + lora_config=None, + speculative_config=None, + compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), + ) + + +def _make_capture_manager(block_table: torch.Tensor): + block_tables = SimpleNamespace( + input_block_tables=(block_table,), + num_kv_cache_groups=1, + kernel_block_sizes=(2,), + blocks_per_kv_block=(1,), + ) + manager = PCPManager( + pcp_world_size=2, + pcp_rank=0, + device=torch.device("cpu"), + max_num_reqs=4, + max_num_tokens=8, + block_tables=block_tables, + ) + return manager, block_tables + + +@pytest.mark.parametrize( + "cudagraph_mode", + [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], +) +def test_validate_config_accepts_decode_only_full_graphs(cudagraph_mode): + PCPManager.validate_config(_make_config(cudagraph_mode), supports_mm_inputs=False) + + +def test_validate_config_rejects_full_graph_for_prefills(): + with pytest.raises(NotImplementedError, match="decode-only routines"): + PCPManager.validate_config( + _make_config(CUDAGraphMode.FULL), supports_mm_inputs=False + ) + + def test_replicated_decode_piecewise_graph_padding(monkeypatch): manager = PCPManager( pcp_world_size=2, @@ -114,6 +170,89 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): ) +@pytest.mark.parametrize( + ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"), + [ + (CUDAGraphMode.NONE, 4, None, None), + (CUDAGraphMode.PIECEWISE, None, 8, None), + (CUDAGraphMode.FULL, 4, 8, 4), + ], +) +def test_partition_padding_is_derived_from_batch_descriptor( + cg_mode, num_reqs, expected_tokens, expected_reqs +): + manager = MagicMock() + input_batch = MagicMock() + manager.partition_batch.return_value = input_batch + batch_desc = BatchExecutionDescriptor( + cg_mode=cg_mode, + num_tokens=8, + num_reqs=num_reqs, + ) + + result = pcp_manager_module.maybe_partition_pcp_batch( + manager, + input_batch, + batch_desc, + ) + + assert result is input_batch + manager.partition_batch.assert_called_once_with( + input_batch, + padded_num_tokens=expected_tokens, + padded_num_reqs=expected_reqs, + ) + + +def test_capture_uses_pcp_persistent_inputs(): + manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32)) + + dummy_batch = InputBatch.make_dummy( + num_reqs=4, + num_tokens=4, + input_buffers=InputBuffers(4, 8, torch.device("cpu")), + max_query_len=1, + ) + dummy_batch.input_ids.copy_(torch.arange(4)) + dummy_batch.positions.fill_(7) + input_batch = manager.prepare_inputs_to_capture(dummy_batch) + + assert input_batch is not dummy_batch + assert input_batch.req_ids == dummy_batch.req_ids + for name in ("input_ids", "positions", "is_padding", "query_start_loc", "seq_lens"): + actual = getattr(input_batch, name) + torch.testing.assert_close(actual, getattr(dummy_batch, name)) + assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr() + + +def test_dummy_context_updates_pcp_local_block_tables(): + global_block_table = torch.full((4, 4), -1, dtype=torch.int32) + manager, block_tables = _make_capture_manager(global_block_table) + input_batch = InputBatch.make_dummy( + num_reqs=2, + num_tokens=2, + input_buffers=manager.input_buffers, + max_query_len=1, + ) + input_batch = manager.prepare_inputs_to_capture(input_batch) + local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs) + + set_dummy_context( + input_batch, + block_tables, + context_len=3, + num_kv_blocks=16, + max_model_len=16, + input_block_tables=local_block_tables, + ) + + torch.testing.assert_close( + local_block_tables[0][:2, :2], + torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + ) + assert torch.all(global_block_table == -1) + + def _rank_rows( pcp_rank: int, pcp_world_size: int, @@ -331,10 +470,22 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2] global_batch.dcp_local_seq_lens.fill_(-1) - local_batch = manager.partition_batch(global_batch, padded_num_tokens=2) + local_batch = manager.partition_batch( + global_batch, + padded_num_tokens=4, + padded_num_reqs=4, + ) assert local_batch.dcp_local_seq_lens is None - assert local_batch.seq_lens.tolist() == [17, 25] + assert local_batch.seq_lens.tolist() == [17, 25, 0, 0] + assert local_batch.dcp_local_seq_lens_cpu_upper_bound is not None + expected_dcp_upper_bound = get_dcp_local_seq_lens( + torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 + ) + assert torch.equal( + local_batch.dcp_local_seq_lens_cpu_upper_bound, + torch.cat([expected_dcp_upper_bound, torch.zeros(2, dtype=torch.int32)]), + ) # What execute_model does next: derive DCP metadata from the final batch # on the PCP-owned buffers. @@ -348,7 +499,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): num_reqs_padded=local_batch.num_reqs_after_padding, ) expected = get_dcp_local_seq_lens( - torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 + torch.tensor([17, 25, 0, 0], dtype=torch.int32), 2, 0, 1 ) assert local_batch.dcp_local_seq_lens is not None assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index d197dbbc6d..9657c28fe6 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -777,11 +777,12 @@ def prepare_inputs_to_capture( input_batch = InputBatch.make_dummy( num_reqs, num_tokens, input_buffers, max_query_len=max_query_len ) - input_block_tables = block_tables.get_dummy_block_tables(num_reqs) - slot_mapping_provider: BlockTables | PCPManager = block_tables if pcp_manager is not None: - slot_mapping_provider = pcp_manager - slot_mappings = slot_mapping_provider.get_dummy_slot_mappings(num_tokens) + input_batch = pcp_manager.prepare_inputs_to_capture(input_batch) + + block_table_provider = pcp_manager or block_tables + input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs) + slot_mappings = block_table_provider.get_dummy_slot_mappings(num_tokens) slot_mappings_by_layer = build_slot_mappings_by_layer( slot_mappings, kv_cache_config ) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 5cbd642d25..d3c65bd06f 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence from dataclasses import dataclass from typing import TYPE_CHECKING @@ -217,9 +218,12 @@ def set_dummy_context( context_len: int, num_kv_blocks: int, max_model_len: int, + input_block_tables: Sequence[torch.Tensor] | None = None, ) -> None: """Give each dummy request context_len of context, used when profiling step cost.""" - if not block_tables.input_block_tables: + if input_block_tables is None: + input_block_tables = block_tables.input_block_tables + if not input_block_tables: # Attention-free models have no KV context to fabricate. return num_reqs = input_batch.num_reqs @@ -241,7 +245,7 @@ def set_dummy_context( seq_len = context_len + query_len for block_table, block_size, bpk in zip( - block_tables.input_block_tables, + input_block_tables, block_tables.kernel_block_sizes, block_tables.blocks_per_kv_block, ): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 761893303a..25395c3330 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1468,7 +1468,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): input_batch = pcp.maybe_partition_pcp_batch( self.pcp_manager, input_batch, - padded_num_tokens=batch_desc.num_tokens, + batch_desc, ) return input_batch @@ -1496,7 +1496,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): def prepare_dummy_attn( self, input_batch: InputBatch, valid_state_slots: bool = False ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: - block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs) + block_table_provider = self.pcp_manager or self.block_tables + block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs) if valid_state_slots: state_slots = torch.arange( 1, @@ -1745,25 +1746,29 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.input_buffers, max_query_len=batch_desc.max_query_len, ) - if not skip_attn_for_dummy_run: - block_tables, slot_mappings = self.prepare_dummy_attn( - input_batch, valid_dummy_state_slots - ) - if context_len: - set_dummy_context( - input_batch, - self.block_tables, - context_len, - self.kv_cache_config.num_blocks, - self.max_model_len, - ) - else: + if self.pcp_manager is not None: + input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch) + if skip_attn_for_dummy_run: assert batch_desc.cg_mode != CUDAGraphMode.FULL, ( "Attention metadata must be prepared for dummy runs when using " "FULL cudagraph mode." ) block_tables = None slot_mappings = None + else: + block_tables, slot_mappings = self.prepare_dummy_attn( + input_batch, valid_dummy_state_slots + ) + if not skip_attn_for_dummy_run and context_len: + assert block_tables is not None + set_dummy_context( + input_batch, + self.block_tables, + context_len, + self.kv_cache_config.num_blocks, + self.max_model_len, + input_block_tables=block_tables, + ) attn_metadata = None slot_mappings_by_layer = None diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 3a51532f57..adfbf2f9f0 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterator from dataclasses import dataclass, replace +from typing import TYPE_CHECKING import numpy as np import torch @@ -17,6 +18,9 @@ from vllm.v1.worker.gpu.input_batch import ( InputBuffers, ) +if TYPE_CHECKING: + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + logger = init_logger(__name__) @@ -169,8 +173,14 @@ class PCPManager: "MRV2 sparse MLA PCP does not support CUDA graphs yet. " "Set -cc.cudagraph_mode=NONE." ) - if cudagraph_mode.has_full_cudagraphs(): - raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") + if ( + cudagraph_mode.has_full_cudagraphs() + and not cudagraph_mode.separate_routine() + ): + raise NotImplementedError( + "MRV2 PCP supports full CUDA graphs for decode-only routines. " + "Use FULL_DECODE_ONLY, FULL_AND_PIECEWISE, PIECEWISE, or NONE." + ) if ( parallel_config.decode_context_parallel_size > 1 and parallel_config.dcp_comm_backend != "ag_rs" @@ -375,6 +385,22 @@ class PCPManager: for rank in range(self.pcp_world_size) ) + @staticmethod + def _resolve_num_reqs_after_padding( + input_batch: InputBatch, + padded_num_reqs: int | None, + num_local_reqs: int, + ) -> int: + if padded_num_reqs is None: + return num_local_reqs + if input_batch.has_prefill: + raise RuntimeError("PCP FULL graphs require a decode-only batch.") + assert padded_num_reqs >= num_local_reqs, ( + "PCP graph request capacity must cover the rank-local batch: " + f"{padded_num_reqs} < {num_local_reqs}." + ) + return padded_num_reqs + @property def input_buffers(self) -> InputBuffers: assert self._input_buffers is not None @@ -384,6 +410,7 @@ class PCPManager: self, input_batch: InputBatch, padded_num_tokens: int | None = None, + padded_num_reqs: int | None = None, ) -> InputBatch: assert self._input_buffers is not None input_buffers = self._input_buffers @@ -414,10 +441,15 @@ class PCPManager: ] num_local_reqs = len(local_segments) - if num_local_reqs > input_buffers.max_num_reqs: + num_reqs_after_padding = self._resolve_num_reqs_after_padding( + global_batch, + padded_num_reqs, + num_local_reqs, + ) + if num_reqs_after_padding > input_buffers.max_num_reqs: raise RuntimeError( - "PCP local request count exceeds the MRV2 input buffer size: " - f"{num_local_reqs} > {input_buffers.max_num_reqs}." + "PCP padded local request count exceeds the MRV2 input buffer size: " + f"{num_reqs_after_padding} > {input_buffers.max_num_reqs}." ) local_to_global_batch_req_idx_np = np.fromiter( @@ -505,21 +537,25 @@ class PCPManager: np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out) local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc) - local_query_start_loc = input_buffers.query_start_loc[: num_local_reqs + 1] + local_query_start_loc = input_buffers.query_start_loc[ + : num_reqs_after_padding + 1 + ] local_to_global_req_idx = async_copy_to_gpu( local_to_global_req_idx_np, device=self.device ) - seq_lens = input_buffers.seq_lens[:num_local_reqs] + seq_lens = input_buffers.seq_lens[:num_reqs_after_padding] + real_seq_lens = seq_lens[:num_local_reqs] if num_local_tokens > 0: local_end_positions = torch.index_select( input_buffers.positions, 0, - local_query_start_loc[1:] - 1, + local_query_start_loc[1 : num_local_reqs + 1] - 1, ) - seq_lens.copy_(local_end_positions + 1) + real_seq_lens.copy_(local_end_positions + 1) else: - seq_lens.zero_() + real_seq_lens.zero_() + seq_lens[num_local_reqs:].zero_() is_padding = input_buffers.is_padding[:num_local_tokens_padded] is_padding[:num_local_tokens].fill_(False) is_padding[num_local_tokens:].fill_(True) @@ -552,11 +588,15 @@ class PCPManager: local_num_computed_prefill_tokens_np = np.minimum( local_start_pos_np, local_prefill_len_np ) - local_is_prefilling_np = ( + real_local_is_prefilling_np = ( local_num_computed_prefill_tokens_np < local_prefill_len_np ) - seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32) - seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens + local_is_prefilling_np = np.zeros(num_reqs_after_padding, dtype=np.bool_) + local_is_prefilling_np[:num_local_reqs] = real_local_is_prefilling_np + seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_after_padding, dtype=np.int32) + seq_lens_cpu_upper_bound_np[:num_local_reqs] = ( + local_start_pos_np + local_num_scheduled_tokens + ) dcp_local_seq_lens_cpu_upper_bound = None if self.dcp_world_size > 1: # The largest DCP shard of each row's whole request, identical on @@ -564,18 +604,24 @@ class PCPManager: request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[ local_to_global_batch_req_idx_np ] - dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( + real_dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( torch.from_numpy(request_seq_lens.astype(np.int32)), self.dcp_world_size, 0, self.cp_interleave, ) + dcp_local_seq_lens_cpu_upper_bound = torch.zeros( + num_reqs_after_padding, dtype=torch.int32 + ) + dcp_local_seq_lens_cpu_upper_bound[:num_local_reqs].copy_( + real_dcp_local_seq_lens_cpu_upper_bound + ) self._local_batch = replace( input_batch, req_ids=local_req_ids, num_reqs=num_local_reqs, - num_reqs_after_padding=num_local_reqs, + num_reqs_after_padding=num_reqs_after_padding, idx_mapping=local_to_global_req_idx, idx_mapping_np=local_to_global_req_idx_np, expanded_idx_mapping=local_to_global_req_idx, @@ -588,7 +634,7 @@ class PCPManager: num_draft_tokens=0, num_draft_tokens_per_req=None, query_start_loc=local_query_start_loc, - query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1], + query_start_loc_np=local_query_start_loc_np[: num_reqs_after_padding + 1], seq_lens=seq_lens, seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np), dcp_local_seq_lens=None, @@ -608,6 +654,31 @@ class PCPManager: ) return self._local_batch + def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch: + """Stage a capture or dummy batch in persistent PCP input buffers.""" + input_buffers = self.input_buffers + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + input_batch = replace( + input_batch, + input_ids=input_buffers.input_ids[:num_tokens].copy_(input_batch.input_ids), + positions=input_buffers.positions[:num_tokens].copy_(input_batch.positions), + is_padding=input_buffers.is_padding[:num_tokens].copy_( + input_batch.is_padding + ), + query_start_loc=input_buffers.query_start_loc[: num_reqs + 1].copy_( + input_batch.query_start_loc + ), + seq_lens=input_buffers.seq_lens[:num_reqs].copy_(input_batch.seq_lens), + ) + return input_batch + + def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]: + assert self._local_block_tables is not None + return tuple( + block_table[:num_reqs].zero_() for block_table in self._local_block_tables + ) + def prepare_attn( self, input_batch: InputBatch ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: @@ -728,13 +799,22 @@ class PCPManager: def maybe_partition_pcp_batch( manager: PCPManager | None, input_batch: InputBatch, - padded_num_tokens: int | None = None, + batch_desc: "BatchExecutionDescriptor", ) -> InputBatch: if manager is None: return input_batch + + padded_num_tokens = None + padded_num_reqs = None + if batch_desc.cg_mode != CUDAGraphMode.NONE: + padded_num_tokens = batch_desc.num_tokens + if batch_desc.cg_mode == CUDAGraphMode.FULL: + padded_num_reqs = batch_desc.num_reqs + return manager.partition_batch( input_batch, padded_num_tokens=padded_num_tokens, + padded_num_reqs=padded_num_reqs, ) From 2c2cbd84fdf199901f3f4e6b0e708ee0664aadd6 Mon Sep 17 00:00:00 2001 From: Yixin Dong Date: Tue, 15 Sep 2026 04:36:16 -0400 Subject: [PATCH 090/139] [Frontend] Use XGrammar schema constraints for DeepSeek V4.1 (#56408) Signed-off-by: Ubospica Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../test_structural_tag_registry.py | 80 ++++++++++++++++++- vllm/tool_parsers/structural_tag_registry.py | 14 +++- 2 files changed, 90 insertions(+), 4 deletions(-) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index b77b4efaff..88e1f3f383 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -5,7 +5,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock import pytest -from xgrammar import Grammar, StructuralTag +from xgrammar import Grammar, StructuralTag, builtin_structural_tag from xgrammar.testing import _is_grammar_accept_string from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -81,10 +81,14 @@ def test_supported_structural_tag_models_include_vllm_builtins(): @pytest.mark.parametrize("choice", ["auto", "required", "get_weather"]) -def test_deepseek_v41_format_constraints_ignore_parameter_schema( +def test_deepseek_v41_older_xgrammar_keeps_format_constraints( choice, sample_tools_strict, + monkeypatch, ): + monkeypatch.setattr( + builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None, raising=False + ) tool_choice = ( ChatCompletionNamedToolChoiceParam( function=ChatCompletionNamedFunction(name=choice) @@ -128,7 +132,11 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools): reasoning=False, ) grammar = Grammar.from_structural_tag(tag) - call = '<｜DSML｜ invoke name="get_weather">\n\n' + call = ( + '<｜DSML｜ invoke name="get_weather">\n' + '<｜DSML｜ parameter name="city" string="true">Paris\n' + "\n" + ) assert _is_grammar_accept_string( grammar, "\n\n<｜DSML｜ calls>\n" + call + "" ) @@ -141,6 +149,72 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools): ) +@pytest.mark.skipif( + not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"), + reason="Requires XGrammar's DeepSeek V4.1 builtin", +) +@pytest.mark.parametrize("choice", ["auto", "required", "get_weather"]) +@pytest.mark.parametrize("reasoning", [False, True]) +def test_deepseek_v41_constrains_parameters_after_reasoning( + sample_tools_strict, choice, reasoning +): + """Strict calls must enforce the schema without requiring another think close.""" + tool_choice = ( + ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name=choice) + ) + if choice == "get_weather" + else choice + ) + sample_tools_strict[0].function.parameters["additionalProperties"] = False + tag = get_model_structural_tag( + "deepseek_v41", sample_tools_strict, tool_choice, reasoning=reasoning + ) + grammar = Grammar.from_structural_tag(tag) + begin = '\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name="get_weather">\n' + end = "\n" + parameter = ( + '<｜DSML｜ parameter name="city" string="true">Paris\n' + ) + assert _is_grammar_accept_string(grammar, begin + parameter + end) + for invalid in ( + "", # missing required city + parameter * 2, + parameter.replace('name="city"', 'name="unknown"'), + parameter.replace('string="true">Paris', 'string="false">42'), + parameter.replace("｜ parameter", "｜parameter"), + '{"city":"Paris"}', # the encoder emits DSML parameters, not a JSON body + ): + assert not _is_grammar_accept_string(grammar, begin + invalid + end) + assert not _is_grammar_accept_string( + grammar, "reason" + begin + parameter + end + ) + assert _is_grammar_accept_string(grammar, "Hello") == (choice == "auto") + + +@pytest.mark.skipif( + not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"), + reason="Requires XGrammar's DeepSeek V4.1 builtin", +) +def test_deepseek_v41_non_strict_parallel_calls_keep_typed_dsml(sample_tools): + sample_tools[0].function.strict = False + tag = get_model_structural_tag( + "deepseek_v41", sample_tools, "required", reasoning=False + ) + grammar = Grammar.from_structural_tag(tag) + call = ( + '<｜DSML｜ invoke name="get_weather">\n' + '<｜DSML｜ parameter name="extra" string="false">' + '{"nested":[true,null,1.5]}\n' + "\n" + ) + output = "\n\n<｜DSML｜ calls>\n" + call * 2 + "" + assert _is_grammar_accept_string(grammar, output) + assert not _is_grammar_accept_string( + grammar, output.replace('{"nested":[true,null,1.5]}', "invalid") + ) + + @pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) def test_get_model_structural_tag_supports_all_xgrammar_builtins( model: str, diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 34e57840be..989250e2ea 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -10,6 +10,7 @@ from openai.types.responses.tool import Tool as ResponsesTool from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed from openai.types.responses.tool_choice_function import ToolChoiceFunction from xgrammar import StructuralTag, normalize_tool_choice +from xgrammar import builtin_structural_tag as xgrammar_builtin_structural_tag from xgrammar import get_model_structural_tag as get_xgrammar_model_structural_tag from xgrammar.openai_tool_call_schema import ( BuiltinToolParam, @@ -242,9 +243,20 @@ def get_deepseek_v41_structural_tag( token_suffix: str = "", ) -> StructuralTag: # Serving enables this visible-text grammar after the reasoning boundary. + builder = getattr( + xgrammar_builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None + ) + if builder is not None: + return builder( + tools=tools, + builtin_tools=builtin_tools, + tool_choice=tool_choice, + reasoning="disabled", + ) + del builtin_tools, reasoning, token_suffix - # TODO: lower parameter schemas into DSML constraints. This builder constrains + # Compatibility with xgrammar releases without the V4.1 builtin. This constrains # tool names and DSML/value syntax only. Parameter names, presence, uniqueness # and value schemas remain unconstrained, including for strict=true. The request # layer still uses strict to decide whether auto tool choice activates a grammar. From 2983fc62869de7943bd68ea58730eec7a39ff597 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 02:00:58 -0700 Subject: [PATCH 091/139] [CI] Fix NIXL push worker test stub after failure deferral (#56955) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- tests/v1/kv_connector/unit/test_nixl_push_connector.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 423b0e9dec..913ed0dcba 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -351,6 +351,7 @@ class _StubWriterWorker(NixlPushConnectorWorker): # Base worker fields touched by start_load_kv / _get_new_notifs. w._recving_metadata = {} w._pending_recv_notifs = {} + w._failed_recv_reqs = queue.Queue() w._recv_failures = set() w._recving_transfers = defaultdict(list) w._reqs_to_process = set() @@ -455,12 +456,14 @@ class TestPushWriterRegSend: list(group) for group in rd["local_block_ids"] ] - def test_fails_request_when_engine_not_handshaked(self): + def test_queues_failed_request_when_engine_not_handshaked(self): w, failures = self._reg_send_worker() w._do_send_reg_notif("req-1", _registration_data("req-1")) - assert failures == [("req-1", None)] + assert failures == [] + assert w._failed_recv_reqs.get_nowait() == "req-1" + assert w._failed_recv_reqs.empty() w.nixl_wrapper.send_notif.assert_not_called() def test_agent_read_is_serialized_with_handshake_lock(self): @@ -886,7 +889,6 @@ class TestPushWriterNotifs: w.nixl_wrapper = MagicMock() w.xfer_stats = MagicMock() w._log_failure = MagicMock() # type: ignore[method-assign] - w._failed_recv_reqs = queue.Queue() w._invalid_block_ids = queue.Queue() w._pending_recv_notifs = {} w._replicated_pcp_done_sending = set() From abb4cb9fc5f1b1039876ed8a657b025f65643d11 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Tue, 15 Sep 2026 05:58:42 -0400 Subject: [PATCH 092/139] =?UTF-8?q?[CI]=20Shard=20V1=20KV=20Connectors=201?= =?UTF-8?q?=E2=86=924=20(#56324)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Thang Nguyen Co-authored-by: Thang Nguyen Co-authored-by: Kimi Code --- .buildkite/test_areas/misc.yaml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index ed82c7d71d..940447bf1b 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -142,19 +142,20 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors" +- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors Shard %N" device: h200_35gb key: v1-kv-connectors timeout_in_minutes: 45 + parallelism: 4 source_file_dependencies: *v1-core-kv-metrics-deps commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit + - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB mirror: amd: - label: ":amd: (MI300) V1 KV Connectors" + label: ":amd: (MI300) V1 KV Connectors Shard %N" dind: false device: mi300_1 timeout_in_minutes: 45 From 073f883b56a5b585366da874473db100bc6fff7e Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Tue, 15 Sep 2026 02:58:57 -0700 Subject: [PATCH 093/139] [Perf][DSpark] Collapse DeepSeek-V4.1 draft states before SP all-gather (#56903) Signed-off-by: Juntian Liu Co-authored-by: Codex Co-authored-by: Woosuk Kwon --- tests/kernels/test_mhc_kernels.py | 21 +++++++++++++++++++++ vllm/models/deepseek_v41/nvidia/dspark.py | 5 ++--- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index f805fe39a5..ff72e2a96c 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -147,6 +147,27 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("num_tokens", [1, 7, 32]) +@pytest.mark.parametrize("tp_size", [4, 8]) +def test_hc_collapse_before_gather_is_bitwise_equal(num_tokens, tp_size): + """Simulate SP on one GPU, including padding and ranks without valid tokens.""" + set_random_seed(0) + padded_tokens = num_tokens + (-num_tokens) % tp_size + streams = torch.randn(padded_tokens, 4, 5120, dtype=torch.bfloat16, device=DEVICE) + pre_mix = torch.rand(padded_tokens, 4, device=DEVICE) + stream_shards = streams.chunk(tp_size) + mix_shards = pre_mix.chunk(tp_size) + before = hc_collapse_triton( + torch.cat(stream_shards)[:num_tokens], torch.cat(mix_shards)[:num_tokens] + ) + after = torch.cat( + [hc_collapse_triton(x, pre) for x, pre in zip(stream_shards, mix_shards)] + )[:num_tokens] + assert after.shape == (num_tokens, 5120) + torch.testing.assert_close(after, before, atol=0, rtol=0) + + @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize( "num_tokens,sinkhorn_repeat,norm_eps", diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 88646404f4..c19dca8005 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -217,15 +217,14 @@ class DSparkDeepseekV4Model(nn.Module): residual, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) - if self.use_sequence_parallel: - hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] - pre_mix = sp_all_gather(pre_mix)[:full_num_tokens] # Collapse the hc copies with the pre-mix from the last layer's FFN # mixes — the mix the reference forward_head applies (hc_pre with the # last block's ffn pre-mix). Return the PRE-norm head hidden; # compute_logits applies self.norm. assert pre_mix is not None hidden_states = hc_collapse_triton(hidden_states, pre_mix) + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] return hidden_states From 142020c6c2c0d4fec4d36df28290e3d9b654a385 Mon Sep 17 00:00:00 2001 From: Yeonwoo Sung Date: Tue, 15 Sep 2026 19:03:26 +0900 Subject: [PATCH 094/139] [Bugfix][Rust Frontend] Honor `add_generation_prompt` in DeepSeek V3.2/V4/V4.1 renderers (#56593) Co-authored-by: Bugen Zhao Signed-off-by: YeonwooSung Signed-off-by: Bugen Zhao --- rust/src/chat/src/renderer/deepseek.rs | 47 ++++++---- .../src/renderer/deepseek_v32/encoding.rs | 18 +++- .../chat/src/renderer/deepseek_v32/tests.rs | 38 ++++++++ .../chat/src/renderer/deepseek_v4/tests.rs | 91 +++++++++++++++++++ .../chat/src/renderer/deepseek_v41/tests.rs | 47 +++++++++- 5 files changed, 217 insertions(+), 24 deletions(-) diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs index 5925c5fd16..01576a7c43 100644 --- a/rust/src/chat/src/renderer/deepseek.rs +++ b/rust/src/chat/src/renderer/deepseek.rs @@ -224,21 +224,30 @@ pub(super) fn render_request_with_media_order( } } - if (is_user_like_entry(message, current_render_index as usize, dialect) - || (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. }))) - && next_rendered_entry_is_assistant_or_end( + if is_user_like_entry(message, current_render_index as usize, dialect) + || (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. })) + { + // Skip the opener only at the last rendered entry when + // add_generation_prompt is off. Inspect the next rendered + // entry, not the raw index: V4 merges trailing tool results. + let should_write_transition = match next_rendered_entry( request.messages.as_slice(), message_index, drop_historical_developers, last_user_like_message_index, - ) - { - write_assistant_transition( - &mut out, - thinking_mode, - drop_thinking, - current_render_index >= last_user_render_index, - ); + ) { + Some(ChatMessage::Assistant { .. }) => true, + None => request.chat_options.add_generation_prompt(), + Some(_) => false, + }; + if should_write_transition { + write_assistant_transition( + &mut out, + thinking_mode, + drop_thinking, + current_render_index >= last_user_render_index, + ); + } } } @@ -420,14 +429,17 @@ fn is_historical_developer( && last_user_like_message_index.is_some_and(|last_index| message_index < last_index) } -/// Return whether the next rendered entry is assistant, or there is no next -/// entry. -fn next_rendered_entry_is_assistant_or_end( +/// Return the next rendered message after `message_index`. +/// +/// Skips merged user-content tails and dropped historical developers so the +/// caller can tell a following assistant turn apart from the end of the +/// conversation. +fn next_rendered_entry( messages: &[ChatMessage], message_index: usize, drop_historical_developers: bool, last_user_like_message_index: Option, -) -> bool { +) -> Option<&ChatMessage> { let mut next_index = message_index + 1; while next_index < messages.len() && (is_following_user_content(messages, next_index) @@ -441,10 +453,7 @@ fn next_rendered_entry_is_assistant_or_end( next_index += 1; } - messages - .get(next_index) - .map(|message| matches!(message, ChatMessage::Assistant { .. })) - .unwrap_or(true) + messages.get(next_index) } /// Render the tool preamble shown to the model for one DeepSeek dialect. diff --git a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs index 3ccf4d11db..4fc5f26530 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs @@ -61,6 +61,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { find_last_user_render_index(request.messages.as_slice(), render_offset); let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice()); let continue_final_message = request.chat_options.continue_final_message(); + let add_generation_prompt = request.chat_options.add_generation_prompt(); let mut prompt = String::from(BOS_TOKEN); if !initial_tools.is_empty() { @@ -79,6 +80,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { thinking_mode, drop_thinking, continue_final_message, + add_generation_prompt, )?; } @@ -110,11 +112,13 @@ fn render_message( thinking_mode: ThinkingMode, drop_thinking: bool, continue_final_message: bool, + add_generation_prompt: bool, ) -> Result<()> { let render_index = message_index as isize + render_offset; let opens_thinking = render_index == last_user_render_index; let after_last_user_turn = render_index > last_user_render_index; let after_or_at_last_user_turn = render_index >= last_user_render_index; + let add_generation_prompt = message_index + 1 != messages.len() || add_generation_prompt; match message { ChatMessage::System { content } => render_system_message(out, Some(content), &[]), @@ -123,11 +127,13 @@ fn render_message( content, tools.as_deref().unwrap_or(&[]), thinking_mode == ThinkingMode::Thinking && opens_thinking, + add_generation_prompt, ), ChatMessage::User { content } => render_user_message( out, content, thinking_mode == ThinkingMode::Thinking && opens_thinking, + add_generation_prompt, ), ChatMessage::Assistant { content } => render_assistant_message( out, @@ -195,6 +201,7 @@ fn render_developer_message( content: &ChatContent, tools: &[ChatTool], opens_thinking: bool, + add_generation_prompt: bool, ) -> Result<()> { if content.is_empty() { return Err(Error::ChatTemplate( @@ -209,7 +216,7 @@ fn render_developer_message( } out.push_str("\n\n# The user's message is: "); write_chat_content(out, content)?; - write_user_like_suffix(out, opens_thinking); + write_user_like_suffix(out, opens_thinking, add_generation_prompt); Ok(()) } @@ -219,17 +226,20 @@ fn render_user_message( out: &mut String, content: &ChatContent, opens_thinking: bool, + add_generation_prompt: bool, ) -> Result<()> { out.push_str("<｜User｜>"); write_chat_content(out, content)?; - write_user_like_suffix(out, opens_thinking); + write_user_like_suffix(out, opens_thinking, add_generation_prompt); Ok(()) } /// Shared trailing wrapper used by both real user turns and native developer /// turns after their content has already been written. -// TODO: respect `add_generation_prompt` option -fn write_user_like_suffix(out: &mut String, opens_thinking: bool) { +fn write_user_like_suffix(out: &mut String, opens_thinking: bool, add_generation_prompt: bool) { + if !add_generation_prompt { + return; + } out.push_str("<｜Assistant｜>"); if opens_thinking { out.push_str(THINKING_START_TOKEN); diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 3f93e3653e..56c03a92ca 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -281,6 +281,44 @@ fn assistant_after_last_user_requires_reasoning_or_tool_calls() { .assert_eq(&error.to_report_string()); } +#[test] +fn last_user_turn_appends_generation_prompt_by_default() { + let request = thinking_request(vec![ChatMessage::user("Hello")]); + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜>")); +} + +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::user("Hello")]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello")); + assert!(!rendered.contains("<｜Assistant｜>")); + assert!(!rendered.contains("")); +} + +#[test] +fn user_then_assistant_still_writes_suffix_without_generation_prompt() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant_text("Hi there."), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜>Hi there.<｜end▁of▁sentence｜>"] + .assert_eq(&rendered); +} + #[test] fn continue_final_assistant_omits_final_eos() { let mut request = ChatRequest { diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 172b67463c..9404b54296 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -633,3 +633,94 @@ fn continue_final_assistant_omits_final_eos() { expect!["<｜begin▁of▁sentence｜><｜User｜>write<｜Assistant｜>partial answer"] .assert_eq(&rendered); } + +#[test] +fn last_user_turn_appends_generation_prompt_by_default() { + let request = thinking_request(vec![ChatMessage::user("Hello")]); + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜>")); +} + +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::user("Hello")]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello")); + assert!(!rendered.contains("<｜Assistant｜>")); + assert!(!rendered.contains("")); +} + +#[test] +fn user_then_assistant_still_writes_suffix_without_generation_prompt() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant_text("Hi there."), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), Value::Bool(false)); + + let rendered = render_request(&request); + + expect![ + "<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜>Hi there.<｜end▁of▁sentence｜>" + ] + .assert_eq(&rendered); +} + +#[test] +fn trailing_system_omits_generation_prompt_when_disabled() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::system("rules"), + ChatMessage::user("question"), + ChatMessage::system("final rules"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("enable_thinking".to_string(), Value::Bool(false)); + request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜>rules<｜User｜>questionfinal rules"].assert_eq(&rendered); +} + +#[test] +fn last_merged_tool_result_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ + ChatMessage::user("question"), + ChatMessage::tool_response("result", "call-1"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>question\n\nresult"] + .assert_eq(&rendered); + assert!(!rendered.contains("<｜Assistant｜>")); +} + +#[test] +fn last_developer_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::developer("latest instruction", None)]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>latest instruction"].assert_eq(&rendered); +} diff --git a/rust/src/chat/src/renderer/deepseek_v41/tests.rs b/rust/src/chat/src/renderer/deepseek_v41/tests.rs index 3babb5c269..3d621c2246 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/tests.rs @@ -10,7 +10,9 @@ use super::DeepSeekV41ChatRenderer; use crate::ChatRenderer; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; -use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort}; +use crate::request::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort, +}; fn render(request: &ChatRequest) -> String { DeepSeekV41ChatRenderer::new() @@ -281,6 +283,49 @@ fn mid_system_advances_last_user_boundary_and_drops_old_reasoning() { .assert_eq(&render(&request)); } +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = request(); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 50 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question"] + .assert_eq(&render(&request)); +} + +#[test] +fn last_mid_system_omits_generation_prompt_when_disabled() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::system("sys"), + ChatMessage::user("q1"), + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::Reasoning { text: "r1".into() }, + AssistantContentBlock::Text { text: "a1".into() }, + ]), + ChatMessage::developer("old policy", None), + ChatMessage::system("mid sys"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("reasoning_effort".into(), json!(88)); + + expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 88 (range 1-100, the higher the value, the more thorough the reasoning)\n\nsys<｜User｜>q1<｜Assistant｜>a1<｜end▁of▁sentence｜><｜System｜>mid sys"] + .assert_eq(&render(&request)); +} + +#[test] +fn chat_mode_omits_generation_prompt_when_disabled() { + let mut request = request(); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request.chat_options.template_kwargs.insert("thinking".into(), json!(false)); + + expect!["<｜begin▁of▁sentence｜><｜User｜>question"].assert_eq(&render(&request)); +} + #[test] fn joins_text_parts_with_reference_separator_and_preserves_literal_tags() { let mut request = ChatRequest { From 24bfbd5d4d530a2657a8925e5d12e86b1ccb9cca Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Tue, 15 Sep 2026 06:07:58 -0400 Subject: [PATCH 095/139] [Refactor] Derive is_reasoning_end from the engine grammar (#56200) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .../engine/test_gemma4_streaming_reasoning.py | 33 +++++++++++++++ tests/parser/engine/test_qwen3_reasoning.py | 42 +++++++++++++++++-- vllm/parser/deepseek_v32.py | 1 + vllm/parser/deepseek_v4.py | 1 + vllm/parser/engine/parser_engine.py | 36 +++++++++------- vllm/parser/engine/parser_engine_config.py | 2 + vllm/parser/gemma4.py | 34 ++++----------- vllm/parser/glm47_moe.py | 6 +-- vllm/parser/inkling.py | 1 + vllm/parser/kimi_k2.py | 19 +-------- vllm/parser/minimax_m2.py | 1 + vllm/parser/mistral.py | 19 +-------- vllm/parser/qwen3.py | 33 ++------------- 13 files changed, 111 insertions(+), 117 deletions(-) diff --git a/tests/parser/engine/test_gemma4_streaming_reasoning.py b/tests/parser/engine/test_gemma4_streaming_reasoning.py index 74224c15b7..ae21d69fc3 100644 --- a/tests/parser/engine/test_gemma4_streaming_reasoning.py +++ b/tests/parser/engine/test_gemma4_streaming_reasoning.py @@ -1616,3 +1616,36 @@ class TestCommaInStringValueRegression: assert result.tools_called is True args = json.loads(result.tool_calls[0].function.arguments) assert args["destination"] == "456 Oakwood Avenue, Rivermist, 83214" + + +class TestGemma4IsReasoningEnd: + @staticmethod + def _parser(thinking: bool): + return Gemma4Parser( + _make_tokenizer(_PLAIN_ANSWER_TOKENS), + chat_template_kwargs={"enable_thinking": thinking}, + ) + + @pytest.mark.parametrize( + ("ids", "thinking", "ended"), + [ + ([CHANNEL_START_ID, 3000, CHANNEL_END_ID], True, True), + ([CHANNEL_START_ID, 3000, TOOL_CALL_START_ID], True, True), + ([CHANNEL_START_ID, 3000], True, False), + ([NEW_TURN_ID, 9100], True, False), + ([CHANNEL_END_ID, NEW_TURN_ID, 9100], True, False), + ([NEW_TURN_ID, 9100], False, True), + # The bundled template's thinking-off prompt: closed channel. + ([NEW_TURN_ID, 9100, CHANNEL_START_ID, 3000, CHANNEL_END_ID], False, True), + # A marker still wins with thinking off: the grammar honours it. + ([CHANNEL_START_ID, 3000], False, False), + # A tool-call opener after the channel closed: still ended. + ( + [CHANNEL_START_ID, 3000, CHANNEL_END_ID, TOOL_CALL_START_ID, 9200], + True, + True, + ), + ], + ) + def test_is_reasoning_end(self, ids, thinking, ended): + assert self._parser(thinking).is_reasoning_end(ids) is ended diff --git a/tests/parser/engine/test_qwen3_reasoning.py b/tests/parser/engine/test_qwen3_reasoning.py index 37e6b5d130..3f285e4d03 100644 --- a/tests/parser/engine/test_qwen3_reasoning.py +++ b/tests/parser/engine/test_qwen3_reasoning.py @@ -180,9 +180,8 @@ class TestIsReasoningEnd: """Tool examples before the generation must not end reasoning.""" assert not parser.is_reasoning_end([_TOOL_CALL_ID, _TEXT_ID, _THINK_START_ID]) - def test_paired_tool_call_not_end(self, parser): - """Paired ... (from template) is NOT end.""" - assert not parser.is_reasoning_end( + def test_completed_tool_call_after_open_think_is_end(self, parser): + assert parser.is_reasoning_end( [_THINK_START_ID, 1, _TOOL_CALL_ID, 2, _TOOL_CALL_END_ID] ) @@ -196,6 +195,43 @@ class TestIsReasoningEnd: assert not parser.is_reasoning_end([]) +class TestIsReasoningEndDerivedFromGrammar: + @pytest.mark.parametrize( + ("ids", "thinking", "ended"), + [ + # Explicit close, and the implicit close on . + ([_THINK_START_ID, 1, _THINK_END_ID], True, True), + ([_THINK_START_ID, 1, _TOOL_CALL_ID], True, True), + # Still inside the block. + ([_THINK_START_ID, 1], True, False), + ([_THINK_END_ID, _THINK_START_ID, 1], True, False), + # Uncommitted: a bare turn, an empty sequence. Thinking on starts + # in REASONING, so reasoning is expected. + ([_IM_END_ID, _IM_START_ID, 1], True, False), + ([], True, False), + # Thinking off starts in CONTENT with no way into REASONING, so + # a the template left behind is inert, and a silent tail + # is ended because no block is expected. + ([_THINK_START_ID, 1], False, True), + ([_THINK_END_ID, _THINK_START_ID, 1], False, True), + ([_IM_END_ID, _IM_START_ID, 1], False, True), + ([], False, True), + ], + ) + def test_is_reasoning_end(self, mock_tokenizer, ids, thinking, ended): + parser = Qwen3Parser( + mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking} + ) + assert parser.is_reasoning_end(ids) is ended + + @pytest.mark.parametrize("thinking", [True, False]) + def test_wait_for_reasoning_follows_thinking(self, mock_tokenizer, thinking): + parser = Qwen3Parser( + mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking} + ) + assert parser.parser_engine_config.wait_for_reasoning is thinking + + class TestIsReasoningEndTurnBoundaries: """Reasoning markers from earlier turns must not mark the new turn's reasoning as finished (think-marker injection / interleaved-thinking diff --git a/vllm/parser/deepseek_v32.py b/vllm/parser/deepseek_v32.py index f76c954a93..de6634cf73 100644 --- a/vllm/parser/deepseek_v32.py +++ b/vllm/parser/deepseek_v32.py @@ -53,6 +53,7 @@ def deepseek_v32_config() -> ParserEngineConfig: return ParserEngineConfig( name="deepseek_v32", initial_state=ParserState.CONTENT, + wait_for_reasoning=False, terminals={ "TOOL_START": DSML_FUNC_START, "TOOL_END": DSML_FUNC_END, diff --git a/vllm/parser/deepseek_v4.py b/vllm/parser/deepseek_v4.py index 13eb83e6b8..79ddd93f6b 100644 --- a/vllm/parser/deepseek_v4.py +++ b/vllm/parser/deepseek_v4.py @@ -140,6 +140,7 @@ def deepseek_v4_config(thinking: bool = False) -> ParserEngineConfig: return ParserEngineConfig( name="deepseek_v4", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ "THINK_START": DSML_THINK_START, "THINK_END": DSML_THINK_END, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 4817dedb46..f873e36342 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -648,24 +648,28 @@ class ParserEngine(Parser): return None def is_reasoning_end(self, input_ids: list[int]) -> bool: - end_id = self._reasoning_end_token_id + config = self.parser_engine_config + wait_for_reasoning = config.wait_for_reasoning + if wait_for_reasoning is None: + wait_for_reasoning = config.initial_state is ParserState.REASONING + start_transition = config.transitions.get((config.initial_state, "THINK_START")) + start_opens_reasoning = ( + start_transition is not None + and start_transition.next_state is ParserState.REASONING + ) + end_ids = self._reasoning_end_token_ids start_id = self._reasoning_start_token_id - if end_id is not None: - if not input_ids: - return self.parser_engine_config.initial_state != ParserState.REASONING - boundary_ids = self._turn_boundary_token_ids - for i in range(len(input_ids) - 1, -1, -1): - token_id = input_ids[i] - if token_id == end_id: - return True - if start_id is not None and token_id == start_id: + boundary_ids = self._turn_boundary_token_ids + for token_id in reversed(input_ids): + if token_id in end_ids: + return True + if token_id == start_id: + if start_opens_reasoning: return False - if token_id in boundary_ids: - return ( - self.parser_engine_config.initial_state != ParserState.REASONING - ) - return False - return self._reasoning_ended + break + if token_id in boundary_ids: + break + return not wait_for_reasoning def extract_content_ids(self, input_ids: list[int]) -> list[int]: end_id = self._reasoning_end_token_id diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py index 5c5cd2b467..b8d7dc6e46 100644 --- a/vllm/parser/engine/parser_engine_config.py +++ b/vllm/parser/engine/parser_engine_config.py @@ -74,6 +74,8 @@ class ParserEngineConfig: initial_state: ParserState = ParserState.CONTENT + wait_for_reasoning: bool | None = None + arg_converter: Callable[[str, bool], str] | None = None stream_arg_deltas: bool = True diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index 51c4696636..e3baaf488a 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -293,10 +293,14 @@ def _gemma4_arg_converter(raw_args: str, partial: bool) -> str: @functools.cache -def gemma4_config() -> ParserEngineConfig: +def gemma4_config(thinking: bool = False) -> ParserEngineConfig: return ParserEngineConfig( name="gemma4", initial_state=ParserState.CONTENT, + # Decides a silent tail: the E-series HF templates write no marker + # with thinking off, and no template does after a tool response. + wait_for_reasoning=thinking, + turn_boundary_tokens=frozenset({"<|turn>", "<|tool_response>"}), terminals={ "THINK_START": CHANNEL_START, "THINK_END": CHANNEL_END, @@ -417,11 +421,11 @@ class Gemma4Parser(ParserEngine): **kwargs, ) -> None: chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} - self._thinking_enabled = chat_kwargs.get("enable_thinking", False) + thinking = chat_kwargs.get("enable_thinking", False) super().__init__( tokenizer, tools, - parser_engine_config=gemma4_config(), + parser_engine_config=gemma4_config(thinking=thinking), **kwargs, ) vocab = self.vocab @@ -472,30 +476,6 @@ class Gemma4Parser(ParserEngine): return delta_text, delta_token_ids - def is_reasoning_end(self, input_ids: list[int]) -> bool: - end_id = self._reasoning_end_token_id - start_id = self._reasoning_start_token_id - tool_call_id = self._tool_call_token_id - new_turn_id = self._new_turn_token_id - tool_response_id = self._tool_response_token_id - - if end_id is not None and not input_ids: - return self.parser_engine_config.initial_state != ParserState.REASONING - - for i in range(len(input_ids) - 1, -1, -1): - tid = input_ids[i] - if start_id is not None and tid == start_id: - return False - if tool_call_id is not None and tid == tool_call_id: - return True - if new_turn_id is not None and tid == new_turn_id: - return not self._thinking_enabled - if tool_response_id is not None and tid == tool_response_id: - return not self._thinking_enabled - if end_id is not None and tid == end_id: - return True - return True - def _prompt_ends_in_open_reasoning(self, prompt_token_ids: Sequence[int]) -> bool: """Whether the prompt tail is inside an open ``<|channel>`` block. diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 8aa4feef25..4a804a4c9c 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -123,6 +123,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: return ParserEngineConfig( name="glm47_moe", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ **reasoning_terminals, "TOOL_START": TOOL_CALL_START, @@ -206,11 +207,6 @@ class Glm47MoeParser(ParserEngine): self._tool_slots[idx].name = self._tool_slots[idx].name.strip() super()._handle_tool_end(event, deltas) - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if not self.thinking_enabled: - return True - return super().is_reasoning_end(input_ids) - def extract_content_ids(self, input_ids: list[int]) -> list[int]: if not self.thinking_enabled: return input_ids diff --git a/vllm/parser/inkling.py b/vllm/parser/inkling.py index 16883b58aa..2081fb2784 100644 --- a/vllm/parser/inkling.py +++ b/vllm/parser/inkling.py @@ -282,6 +282,7 @@ def inkling_config() -> ParserEngineConfig: # `<|message_model|>`. Non-streaming parsing receives only the generated # suffix, so begin in the corresponding message-header state as well. initial_state=ParserState.MESSAGE_HEADER, + wait_for_reasoning=True, terminals=terminals, # Inkling content-kind markers are the grammar. When the engine is # used through DelegatingParser, the reasoning pass can hand the tool diff --git a/vllm/parser/kimi_k2.py b/vllm/parser/kimi_k2.py index 1e1067a8c2..7c5f30aafe 100644 --- a/vllm/parser/kimi_k2.py +++ b/vllm/parser/kimi_k2.py @@ -81,6 +81,7 @@ def kimi_k2_config(thinking: bool = True) -> ParserEngineConfig: return ParserEngineConfig( name="kimi_k2", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ **reasoning_terminals, "TOOL_SECTION_START": TOOL_SECTION_START, @@ -236,24 +237,6 @@ class KimiK2Parser(ParserEngine): def _extract_args_json(self, raw_args: str, func_name: str) -> str: return raw_args.strip() or "{}" - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if not self.thinking_enabled: - return True - - start_id = self._start_token_id - end_id = self._end_token_id - tool_section_id = self._tool_section_start_token_id - - for i in range(len(input_ids) - 1, -1, -1): - token_id = input_ids[i] - if start_id is not None and token_id == start_id: - return False - if end_id is not None and token_id == end_id: - return True - if tool_section_id is not None and token_id == tool_section_id: - return True - return False - def extract_content_ids(self, input_ids: list[int]) -> list[int]: if not self.thinking_enabled: return input_ids diff --git a/vllm/parser/minimax_m2.py b/vllm/parser/minimax_m2.py index ab1b7eb114..c96f852b9d 100644 --- a/vllm/parser/minimax_m2.py +++ b/vllm/parser/minimax_m2.py @@ -95,6 +95,7 @@ def minimax_m2_config() -> ParserEngineConfig: return ParserEngineConfig( name="minimax_m2", initial_state=ParserState.REASONING, + wait_for_reasoning=True, terminals={ "THINK_START": THINK_START, "THINK_END": THINK_END, diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index fc9b8fbefe..ce516eea14 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -193,6 +193,7 @@ def mistral_config( return ParserEngineConfig( name=name, initial_state=ParserState.CONTENT, + wait_for_reasoning=reasoning_encoding != "none", terminals={ **reasoning_terminals, "TOOL_CALLS": _TOOL_CALLS, @@ -891,24 +892,6 @@ class MistralParser(ParserEngine): return (delta_to_be_parsed, "") - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if self._reasoning_encoding == "none": - return True - if super().is_reasoning_end(input_ids): - return True - # [TOOL_CALLS] acts as an implicit reasoning-end marker - if self.bot_token_id is not None: - reasoning_start_id = self._reasoning_start_token_id - for i in range(len(input_ids) - 1, -1, -1): - if ( - reasoning_start_id is not None - and input_ids[i] == reasoning_start_id - ): - return False - if input_ids[i] == self.bot_token_id: - return True - return False - def extract_reasoning( self, model_output: str, diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py index 8edf42e263..a11551a209 100644 --- a/vllm/parser/qwen3.py +++ b/vllm/parser/qwen3.py @@ -100,6 +100,7 @@ def qwen3_config( return ParserEngineConfig( name=name, initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, turn_boundary_tokens=turn_boundary_tokens, terminals={ # Reasoning terminals @@ -205,8 +206,8 @@ class Qwen3Parser(ParserEngine): """Qwen3 parser: ````/```` reasoning + ```` XML tool calls in a single engine. - - ```` as implicit reasoning end - - Unpaired ```` token ID detection for ``is_reasoning_end`` + - ```` as implicit reasoning end (a grammar transition, so it + also feeds ``is_reasoning_end`` and the structured-output gate) Subclasses that share the grammar but differ only in the four wrapper token strings (reasoning + tool-call) override the class attributes @@ -245,9 +246,6 @@ class Qwen3Parser(ParserEngine): tools, **kwargs, ) - vocab = self.vocab - self._tool_call_token_id: int | None = vocab.get(self.TOOL_START) - self._tool_call_end_token_id: int | None = vocab.get(self.TOOL_END) def extract_reasoning( self, @@ -257,28 +255,3 @@ class Qwen3Parser(ParserEngine): if not self.thinking_enabled: return None, model_output return super().extract_reasoning(model_output, request) - - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if super().is_reasoning_end(input_ids): - return True - tool_call_id = self._tool_call_token_id - tool_call_end_id = self._tool_call_end_token_id - reasoning_start_id = self._reasoning_start_token_id - boundary_ids = self._turn_boundary_token_ids - if tool_call_id is not None: - for i in range(len(input_ids) - 1, -1, -1): - if ( - reasoning_start_id is not None - and input_ids[i] == reasoning_start_id - ): - return False - if input_ids[i] in boundary_ids: - return False - if input_ids[i] == tool_call_id: - if tool_call_end_id is not None and any( - input_ids[j] == tool_call_end_id - for j in range(i + 1, len(input_ids)) - ): - continue - return True - return False From 5be911ee871604895c5a950ac1a5675c71896fd2 Mon Sep 17 00:00:00 2001 From: ALEX LILUZ <49665315+alexliluz@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:31:25 +0800 Subject: [PATCH 096/139] [Bugfix] Load reasoning parser plugins before headless engine config (#53124) Signed-off-by: alexliluz <49665315+alexliluz@users.noreply.github.com> Co-authored-by: Jiangyun Zhu --- tests/entrypoints/cli/test_serve.py | 42 +++++++++++++++++++++++++++++ vllm/entrypoints/cli/serve.py | 4 +++ 2 files changed, 46 insertions(+) create mode 100644 tests/entrypoints/cli/test_serve.py diff --git a/tests/entrypoints/cli/test_serve.py b/tests/entrypoints/cli/test_serve.py new file mode 100644 index 0000000000..9ba9ac7455 --- /dev/null +++ b/tests/entrypoints/cli/test_serve.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the ``vllm serve`` CLI subcommand.""" + +import argparse +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.entrypoints.cli.serve import run_headless + + +class _StopAfterEngineConfig(Exception): + pass + + +def test_headless_imports_reasoning_parser_plugin_before_engine_config(): + plugin_path = "/tmp/custom_reasoning_parser.py" + args = argparse.Namespace( + api_server_count=0, + reasoning_parser_plugin=plugin_path, + ) + engine_args = MagicMock() + + with ( + patch( + "vllm.entrypoints.cli.serve.vllm.AsyncEngineArgs.from_cli_args", + return_value=engine_args, + ), + patch( + "vllm.reasoning.ReasoningParserManager.import_reasoning_parser" + ) as import_plugin, + ): + + def create_engine_config(**kwargs): + import_plugin.assert_called_once_with(plugin_path) + raise _StopAfterEngineConfig + + engine_args.create_engine_config.side_effect = create_engine_config + + with pytest.raises(_StopAfterEngineConfig): + run_headless(args) diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index e43b273e5e..6a8492f4c4 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -18,6 +18,7 @@ from vllm.entrypoints.launchers.cli_args import ( from vllm.entrypoints.launchers.dp_supervisor import run_dp_supervisor from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger +from vllm.reasoning import ReasoningParserManager from vllm.usage.usage_lib import UsageContext from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.network_utils import get_tcp_uri @@ -179,6 +180,9 @@ def run_headless(args: argparse.Namespace): if args.api_server_count > 1: raise ValueError("api_server_count can't be set in headless mode") + if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: + ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + # Create the EngineConfig. engine_args = vllm.AsyncEngineArgs.from_cli_args(args) usage_context = UsageContext.OPENAI_API_SERVER From bb5507741155a698750bc422cda1f7c57f7d114d Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Tue, 15 Sep 2026 04:01:06 -0700 Subject: [PATCH 097/139] [Perf][Kernel] Integrate Mega-mHC from DeepGEMM for DeepSeek V4.1 (reopen of #56255) (#56962) Signed-off-by: Yongye Zhu Signed-off-by: Juntian Liu Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Thien Tran --- tests/kernels/quantization/test_block_fp8.py | 7 +- tests/kernels/test_mhc_kernels.py | 105 +++++++++++- vllm/models/deepseek_v41/nvidia/model.py | 72 ++++---- .../deepseek_v41/nvidia/ops/__init__.py | 2 + .../deepseek_v41/nvidia/ops/mega_mhc.py | 156 ++++++++++++++++++ vllm/utils/deep_gemm.py | 15 +- 6 files changed, 316 insertions(+), 41 deletions(-) create mode 100644 vllm/models/deepseek_v41/nvidia/ops/__init__.py create mode 100644 vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index fd530e7ffb..d02527aba0 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -27,6 +27,7 @@ from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( fp8_gemm_nt, get_tma_aligned_size, + is_deep_gemm_e8m0_used, per_block_cast_to_fp8, should_use_deepgemm_for_fp8_linear, ) @@ -291,7 +292,11 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed): A_fp8, As_fp8 = per_token_group_quant_fp8( A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True ) - B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size) + B_fp8, Bs_fp8 = per_block_cast_to_fp8( + B_fp32, + block_size=block_size, + use_ue8m0=is_deep_gemm_e8m0_used(), + ) As = As_fp8.to(torch.float32) Bs = Bs_fp8.to(torch.float32) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index ff72e2a96c..5fab150d9f 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -11,6 +11,7 @@ import vllm.model_executor.kernels.mhc # noqa: F401 import vllm.model_executor.layers.mhc as mhc_layers from vllm.model_executor.kernels.mhc.tilelang import ( _torch_hc_prenorm_gemm, + mhc_post_tilelang, mhc_pre_delayed_tilelang, ) from vllm.model_executor.kernels.mhc.tilelang_kernels import ( @@ -39,6 +40,10 @@ from vllm.models.deepseek_v4.nvidia.model import ( from vllm.models.deepseek_v41.nvidia.model import ( DeepseekV4DecoderLayer as DeepseekV41DecoderLayer, ) +from vllm.models.deepseek_v41.nvidia.ops.mega_mhc import ( + is_mega_mhc_supported, + mhc_shifted_post_pre_deep_gemm, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -495,7 +500,7 @@ def test_deepseek_v41_decoder_mixes_match_torch( reference, ) monkeypatch.setattr( - "vllm.models.deepseek_v41.nvidia.model.mhc_fused_post_pre_delayed_tilelang", + "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre", fused_reference, ) expected = decoder(x, positions, None, **kwargs) @@ -580,6 +585,104 @@ def test_mhc_pre_delayed_custom_op_supports_compile(carried): ) +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="DeepGEMM Mega mHC requires SM100-family CUDA", +) +@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 1024]) +@pytest.mark.parametrize("hidden_size", [5120, 7168]) +def test_deep_gemm_mega_mhc_correctness(num_tokens, hidden_size): + set_random_seed(0) + hc_mult = 4 + if not is_mega_mhc_supported(hidden_size, 4): + pytest.skip("DeepGEMM Mega mHC is not available") + + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) + residual = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + previous_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) + post_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE) + res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE) + fn = torch.randn(24, hc_mult * hidden_size, device=DEVICE) * 0.01 + scale = torch.randn(3, device=DEVICE) * 0.1 + base = torch.randn(24, device=DEVICE) * 0.1 + norm_weight = torch.empty( + hidden_size, dtype=torch.bfloat16, device=DEVICE + ).uniform_(0.9, 1.1) + + expected_residual = ( + mhc_post_tilelang(x, residual, post_mix, res_mix) if num_tokens else residual + ) + ( + expected_post_mix, + expected_res_mix, + expected_y_bf16, + expected_previous_mix, + ) = mhc_pre_delayed_tilelang( + expected_residual, + fn, + scale, + base, + 2e-5, + 3e-4, + 2e-6, + 1.25, + 10, + pre_mix=previous_mix, + norm_weight=norm_weight, + norm_eps=7e-6, + ) + args = ( + x, + residual, + previous_mix, + post_mix, + res_mix, + fn, + scale, + base, + 2e-5, + 3e-4, + 1.25, + 2e-6, + 10, + norm_weight, + 7e-6, + ) + actual = mhc_shifted_post_pre_deep_gemm(*args) + expected = ( + expected_residual, + expected_post_mix, + expected_res_mix, + expected_y_bf16, + expected_previous_mix, + ) + for i, (result, ref) in enumerate(zip(actual, expected, strict=True)): + # Match the existing mHC BF16 tolerance: post rounding differences + # can be amplified by cancellation in the carried collapse. + atol, rtol = (1.6e-2, 1e-2) if i in (0, 3) else (1e-6, 1e-3) + torch.testing.assert_close(result, ref, atol=atol, rtol=rtol) + + if num_tokens: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + # DeepGEMM initializes barriers per stream before capture. + mhc_shifted_post_pre_deep_gemm(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = mhc_shifted_post_pre_deep_gemm(*args) + for multiplier in (0.5, -1.0): + x.mul_(multiplier) + graph.replay() + eager = mhc_shifted_post_pre_deep_gemm(*args) + for result, ref in zip(captured, eager, strict=True): + torch.testing.assert_close(result, ref, atol=0, rtol=0) + torch.cuda.current_stream().wait_stream(stream) + + def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor: x = x.softmax(-1) + eps x = x / (x.sum(-2, keepdim=True) + eps) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index f3530544da..9c9719e977 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -20,7 +20,6 @@ from vllm.distributed import ( from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( - mhc_fused_post_pre_delayed_tilelang, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -79,6 +78,7 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask from .engram import Engram, gather_engram_hashes +from .ops.mega_mhc import mhc_shifted_post_pre if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -403,25 +403,23 @@ class DeepseekV4DecoderLayer(nn.Module): else: # The collapse already reads the post-mapped streams, so the mean # aux consumers want comes out of the same kernel. - residual, post_mix, res_mix, x, attn_pre, aux = ( - mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=pre_mix, - norm_weight=self.attn_norm.weight, - norm_eps=self.attn_norm.variance_epsilon, - capture_aux=capture_previous_aux, - ) + residual, post_mix, res_mix, x, attn_pre, aux = mhc_shifted_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + norm_weight=self.attn_norm.weight, + norm_eps=self.attn_norm.variance_epsilon, + capture_aux=capture_previous_aux, ) if capture_previous_aux: previous_aux = aux @@ -433,24 +431,22 @@ class DeepseekV4DecoderLayer(nn.Module): if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual, post_mix, res_mix, x, ffn_pre, _ = ( - mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=attn_pre, - norm_weight=self.ffn_norm.weight, - norm_eps=self.ffn_norm.variance_epsilon, - ) + residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=attn_pre, + norm_weight=self.ffn_norm.weight, + norm_eps=self.ffn_norm.variance_epsilon, ) x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix, ffn_pre, previous_aux diff --git a/vllm/models/deepseek_v41/nvidia/ops/__init__.py b/vllm/models/deepseek_v41/nvidia/ops/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py new file mode 100644 index 0000000000..f2c7490261 --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools + +import torch + +from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_fused_post_pre_delayed_tilelang, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import ( + _import_deep_gemm, + is_deep_gemm_supported, + mega_mhc, +) + + +@functools.cache +def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool: + """Return whether DeepGEMM can run the DSV4.1 shifted mHC kernel.""" + if not ( + is_deep_gemm_supported() + and current_platform.is_device_capability_family(100) + and hidden_size > 0 + and hidden_size % 1024 == 0 + and hc_mult == 4 + ): + return False + deep_gemm = _import_deep_gemm() + return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None)) + + +def mhc_shifted_post_pre_deep_gemm( + x: torch.Tensor, + residual: torch.Tensor, + shifted_prev_mix: torch.Tensor, + post_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + mix_scales: torch.Tensor, + mix_bases: torch.Tensor, + hc_norm_eps: float, + hc_pre_eps: float, + hc_post_scale: float, + sinkhorn_eps: float, + num_sinkhorn_iters: int, + rmsnorm_weight: torch.Tensor, + rmsnorm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Run DSV4.1 shifted post, next pre, and BF16 RMSNorm with Mega mHC.""" + num_tokens, hidden_size = x.shape + hc_mult = residual.shape[1] + new_residual = torch.empty_like(residual) + new_prev_mix = x.new_empty(num_tokens, hc_mult, 1, dtype=torch.float32) + new_post_mix = torch.empty_like(post_mix) + new_comb_res_mix = torch.empty_like(comb_res_mix) + y_bf16 = x.new_empty(num_tokens, hidden_size, dtype=torch.bfloat16) + mega_mhc( + x=x, + residual=residual, + shifted_prev_mix=shifted_prev_mix.unsqueeze(-1), + post_mix=post_mix, + comb_res_mix=comb_res_mix, + fn=fn, + mix_scales=mix_scales, + mix_bases=mix_bases, + hc_mult=hc_mult, + hc_norm_eps=hc_norm_eps, + hc_pre_eps=hc_pre_eps, + hc_post_scale=hc_post_scale, + sinkhorn_eps=sinkhorn_eps, + num_sinkhorn_iters=num_sinkhorn_iters, + rmsnorm_weight=rmsnorm_weight, + rmsnorm_eps=rmsnorm_eps, + rmsnorm_scale=1.0, + new_residual=new_residual, + new_prev_mix=new_prev_mix, + new_post_mix=new_post_mix, + new_comb_res_mix=new_comb_res_mix, + y_bf16=y_bf16, + ) + return ( + new_residual, + new_post_mix, + new_comb_res_mix, + y_bf16, + new_prev_mix.squeeze(-1), + ) + + +def mhc_shifted_post_pre( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Use Mega-mHC for carried mixing, retaining TileLang for aux capture.""" + if ( + pre_mix is not None + and norm_weight is not None + and not capture_aux + and x.shape[0] <= 1 << 20 + and is_mega_mhc_supported(x.shape[1], residual.shape[1]) + ): + outputs = mhc_shifted_post_pre_deep_gemm( + x, + residual, + pre_mix, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_post_mult_value, + hc_sinkhorn_eps, + sinkhorn_repeat, + norm_weight, + norm_eps, + ) + return *outputs, x.new_empty(0, x.shape[1]) + + return mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + capture_aux=capture_aux, + ) diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index d515d5131e..5d05c29a84 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -159,6 +159,7 @@ _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None +_mega_mhc_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = ( @@ -231,7 +232,7 @@ def _lazy_init() -> None: global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl global _get_paged_mqa_logits_metadata_impl - global _tf32_hc_prenorm_gemm_impl + global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl global _get_theoretical_mk_alignment_for_contiguous_layout_impl @@ -251,6 +252,7 @@ def _lazy_init() -> None: or _fp8_fp4_paged_mqa_logits_impl is not None or _get_paged_mqa_logits_metadata_impl is not None or _tf32_hc_prenorm_gemm_impl is not None + or _mega_mhc_impl is not None or _get_mk_alignment_for_contiguous_layout_impl is not None or _transform_sf_into_required_layout_impl is not None or _pack_ue8m0_to_int_impl is not None @@ -294,6 +296,7 @@ def _lazy_init() -> None: _dg, "get_paged_mqa_logits_metadata", None ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) + _mega_mhc_impl = getattr(_dg, "mega_mhc", None) _get_mn_major_tma_aligned_tensor_impl = getattr( _dg, "get_mn_major_tma_aligned_tensor", None ) @@ -704,6 +707,15 @@ def tf32_hc_prenorm_gemm( ) +def mega_mhc(*args: Any, **kwargs: Any) -> None: + """Run DeepGEMM Mega mHC with caller-owned output tensors.""" + _lazy_init() + if _mega_mhc_impl is None: + _missing() + return + _mega_mhc_impl(*args, **kwargs) + + def _ceil_to_ue8m0(x: torch.Tensor): return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) @@ -797,6 +809,7 @@ __all__ = [ "per_block_cast_to_fp8", "is_deep_gemm_e8m0_used", "is_deep_gemm_supported", + "mega_mhc", "get_num_sms", "set_num_sms", "should_use_deepgemm_for_fp8_linear", From 78eaaad2c7a7f8c7864d9c9c72772e3f9f2669d0 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 15 Sep 2026 20:12:15 +0800 Subject: [PATCH 098/139] [Rust Frontend] Add iteration token histogram (#56990) Co-authored-by: Codex Signed-off-by: Bugen Zhao --- rust/src/engine-core-client/src/client/imp.rs | 26 +++++++++++++++++- rust/src/engine-core-client/src/metrics.rs | 21 +++++++++++++-- rust/src/metrics/src/request.rs | 27 ++++++++++++++----- rust/src/mock-engine/src/engine.rs | 19 +++++++++---- rust/src/mock-engine/src/tests.rs | 11 ++++++++ 5 files changed, 90 insertions(+), 14 deletions(-) diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 117473ce47..7318fa01b8 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -19,7 +19,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; -use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder}; +use crate::metrics::{IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder}; use crate::protocol::encode_msgpack; use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs}; use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; @@ -47,6 +47,7 @@ pub(crate) struct ClientInner { /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created. msgpack_zero_copy_threshold: usize, scheduler_stats_recorder: SchedulerStatsRecorder, + iteration_metrics: BTreeMap, request_reg: Mutex, utility_reg: Mutex, health_error: ArcSwapOption, @@ -64,12 +65,21 @@ impl ClientInner { ) -> Self { let scheduler_stats_recorder = SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines); + let iteration_metrics = engines + .iter() + .filter_map(|engine| { + let engine = engine.engine_id.engine_index()?; + let handles = IterationMetricHandles::new(&METRICS.request, &model_name, engine); + Some((engine, handles)) + }) + .collect(); Self { input_send, handle, model_name, msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(), scheduler_stats_recorder, + iteration_metrics, request_reg: Mutex::new(RequestRegistry::new(engines)), utility_reg: Mutex::new(UtilityRegistry::default()), health_error: ArcSwapOption::empty(), @@ -415,6 +425,8 @@ pub(crate) async fn run_output_dispatcher_loop( match outputs { EngineCoreOutputs::RequestBatch(batch) => { + let has_outputs = !batch.outputs.is_empty(); + let mut iteration_tokens = 0_u64; let senders = inner.take_senders_for_outputs(&batch.outputs); for (output, sender) in batch.outputs.into_iter().zip(senders) { let request_id = output.request_id.clone(); @@ -423,6 +435,12 @@ pub(crate) async fn run_output_dispatcher_loop( continue; }; + iteration_tokens += output.new_token_ids.len() as u64; + if let Some(prefill) = &output.prefill_stats { + // The engine emits prefill_stats once, on the first output. + iteration_tokens += u64::from(prefill.num_computed_tokens); + } + let wrapped_output = EngineCoreStreamOutput { engine_index: batch.engine_index, timestamp: batch.timestamp, @@ -433,6 +451,12 @@ pub(crate) async fn run_output_dispatcher_loop( } } + if has_outputs + && let Some(handles) = inner.iteration_metrics.get(&batch.engine_index) + { + handles.iteration_tokens_total.observe(iteration_tokens as f64); + } + // The sender for normally-finished requests should have already been removed // from the registry when their final output was dispatched // above. This serves as a safety net to capture any diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index e15258d608..b392d52601 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -8,8 +8,8 @@ use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, - MooncakeOperationLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, - WaitingReasonLabels, + MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics, + U64Counter, U64Gauge, WaitingReasonLabels, }; use crate::protocol::stats::{ @@ -20,6 +20,23 @@ use crate::transport::ConnectedEngine; const WAITING_REASON_CAPACITY: &str = "capacity"; const WAITING_REASON_DEFERRED: &str = "deferred"; +/// Cached output-batch metric handles for one model and engine index. +pub(crate) struct IterationMetricHandles { + pub iteration_tokens_total: HistogramMetric, +} + +impl IterationMetricHandles { + pub(crate) fn new(metrics: &RequestMetrics, model_name: &str, engine: u32) -> Self { + let labels = EngineLabels { + model_name: model_name.to_string(), + engine, + }; + Self { + iteration_tokens_total: metrics.iteration_tokens_total.get_or_create_owned(&labels), + } + } +} + /// Cached scheduler-stats metric handles for all engines connected to one /// frontend client. pub(crate) struct SchedulerStatsRecorder { diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index c1135781b2..691c789936 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -21,6 +21,9 @@ const REQUEST_LATENCY_BUCKETS: [f64; 21] = [ 480.0, 960.0, 1920.0, 7680.0, ]; const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0]; +const ITERATION_TOKENS_BUCKETS: [f64; 13] = [ + 1.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, 16384.0, +]; fn build_1_2_5_buckets(max_value: u32) -> Vec { let mut buckets = Vec::new(); @@ -66,6 +69,10 @@ fn request_params_n_histogram() -> Histogram { Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied()) } +fn iteration_tokens_histogram() -> Histogram { + Histogram::new(ITERATION_TOKENS_BUCKETS.iter().copied()) +} + #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] pub struct FinishedReasonLabels { pub model_name: String, @@ -83,8 +90,12 @@ pub struct PromptTokenSourceLabels { pub(crate) type FinishedReasonCounterFamily = Family; pub(crate) type PromptTokenSourceCounterFamily = Family; -/// Request-lifecycle Prometheus families exported from the `llm` layer. +/// Request and output-batch Prometheus families. pub struct RequestMetrics { + /// Computed prompt and generated tokens per nonempty engine output batch, + /// recorded by the engine-core client for requests tracked by this frontend. + pub iteration_tokens_total: HistogramFamily, + // Request-derived counters. pub num_preemptions: Family, pub prompt_tokens: Family, @@ -92,11 +103,6 @@ pub struct RequestMetrics { pub prompt_tokens_cached: Family, pub generation_tokens: Family, - // We intentionally don't support iteration-level histograms for now, since it seems to make - // more sense if the engine maintains these metrics and frontend simply forwards. - // - // pub iteration_tokens_total: HistogramFamily, - // Request lifecycle counters and histograms. pub request_success: FinishedReasonCounterFamily, pub request_prompt_tokens: HistogramFamily, @@ -118,6 +124,14 @@ pub struct RequestMetrics { impl RequestMetrics { /// Register the request-oriented metric families into the shared registry. pub(crate) fn register(registry: &mut Registry) -> Self { + let iteration_tokens_total = + Family::new_with_constructor(iteration_tokens_histogram as fn() -> Histogram); + registry.register( + "vllm:iteration_tokens_total", + "Histogram of computed prompt and generated tokens per nonempty engine output batch.", + iteration_tokens_total.clone(), + ); + // Request-derived counters. let num_preemptions = Family::default(); registry.register( @@ -276,6 +290,7 @@ impl RequestMetrics { ); Self { + iteration_tokens_total, num_preemptions, prompt_tokens, prompt_tokens_by_source, diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index c216829159..38f9b21be1 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/rust/src/mock-engine/src/engine.rs @@ -19,6 +19,7 @@ use vllm_engine_core_client::protocol::output::{ UtilityCallOutput, }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; +use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::protocol::utility::{ EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope, }; @@ -187,6 +188,11 @@ impl ActiveRequest { /// Advance this request by one mock engine step. fn step(&mut self, opt: &Opt) -> EngineCoreOutput { + let prefill_stats = (self.generated == 0).then(|| PrefillStats { + num_prompt_tokens: self.prompt_len as u32, + num_computed_tokens: self.prompt_len as u32, + ..Default::default() + }); let remaining = self.max_tokens - self.generated; let chunk_len = remaining.min(opt.output_token_chunk_size); let mut new_token_ids = Vec::with_capacity(chunk_len); @@ -196,11 +202,14 @@ impl ActiveRequest { self.generated += chunk_len; let finished = self.generated >= self.max_tokens; - request_output( - self.request_id.clone(), - new_token_ids, - finished.then_some(EngineCoreFinishReason::Length), - ) + EngineCoreOutput { + prefill_stats, + ..request_output( + self.request_id.clone(), + new_token_ids, + finished.then_some(EngineCoreFinishReason::Length), + ) + } } } diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index 9bfee620ec..bee710fd0c 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -10,6 +10,7 @@ use tokio_util::sync::CancellationToken; use vllm_engine_core_client::protocol::output::EngineCoreFinishReason; use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::test_utils::IpcNamespace; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode}; @@ -119,12 +120,22 @@ async fn chunk_size_one_outputs_one_token_per_update() { let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 1); assert_eq!(first.finish_reason, None); + assert_eq!( + first.prefill_stats, + Some(PrefillStats { + num_prompt_tokens: 3, + num_computed_tokens: 3, + ..Default::default() + }) + ); let second = stream.next().await.expect("second").expect("second ok"); assert_eq!(second.new_token_ids.len(), 1); assert_eq!(second.finish_reason, None); + assert!(second.prefill_stats.is_none()); let third = stream.next().await.expect("third").expect("third ok"); assert_eq!(third.new_token_ids.len(), 1); assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length)); + assert!(third.prefill_stats.is_none()); assert!(stream.next().await.is_none()); shutdown_mock(client, shutdown, task).await; From 64856080b05054fc62754f37234f7b483b753445 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 15 Sep 2026 08:21:05 -0500 Subject: [PATCH 099/139] [Bugfix] Measure complete pooling responses (#56760) Signed-off-by: Taneem Ibrahim --- vllm/benchmarks/lib/endpoint_request_func.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index b55c5f4e16..54c6ce63b0 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -606,6 +606,7 @@ async def _run_pooling_request( headers: dict[str, Any], pbar: tqdm | None = None, num_input_sequences: int = 1, + prompt_len: int = 0, ) -> RequestFuncOutput: output = RequestFuncOutput(num_input_sequences=num_input_sequences) st = time.perf_counter() @@ -613,11 +614,19 @@ async def _run_pooling_request( try: async with session.post(url=api_url, headers=headers, json=payload) as response: if response.status == 200: + encoding_format = payload.get("encoding_format", "float") + if encoding_format in ("bytes", "bytes_only"): + async for _ in response.content.iter_any(): + pass + else: + await response.read() output.ttft = output.latency = time.perf_counter() - st - if payload.get("encoding_format", "float") == "bytes": + if encoding_format == "bytes": metadata = json.loads(response.headers["metadata"]) usage = metadata.get("usage", {}) + elif encoding_format == "bytes_only": + usage = {"prompt_tokens": prompt_len} else: data = await response.json() usage = data.get("usage", {}) @@ -672,6 +681,7 @@ async def async_request_openai_embeddings( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) @@ -709,6 +719,7 @@ async def async_request_vllm_rerank( headers=headers, pbar=pbar, num_input_sequences=len(request_func_input.prompt) - 1, + prompt_len=request_func_input.prompt_len, ) @@ -743,6 +754,7 @@ async def async_request_openai_embeddings_chat( payload=payload, headers=headers, pbar=pbar, + prompt_len=request_func_input.prompt_len, ) @@ -847,6 +859,7 @@ async def async_request_infinity_embeddings( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) @@ -896,6 +909,7 @@ async def async_request_vllm_pooling( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) From e6960af33b379d502f409e3e2241bbf2b2c2f68d Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Tue, 15 Sep 2026 09:48:49 -0400 Subject: [PATCH 100/139] [Refactor] StructuredOutputManager x Speculative Decoding Refactor (#48200) --- tests/parser/engine/test_deepseek_v4.py | 2 +- tests/parser/engine/test_parser_engine.py | 4 +- tests/v1/core/test_async_scheduler.py | 6 +- tests/v1/core/test_scheduler.py | 28 +- .../spec_decode/test_mtp_structured_output.py | 359 +---------- .../test_reasoning_structured_output.py | 557 ------------------ .../test_structured_output_manager.py | 414 +++++++++++++ vllm/parser/abstract_parser.py | 8 +- vllm/parser/engine/parser_engine.py | 2 +- vllm/v1/core/sched/scheduler.py | 53 +- vllm/v1/structured_output/__init__.py | 349 +++++------ vllm/v1/structured_output/request.py | 6 - 12 files changed, 596 insertions(+), 1192 deletions(-) delete mode 100644 tests/v1/structured_output/test_reasoning_structured_output.py create mode 100644 tests/v1/structured_output/test_structured_output_manager.py diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 1c1029b51a..7cbf9791a7 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -574,7 +574,7 @@ class TestImplicitReasoningEnd: # the text lexer. assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID} assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1 - assert thinking_parser.find_reasoning_end_offset([7, 8]) is None + assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2 def test_streaming_reasoning_implicit_end(self, thinking_parser): chunks = [ diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index f27ce52cb3..8fd7c4d262 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -185,8 +185,8 @@ class TestReasoningEndTokenIds: def test_find_reasoning_end_offset_returns_first_match(self): engine = _make_engine() assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1 - assert engine.find_reasoning_end_offset([5, 6]) is None - assert engine.find_reasoning_end_offset([]) is None + assert engine.find_reasoning_end_offset([5, 6]) == 2 + assert engine.find_reasoning_end_offset([]) == 0 # Rejected-draft placeholders never match. assert engine.find_reasoning_end_offset([-1, 201]) == 1 diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 7ff63d019c..7e27a4517e 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -334,7 +334,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): request = create_requests(num_requests=1, num_tokens=1)[0] request.structured_output_request = Mock() request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) - request.structured_output_request.grammar.accept_tokens.return_value = False request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens request.num_output_placeholders = 1 @@ -343,10 +342,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.connector = None scheduler.ec_connector = None scheduler.structured_output_manager = Mock() - scheduler.structured_output_manager.should_advance.return_value = True - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( - lambda request, new_token_ids: new_token_ids - ) + scheduler.structured_output_manager.accept_tokens.return_value = False scheduler.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 6ca6a8cf87..061c825fc8 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -3674,7 +3674,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): ) request.structured_output_request = Mock() request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) - request.structured_output_request.grammar.accept_tokens.return_value = False request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens @@ -3682,10 +3681,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.connector = None scheduler.ec_connector = None scheduler.structured_output_manager = Mock() - scheduler.structured_output_manager.should_advance.return_value = True - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( - lambda request, new_token_ids: new_token_ids - ) + scheduler.structured_output_manager.accept_tokens.return_value = False scheduler.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() @@ -3734,8 +3730,8 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): ) engine_core_outputs = scheduler.update_from_output(output, model_runner_output) - request.structured_output_request.grammar.accept_tokens.assert_called_once_with( - request.request_id, [123] + scheduler.structured_output_manager.accept_tokens.assert_called_once_with( + request, [123] ) assert request.resumable is False assert request.status == RequestStatus.FINISHED_ERROR @@ -6573,16 +6569,15 @@ def _decode_ready_request(scheduler): return request -def test_update_draft_token_ids_strips_ngram_padding(monkeypatch): - """ngram_gpu pads unfilled draft slots with -1; update_draft_token_ids must +def test_update_draft_token_ids_strips_ngram_padding(): + """ngram_gpu pads unfilled draft slots with -1; manager validation must strip them before grammar.validate_tokens (which otherwise raises).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) grammar = _RecordingGrammar() - request.structured_output_request = SimpleNamespace(grammar=grammar) - monkeypatch.setattr( - scheduler.structured_output_manager, "should_advance", lambda req: True + request.structured_output_request = SimpleNamespace( + grammar=grammar, reasoning_ended=True ) scheduler.update_draft_token_ids( @@ -6593,16 +6588,15 @@ def test_update_draft_token_ids_strips_ngram_padding(monkeypatch): assert request.spec_token_ids == [10, 11] -def test_update_draft_token_ids_in_output_strips_padding(monkeypatch): +def test_update_draft_token_ids_in_output_strips_padding(): """Same guard on the output path; the -1 pad-back for the rejected count - is preserved (only the input to validate_tokens is stripped).""" + is preserved (only the input to manager validation is stripped).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) grammar = _RecordingGrammar() - request.structured_output_request = SimpleNamespace(grammar=grammar) - monkeypatch.setattr( - scheduler.structured_output_manager, "should_advance", lambda req: True + request.structured_output_request = SimpleNamespace( + grammar=grammar, reasoning_ended=True ) scheduler_output = SimpleNamespace( diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 330e10438f..4e12a6f49d 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""grammar_bitmask under spec-decode draft padding (#44006).""" +"""Backend and engine-reasoner checks that are not manager-flow tests.""" import pytest from transformers import AutoTokenizer @@ -49,221 +49,6 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'): return tokenizer, manager, request, prompt -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_with_padded_invalid_drafts(backend): - """Bitmask handles -1 padded drafts and returns N+1 rows.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - valid_drafts = [tokenizer.encode(":")[0], tokenizer.encode(' "')[0]] - padded = valid_drafts + [-1, -1] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: padded}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(padded) + 1 - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_when_grammar_terminates_mid_window(backend): - """Drafts following an EOS that terminates the grammar are a no-op.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - eos = tokenizer.eos_token_id - drafts = [eos] + [tokenizer.encode(" ")[0]] * (NUM_SPEC_TOKENS - 1) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1 - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_idempotent_across_calls(backend): - """Repeated calls with the same input return the same bitmask.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - drafts = [tokenizer.encode(":")[0], -1, -1, -1] - - first = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - second = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert first is not None and second is not None - assert (first == second).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bonus_position_constrained_after_invalid_drafts(backend): - """Regression for #44006: bonus row stays constrained after -1 padding.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - valid = tokenizer.encode(":")[0] - drafts = [valid, -1, -1, -1] - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_constrained_when_reasoning_ends_midwindow(backend): - """Drafts after a mid-window reasoning-end marker stay constrained.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - marker = tokenizer.encode("\n")[0] - - class StubReasoner: - def __init__(self, *_, **__): - self.end_token_id = marker - - def is_reasoning_end(self, input_ids): - return marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - manager.reasoner_cls = StubReasoner - request.structured_output_request.reasoner = StubReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode(",")[0] - drafts = [pre, marker, post] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - assert (bitmask[0] == -1).all() - assert (bitmask[1] == -1).all() - assert not (bitmask[2] == -1).all() - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_post_reasoning_end_drafts_skip_grammar_advance(backend, caplog): - """Post-marker drafts predate the bitmask and may be grammar-invalid; - grammar_bitmask must skip the grammar advance instead of asserting. - """ - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str="{" - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - assert not grammar.is_terminated() - - marker = tokenizer.encode("\n")[0] - - class StubReasoner: - def __init__(self, *_, **__): - self.end_token_id = marker - - def is_reasoning_end(self, input_ids): - return marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - manager.reasoner_cls = StubReasoner - request.structured_output_request.reasoner = StubReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - # A token that the JSON grammar would reject as the first post-marker - # token; without the fix grammar.accept_tokens fires the assertion. - invalid_post = tokenizer.encode("z")[0] - drafts = [pre, marker, invalid_post] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - # Post-marker position is still bitmask-constrained. - assert not (bitmask[2] == -1).all() - # Grammar must not have advanced through the unvalidated draft. - assert not grammar.is_terminated() - assert "Failed to advance FSM" not in caplog.text - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_validate_tokens_then_bitmask_round_trip(backend): - """validate_tokens -> pad with -1 -> grammar_bitmask must not assert.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - raw_drafts = [tokenizer.encode(",")[0], 99999, 12345, 67890] - accepted = grammar.validate_tokens(raw_drafts) - assert len(accepted) <= len(raw_drafts) - - padded = accepted + [-1] * (len(raw_drafts) - len(accepted)) - assert len(padded) == len(raw_drafts) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: padded}, - ) - assert bitmask is not None - assert bitmask.shape[0] == len(padded) + 1 - assert not grammar.is_terminated() - - def test_xgrammar_accept_tokens_stops_at_termination(capfd): """Tokens after a terminating EOS do not reach the matcher.""" tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") @@ -302,7 +87,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert grammar.validate_tokens([eos, trailing]) == [eos] assert "trying to accept new token" not in capfd.readouterr().err - # Check matcher state directly to verify validation rolled it back. assert not grammar.matcher.is_terminated() assert grammar.accept_tokens(request.request_id, [eos]) @@ -312,87 +96,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert "trying to accept new token" not in capfd.readouterr().err -class _MarkerReasoner: - """Stub reasoner whose reasoning-end marker is a single fixed token.""" - - def __init__(self, marker: int): - self.marker = marker - - def is_reasoning_end(self, input_ids): - return self.marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return self.marker in list(delta_ids) - - -def _setup_boundary_request(backend: str): - """Request with a structural-tag key and reasoning not yet ended.""" - from vllm.v1.structured_output.backend_types import StructuredOutputOptions - - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - marker = tokenizer.encode("\n")[0] - structured_req = request.structured_output_request - # The grammar itself is JSON (cheap to build); only the key kind matters - # for the should_advance structural-tag branch, so pre-seed the cached - # property. - structured_req.__dict__["structured_output_key"] = ( - StructuredOutputOptions.STRUCTURAL_TAG, - "", - ) - manager.reasoner_cls = _MarkerReasoner - structured_req.reasoner = _MarkerReasoner(marker) - structured_req.reasoning_ended = False - return tokenizer, manager, request, prompt, marker - - -def test_should_advance_records_reasoning_end_index(): - """Regression for #44006 on post-#42452 main: the boundary step must - record where reasoning ends so the scheduler can trim before advancing. - """ - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar") - structured_req = request.structured_output_request - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode("{")[0] - request.append_output_token_ids([pre, marker, post]) - - assert manager.should_advance(request) - assert structured_req.reasoning_ended - # Marker sits at absolute index len(prompt) + 1. - assert structured_req.reasoning_end_token_index == len(prompt) + 1 - - -def test_trim_reasoning_for_advance(): - """trim drops the marker and everything before it; later steps and - requests without a recorded boundary pass through unchanged. - """ - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar") - structured_req = request.structured_output_request - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode("{")[0] - - # No boundary recorded yet: pass-through. - assert manager.trim_reasoning_for_advance(request, [pre]) == [pre] - - # Boundary step: marker mid-step keeps only the suffix. - step_tokens = [pre, marker, post] - request.append_output_token_ids(step_tokens) - assert manager.should_advance(request) - assert manager.trim_reasoning_for_advance(request, step_tokens) == [post] - - # Boundary step variant: marker last (the #44006 crash shape - # [198, ]) trims to empty -> scheduler skips accept_tokens. - structured_req.reasoning_end_token_index = len(request.all_token_ids) - 1 - assert manager.trim_reasoning_for_advance(request, step_tokens) == [] - - # Later steps: tokens are past the boundary, returned unchanged. - structured_req.reasoning_end_token_index = len(prompt) + 1 - next_step = [post, post] - request.append_output_token_ids(next_step) - assert manager.trim_reasoning_for_advance(request, next_step) == next_step - - class _EngineReasonerStub(ParserEngineReasoningAdapter): """Adapter-typed reasoner with a fixed end-token set and no real engine.""" @@ -409,7 +112,7 @@ class _EngineReasonerStub(ParserEngineReasoningAdapter): for offset, token in enumerate(token_ids): if token in self._end_token_ids: return offset - return None + return len(token_ids) def is_reasoning_end(self, input_ids): return any(token in self._end_token_ids for token in input_ids) @@ -449,61 +152,3 @@ def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend): assert not (bitmask[2] == -1).all() assert reasoner.windows == [[pre, marker, post]] assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_legacy_order_sensitive_predicate_probed_per_token(backend): - """Legacy predicates are probed one draft at a time, never on the whole - window. A predicate that reports which marker is *newest* (KimiK3) would - otherwise answer False for a window that closes and then reopens - reasoning, leaving the post-marker drafts and the bonus row - unconstrained. - """ - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - close = tokenizer.encode("\n")[0] - reopen = tokenizer.encode("#")[0] - assert close != reopen - - class NewestMarkerReasoner: - def __init__(self, *_, **__): - pass - - def is_reasoning_end(self, input_ids): - return self.is_reasoning_end_streaming(input_ids, input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - for token in reversed(list(delta_ids)): - if token == close: - return True - if token == reopen: - return False - return False - - manager.reasoner_cls = NewestMarkerReasoner - request.structured_output_request.reasoner = NewestMarkerReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode(",")[0] - drafts = [pre, close, post, reopen] - # Sanity: the whole window says "reopened", so only per-token probing - # can find the close marker. - assert not NewestMarkerReasoner().is_reasoning_end_streaming([], drafts) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - assert (bitmask[0] == -1).all() - assert (bitmask[1] == -1).all() - assert not (bitmask[2] == -1).all() - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py deleted file mode 100644 index a57e5da72b..0000000000 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ /dev/null @@ -1,557 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Unit tests for reasoning-aware structured output functionality (PR #25515).""" - -from collections.abc import Sequence -from unittest.mock import Mock - -import pytest - -from vllm.config import ModelConfig, SchedulerConfig, VllmConfig -from vllm.v1.request import Request -from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.structured_output.backend_types import StructuredOutputOptions - - -class UntouchableTokenIds(Sequence[int]): - """A token history that fails the test if anything reads it.""" - - def __init__(self, length: int) -> None: - self._length = length - - def __len__(self) -> int: - return self._length - - def __getitem__(self, item): - raise AssertionError("token history must not be read") - - def __iter__(self): - raise AssertionError("token history must not be iterated") - - -def make_qwen3_reasoner(): - """A real engine-based reasoning adapter over a tiny mock vocabulary.""" - from tests.parser.engine.conftest import make_mock_tokenizer - from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter - - tokenizer = make_mock_tokenizer( - {"": 200, "": 201, "": 202, "": 203} - ) - return Qwen3ParserReasoningAdapter(tokenizer) - - -class MockReasoner: - def __init__(self, tokenizer): - self.is_reasoning_end = Mock(return_value=False) - self.is_reasoning_end_streaming = Mock(return_value=False) - - -class TestReasoningStructuredOutput: - """Test reasoning-aware structured output functionality.""" - - @pytest.fixture - def mock_model_config(self): - """Create a mock ModelConfig.""" - config = Mock(spec=ModelConfig) - config.skip_tokenizer_init = True # Skip tokenizer init to avoid network calls - config.get_vocab_size = Mock(return_value=50000) - # Add missing runner_type attribute that tokenizer initialization expects - config.runner_type = "generate" - # Add other attributes that tokenizer initialization might need - config.tokenizer = "test-tokenizer" - config.tokenizer_mode = "auto" - config.trust_remote_code = False - config.tokenizer_revision = None - return config - - @pytest.fixture - def mock_scheduler_config(self): - """Create a mock SchedulerConfig.""" - config = Mock(spec=SchedulerConfig) - config.max_num_seqs = 128 - return config - - @pytest.fixture - def mock_vllm_config(self, mock_model_config, mock_scheduler_config): - """Create a mock VllmConfig.""" - config = Mock(spec=VllmConfig) - config.model_config = mock_model_config - config.scheduler_config = mock_scheduler_config - config.structured_outputs_config = Mock() - config.structured_outputs_config.reasoning_parser = None - config.structured_outputs_config.enable_in_reasoning = False - config.speculative_config = None - return config - - @pytest.fixture - def mock_request_with_structured_output(self): - """Create a mock request with structured output.""" - request = Mock(spec=Request) - request.structured_output_request = Mock() - request.structured_output_request.reasoning_ended = None - request.structured_output_request.grammar = Mock() - request.structured_output_request.reasoning_parser_kwargs = None - request.structured_output_request.reasoner = None - request.structured_output_request.grammar.is_terminated = Mock( - return_value=False - ) - request.use_structured_output = True - request.prompt_token_ids = [1, 2, 3, 4, 5] - request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8] - request.num_computed_tokens = 5 - request.num_output_placeholders = 0 - request.request_id = "mock_req" - return request - - @pytest.fixture - def manager_with_reasoner(self, mock_vllm_config): - manager = StructuredOutputManager(mock_vllm_config) - manager.reasoner_cls = MockReasoner - manager.tokenizer = Mock() - return manager - - def test_should_fill_bitmask_with_enable_in_reasoning( - self, manager_with_reasoner, mock_request_with_structured_output - ): - """Test should_fill_bitmask when enable_in_reasoning is True.""" - manager_with_reasoner.enable_in_reasoning = True - - # Should always return True when enable_in_reasoning is enabled - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_fill_bitmask_reasoning_already_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """An active grammar does not need a request-local reasoner.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = True - - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - - assert result is True - assert structured_req.reasoner is None - - def test_grammar_bitmask_skips_reasoner_when_already_active( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Bitmask generation skips reasoning-boundary detection.""" - manager_with_reasoner.vllm_config.num_speculative_tokens = 1 - manager_with_reasoner._grammar_bitmask = Mock() - manager_with_reasoner._grammar_bitmask.shape = (1, 1) - expected_bitmask = Mock() - manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask - manager_with_reasoner._fill_bitmasks = Mock() - - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = True - request_id = mock_request_with_structured_output.request_id - - bitmask = manager_with_reasoner.grammar_bitmask( - requests={request_id: mock_request_with_structured_output}, - structured_output_request_ids=[request_id], - scheduled_spec_decode_tokens={}, - ) - - assert bitmask is expected_bitmask - assert structured_req.reasoner is None - - def test_should_fill_bitmask_without_enable_in_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_fill_bitmask when enable_in_reasoning is False.""" - # Keep enable_in_reasoning as False (default) - config = manager_with_reasoner.vllm_config.structured_outputs_config - assert config.enable_in_reasoning is False - - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - - # Should set reasoning_ended and return its value - assert ( - mock_request_with_structured_output.structured_output_request.reasoning_ended - is False - ) - assert result is False - - def test_should_fill_bitmask_no_reasoner( - self, mock_vllm_config, mock_request_with_structured_output - ): - """Test should_fill_bitmask when no reasoner is configured.""" - manager = StructuredOutputManager(mock_vllm_config) - - result = manager.should_fill_bitmask(mock_request_with_structured_output) - - # Should default to True when no reasoner - assert result is True - - def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs( - self, mock_vllm_config, mock_request_with_structured_output - ): - """Test request-level parser kwargs override the default reasoner.""" - - class KwargReasoner: - def __init__(self, tokenizer, chat_template_kwargs=None): - self.chat_template_kwargs = chat_template_kwargs or {} - - def is_reasoning_end(self, input_ids): - return not self.chat_template_kwargs.get("enable_thinking", False) - - manager = StructuredOutputManager(mock_vllm_config) - manager.reasoner_cls = KwargReasoner - manager.tokenizer = Mock() - - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_parser_kwargs = { - "chat_template_kwargs": {"enable_thinking": True} - } - - result = manager.should_fill_bitmask(mock_request_with_structured_output) - - assert result is False - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is not None - ) - - def test_should_advance_with_enable_in_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when enable_in_reasoning is True.""" - # Enable enable_in_reasoning - manager_with_reasoner.enable_in_reasoning = True - - # Should always return True when enable_in_reasoning is enabled - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_advance_reasoning_not_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning has not ended.""" - # Set reasoning as not ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = False - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # Should return False since reasoning hasn't ended - assert result is False - - def test_should_advance_reasoning_just_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning ends in current step.""" - # Set reasoning as not ended initially, but ends in this step - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = True - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoner = reasoner - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # The scheduler trims the reasoning prefix before advancing the grammar. - assert ( - mock_request_with_structured_output.structured_output_request.reasoning_ended - is True - ) - assert result is True - - def test_should_advance_reasoning_already_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning has already ended.""" - # Set reasoning as already ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = True - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # Should return True since reasoning has ended - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_advance_uses_new_token_ids_when_provided( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Regression for #43388: when caller passes new_token_ids, the - reasoner sees the exact multi-token delta rather than the - placeholder-derived window. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - - end_token_id = 248069 - - reasoner = MockReasoner(tokenizer=Mock()) - # Detection mirrors the real Qwen3 parser: end token in the delta. - reasoner.is_reasoning_end_streaming = Mock( - side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids) - ) - structured_req.reasoner = reasoner - - # Scenario from #43388: async + spec decode K=4, 4 tokens accepted - # but only 1 placeholder remains (some drafts were rejected). - # The placeholder math would yield delta=[271] and miss . - # Passing new_token_ids must override that. - new_token_ids = [9, 198, end_token_id, 271] - mock_request_with_structured_output.all_token_ids = [ - 1, - 2, - 3, - 4, - 5, - ] + new_token_ids - mock_request_with_structured_output.num_computed_tokens = 9 - mock_request_with_structured_output.num_output_placeholders = 1 - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output, - new_token_ids=new_token_ids, - ) - - # First call to is_reasoning_end_streaming was with the full - # new_token_ids (not the truncated placeholder window). - first_call = reasoner.is_reasoning_end_streaming.call_args_list[0] - _, called_delta = first_call.args - assert list(called_delta) == new_token_ids - - assert structured_req.reasoning_ended is True - assert result is True - - def test_should_advance_without_new_token_ids_falls_back( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Backward compat: callers that don't pass new_token_ids keep - the original placeholder-derived delta window. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = False - structured_req.reasoner = reasoner - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] - mock_request_with_structured_output.num_computed_tokens = 5 - mock_request_with_structured_output.num_output_placeholders = 2 - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # placeholder window: start = 5 - 2 = 3, delta = [4, 5] - _, called_delta = reasoner.is_reasoning_end_streaming.call_args[0] - assert list(called_delta) == [4, 5] - assert result is False - - def test_should_advance_whole_delta_fallback_pins_last_sequence_index( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """When only the whole-delta predicate fires, the boundary is the - last index of the sequence, even if the placeholder window - overshoots it and the delta is empty. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - # Base-default shape: rescans the whole sequence, ignores the delta. - reasoner.is_reasoning_end_streaming.return_value = True - structured_req.reasoner = reasoner - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] - # placeholder window: start = 9 - 1 = 8, past the end -> delta = [] - mock_request_with_structured_output.num_computed_tokens = 9 - mock_request_with_structured_output.num_output_placeholders = 1 - - assert manager_with_reasoner.should_advance(mock_request_with_structured_output) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 4 - - def test_should_advance_trims_reasoning_prefix_for_json( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """JSON uses the common trim-then-advance path at the boundary.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.structured_output_key = ( - StructuredOutputOptions.JSON_OBJECT, - "{}", - ) - - marker = 248069 - - class MarkerReasoner: - def __init__(self, *_, **__): - pass - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - structured_req.reasoner = MarkerReasoner() - - new_token_ids = [9, 198, marker, 271, 5005] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output, - new_token_ids=new_token_ids, - ) - - structured_req.grammar.accept_tokens.assert_not_called() - assert structured_req.reasoning_ended is True - assert result is True - assert structured_req.reasoning_end_token_index == 5 - assert manager_with_reasoner.trim_reasoning_for_advance( - mock_request_with_structured_output, new_token_ids - ) == [271, 5005] - - def test_should_advance_engine_adapter_scans_only_the_delta( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Engine-based parsers locate the boundary from the delta alone.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = make_qwen3_reasoner() - assert reasoner.reasoning_end_token_ids == {201, 202} - reasoner.is_reasoning_end_streaming = Mock( - side_effect=AssertionError("engine path must not rescan the prefix") - ) - structured_req.reasoner = reasoner - - new_token_ids = [9, 198, 201, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 5 - assert manager_with_reasoner.trim_reasoning_for_advance( - mock_request_with_structured_output, new_token_ids - ) == [271] - - def test_should_advance_engine_adapter_tool_start_ends_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 202, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_end_token_index == 4 - - def test_should_advance_engine_adapter_no_boundary( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 198, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert not manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_ended is False - - def test_should_advance_engine_adapter_never_reads_history( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Per-step cost must not depend on how much has been generated.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 201] - total = (1 << 20) + len(new_token_ids) - mock_request_with_structured_output.all_token_ids = UntouchableTokenIds(total) - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_end_token_index == total - 1 - - def test_should_advance_engine_adapter_without_new_tokens_persists( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Draft-validation callers persist the boundary like the legacy path.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5, 201, 7] - # placeholder window: start = 7 - 2 = 5, delta = [201, 7] - mock_request_with_structured_output.num_computed_tokens = 7 - mock_request_with_structured_output.num_output_placeholders = 2 - - assert manager_with_reasoner.should_advance(mock_request_with_structured_output) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 5 diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py new file mode 100644 index 0000000000..17bda09b3e --- /dev/null +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import time +from dataclasses import dataclass + +import pytest +from transformers import AutoTokenizer + +from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig +from vllm.config.model import ModelConfig +from vllm.config.speculative import SpeculativeConfig +from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.request import Request +from vllm.v1.structured_output import StructuredOutputManager + +TOKENIZER = "gpt2" +THINK_END = "\n" # reasoning-end marker (single GPT-2 token) +EOS = "<|eos|>" # resolved to tokenizer.eos_token_id +JSON_SCHEMA = '{"type": "object"}' +BACKENDS = ("xgrammar", "guidance") +MAX_WAIT_SECONDS = 5 +NUM_SPEC_TOKENS = 8 + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(TOKENIZER) + + +@dataclass(frozen=True) +class FlowCase: + """One validate -> bitmask -> accept -> post-accept bitmask scenario.""" + + raw_drafts: tuple[str, ...] + expected_row_pattern: str + expected_reasoning: bool | None + expect_terminated: bool + prefix: str = "" + expected_validated: tuple[str, ...] | None = None + reasoning_ended: bool | None = False + xfail_guidance: str | None = None + + +class MockReasoner: + def __init__(self, tokenizer, marker: int | None = None): + self.marker = marker + + def is_reasoning_end(self, input_ids): + if self.marker is None: + return True + return self.marker in list(input_ids) + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return self.is_reasoning_end(delta_ids) + + +def _single_token(tokenizer, text: str) -> int: + token_ids = tokenizer.encode(text) + assert len(token_ids) == 1, (text, token_ids) + return token_ids[0] + + +def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]: + """Convert literal text (or EOS sentinel) to token IDs.""" + return [ + tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t) + for t in texts + ] + + +def _wait_for_grammar(request: Request) -> None: + structured_req = request.structured_output_request + assert structured_req is not None + + deadline = time.time() + MAX_WAIT_SECONDS + while not structured_req._check_grammar_completion(): + if time.time() > deadline: + pytest.fail("Grammar compilation timed out") + time.sleep(0.01) + + +def _build_harness( + tokenizer, + backend: str, + prefix: str = "", + use_reasoner: bool = True, + reasoning_ended: bool | None = None, + enable_in_reasoning: bool = False, + reasoning_parser_kwargs: dict | None = None, +) -> tuple[StructuredOutputManager, Request]: + vllm_config = VllmConfig( + model_config=ModelConfig(tokenizer=TOKENIZER), + device_config=DeviceConfig(device="cpu"), + structured_outputs_config=StructuredOutputsConfig( + backend=backend, + enable_in_reasoning=enable_in_reasoning, + ), + speculative_config=SpeculativeConfig( + model="[ngram]", + num_speculative_tokens=NUM_SPEC_TOKENS, + ), + ) + manager = StructuredOutputManager(vllm_config) + if use_reasoner: + manager.reasoner_cls = MockReasoner + + sampling_params = SamplingParams( + structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) + ) + sampling_params.structured_outputs._backend = backend + sampling_params.update_from_generation_config({}, tokenizer.eos_token_id) + + prompt_ids = tokenizer.encode(prefix) if prefix else [] + request = Request( + request_id=f"{backend}-flow", + prompt_token_ids=prompt_ids, + sampling_params=sampling_params, + pooling_params=None, + reasoning_ended=reasoning_ended, + reasoning_parser_kwargs=reasoning_parser_kwargs, + ) + manager.grammar_init(request) + _wait_for_grammar(request) + + structured_req = request.structured_output_request + assert request.prompt_token_ids is not None + assert structured_req is not None + assert structured_req.grammar is not None + if prompt_ids: + assert structured_req.grammar.accept_tokens( + request.request_id, + prompt_ids, + ) + return manager, request + + +def _row_pattern(bitmask) -> str: + return "".join("U" if (row == -1).all() else "C" for row in bitmask) + + +def _run_real_flow( + manager: StructuredOutputManager, + request: Request, + raw_drafts: list[int], + expected_validated: list[int], + expected_row_pattern: str, + expected_reasoning: bool | None, + expect_terminated: bool, +) -> None: + assert len(expected_validated) <= len(raw_drafts) + + validated = manager.validate_tokens(request, list(raw_drafts)) + assert validated == expected_validated + + padded = validated + [-1] * (len(raw_drafts) - len(validated)) + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: padded}, + ) + assert bitmask is not None + assert bitmask.shape[0] == len(raw_drafts) + 1 + assert _row_pattern(bitmask) == expected_row_pattern + + # grammar_bitmask() must rollback any speculative state it advanced. + assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated + + structured_req = request.structured_output_request + assert structured_req is not None + grammar = structured_req.grammar + assert grammar is not None + assert not grammar.is_terminated() + # `bitmask[i]` is the grammar state after the first `i` scheduled tokens. + # These tests commit the validated prefix, so the matching post-accept + # state is already present in the scheduled bitmask. + expected_post_accept_row = bitmask[len(expected_validated)].copy() + + # Mirror the scheduler flow: sampled tokens are appended before + # StructuredOutputManager.accept_tokens() trims them to the + # grammar-constrained suffix. + request.append_output_token_ids(list(expected_validated)) + assert manager.accept_tokens(request, list(expected_validated)) + + post_accept_bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert post_accept_bitmask is not None + assert post_accept_bitmask.shape[0] == 1 + assert (post_accept_bitmask[0] == expected_post_accept_row).all() + assert structured_req.reasoning_ended == expected_reasoning + assert grammar.is_terminated() is expect_terminated + + +FLOW_CASES = [ + pytest.param( + FlowCase( + raw_drafts=(" ", "z", " "), + expected_row_pattern="UUUU", + expected_reasoning=False, + expect_terminated=False, + ), + id="inactive", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", ' "', "b", '"}'), + expected_row_pattern="CCCCC", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_all_valid", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=("z", ":", ' "', "b"), + expected_validated=(), + expected_row_pattern="CUUUU", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_first_token_invalid", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", "z", ' "', "b"), + expected_validated=(":",), + expected_row_pattern="CCUUU", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_later_token_invalid", + ), + pytest.param( + FlowCase( + prefix='{"a": "b"}', + raw_drafts=(EOS, " ", " "), + expected_validated=(EOS,), + expected_row_pattern="CUUU", + expected_reasoning=True, + expect_terminated=True, + reasoning_ended=True, + xfail_guidance=( + "guidance validate_tokens rejects EOS on an already-complete " + "object even though accept_tokens would accept it" + ), + ), + id="active_first_token_terminates", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", ' "', "b", '"}', EOS, " "), + expected_validated=(":", ' "', "b", '"}', EOS), + expected_row_pattern="CCCCCUU", + expected_reasoning=True, + expect_terminated=True, + reasoning_ended=True, + ), + id="active_later_token_terminates", + ), + pytest.param( + FlowCase( + raw_drafts=(THINK_END, "{", ' "', "b"), + expected_row_pattern="UCCCC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_first_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", ' "'), + expected_row_pattern="UUCCC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_middle_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", " ", THINK_END), + expected_row_pattern="UUUC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_last_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "z", "{"), + expected_validated=(" ", THINK_END), + expected_row_pattern="UUCUU", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_first_token_invalid", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", "z"), + expected_validated=(" ", THINK_END, "{"), + expected_row_pattern="UUCCU", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_later_token_invalid", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", "}", EOS, " "), + expected_validated=(" ", THINK_END, "{", "}", EOS), + expected_row_pattern="UUCCCUU", + expected_reasoning=True, + expect_terminated=True, + ), + id="becomes_active_terminates", + ), +] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("case", FLOW_CASES) +def test_real_flow( + tokenizer, + backend: str, + case: FlowCase, +): + if backend == "guidance" and case.xfail_guidance: + pytest.xfail(case.xfail_guidance) + + reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)} + manager, request = _build_harness( + tokenizer, + backend, + prefix=case.prefix, + reasoning_ended=case.reasoning_ended, + reasoning_parser_kwargs=reasoner_kwargs, + ) + + raw_drafts = _to_token_ids(tokenizer, case.raw_drafts) + expected_texts = ( + case.raw_drafts if case.expected_validated is None else case.expected_validated + ) + expected_validated = _to_token_ids(tokenizer, expected_texts) + + _run_real_flow( + manager, + request, + raw_drafts=raw_drafts, + expected_validated=expected_validated, + expected_row_pattern=case.expected_row_pattern, + expected_reasoning=case.expected_reasoning, + expect_terminated=case.expect_terminated, + ) + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + ("use_reasoner", "reasoning_ended", "enable_in_reasoning"), + [ + pytest.param(True, False, True, id="enable_in_reasoning"), + pytest.param(True, True, False, id="reasoning_ended"), + pytest.param(False, False, False, id="no_reasoner"), + pytest.param(True, None, False, id="inferred"), + ], +) +def test_initial_constraint_activation( + tokenizer, + backend: str, + use_reasoner: bool, + reasoning_ended: bool | None, + enable_in_reasoning: bool, +): + manager, request = _build_harness( + tokenizer, + backend, + use_reasoner=use_reasoner, + reasoning_ended=reasoning_ended, + enable_in_reasoning=enable_in_reasoning, + reasoning_parser_kwargs={ + "marker": _single_token(tokenizer, THINK_END) + if reasoning_ended is not None + else None + }, + ) + + # "{" is valid JSON start; "z" is not. Truncation proves grammar is active. + open_brace = _single_token(tokenizer, "{") + z = _single_token(tokenizer, "z") + assert manager.validate_tokens(request, [open_brace, z]) == [open_brace] + + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert bitmask is not None + request.append_output_token_ids([open_brace]) + assert manager.accept_tokens(request, [open_brace]) + + # Short-circuit paths must not construct a request-local reasoner (#52573). + if enable_in_reasoning or reasoning_ended is True: + structured_req = request.structured_output_request + assert structured_req is not None + assert structured_req.reasoner is None diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index cf83bb996e..f099d7752a 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -197,6 +197,8 @@ class Parser: Used by structured engines like `xgrammar` to check if the reasoning content ends in the model output. + Must be a pure function of the input_ids. + Args: input_ids: The token IDs of the model output. @@ -210,6 +212,8 @@ class Parser: """ Check if the reasoning content ends during a decode step. + Must be a pure function of the input_ids and delta_ids. + Args: input_ids: The entire model output token IDs. delta_ids: The last few computed tokens at the current decode step. @@ -228,8 +232,8 @@ class Parser: Returns: The offset within ``token_ids`` of the last reasoning token, or - ``None`` when reasoning does not end in this window. Parsers that - cannot answer from a window alone always return ``None``; callers + ``len(token_ids)`` when reasoning does not end in this window. Parsers + that cannot answer from a window alone always return ``None``; callers fall back to :meth:`is_reasoning_end_streaming`. """ return None diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index f873e36342..fa306a5da4 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -645,7 +645,7 @@ class ParserEngine(Parser): for offset, token_id in enumerate(token_ids): if token_id in end_ids: return offset - return None + return len(token_ids) def is_reasoning_end(self, input_ids: list[int]) -> bool: config = self.parser_engine_config diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 3af527124d..daf2b316d0 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -70,8 +70,7 @@ from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats -from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager -from vllm.v1.structured_output.utils import strip_speculative_padding +from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import record_function_or_nullcontext logger = init_logger(__name__) @@ -2052,33 +2051,18 @@ class Scheduler(SchedulerInterface): request.status = RequestStatus.FINISHED_STOPPED stopped = True - if new_token_ids and self.structured_output_manager.should_advance( - request, new_token_ids=new_token_ids + if new_token_ids and not self.structured_output_manager.accept_tokens( + request, new_token_ids ): - struct_output_request = request.structured_output_request - assert struct_output_request is not None - grammar = struct_output_request.grammar - assert isinstance(grammar, StructuredOutputGrammar) - # new_token_ids can be a mixed block of reasoning content, then - # the reasoning end marker, then the start of the grammar content. - # Trim the reasoning content so the grammar only sees grammar content. - advance_token_ids = ( - self.structured_output_manager.trim_reasoning_for_advance( - request, new_token_ids - ) + logger.error( + "Unexpected: grammar rejected tokens %s for request %s. " + "Terminating request.", + new_token_ids, + req_id, ) - if advance_token_ids and not grammar.accept_tokens( - req_id, advance_token_ids - ): - logger.error( - "Unexpected: grammar rejected tokens %s for request %s. " - "Terminating request.", - advance_token_ids, - req_id, - ) - request.status = RequestStatus.FINISHED_ERROR - request.resumable = False - stopped = True + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True routed_experts = None if ( @@ -2434,11 +2418,9 @@ class Scheduler(SchedulerInterface): continue # Add newly generated spec token ids to the request. - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - spec_token_ids = strip_speculative_padding(spec_token_ids) - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] - request.spec_token_ids = spec_token_ids + request.spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) def update_draft_token_ids_in_output( self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput @@ -2464,10 +2446,9 @@ class Scheduler(SchedulerInterface): # (needed for chunked prefill case for example). del spec_token_ids[orig_num_spec_tokens:] # Filter out spec tokens which do not adhere to the grammar. - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - spec_token_ids = strip_speculative_padding(spec_token_ids) - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] + spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) # Pad to original number of spec tokens. num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids) if num_invalid_tokens: diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index ca6adc736d..7fd311c873 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import itertools import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor @@ -18,6 +17,7 @@ from vllm.v1.structured_output.backend_types import ( StructuredOutputGrammar, ) from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend +from vllm.v1.structured_output.utils import strip_speculative_padding if TYPE_CHECKING: import numpy as np @@ -217,6 +217,100 @@ class StructuredOutputManager: ) -> Future: return self.executor_for_fillmask.submit(self._fill_bitmasks, batch) + def _get_constraint_start( + self, + request: "Request", + spec_tokens: Sequence[int], + spec_tokens_committed: bool = False, + ) -> int: + """Return the index into `spec_tokens` where tokens should start being + constrained by the grammar. + Assumes `spec_tokens` are stripped of -1 padding tokens. + Returns `len(spec_tokens) + 1` if no token should be constrained after + accepting all `spec_tokens`. + + `spec_tokens_committed` is True if spec_tokens is already in + `request.all_token_ids`. + """ + if self.enable_in_reasoning: + return 0 + + structured_req = request.structured_output_request + assert structured_req is not None + if structured_req.reasoning_ended: + return 0 + + reasoner = self._get_reasoner(request) + if reasoner is None: + return 0 + + if structured_req.reasoning_ended is None: + # This should be removed here, but since `openai_gptoss` + # is an independent code path, it is kept for now. + # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, + # it can be removed. + if reasoner.is_reasoning_end(request.prompt_token_ids or []): + return 0 + structured_req.reasoning_ended = False + + num_spec_tokens = len(spec_tokens) + if num_spec_tokens <= 0: + return num_spec_tokens + 1 + + # Use `find_reasoning_end_offset` to find constraint start if supported + if ( + isinstance(reasoner, ParserEngineReasoningAdapter) + and reasoner.reasoning_end_token_ids + ): + offset = reasoner.find_reasoning_end_offset(spec_tokens) + if offset is not None: + return offset + 1 + + # Fallback to `find_reasoning_end_offset` + # TODO: Build a read-only Sequence view over all_token_ids + # instead of copying the entire all_token_ids into input_ids. + if spec_tokens_committed: + if not spec_tokens or not reasoner.is_reasoning_end_streaming( + request.all_token_ids, spec_tokens + ): + return num_spec_tokens + 1 + + input_ids = request.all_token_ids.copy() + delta_ids = list(spec_tokens) + else: + input_ids = request.all_token_ids.copy() + input_ids.extend(spec_tokens) + if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens): + return num_spec_tokens + 1 + delta_ids = list(spec_tokens) + + for i in range(num_spec_tokens - 1, 0, -1): + input_ids.pop() + delta_ids.pop() + if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids): + return i + 1 + return 1 + + def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]: + """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`.""" + if not request.use_structured_output: + return spec_tokens + + spec_tokens = strip_speculative_padding(spec_tokens) + constraint_start = self._get_constraint_start(request, spec_tokens) + if constraint_start >= len(spec_tokens): + return spec_tokens + + structured_req = request.structured_output_request + if TYPE_CHECKING: + assert structured_req is not None + grammar = structured_req.grammar + if TYPE_CHECKING: + assert isinstance(grammar, StructuredOutputGrammar) + prefix = spec_tokens[:constraint_start] + validated = grammar.validate_tokens(spec_tokens[constraint_start:]) + return prefix + validated + def grammar_bitmask( self, requests: dict[str, "Request"], @@ -264,7 +358,7 @@ class StructuredOutputManager: if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self.should_fill_bitmask(request) + apply_bitmask = self._get_constraint_start(request, ()) == 0 batch.append((grammar, cumulative_index, apply_bitmask)) if len(batch) == self.fill_bitmask_parallel_batch_size: promises.append(self._async_submit_fill_bitmask(batch)) @@ -288,76 +382,47 @@ class StructuredOutputManager: grammar = structured_output_request.grammar if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self.should_fill_bitmask(request) - - reasoner = None if apply_bitmask else self._get_reasoner(request) + req_tokens = scheduled_spec_decode_tokens.get(req_id, list()) + constraint_start = self._get_constraint_start( + request, strip_speculative_padding(req_tokens) + ) state_advancements = 0 - post_reasoning_end_in_window = False - req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) - # Locate the reasoning-end marker in the draft window once. - # Rejected-draft padding (-1) is trailing, so the window is - # cut there and the reasoner never sees a placeholder. - window_start = len(request.all_token_ids) - reasoning_end_index: int | None = None - if reasoner is not None and req_tokens: - window = req_tokens - if -1 in window: - window = window[: window.index(-1)] - reasoning_end_index = self._find_reasoning_end_index( - reasoner, - request.all_token_ids, - window_start, - window, - delta_appended=False, - ) + seen_padding = False + failed = False for i, token in enumerate(req_tokens): + apply_bitmask = ( + not failed and not seen_padding and i >= constraint_start + ) self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) - advance_grammar = apply_bitmask if token == -1: - apply_bitmask = False - advance_grammar = False - elif not apply_bitmask and window_start + i == reasoning_end_index: - # Reasoning ended mid-window. Constrain the rest - # of the window via bitmask. Skip grammar advance - # through the marker (it is reasoning content); - # try to advance through subsequent drafts so the - # next bitmask row reflects the post-advance state, - # but tolerate rejection since those drafts predate - # the bitmask and are not guaranteed valid. - apply_bitmask = True - advance_grammar = False - post_reasoning_end_in_window = True - if advance_grammar and not grammar.is_terminated(): - if post_reasoning_end_in_window: - accepted = bool(grammar.validate_tokens([token])) - if accepted: - accepted = grammar.accept_tokens(req_id, [token]) - else: - accepted = grammar.accept_tokens(req_id, [token]) - if accepted: + seen_padding = True + elif apply_bitmask: + if not grammar.is_terminated() and grammar.accept_tokens( + req_id, [token] + ): state_advancements += 1 - elif not post_reasoning_end_in_window: - raise AssertionError( - (token, req_id, scheduled_spec_decode_tokens) + else: + failed = True + logger.error( + "Unexpected: grammar terminated or rejected draft " + "token %s for request %s during bitmask fill.", + token, + req_id, ) cumulative_index += 1 + # Diffusion LLMs don't sample a bonus token after the # scheduled positions, so skip its bitmask in that case. if not (self.vllm_config.model_config.is_diffusion and req_tokens): - # bonus_apply must be True when the bonus-row position - # should be grammar-constrained. Two triggers: - # - should_fill_bitmask(request): reasoning was already - # over at step start (or no reasoner / - # enable_in_reasoning). - # - apply_bitmask: reasoning ended mid-window in this - # call and was flipped True after the marker; - # should_fill_bitmask still returns False here because - # reasoning_ended is only persisted later by - # should_advance. - bonus_apply = self.should_fill_bitmask(request) or apply_bitmask + bonus_apply = ( + not failed + and not seen_padding + and constraint_start <= len(req_tokens) + ) self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),)) cumulative_index += 1 + if state_advancements > 0: grammar.rollback(state_advancements) @@ -370,161 +435,29 @@ class StructuredOutputManager: # and deserialization when sending this to the GPU workers. return bitmask_tensor.numpy() - def should_fill_bitmask(self, request: "Request") -> bool: - # NOTE (Hanchen) if enable_in_reasoning is True, it means that - # the model needs to be constrained in reasoning. So we should always - # enable the bitmask filling. - structured_req = request.structured_output_request - if self.enable_in_reasoning or ( - structured_req is not None and structured_req.reasoning_ended is True - ): + def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool: + """Advance grammar with accepted tokens. Returns False on rejection.""" + if not request.use_structured_output or not new_token_ids: return True - reasoner = self._get_reasoner(request) - if reasoner is not None: - assert structured_req is not None - if structured_req.reasoning_ended is None: - # This should be removed here, but since `openai_gptoss` - # is an independent code path, it is kept for now. - # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, - # it can be removed. - structured_req.reasoning_ended = reasoner.is_reasoning_end( - request.prompt_token_ids or [] - ) - return structured_req.reasoning_ended - return True - - def should_advance( - self, - request: "Request", - new_token_ids: list[int] | None = None, - ) -> bool: - if not request.use_structured_output: - return False - - # To determine whether we can advance the FSM. - # Supports thinking usage where we skip the reasoning components. + structured_req = request.structured_output_request if TYPE_CHECKING: - assert request.structured_output_request is not None - assert request.structured_output_request.grammar is not None - structured_req = request.structured_output_request - if self.enable_in_reasoning or ( - structured_req is not None and structured_req.reasoning_ended is True - ): - return True + assert structured_req is not None + grammar = structured_req.grammar + if TYPE_CHECKING: + assert isinstance(grammar, StructuredOutputGrammar) - # by default, we should always advance - # for cases that don't use thinking mode. - reasoner = self._get_reasoner(request) - if reasoner is None: - return True - - assert structured_req is not None - - # Check if reasoning ends in *this* step. - # When the caller passes new_token_ids (the tokens that were just - # appended this step), use it directly as the delta window. The - # placeholder-derived fallback assumes num_output_placeholders == - # len(new_token_ids), which breaks under async scheduling + spec - # decode when some drafts are rejected (#43388): the placeholder - # count remains > 0 after the step and the computed delta window - # starts past the reasoning-end marker. - all_token_ids = request.all_token_ids - if new_token_ids: - # The tokens were already appended this step, so the step window - # starts exactly len(new_token_ids) from the end. - start = len(all_token_ids) - len(new_token_ids) - delta_ids: Sequence[int] = new_token_ids - else: - # Draft-validation callers arrive here after update_from_output - # has already persisted any boundary found this step, so this - # window only needs to cover tokens not examined yet. - delta_from = request.num_computed_tokens - request.num_output_placeholders - start = ( - delta_from - if delta_from >= 0 - else max(len(all_token_ids) + delta_from, 0) - ) - delta_ids = all_token_ids[start:] - end_index = self._find_reasoning_end_index( - reasoner, all_token_ids, start, delta_ids, delta_appended=True + constraint_start = self._get_constraint_start( + request, new_token_ids, spec_tokens_committed=True ) - if end_index is not None: - structured_req.reasoning_ended = True - # Record the boundary so the scheduler can exclude reasoning tokens. - structured_req.reasoning_end_token_index = end_index + # Early return only when the constraint hasn't started. + # Otherwise, latch structured_req.reasoning_ended. + if constraint_start > len(new_token_ids): return True - - return False - - @staticmethod - def _find_reasoning_end_index( - reasoner: "ReasoningParser", - all_token_ids: Sequence[int], - start: int, - delta_ids: Sequence[int], - *, - delta_appended: bool, - ) -> int | None: - """Locates the token that ends reasoning within ``delta_ids``. - - ``delta_ids`` occupy positions ``start..``; ``delta_appended`` says - whether they are already in ``all_token_ids`` (accepted tokens) or - not (speculative drafts). - - Legacy parsers are probed one token at a time. Accepted tokens get a - whole-delta check as well, and a marker seen only that way reports - the final token. Drafts skip it: some predicates are order-sensitive - over a multi-token window and would miss the marker. - - Returns: - Absolute index of the last reasoning token, or ``None``. - """ - if ( - isinstance(reasoner, ParserEngineReasoningAdapter) - and reasoner.reasoning_end_token_ids - ): - offset = reasoner.find_reasoning_end_offset(delta_ids) - return None if offset is None else start + offset - - if delta_appended and not reasoner.is_reasoning_end_streaming( - all_token_ids, delta_ids - ): - return None - prefix = list(itertools.islice(all_token_ids, start)) - for offset, token in enumerate(delta_ids): - prefix.append(token) - if reasoner.is_reasoning_end_streaming(prefix, [token]): - return start + offset - if delta_appended: - return len(all_token_ids) - 1 - return None - - def trim_reasoning_for_advance( - self, request: "Request", new_token_ids: list[int] - ) -> list[int]: - """Drops reasoning content from tokens about to advance the grammar. - - When reasoning ends mid-step (see should_advance), the step's output - still contains reasoning tokens up to and including the end marker. - Those are not grammar content: feeding them to accept_tokens makes - the grammar reject the marker and kills the request (#44006). - - Returns: - The suffix of ``new_token_ids`` that follows the reasoning-end - marker. Steps fully after the boundary are returned unchanged. - """ - structured_req = request.structured_output_request - if structured_req is None: - return new_token_ids - end_idx = structured_req.reasoning_end_token_index - if end_idx is None: - return new_token_ids - first_idx = len(request.all_token_ids) - len(new_token_ids) - num_reasoning = end_idx + 1 - first_idx - if num_reasoning <= 0: - return new_token_ids - return new_token_ids[num_reasoning:] + structured_req.reasoning_ended = True + return grammar.accept_tokens( + request.request_id, new_token_ids[constraint_start:] + ) def clear_backend(self) -> None: if self.backend is not None: diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py index b0f4f3c3a8..dce0038782 100644 --- a/vllm/v1/structured_output/request.py +++ b/vllm/v1/structured_output/request.py @@ -25,12 +25,6 @@ class StructuredOutputRequest: Future[StructuredOutputGrammar] | StructuredOutputGrammar | Exception | None ) = None reasoning_ended: bool | None = None - # Absolute index into the request's all_token_ids of the last reasoning - # token (the reasoning-end marker). Tokens at or before this index are - # reasoning content and must never be fed to the grammar. Only set when - # reasoning ends in a step whose tokens the scheduler advances immediately - # (structural tags + speculative decoding, see #42452). - reasoning_end_token_index: int | None = None reasoning_parser_kwargs: dict[str, Any] | None = None # Cached per request; do not share reasoning parsers across requests because # their behavior can depend on reasoning_parser_kwargs. From a2e8f1ea94b1332428e2282744cec12ce7c75d88 Mon Sep 17 00:00:00 2001 From: xiangdong <40376367+zxd1997066@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:04:09 +0800 Subject: [PATCH 101/139] [XPU][CI]Skip test_chat_completion_with_tools in Intel GPU CI (#56934) Signed-off-by: zengxian Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/entrypoints_intel.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index efed04c1b3..e8006bac8f 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -134,7 +134,8 @@ steps: 'cd tests && pytest -v -s tool_use --deselect="tests/tool_use/mistral/test_mistral_tool_calls.py::test_tool_call_none_with_tools[ministral-3b]" - --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]" && + --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]" + --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[llama3.2]" && pytest -v -s entrypoints/tool_parsers && pytest -v -s entrypoints/generate && pytest -v -s entrypoints/anthropic && From 4bac767695d2dfecd4098a77f71dacd100061247 Mon Sep 17 00:00:00 2001 From: Junhao Shen Date: Tue, 15 Sep 2026 22:33:51 +0800 Subject: [PATCH 102/139] [Perf][Nemotron] Skip redundant latent-MoE all-reduce at TP>1 (~13% decode win) (#52301) Signed-off-by: Junhao Shen Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Misha Goin Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../layers/fused_moe/runner/moe_runner.py | 9 +++++++++ vllm/model_executor/models/nemotron_h.py | 12 ++++++++++++ 2 files changed, 21 insertions(+) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 2868690ebe..6e63563ab8 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -457,9 +457,18 @@ class MoERunner(MoERunnerInterface): Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm. TP partial routed outputs must be summed in latent space before such transforms are applied. + + A transform that commutes with the TP sum is exempt: if + ``sum_r T(x_r) == T(sum_r x_r)``, applying the transform to the local + partial output and letting the existing late all-reduce sum the + combined result is equivalent, and costs one collective instead of two. + Such a transform opts out by setting ``reduce_commutative = True``. + The default is False, so transforms that do not declare themselves + keep being reduced early. """ if ( self.routed_output_transform is not None + and not getattr(self.routed_output_transform, "reduce_commutative", False) and not self.moe_config.is_sequence_parallel and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) and not fused_output_is_reduced diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index c1bc4c6823..ea8a28ebd9 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -49,6 +49,7 @@ from vllm.model_executor.layers.linear import ( QKVParallelLinear, ReplicatedLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 @@ -204,6 +205,17 @@ class NemotronHMoE(nn.Module): disable_tp=self.is_sequence_parallel, prefix=f"{prefix}.fc2_latent_proj", ) + # A bias-free, unquantized linear commutes with the TP sum + # (sum_r W x_r == W sum_r x_r), so one reduce after the transform + # suffices. Test the layer, not the model-wide `quant_config`: ModelOpt + # excludes the latent projections, so quantized checkpoints still get + # an UnquantizedLinearMethod here. + self.fc2_latent_proj.reduce_commutative = ( + not config.mlp_bias + and isinstance( + self.fc2_latent_proj.quant_method, UnquantizedLinearMethod + ) + ) else: self.fc1_latent_proj = None self.fc2_latent_proj = None From 0136df94b0d75732be6c66f620f51289c888e41d Mon Sep 17 00:00:00 2001 From: zack <51604064+luoyuctl@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:35:47 +0800 Subject: [PATCH 103/139] [Bugfix][Spec Decode][MoE] Avoid uninitialized EPLB state in DeepSeek V4.1 DSpark drafter (#56387) Signed-off-by: zack <51604064+luoyuctl@users.noreply.github.com> --- tests/v1/worker/test_dspark_utils.py | 51 +++++++++++++++++++ .../v1/worker/gpu/spec_decode/dspark/utils.py | 32 +++++++++--- 2 files changed, 77 insertions(+), 6 deletions(-) create mode 100644 tests/v1/worker/test_dspark_utils.py diff --git a/tests/v1/worker/test_dspark_utils.py b/tests/v1/worker/test_dspark_utils.py new file mode 100644 index 0000000000..94655f192c --- /dev/null +++ b/tests/v1/worker/test_dspark_utils.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass, field + +from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config + + +@dataclass +class _FakeEPLBConfig: + num_redundant_experts: int = 0 + + +@dataclass +class _FakeParallelConfig: + pipeline_parallel_size: int = 2 + tensor_parallel_size: int = 8 + enable_eplb: bool = True + eplb_config: _FakeEPLBConfig = field( + default_factory=lambda: _FakeEPLBConfig(num_redundant_experts=32) + ) + enable_elastic_ep: bool = True + + def __post_init__(self) -> None: + if not self.enable_eplb and self.eplb_config.num_redundant_experts: + raise ValueError("redundant experts require EPLB") + if self.enable_elastic_ep and not self.enable_eplb: + raise ValueError("elastic EP requires EPLB") + + +def test_dspark_parallel_config_disables_eplb_atomically(): + target_config = _FakeParallelConfig() + + draft_config = _get_dspark_parallel_config( + target_config, + tensor_parallel_size=4, + ) + + assert target_config.pipeline_parallel_size == 2 + assert target_config.tensor_parallel_size == 8 + assert target_config.enable_eplb + assert target_config.eplb_config.num_redundant_experts == 32 + assert target_config.enable_elastic_ep + + assert draft_config is not target_config + assert draft_config.pipeline_parallel_size == 1 + assert draft_config.tensor_parallel_size == 4 + assert not draft_config.enable_eplb + assert draft_config.eplb_config.num_redundant_experts == 0 + assert not draft_config.enable_elastic_ep + assert draft_config.eplb_config is not target_config.eplb_config diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index db1e17a9f2..bf1a5761b2 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -3,7 +3,7 @@ import torch.nn as nn -from vllm.config import ModelConfig, VllmConfig, replace +from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace from vllm.logger import init_logger from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config @@ -31,6 +31,29 @@ def _resolve_dspark_attention_backend( return None +def _get_dspark_parallel_config( + parallel_config: ParallelConfig, + tensor_parallel_size: int, +) -> ParallelConfig: + if parallel_config.enable_eplb: + logger.warning_once( + "EPLB is disabled for the DSpark draft model. EPLB remains enabled " + "for the target model." + ) + + return replace( + parallel_config, + pipeline_parallel_size=1, + tensor_parallel_size=tensor_parallel_size, + enable_eplb=False, + eplb_config=replace( + parallel_config.eplb_config, + num_redundant_experts=0, + ), + enable_elastic_ep=False, + ) + + def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: speculative_config = vllm_config.speculative_config assert speculative_config is not None @@ -54,12 +77,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_vllm_config = replace( vllm_config, - parallel_config=replace( + parallel_config=_get_dspark_parallel_config( vllm_config.parallel_config, - pipeline_parallel_size=1, - tensor_parallel_size=( - speculative_config.draft_parallel_config.tensor_parallel_size - ), + speculative_config.draft_parallel_config.tensor_parallel_size, ), attention_config=replace( vllm_config.attention_config, From f0fbc98c488b40a07d3b269aeba2bda3c4c3b03b Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 15 Sep 2026 22:50:26 +0800 Subject: [PATCH 104/139] [Frontend] Share max_num_queued_reqs across API server processes (#54746) Signed-off-by: chaunceyjiang --- .../test_api_server_process_manager.py | 39 ++++++++++++++ tests/v1/engine/test_admission_control.py | 39 ++++++++++++++ vllm/v1/engine/admission_control.py | 52 +++++++++++++++++++ vllm/v1/engine/async_llm.py | 24 +++++++-- vllm/v1/engine/output_processor.py | 14 ++++- vllm/v1/utils.py | 11 ++++ 6 files changed, 173 insertions(+), 6 deletions(-) create mode 100644 vllm/v1/engine/admission_control.py diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 40fefc2374..668f115448 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -5,6 +5,7 @@ import multiprocessing import socket import threading import time +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -14,6 +15,7 @@ from tests.entrypoints.launchers.api_server._api_server_spawn_workers import ( exit_before_report_worker, ) from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.utils import ( APIServerProcessManager, get_engine_client_zmq_addr, @@ -32,6 +34,18 @@ def mock_run_api_server_worker(listen_address, sock, args, client_config=None): print("Mock worker completed successfully") +def update_admission_stats_worker(listen_address, sock, args, client_config): + stats = SharedAdmissionStats( + client_config, + client_config["client_count"], + client_config["client_index"], + ) + value = client_config["client_index"] + 1 + stats.set_num_requests(value) + args.barrier.wait() + args.results.put(stats.get_num_requests()) + + # Module-level stub for the gather_actual_addresses test. Must be # importable by `multiprocessing.spawn` (no closures, no nesting). def defer_addresses_stub_worker(listen_address, sock, args, client_config): @@ -131,6 +145,31 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): assert not proc.is_alive() +@pytest.mark.skip_global_cleanup +def test_api_server_processes_share_admission_stats(api_server_args): + args = api_server_args.copy() + spawn_context = multiprocessing.get_context("spawn") + worker_args = SimpleNamespace( + barrier=spawn_context.Barrier(args["num_servers"]), + results=spawn_context.Queue(), + max_num_queued_reqs=10, + ) + args["args"] = worker_args + args["target_server_fn"] = update_admission_stats_worker + manager = APIServerProcessManager(**args) + + try: + for process in manager.processes: + process.join(timeout=30) + assert process.exitcode == 0 + + assert [worker_args.results.get(timeout=5) for _ in range(3)] == [6] * 3 + finally: + manager.shutdown() + worker_args.results.close() + worker_args.results.join_thread() + + @patch("vllm.v1.utils.run_api_server_worker_proc", mock_run_api_server_worker) def test_wait_for_completion_or_failure(api_server_args): """Test that wait_for_completion_or_failure works with failures.""" diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py index a03718ab5a..219850f5d4 100644 --- a/tests/v1/engine/test_admission_control.py +++ b/tests/v1/engine/test_admission_control.py @@ -14,6 +14,7 @@ These tests cover: import argparse import asyncio +import multiprocessing from http import HTTPStatus from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -34,6 +35,7 @@ from vllm.exceptions import ( from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.utils.argparse_utils import human_readable_int +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.output_processor import OutputProcessor @@ -66,6 +68,7 @@ def _make_async_llm( llm.output_processor = MagicMock() llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens + llm.admission_stats = None return llm @@ -84,6 +87,42 @@ def _make_scheduler_config(**kwargs) -> SchedulerConfig: ) +def _make_shared_stats(client_count: int = 2) -> list[SharedAdmissionStats]: + counters = multiprocessing.RawArray( + "q", SharedAdmissionStats.num_counters(client_count) + ) + return [ + SharedAdmissionStats({"mp_admission_counters": counters}, client_count, index) + for index in range(client_count) + ] + + +@pytest.mark.skip_global_cleanup +def test_shared_admission_stats_aggregate_api_servers(): + server_0, server_1 = _make_shared_stats() + + server_0.set_num_requests(2) + assert server_0.get_num_requests() == 2 + server_1.set_num_requests(3) + assert server_0.get_num_requests() == server_1.get_num_requests() == 5 + + server_0.set_num_requests(1) + assert server_0.get_num_requests() == 4 + + +@pytest.mark.skip_global_cleanup +def test_shared_request_snapshot_enforces_global_limit(): + server_0, server_1 = _make_shared_stats() + server_0.set_num_requests(2) + llm = _make_async_llm(max_num_queued_reqs=2) + llm.admission_stats = server_1 + + with pytest.raises(QueueOverflowError): + llm.check_admission() + + assert server_0.get_num_requests() == 2 + + # --------------------------------------------------------------------------- # Exception classes # --------------------------------------------------------------------------- diff --git a/vllm/v1/engine/admission_control.py b/vllm/v1/engine/admission_control.py new file mode 100644 index 0000000000..b72d50f346 --- /dev/null +++ b/vllm/v1/engine/admission_control.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any, Protocol, cast + + +class AdmissionCounterBuffer(Protocol): + def __getitem__(self, index: int) -> int: ... + + def __setitem__(self, index: int, value: int) -> None: ... + + +class SharedAdmissionStats: + """Lock-free admission stats shared by API server processes. + + Each process writes exclusively to its own cache-line-sized slot. Readers + aggregate the slots, avoiding both inter-process locks and cache-line + contention on the request path. + """ + + # Space per-process counters 64 bytes apart to avoid false sharing on CPUs + # with 64-byte cache lines: one process's writes would otherwise invalidate + # the cache line containing another process's counter. + CACHE_LINE_BYTES = 64 + COUNTER_BYTES = 8 + COUNTERS_PER_SLOT = CACHE_LINE_BYTES // COUNTER_BYTES + + def __init__( + self, + client_addresses: dict[str, Any], + client_count: int = 1, + client_index: int = 0, + ): + self._counters = cast( + AdmissionCounterBuffer, + client_addresses.pop("mp_admission_counters"), + ) + self._client_count = client_count + self._slot_offset = client_index * self.COUNTERS_PER_SLOT + + @classmethod + def num_counters(cls, client_count: int) -> int: + return client_count * cls.COUNTERS_PER_SLOT + + def set_num_requests(self, count: int) -> None: + self._counters[self._slot_offset] = count + + def get_num_requests(self) -> int: + return sum( + self._counters[client_index * self.COUNTERS_PER_SLOT] + for client_index in range(self._client_count) + ) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 91618c4093..bf3f5f420a 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -44,6 +44,7 @@ from vllm.usage.usage_lib import UsageContext from vllm.utils.async_utils import cancel_task_threadsafe from vllm.utils.collection_utils import as_list from vllm.v1.engine import EngineCoreRequest, PauseMode +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.core_client import EngineCoreClient from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm.v1.engine.input_processor import InputProcessor @@ -144,12 +145,20 @@ class AsyncLLM(EngineClient): # Convert EngineInput --> EngineCoreRequest. self.input_processor = InputProcessor(self.vllm_config, renderer) + self.admission_stats = ( + SharedAdmissionStats(client_addresses, client_count, client_index) + if client_addresses is not None + and "mp_admission_counters" in client_addresses + else None + ) + # Converts EngineCoreOutputs --> RequestOutput. self.output_processor = OutputProcessor( renderer.tokenizer, log_stats=self.log_stats, stream_interval=self.vllm_config.scheduler_config.stream_interval, tracing_enabled=tracing_endpoint is not None, + admission_stats=self.admission_stats, ) # EngineCore (starts the engine in background process). @@ -291,8 +300,9 @@ class AsyncLLM(EngineClient): Both limits return HTTP 503 (Service Unavailable) so that load balancers and client SDKs retry on a different instance. - - ``max_num_queued_reqs``: hard cap on the number of unfinished requests - (waiting + running). A request with ``n > 1`` counts as ``n`` slots. + - ``max_num_queued_reqs``: approximate global cap on the number of + unfinished requests (waiting + running). A request with ``n > 1`` + counts as ``n`` slots. - ``max_num_queued_tokens``: TTFT QoS — cap on the total prompt tokens of requests still in prefill. @@ -312,13 +322,17 @@ class AsyncLLM(EngineClient): """ max_num_reqs = self.scheduler_config.max_num_queued_reqs if max_num_reqs is not None: - current = self.get_num_unfinished_requests() - if current + n > max_num_reqs: + current_requests = ( + self.admission_stats.get_num_requests() + if self.admission_stats is not None + else self.get_num_unfinished_requests() + ) + if current_requests + n > max_num_reqs: logger.info( "Request queue full - rejecting request %s " "(current=%d, n=%d, max=%d).", request_id, - current, + current_requests, n, max_num_reqs, ) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 419e3aa8e9..257287243d 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -5,7 +5,7 @@ import asyncio from collections import defaultdict, deque from collections.abc import Iterable from dataclasses import dataclass -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import numpy as np import torch @@ -41,6 +41,9 @@ from vllm.v1.metrics.stats import ( ) from vllm.v1.outputs import SamplingMaskLists +if TYPE_CHECKING: + from vllm.v1.engine.admission_control import SharedAdmissionStats + # shared empty CPU tensor used as a placeholder pooling output EMPTY_CPU_TENSOR = torch.empty(0, device="cpu") @@ -455,6 +458,7 @@ class OutputProcessor: log_stats: bool, stream_interval: int = 1, tracing_enabled: bool = False, + admission_stats: "SharedAdmissionStats | None" = None, ): self.log_stats = log_stats self.tokenizer = tokenizer @@ -464,6 +468,7 @@ class OutputProcessor: self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list) self.lora_states = LoRARequestStates(log_stats) self.tracing_enabled = tracing_enabled + self.admission_stats = admission_stats def get_num_unfinished_requests(self): return len(self.request_states) @@ -553,6 +558,7 @@ class OutputProcessor: child_reqs = self.abort_requests(child_reqs, internal=True) request_ids_to_abort.extend(child_reqs) self.parent_requests.pop(request_id, None) + self._update_admission_stats() return request_ids_to_abort def add_request( @@ -585,6 +591,7 @@ class OutputProcessor: # Track the external_req_id -> [internal_req_id, ...] mapping self.external_req_ids[req_state.external_req_id].append(request_id) + self._update_admission_stats() def _update_streaming_request_state( self, req_state: RequestState, request: EngineCoreRequest, prompt: str | None @@ -763,6 +770,11 @@ class OutputProcessor: parent_req = req_state.parent_req if parent_req and not parent_req.child_requests: self.parent_requests.pop(parent_req.request_id, None) + self._update_admission_stats() + + def _update_admission_stats(self) -> None: + if self.admission_stats is not None: + self.admission_stats.set_num_requests(self.get_num_unfinished_requests()) def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None): self.lora_states.update_scheduler_stats(scheduler_stats) diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index b0411dc9e8..da17110d80 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -209,6 +209,15 @@ class APIServerProcessManager: self.processes: list[BaseProcess] = [] self._address_pipes: list[connection.Connection] = [] + admission_counters = None + if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None: + from vllm.v1.engine.admission_control import SharedAdmissionStats + + admission_counters = spawn_context.RawArray( + "q", + SharedAdmissionStats.num_counters(num_servers), + ) + for i, in_addr, out_addr in zip( range(num_servers), input_addresses, output_addresses ): @@ -218,6 +227,8 @@ class APIServerProcessManager: "client_count": num_servers, "client_index": i, } + if admission_counters is not None: + client_config["mp_admission_counters"] = admission_counters if stats_update_address is not None: client_config["stats_update_address"] = stats_update_address if tensor_queue is not None: From 1e47ec00d2e6ad911969508068b72f7d1aae1968 Mon Sep 17 00:00:00 2001 From: Abhinav Verma Date: Tue, 15 Sep 2026 20:27:32 +0530 Subject: [PATCH 105/139] [Bugfix][Gemma 4] Don't read fft_length when profiling unified audio (#56721) Signed-off-by: Abhinav Verma Co-authored-by: Claude Opus 5 --- .../processing/test_gemma4_unified.py | 44 ++++++++++++++++++- vllm/model_executor/models/gemma4_mm.py | 25 ++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/tests/models/multimodal/processing/test_gemma4_unified.py b/tests/models/multimodal/processing/test_gemma4_unified.py index 67a81ddb7b..fe544e2338 100644 --- a/tests/models/multimodal/processing/test_gemma4_unified.py +++ b/tests/models/multimodal/processing/test_gemma4_unified.py @@ -6,9 +6,13 @@ from collections.abc import Mapping import pytest import torch from PIL import Image as PILImage +from transformers.models.gemma4 import Gemma4AudioFeatureExtractor from vllm.exceptions import VLLMValidationError -from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs +from vllm.model_executor.models.gemma4_mm import ( + Gemma4ImagePixelInputs, + _dummy_audio_num_samples, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig @@ -204,3 +208,41 @@ def test_limit_mm_per_prompt( mm_items=processor.info.parse_mm_data(mm_data), hf_processor_mm_kwargs={}, ) + + +def test_dummy_audio_length_falls_back_for_the_unified_extractor(): + """The unified audio feature extractor does not define ``fft_length``. + + ``Gemma4DummyInputsBuilder.get_dummy_mm_data`` read that attribute + unconditionally, so memory profiling raised ``AttributeError`` for every + unified Gemma 4 deployment in which audio is the modality with the largest + per-item token count -- which is the case as soon as ``audio_seq_length`` + is raised above the video budget. + """ + + class _UnifiedFeatureExtractor: + sampling_rate = 16000 + + class _UnifiedProcessor: + feature_extractor = _UnifiedFeatureExtractor() + audio_seq_length = 3000 + audio_ms_per_token = 40 + + # 3000 tokens * 40 ms = 120 s at 16 kHz + assert _dummy_audio_num_samples(_UnifiedProcessor()) == 120 * 16000 + + +def test_dummy_audio_length_still_uses_fft_length_for_the_tower_extractor(): + """The mel/tower extractor keeps its existing behaviour.""" + + class _TowerFeatureExtractor(Gemma4AudioFeatureExtractor): + def __init__(self): # noqa: D107 - bypass the real feature-extractor init + self.fft_length = 512 + self.sampling_rate = 16000 + + class _TowerProcessor: + feature_extractor = _TowerFeatureExtractor() + audio_seq_length = 750 + audio_ms_per_token = 40 + + assert _dummy_audio_num_samples(_TowerProcessor()) == 512 diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 2dddc72b21..2fe2752475 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -24,6 +24,7 @@ from PIL import Image as PILImage from torch import nn from transformers import AutoModel, BatchFeature from transformers.models.gemma4 import ( + Gemma4AudioFeatureExtractor, Gemma4Config, Gemma4Processor, Gemma4VisionConfig, @@ -456,6 +457,28 @@ class Gemma4ProcessingInfo(BaseProcessingInfo): # --------------------------------------------------------------------------- +def _dummy_audio_num_samples(processor: Any) -> int: + """Length in samples of a worst-case dummy audio item. + + ``fft_length`` is set in ``Gemma4AudioFeatureExtractor.__init__`` and is not + defined by ``Gemma4UnifiedAudioFeatureExtractor``, which the encoder-free + Gemma 4 variant uses. Reading it unconditionally therefore raises + ``AttributeError`` during memory profiling on that variant, which is reached + whenever audio is the modality with the largest per-item token count. + + For any other extractor, fall back to the audio budget the processor itself + declares -- the same pair of values used to warn about over-long audio + elsewhere in this file. + """ + feature_extractor = processor.feature_extractor + if isinstance(feature_extractor, Gemma4AudioFeatureExtractor): + return feature_extractor.fft_length + + sampling_rate = feature_extractor.sampling_rate + max_duration_s = processor.audio_seq_length * processor.audio_ms_per_token / 1000.0 + return int(max_duration_s * sampling_rate) + + class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_images = mm_counts.get("image", 0) @@ -508,7 +531,7 @@ class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]): } if num_audios > 0: - audio_len = processor.feature_extractor.fft_length + audio_len = _dummy_audio_num_samples(processor) data["audio"] = self._get_dummy_audios( length=audio_len, num_audios=num_audios, From 9446ea168067c141992f86c3a796e8551f37139e Mon Sep 17 00:00:00 2001 From: drslark <96540755+drslark@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:59:02 +0800 Subject: [PATCH 106/139] [Bugfix][Spec Decode] Only create draft_id_to_target_id when draft vocab differs (#53458) Signed-off-by: drslark Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: GLM-5.3 Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py --- vllm/model_executor/models/deepseek_eagle3.py | 12 ++++++++---- vllm/model_executor/models/llama_eagle3.py | 12 ++++++++---- vllm/model_executor/models/qwen3_eagle3.py | 12 ++++++++---- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py index 80c41501e6..48d019cd35 100644 --- a/vllm/model_executor/models/deepseek_eagle3.py +++ b/vllm/model_executor/models/deepseek_eagle3.py @@ -305,10 +305,14 @@ class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None def embed_input_ids( self, diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index c3d956a946..68e21c9dc2 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -335,10 +335,14 @@ class Eagle3LlamaForCausalLM(_Eagle3LlamaForCausalLMBase): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None self.use_parallel_drafting = speculative_config.parallel_drafting diff --git a/vllm/model_executor/models/qwen3_eagle3.py b/vllm/model_executor/models/qwen3_eagle3.py index 39d0173b3c..745c77a977 100644 --- a/vllm/model_executor/models/qwen3_eagle3.py +++ b/vllm/model_executor/models/qwen3_eagle3.py @@ -309,10 +309,14 @@ class Eagle3Qwen3ForCausalLM(Qwen3ForCausalLM): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting From 836bb3839ffefcda8283ea7d41671a89e1a613df Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 15 Sep 2026 23:03:58 +0800 Subject: [PATCH 107/139] [Bugfix] Prevent out-of-bounds access in FlashInfer SM90 sparse MLA mixed batches (#56969) Signed-off-by: chaunceyjiang --- .../test_flashinfer_mla_sparse_sm90.py | 26 ++++++++++++++----- .../mla/flashinfer_mla_sparse_sm90.py | 8 ++++++ 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index b20c163306..8dc082b08b 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -151,32 +151,44 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): assert kwargs == {} -def test_builder_attaches_its_state(monkeypatch): +@pytest.mark.parametrize("use_mha", [False, True]) +@pytest.mark.parametrize("num_decodes", [0, 1]) +def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_decodes): + """MHA prefill rows must not make the MQA kernel read beyond its query.""" builder = object.__new__(FlashInferMLASparseSM90Builder) builder._index_topk = 2048 builder._index_kpool = 4 builder._async_scheduling = False builder.state = FakeState(TOPK) + builder._attention_layer = SimpleNamespace(_use_sparse_mha=lambda _: use_mha) metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata) metadata.state = None + metadata.num_prefills = 1 + metadata.num_decode_tokens = num_decodes monkeypatch.setattr( sm90_mod.FlashInferMLASparseMetadataBuilder, "build", lambda *_args, **_kwargs: metadata, ) cam = SimpleNamespace( - num_reqs=1, - query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), - seq_lens=torch.tensor([1], dtype=torch.int32), - seq_lens_cpu_upper_bound=torch.tensor([1], dtype=torch.int32), + num_reqs=num_decodes + 1, + query_start_loc_cpu=torch.tensor( + [0, 1, 6] if num_decodes else [0, 5], dtype=torch.int32 + ), + seq_lens_cpu_upper_bound=torch.tensor( + [1402, 5] if num_decodes else [5], dtype=torch.int32 + ), positions=None, ) result = builder.build(0, cam) assert result.state is builder.state - assert builder.state.plan_calls[0][0] == 1 - assert builder.state.plan_calls[0][1].tolist() == [1] + expected_lens = [1402] if num_decodes else [] + if not use_mha: + expected_lens += [1, 2, 3, 4, 5] + assert builder.state.plan_calls[0][0] == len(expected_lens) + assert builder.state.plan_calls[0][1].tolist() == expected_lens def test_plan_uses_state_params(monkeypatch): diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 90dc4cd947..830bc52ab9 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py @@ -277,6 +277,7 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): attention_layer = vllm_config.compilation_config.static_forward_context[ layer_names[0] ] + self._attention_layer = attention_layer impl = attention_layer.impl if not isinstance(impl, FlashInferMLASparseSM90Impl): raise TypeError( @@ -376,6 +377,13 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): # Replan every step outside any CUDA graph capture with this step's # exact per-row lengths; captured runs read the refreshed buffers. num_rows, kv_lens = self._kv_lens_host(common_attn_metadata) + # MHA prefills are removed from q before forward_mqa. Plan only the + # decode prefix in that case, using the same routing decision as MLA. + if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha( + metadata + ): + num_rows = metadata.num_decode_tokens + kv_lens = kv_lens[:num_rows] self.state.plan(num_rows, kv_lens) metadata.state = self.state return metadata From 241e9391ed85f6c967432135427717a529776a63 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 09:58:30 -0700 Subject: [PATCH 108/139] [Bugfix][Attention] Stabilize sparse-MLA DCP for GLM PCP evals (#55879) Signed-off-by: khluu Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Tyler Michael Smith Co-authored-by: Jason Yao --- .../GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml | 5 +- .../configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml | 2 + tests/v1/attention/test_dcp_a2a_pack_mask.py | 16 +- .../test_flashinfer_sparse_mla_workspace.py | 166 ++++++++++++++++++ .../v1/attention/test_indexer_dcp_localize.py | 17 ++ .../test_indexer_deepseek_v4_slot_mapping.py | 32 ++++ .../backends/mla/flashinfer_mla_sparse.py | 125 ++++++++++++- vllm/v1/attention/backends/mla/indexer.py | 46 ++++- vllm/v1/attention/backends/utils.py | 8 +- vllm/v1/attention/ops/dcp.py | 7 + 10 files changed, 411 insertions(+), 13 deletions(-) create mode 100644 tests/v1/attention/test_flashinfer_sparse_mla_workspace.py diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml index b9d2675431..8f3e9ce918 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml @@ -3,10 +3,13 @@ accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 max_concurrency: 100 +# Preserve KV-cache headroom after CUDA graph memory profiling and bound the +# DCP-aware FlashInfer workspace to 2.39 GiB at the 4,096-token scheduler cap. server_args: >- --compilation-config '{"cudagraph_mode":"PIECEWISE"}' + --gpu-memory-utilization 0.96 --max-model-len 4096 - --max-num-batched-tokens 32768 + --max-num-batched-tokens 4096 --safetensors-load-strategy prefetch --moe-backend flashinfer_cutlass --prefill-context-parallel-size 4 diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml index 72f7fdf1ae..bc5c11fa53 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml @@ -3,8 +3,10 @@ accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 max_concurrency: 100 +# Leave headroom for FlashInfer CUTLASS autotuning after KV-cache allocation. server_args: >- --compilation-config '{"cudagraph_mode":"NONE"}' + --gpu-memory-utilization 0.85 --max-model-len 4096 --max-num-batched-tokens 32768 --safetensors-load-strategy prefetch diff --git a/tests/v1/attention/test_dcp_a2a_pack_mask.py b/tests/v1/attention/test_dcp_a2a_pack_mask.py index 36032e7a69..cff1228baf 100644 --- a/tests/v1/attention/test_dcp_a2a_pack_mask.py +++ b/tests/v1/attention/test_dcp_a2a_pack_mask.py @@ -13,7 +13,7 @@ from vllm.v1.attention.ops.dcp import ( mask_dcp_empty_shards_, ) -pytestmark = pytest.mark.skipif( +requires_accelerator = pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="needs a CUDA or ROCm device for the Triton pack kernel", ) @@ -44,6 +44,7 @@ def _pack(out, lse, world_size, h_per_rank, head_dim, seq_lens, query_start_loc) @pytest.mark.parametrize("world_size,h_per_rank", [(8, 2), (8, 16), (4, 4), (2, 1)]) @pytest.mark.parametrize("tokens_per_req", [1, 2, 3]) @pytest.mark.parametrize("num_pad_rows", [0, 2]) +@requires_accelerator @pytest.mark.parametrize( "seq_lens_list", [ @@ -97,6 +98,7 @@ def test_fused_mask_matches_eager( @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("seed", range(8)) +@requires_accelerator def test_fused_mask_matches_eager_ragged(dtype, seed): """Non-uniform per-request query lengths. @@ -137,6 +139,7 @@ def test_fused_mask_matches_eager_ragged(dtype, seed): ) +@requires_accelerator def test_mask_disabled_is_unmasked(): """Passing no seq_lens/query_start_loc must leave the LSE untouched.""" torch.manual_seed(0) @@ -149,3 +152,14 @@ def test_mask_disabled_is_unmasked(): packed = _pack(out, lse, world_size, h_per_rank, head_dim, None, None) assert torch.isfinite(packed.float()).all() + + +def test_eager_mask_handles_rank_with_no_local_sequences(): + """Padded graph rows are all empty when a DCP rank has no sequences.""" + lse = torch.randn(8, 4, dtype=torch.float32) + seq_lens = torch.empty(0, dtype=torch.int32) + query_start_loc = torch.tensor([0], dtype=torch.int32) + + mask_dcp_empty_shards_(lse, seq_lens, query_start_loc) + + assert torch.isneginf(lse).all() diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py new file mode 100644 index 0000000000..5b63c2c609 --- /dev/null +++ b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only tests for FlashInfer sparse MLA backend constraints.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + FlashInferMLASparseTRTLLMBackend, + _required_workspace_bytes, + compute_trtllm_sparse_mla_workspace_bytes, +) + +# Reporter config from vllm#50781: GLM-5.2 (64 q heads), TP=8, DCP=8, +# max_num_batched_tokens=16384, FlashInfer 0.6.14. +_REPORTER_HEADS_PER_RANK = 64 // 8 +_REPORTER_DCP = 8 +_REPORTER_MAX_TOKENS = 16384 +# From the reported crash: trtllm-gen requested exactly this many bytes for +# trtllm_gen_softmax_workspace after an 8 MiB counter carve had left +# 404,750,336 of the 413,138,944-byte default buffer. +_REPORTER_OBSERVED_SOFTMAX_BYTES = 1_611_661_312 +_REPORTER_OBSERVED_PRECARVE_BYTES = 8_388_608 + + +def test_reporter_softmax_carve_is_reproduced_byte_exactly(): + # The crashing step scheduled 12288 tokens (a single 12K-token request); + # the softmax slab formula must reproduce the kernel's request exactly: + # 8 * (8 heads/rank * 8 dcp) * 12288 tokens * 256 + 1 MiB guard. + softmax_bytes = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=0, + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=12288, + ) + assert softmax_bytes == _REPORTER_OBSERVED_SOFTMAX_BYTES + + +def test_reporter_config_covers_observed_overflow(): + computed = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert computed >= ( + _REPORTER_OBSERVED_SOFTMAX_BYTES + _REPORTER_OBSERVED_PRECARVE_BYTES + ) + + +def test_non_dcp_size_is_unchanged(): + computed = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size=1, + num_heads_per_rank=128, + max_num_batched_tokens=65536, + ) + assert computed == _DEFAULT_WORKSPACE_BUFFER_SIZE + + +def test_combined_pcp_dcp_rejects_flashinfer_sparse(): + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=4, + decode_context_parallel_size=4, + ), + model_config=None, + ) + with patch("vllm.config.get_current_vllm_config", return_value=config): + reason = FlashInferMLASparseTRTLLMBackend.supports_combination( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + device_capability=SimpleNamespace(major=10), + ) + + assert reason is not None + assert "use FLASHMLA_SPARSE" in reason + + +@pytest.mark.parametrize("pcp_size,dcp_size", [(4, 1), (1, 4)]) +def test_single_context_parallel_mode_allows_flashinfer_sparse( + pcp_size: int, dcp_size: int +): + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=pcp_size, + decode_context_parallel_size=dcp_size, + ), + model_config=None, + ) + with patch("vllm.config.get_current_vllm_config", return_value=config): + reason = FlashInferMLASparseTRTLLMBackend.supports_combination( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + device_capability=SimpleNamespace(major=10), + ) + + assert reason is None + + +def test_default_constant_matches_envs_default(monkeypatch): + from vllm import envs + + monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False) + assert _DEFAULT_WORKSPACE_BUFFER_SIZE == envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + + +def test_env_unset_returns_computed(monkeypatch): + monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False) + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == compute_trtllm_sparse_mla_workspace_bytes( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + _REPORTER_DCP, + _REPORTER_HEADS_PER_RANK, + _REPORTER_MAX_TOKENS, + ) + + +def test_env_override_below_computed_is_respected_with_warning(monkeypatch): + override = 100 * 1024 * 1024 + monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override)) + with patch( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger" + ) as mock_logger: + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == override + mock_logger.warning_once.assert_called_once() + + +def test_env_override_above_computed_is_respected_without_warning(monkeypatch): + override = 8 * 1024 * 1024 * 1024 + monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override)) + with patch( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger" + ) as mock_logger: + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == override + mock_logger.warning_once.assert_not_called() diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index eab4759640..812bce39d8 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -307,6 +307,23 @@ def test_get_dcp_local_seq_lens_preserves_mtp_bounds_shape(): torch.testing.assert_close(actual, expected) +def test_get_dcp_local_seq_lens_reuses_device_rank_tensor(monkeypatch): + seq_lens = torch.tensor([0, 1, 7, 8, 17], dtype=torch.int32) + world = 4 + rank = 2 + rank_tensor = torch.tensor(rank, dtype=torch.int32) + expected = get_dcp_local_seq_lens(seq_lens, world, rank, 2) + + def fail_tensor_construction(*args, **kwargs): + raise AssertionError("rank tensor must be reused") + + with monkeypatch.context() as patch: + patch.setattr(torch, "tensor", fail_tensor_construction) + actual = get_dcp_local_seq_lens(seq_lens, world, rank_tensor, 2) + + torch.testing.assert_close(actual, expected) + + def test_get_dcp_local_seq_lens_must_run_after_decode_expansion(): world = 2 rank = 1 diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 87df1e409c..f3b4b7f5ba 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -179,6 +179,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): ), parallel_config=SimpleNamespace( decode_context_parallel_size=1, + prefill_context_parallel_size=1, cp_kv_cache_interleave_size=1, ), ) @@ -193,6 +194,37 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): } +def test_indexer_warmup_includes_pcp_normalized_dcp_key(monkeypatch): + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.indexer.get_dcp_group", + lambda: SimpleNamespace(rank_in_group=2), + ) + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace(compress_ratios=[32], index_kpool=32) + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=4, + cp_kv_cache_interleave_size=1, + ), + ) + + keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config) + + # Triton's compile key normalizes generic i32 values to 2 and divisible + # i32 values (including zero) to 16. + assert {(key.dcp_rank, key.dcp_world) for key in keys} == {(2, 2), (16, 1)} + assert { + ( + key.input_variant.is_aligned("uncompressed_seq_lens"), + key.input_variant.is_aligned("cu_compressed_seq_lens"), + ) + for key in keys + } == {(False, False), (False, True), (True, False), (True, True)} + + def test_compressed_slot_mapping_warmup_includes_index_kpool(): config = SimpleNamespace( cache_config=SimpleNamespace(block_size=256), diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index ae2e512840..ba10a2e91c 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -113,6 +113,16 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase): from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + if ( + parallel_config.prefill_context_parallel_size > 1 + and parallel_config.decode_context_parallel_size > 1 + ): + return ( + "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; " + "use FLASHMLA_SPARSE, which gathers each DCP KV shard before " + "running the rank-local PCP prefill queries" + ) if kv_cache_dtype == "fp8_ds_mla": return ( "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype" @@ -265,6 +275,31 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS hisparse_supports_multi_token_decode: ClassVar[bool] = True + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + + # Under DCP the workspace must hold the trtllm-gen softmax-stats slab. + # The buffer address is baked into CUDA graphs, so it has to reach its + # final size here, before the first forward/capture (no lazy regrow). + if self.dcp_world_size > 1: + num_q_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + _get_workspace_buffer( + device, + _required_workspace_bytes( + self.dcp_world_size, + num_q_heads, + vllm_config.scheduler_config.max_num_batched_tokens, + ), + ) + def _build_req_id_per_token( self, common_attn_metadata: "CommonAttentionMetadata", @@ -275,14 +310,98 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde # Global workspace buffer (lazily initialized) _fi_sparse_workspace: torch.Tensor | None = None +# trtllm-gen carves a softmax-stats slab from the workspace whenever LSE is +# requested (FlashInfer csrc/trtllm_fmha_kernel_launcher.cu, unchanged from +# v0.6.14 through current main): +# sizeof(float2) * num_qo_heads * batch_size * round_up(max_q_len, 256) +# + 1 MiB guard +# forward_mqa always passes q_len == 1, so each (head, token) pair costs +# round_up(1, 256) == 256 slots. +_TRTLLM_GEN_SOFTMAX_STAT_BYTES = 8 # sizeof(float2) +_TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN = 256 # round_up(q_len=1, 256) +_TRTLLM_GEN_SOFTMAX_GUARD_BYTES = 1024 * 1024 -def _get_workspace_buffer(device: torch.device) -> torch.Tensor: +# Keep in sync with the VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE default in +# vllm/envs.py. +_DEFAULT_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024 + + +def compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes: int, + dcp_world_size: int, + num_heads_per_rank: int, + max_num_batched_tokens: int, +) -> int: + """Workspace bytes needed by the trtllm-gen sparse MLA decode kernel. + + Under DCP the query is all-gathered across the DCP group in the head dim + and the kernel is asked for LSE, so trtllm-gen carves the softmax-stats + slab described above from the workspace before its (batch-independent) + counter and scratch regions. ``base_workspace_bytes`` must stay available + for those regions, so the slab is added on top of it (see #50781). + + Without DCP no LSE is requested, no slab is carved, and the base size is + returned unchanged. + """ + if dcp_world_size <= 1: + return base_workspace_bytes + softmax_bytes = ( + _TRTLLM_GEN_SOFTMAX_STAT_BYTES + * (num_heads_per_rank * dcp_world_size) + * max_num_batched_tokens + * _TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN + + _TRTLLM_GEN_SOFTMAX_GUARD_BYTES + ) + return base_workspace_bytes + softmax_bytes + + +def _required_workspace_bytes( + dcp_world_size: int, + num_heads_per_rank: int, + max_num_batched_tokens: int, +) -> int: + """Resolve the workspace size, honoring an explicit env override.""" + computed = compute_trtllm_sparse_mla_workspace_bytes( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size, + num_heads_per_rank, + max_num_batched_tokens, + ) + if not envs.is_set("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE"): + return computed + env_bytes = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + if env_bytes < computed: + logger.warning_once( + "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d is below the %d bytes " + "computed for the sparse MLA workspace with dcp_world_size=%d, " + "%d heads/rank and max_num_batched_tokens=%d. Respecting the " + "override, but the trtllm-gen kernel may crash with a workspace " + "overflow; set VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d or unset " + "it to use the computed size.", + env_bytes, + computed, + dcp_world_size, + num_heads_per_rank, + max_num_batched_tokens, + computed, + ) + return env_bytes + + +def _get_workspace_buffer( + device: torch.device, min_bytes: int | None = None +) -> torch.Tensor: global _fi_sparse_workspace - if _fi_sparse_workspace is None: + required = ( + min_bytes + if min_bytes is not None + else envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + ) + if _fi_sparse_workspace is None or _fi_sparse_workspace.numel() < required: # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace; # the trtllm-gen path views it as uint8, so int8 is safe for all backends. _fi_sparse_workspace = torch.zeros( - envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, + required, dtype=torch.int8, device=device, ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2321cf8d13..28bd450da1 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -10,7 +10,7 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger -from vllm.model_executor.warmup.jit_warmup import kernel_launcher +from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonPointerInputVariant, @@ -493,6 +493,14 @@ class BuildPrefillChunkMetadataKernel( dcp_world = parallel_config.decode_context_parallel_size dcp_interleave = parallel_config.cp_kv_cache_interleave_size dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0 + dcp_cases = [dict(DCP_RANK=dcp_rank, DCP_WORLD=dcp_world)] + if parallel_config.prefill_context_parallel_size > 1 and dcp_world > 1: + # PCP+DCP prefill constructs a global PCP chunk plan and therefore + # dispatches this kernel with already-localized row starts. The + # runtime key is normalized to rank 0/world 1 in that path; warm it + # explicitly so one PCP rank cannot enter a collective while + # another rank is still compiling the first real inference batch. + dcp_cases.append(dict(DCP_RANK=0, DCP_WORLD=1)) compress_ratios = tuple( dict.fromkeys( max(1, int(ratio)) @@ -506,17 +514,33 @@ class BuildPrefillChunkMetadataKernel( if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios: compress_ratios = compress_ratios + (index_kpool,) return self._trace_dispatch(self.dispatch)( + zip_inputs(*dcp_cases), # Cover Triton's divisible, exact-one, and generic i32 classes. query_slice_start=(0, 1, 2), query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens), - DCP_RANK=dcp_rank, - DCP_WORLD=dcp_world, DCP_INTERLEAVE=dcp_interleave, BLOCK_SIZE=self.BLOCK_SIZE, COMPRESS_RATIO=list(compress_ratios), + # PCP's global cumulative lengths are the second row of one packed + # allocation, so their pointer is not always 16-byte aligned. + # Prefill chunking can independently unalign uncompressed lengths. input_variant=( - TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True), - TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=True, + cu_compressed_seq_lens=True, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=True, + cu_compressed_seq_lens=False, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=False, + cu_compressed_seq_lens=True, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=False, + cu_compressed_seq_lens=False, + ), ), ) @@ -527,7 +551,9 @@ class BuildPrefillChunkMetadataKernel( uncompressed_seq_lens=compile_key.input_variant.pointer( "uncompressed_seq_lens", torch.int32 ), - cu_compressed_seq_lens=int32_ptr, + cu_compressed_seq_lens=compile_key.input_variant.pointer( + "cu_compressed_seq_lens", torch.int32 + ), row_start_cu_compressed_seq_lens=int32_ptr, token_to_seq=int32_ptr, cu_compressed_seq_len_ks=int32_ptr, @@ -836,6 +862,12 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): dtype=torch.int32, device=self.device, ) + # Materialize the rank on device during builder initialization. Creating + # this scalar in build() would introduce a GPU<->CPU sync in the decode + # hot path. + self.dcp_rank_tensor = torch.tensor( + self.dcp_rank, dtype=torch.int32, device=self.device + ) self.expanded_block_table_buffer = torch.zeros( (scheduler_config.max_num_batched_tokens, block_table_width), dtype=torch.int32, @@ -889,7 +921,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): local_seq_lens = get_dcp_local_seq_lens( seq_lens, self.dcp_world_size, - self.dcp_rank, + self.dcp_rank_tensor, self.cp_kv_cache_interleave_size, ) if seq_lens_is_buffer_view: diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index a3adf70d25..e656343a91 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -1094,7 +1094,7 @@ def compute_causal_conv1d_metadata( def get_dcp_local_seq_lens( seq_lens: torch.Tensor, dcp_size: int = 1, - dcp_rank: int | None = None, + dcp_rank: int | torch.Tensor | None = None, cp_kv_cache_interleave_size: int = 1, ) -> torch.Tensor: """While using dcp, kv_cache size stored on each rank may be different, @@ -1112,6 +1112,12 @@ def get_dcp_local_seq_lens( dcp_size, ) seq_lens_tiled = seq_lens_i32.unsqueeze(-1) + elif isinstance(dcp_rank, torch.Tensor): + assert dcp_rank.dtype == torch.int32 + assert dcp_rank.device == seq_lens.device + assert dcp_rank.numel() == 1 + rank_offsets = dcp_rank + seq_lens_tiled = seq_lens_i32 else: rank_offsets = torch.tensor(dcp_rank, dtype=torch.int32, device=seq_lens.device) seq_lens_tiled = seq_lens_i32 diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index bd20c9020d..58ac1e14a0 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -77,6 +77,13 @@ def mask_dcp_empty_shards_( return assert seq_lens is not None and query_start_loc is not None + # A DCP rank can receive no local sequences during CUDA graph warmup even + # though the padded LSE buffer still has rows. In that case every row is an + # empty shard; avoid indexing the empty seq_lens tensor below. + if seq_lens.shape[0] == 0: + lse.fill_(float("-inf")) + return + row_indices = torch.arange( lse.shape[0], device=lse.device, dtype=query_start_loc.dtype ) From 35dc273072ff133743bc5fcfa306bf030fa8730b Mon Sep 17 00:00:00 2001 From: ppalanga Date: Tue, 15 Sep 2026 10:18:09 -0700 Subject: [PATCH 109/139] [ROCm][Spec Decode] Add Aiter MLA decode support non-causal draft block (#55966) Signed-off-by: Poovaiah Palangappa Co-authored-by: Poovaiah Palangappa Co-authored-by: John Qin Co-authored-by: Cursor Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> --- .../test_rocm_aiter_mla_fp8_support.py | 117 +++++++++++++- .../test_rocm_aiter_mla_fp8_decode_routing.py | 5 + .../test_rocm_aiter_mla_mtp_split.py | 98 ++++++++++++ vllm/_aiter_ops.py | 27 +++- .../attention/backends/mla/rocm_aiter_mla.py | 148 +++++++++++++++++- 5 files changed, 390 insertions(+), 5 deletions(-) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py index d67df8ee2b..d5c37724aa 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py @@ -30,7 +30,10 @@ def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None: def _install_fake_aiter_modules( - monkeypatch: pytest.MonkeyPatch, *, supports_fp8: bool + monkeypatch: pytest.MonkeyPatch, + *, + supports_fp8: bool, + supports_causal: bool = False, ) -> None: aiter_mod: Any = types.ModuleType("aiter") mla_mod: Any = types.ModuleType("aiter.mla") @@ -65,6 +68,14 @@ def _install_fake_aiter_modules( mla_decode_fwd = mla_decode_fwd_without_fp8 + if supports_causal: + inner = mla_decode_fwd + + def mla_decode_fwd_with_causal(*args, causal=True, **kwargs): + return inner(*args, **kwargs) + + mla_decode_fwd = mla_decode_fwd_with_causal + mla_mod.mla_decode_fwd = mla_decode_fwd aiter_mod.mla = mla_mod @@ -135,3 +146,107 @@ def test_aiter_mla_fp8_support_result_is_cached(monkeypatch): assert _check_aiter_mla_fp8_support() is True assert _check_aiter_mla_fp8_support() is True assert signature_mock.call_count == 1 + + +def test_missing_causal_arg_fails_closed(monkeypatch): + """Old aiter wheels have no causal=; do not report that as a valid draft path.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=False) + with pytest.raises(RuntimeError, match="causal-only"): + rocm_aiter_ops.mla_decode_supports_non_causal() + with pytest.raises(RuntimeError, match="causal-only"): + AiterMLABackend.supports_non_causal() + + +@pytest.mark.parametrize("has_kernels", [False, True]) +def test_supports_non_causal_requires_arch_kernels(monkeypatch, has_kernels): + """The Python argument alone does not guarantee an arch has the kernel.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla import rocm_aiter_mla + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True) + monkeypatch.setattr( + rocm_aiter_mla, + "_aiter_mla_non_causal_asm_kernels", + lambda: has_kernels, + ) + assert rocm_aiter_ops.mla_decode_supports_non_causal() is True + assert AiterMLABackend.supports_non_causal() is has_kernels + + +@pytest.mark.parametrize("causal", [True, False]) +def test_decode_fwd_forwards_causal_to_aiter(monkeypatch, causal): + """The decode op must pass the block mask through; omitting it is causal-only.""" + seen: list = [] + + def mla_decode_fwd(*args, causal=None, **kwargs): + seen.append(causal) + + mla_mod: Any = types.ModuleType("aiter.mla") + mla_mod.mla_decode_fwd = mla_decode_fwd + aiter_mod: Any = types.ModuleType("aiter") + aiter_mod.mla = mla_mod + monkeypatch.setitem(sys.modules, "aiter", aiter_mod) + monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod) + + import torch + + import vllm._aiter_ops as aiter_ops + + aiter_ops._rocm_aiter_mla_decode_fwd_impl( + torch.zeros(2, 1, 8), + torch.zeros(2, 1, 1, 8), + torch.zeros(2, 1, 8), + torch.zeros(3, dtype=torch.int32), + 1, + causal=causal, + ) + assert seen == [causal] + + +def _aiter_mla_validate( + monkeypatch, + *, + dtype, + kv_cache_dtype, + use_non_causal, +): + from vllm.platforms.interface import DeviceCapability + from vllm.v1.attention.backend import AttentionType + from vllm.v1.attention.backends.mla import rocm_aiter_mla + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True) + monkeypatch.setattr( + rocm_aiter_mla, "_aiter_mla_non_causal_asm_kernels", lambda: True + ) + return AiterMLABackend.validate_configuration( + head_size=576, + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + block_size=1, + use_mla=True, + has_sink=False, + use_sparse=False, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(9, 5), + attn_type=AttentionType.DECODER, + use_non_causal=use_non_causal, + ) + + +def test_non_causal_fp16_auto_cache_is_rejected(monkeypatch): + """Pinned AITER aborts fp16 Q; auto-select must fall through to Triton.""" + import torch + + reasons = _aiter_mla_validate( + monkeypatch, + dtype=torch.float16, + kv_cache_dtype="auto", + use_non_causal=True, + ) + assert any("non-causal fp16" in reason for reason in reasons) diff --git a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py index 06684302cf..a86c7e617b 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py +++ b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py @@ -176,3 +176,8 @@ def test_pad_unpad_round_trip_preserves_head_order(num_heads): assert unpadded.shape == q.shape torch.testing.assert_close(unpadded, q) + + +def test_a_non_causal_block_never_routes_to_gluon(gluon_available): + """A small-head block that normally uses Gluon must use ASM when non-causal.""" + assert not AiterMLAHelper.use_gluon_verify(12, 8, "auto", causal=False) diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index 112cdff324..efffcad927 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -130,6 +130,7 @@ def _builder( _mla_q_dtype=torch.bfloat16, _mla_kv_dtype=torch.bfloat16, decode_attn_out_dtype=torch.bfloat16, + _decode_causal=True, ) # Bound to the stub rather than faked: the verify flatten's per-row view is # part of what _build_decode is being tested for. @@ -155,6 +156,25 @@ def test_backend_declares_uniform_batch_support(): ) +@pytest.mark.parametrize("causal", [True, False]) +def test_build_hands_common_causal_flag_to_aiter_builder(monkeypatch, causal): + seen = [] + + def fake_common_build(self, *args): + seen.append(self._decode_causal) + return SimpleNamespace(decode=None, prefill=None) + + monkeypatch.setattr( + rocm_aiter_mla.MLACommonMetadataBuilder, "build", fake_common_build + ) + builder = object.__new__(AiterMLAMetadataBuilder) + builder._fp8_prefill_enabled = False + + builder.build(0, SimpleNamespace(causal=causal)) + + assert seen == [causal] + + def test_dcp_verify_row_view_is_causal_per_row(): """Row t must cover this rank's share of positions [0, seq_len - qlen + t]. @@ -949,3 +969,81 @@ def test_persistent_metadata_gate_without_gluon_build( assert metadata.has_persistent_metadata is expect_persistent assert get_mla_metadata_v1.called is expect_persistent + + +def _build_non_causal(monkeypatch, *, num_heads, kv_cache_dtype, qlen, mtp_qlen): + """Drive _build_decode with causal=False and hand back the aiter mock.""" + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + builder = _builder( + mtp_decode_qlen=mtp_qlen, + kv_cache_dtype=kv_cache_dtype, + num_heads=num_heads, + ) + builder._decode_causal = False + metadata = AiterMLAMetadataBuilder._build_decode( + builder, + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=torch.tensor([7, 5], dtype=torch.int32), + max_seq_len=7, + query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), + num_decode_tokens=2 * qlen, + dcp_tot_seq_lens_device=None, + ) + return metadata, get_mla_metadata_v1 + + +def test_non_causal_build_hands_the_mask_to_aiter(monkeypatch): + """The schedule carries the mask, so is_causal must follow the block.""" + metadata, get_mla_metadata_v1 = _build_non_causal( + monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=8, mtp_qlen=8 + ) + assert metadata.has_persistent_metadata + # 6th positional arg of get_mla_metadata_v1 is is_causal. + assert get_mla_metadata_v1.call_args.args[5] is False + + +@pytest.mark.parametrize("num_heads", [12, 16, 48]) +def test_a_two_token_fp8_block_is_refused_without_a_fold(monkeypatch, num_heads): + """16-head (and padded-to-16) qlen-2 fp8 has no non-causal kernel. + + 48/80/112 (and 33 padded to 48) fold to H16 while keeping Q2, so they + must be refused the same way as native 16-head qlen 2. + """ + with pytest.raises(ValueError, match="2-token"): + _build_non_causal( + monkeypatch, + num_heads=num_heads, + kv_cache_dtype="fp8", + qlen=2, + mtp_qlen=8, + ) + + +def test_a_two_token_fp8_block_is_allowed_when_folded(monkeypatch): + """32 heads at qlen 2 folds onto the 16-head / 4-token kernel.""" + metadata, get_mla_metadata_v1 = _build_non_causal( + monkeypatch, + num_heads=32, + kv_cache_dtype="fp8", + qlen=2, + mtp_qlen=8, + ) + assert metadata.has_persistent_metadata + assert get_mla_metadata_v1.called + + +def test_a_two_token_bf16_block_is_allowed(monkeypatch): + """The bf16 fold has no such hole, so the guard must not fire there.""" + metadata, _ = _build_non_causal( + monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8 + ) + assert metadata.has_persistent_metadata diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 4cfb3b3884..9bea4f8b3a 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -583,10 +583,11 @@ def _rocm_aiter_mla_decode_fwd_impl( reduce_indptr: torch.Tensor | None = None, reduce_final_map: torch.Tensor | None = None, reduce_partial_map: torch.Tensor | None = None, + causal: bool = True, ) -> None: from aiter.mla import mla_decode_fwd - kwargs: dict[str, float | torch.Tensor | None] = { + kwargs: dict[str, float | torch.Tensor | None | bool] = { "sm_scale": sm_scale, "logit_cap": logit_cap, } @@ -596,6 +597,8 @@ def _rocm_aiter_mla_decode_fwd_impl( kwargs["q_scale"] = q_scale kwargs["kv_scale"] = kv_scale + kwargs["causal"] = causal + if work_meta_data is not None: assert work_indptr is not None, ( "work_indptr must be provided with work_meta_data" @@ -2036,6 +2039,26 @@ class rocm_aiter_ops: def is_mla_enabled(cls) -> bool: return cls._AITER_ENABLED and cls._MLA_ENABLED + @classmethod + @if_aiter_supported + def mla_decode_supports_non_causal(cls) -> bool: + """Whether installed aiter.mla.mla_decode_fwd accepts `causal`. + + Added in aiter v0.1.20. A missing argument is a hard error rather + than a boolean the caller can use to pick another MLA backend. + """ + import inspect + + from aiter.mla import mla_decode_fwd + + if "causal" not in inspect.signature(mla_decode_fwd).parameters: + raise RuntimeError( + "ROCM_AITER_MLA requires aiter.mla.mla_decode_fwd(..., causal=). " + "The installed aiter is causal-only; upgrade aiter rather than " + "falling back to another MLA backend." + ) + return True + @classmethod @if_aiter_supported def is_mha_enabled(cls) -> bool: @@ -2858,6 +2881,7 @@ class rocm_aiter_ops: reduce_indptr: torch.Tensor | None = None, reduce_final_map: torch.Tensor | None = None, reduce_partial_map: torch.Tensor | None = None, + causal: bool = True, ): torch.ops.vllm.rocm_aiter_mla_decode_fwd( q, @@ -2878,6 +2902,7 @@ class rocm_aiter_ops: reduce_indptr=reduce_indptr, reduce_final_map=reduce_final_map, reduce_partial_map=reduce_partial_map, + causal=causal, ) @staticmethod diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index b850deb070..c9ba2b0a56 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -4,7 +4,7 @@ import functools from dataclasses import dataclass from pathlib import Path -from typing import ClassVar, Final +from typing import TYPE_CHECKING, ClassVar, Final import torch @@ -38,6 +38,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_merge import ( ) from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache +if TYPE_CHECKING: + from vllm.platforms.interface import DeviceCapability + logger = init_logger(__name__) @@ -144,6 +147,20 @@ def _aiter_mla_native_h24_supported() -> bool: ) +def _aiter_mla_non_causal_asm_kernels() -> bool: + """Whether this arch ships non-causal MLA decode ASM kernels. + + The Python `causal=` probe cannot see the per-arch manifest. gfx950 has + the kernels; gfx942 and everything else must fall through to another + backend. + """ + try: + from vllm.platforms.rocm import on_gfx950 + except Exception: # noqa: BLE001 + return False + return bool(on_gfx950()) + + @functools.lru_cache(maxsize=1) def _gluon_mla_decode_supported() -> bool: """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build. @@ -252,6 +269,72 @@ class AiterMLABackend(MLACommonBackend): def get_name() -> str: return "ROCM_AITER_MLA" + @classmethod + def supports_non_causal(cls) -> bool: + # causal= on mla_decode_fwd is necessary but not sufficient: pinned + # AITER v0.1.21.post2 still exports that argument on gfx942, whose ASM + # manifest has no non-causal decode kernels. Selecting this backend + # there reaches `cannot get heuristic kernel!` instead of TRITON_MLA. + return bool( + rocm_aiter_ops.mla_decode_supports_non_causal() + and _aiter_mla_non_causal_asm_kernels() + ) + + @classmethod + def validate_configuration( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + use_per_head_quant_scales: bool, + device_capability: "DeviceCapability", + attn_type: str, + has_sliding_window: bool = False, + use_non_causal: bool = False, + use_batch_invariant: bool = False, + use_kv_connector: bool = False, + use_pcp: bool = False, + use_adaptive_verification: bool = False, + use_dcp: bool = False, + use_rswa: bool = False, + ) -> list[str]: + invalid_reasons = super().validate_configuration( + head_size=head_size, + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + block_size=block_size, + use_mla=use_mla, + has_sink=has_sink, + use_sparse=use_sparse, + use_mm_prefix=use_mm_prefix, + use_per_head_quant_scales=use_per_head_quant_scales, + device_capability=device_capability, + attn_type=attn_type, + has_sliding_window=has_sliding_window, + use_non_causal=use_non_causal, + use_batch_invariant=use_batch_invariant, + use_kv_connector=use_kv_connector, + use_pcp=use_pcp, + use_adaptive_verification=use_adaptive_verification, + use_dcp=use_dcp, + use_rswa=use_rswa, + ) + # Arch/signature cannot express this: the backend still advertises + # fp16, so a DSpark draft with kv_cache_dtype="auto" would select + # AITER and abort on fp16 query. Causal fp16 is fine. + if ( + use_non_causal + and dtype == torch.float16 + and kv_cache_dtype in (None, "auto", "float16") + ): + invalid_reasons.append("non-causal fp16 MLA decode not supported") + return invalid_reasons + @staticmethod def get_impl_cls() -> type["AiterMLAImpl"]: return AiterMLAImpl @@ -347,6 +430,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # https://github.com/vllm-project/vllm/issues/22945 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM + # Served by passing the mask to the kernel; _build_decode turns away the + # shapes AITER has no non-causal kernel for. + supports_non_causal_multi_token_decode: ClassVar[bool] = True + # Set from the common metadata every build; a batch is causal unless the + # drafter says otherwise. + _decode_causal: bool = True @staticmethod def _uniform_padded_mtp_qo_len( @@ -919,6 +1008,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): num_decode_tokens: int, dcp_tot_seq_lens_device: torch.Tensor | None, ) -> AiterMLADecodeMetadata: + causal = self._decode_causal device = self.device num_reqs = seq_lens_device.size(0) qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] @@ -970,6 +1060,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): int(max_qo_len), self._kv_cache_dtype_str, self.dcp_world_size, + causal, ) use_segmented_dcp_verify = ( self._supports_segmented_dcp_verify and max_qo_len > 1 @@ -1047,14 +1138,36 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # A padded rank has no bf16 persistent kernel past qlen 4 where the # gfx950 fold is absent; the non-persistent entry covers it. fp8 # keeps the schedule -- its fold rejects non-persistent outright. + # A non-causal block keeps it too: what the fold drops past qlen 4 + # is the block's causal staircase, which a non-causal block does + # not have, and the schedule is the only thing carrying its mask. and ( - self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS + not causal + or self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN or is_quantized_kv_cache(self._kv_cache_dtype_str) ) and max_qo_len >= 1 and max_qo_len <= self._mtp_decode_qlen ) + if ( + not causal + and max_qo_len == 2 + and is_quantized_kv_cache(self._kv_cache_dtype_str) + and not AiterMLAHelper.has_fp8_non_causal_qlen2_kernel( + self._decode_num_heads + ) + ): + # AITER's fp8 dispatch folds (gqa 16, qlen 3 or 4) onto the + # qseqlen-4 kernel but never lists 2. Only 32/64/96/128 heads + # at qlen 2 fold onto that 16-head / 4-token non-causal kernel; + # 16-head qlen 2, and padded 48/80/112, keep Q2 and abort the + # process rather than raising. The bf16 fold has no such hole. + raise ValueError( + "AITER has no non-causal fp8 MLA kernel for this 2-token " + "query block. Pin the draft to TRITON_MLA for this " + "speculative config." + ) if use_persistent_metadata: from aiter import get_mla_metadata_v1 @@ -1067,7 +1180,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): paged_kv_last_page_len, self._num_attention_heads, 1, - True, + causal, self._mla_work_meta_data, self._mla_work_info_set, self._mla_work_indptr, @@ -1137,6 +1250,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> AiterMLAMetadata: + # The common MLA builder owns the causal flag but its _build_decode + # contract is shared with CUDA backends. Keep this ROCm-specific state + # on the per-ubatch Aiter builder instead of changing that shared API; + # super().build invokes _build_decode synchronously. + self._decode_causal = common_attn_metadata.causal attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) @@ -1226,6 +1344,10 @@ class AiterMLAHelper: # fold that reaches a persistent one is gfx950-only. _ASM_PADDED_MAX_PS_QLEN: Final = 4 _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = () + # Pinned AITER v0.1.21.post2 folds these fp8 qlen-2 counts onto the + # non-causal 16-head / 4-token kernel. Other multiples of 16 (48/80/112) + # fold to H16 while keeping Q2, which has no non-causal fp8 entry. + _AITER_FP8_NON_CAUSAL_QLEN2_HEADS: ClassVar[tuple[int, ...]] = (32, 64, 96, 128) @staticmethod def qo_indptr_for_uniform_qlen( @@ -1263,6 +1385,18 @@ class AiterMLAHelper: ) ) + @staticmethod + def has_fp8_non_causal_qlen2_kernel(num_heads: int) -> bool: + """Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel. + + Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the + non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2 + block has no matching entry, and padded 48/80/112 keep Q2 after the + H16 fold. + """ + kernel_heads = AiterMLAHelper.get_actual_mla_num_heads(num_heads) + return kernel_heads in AiterMLAHelper._AITER_FP8_NON_CAUSAL_QLEN2_HEADS + @staticmethod def get_actual_mla_num_heads(num_heads: int) -> int: if num_heads == 24 and _aiter_mla_native_h24_supported(): @@ -1364,6 +1498,7 @@ class AiterMLAHelper: max_qo_len: int, kv_cache_dtype: str, dcp_world_size: int = 1, + causal: bool = True, ) -> bool: """Whether a small-head multi-token verify uses native Gluon MTP. @@ -1376,7 +1511,13 @@ class AiterMLAHelper: DCP verify is excluded: its per-row causal windows are served by the segmented path, which Gluon's MTP entry cannot express. + + Gluon masks the block causally with no way to turn it off, so a + non-causal block takes the padded asm decode whatever the head count + -- padding is what gives it a kernel there. """ + if not causal: + return False if max_qo_len <= 1 or dcp_world_size > 1: return False if is_quantized_kv_cache(kv_cache_dtype): @@ -1974,6 +2115,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): decode.paged_kv_indptr, decode.paged_kv_indices, decode.paged_kv_last_page_len, + causal=attn_metadata.causal, **mla_kwargs, ) From 2f46a1d9c6ab33c4ea8196c375ff85b661982bff Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 11:19:14 -0700 Subject: [PATCH 110/139] [CI] Keep Qwen3 Omni DSpark config fixture complete (#57044) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> --- tests/model_executor/test_qwen3_omni.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 42accbef85..30d9909a5a 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn from transformers import PretrainedConfig +from vllm.config import ParallelConfig from vllm.multimodal.processing import InputProcessingContext @@ -371,7 +372,7 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary(): attention_backend=None, kv_cache_dtype=None, ), - parallel_config=SimpleNamespace(tensor_parallel_size=1), + parallel_config=ParallelConfig(), attention_config=SimpleNamespace(backend=None), cache_config=SimpleNamespace(), load_config=SimpleNamespace(), From dffbb714e4e8e4b95ccc888df98d47c7d2cef78d Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:34:24 -0400 Subject: [PATCH 111/139] [ROCm][Perf] Optimize DSV4.1 K=512 decode top-k on gfx950 (#56743) Signed-off-by: fai --- csrc/libtorch_stable/sampler.cu | 263 ++++++++++++++---- tests/kernels/test_top_k_per_row.py | 189 +++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 12 + 3 files changed, 406 insertions(+), 58 deletions(-) diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index ada1c4ba39..b8d64f7af1 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -154,7 +154,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads, } template __device__ bool processHistogramStep( const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern, @@ -180,10 +181,15 @@ __device__ bool processHistogramStep( << patternShift; } + int negativeInfinityCount = 0; auto distributeToBins = [&](float logit, int /* idx */ = 0) { if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - atomicAdd(&smemFinal.histo.data[binIdx], 1); + if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) { + ++negativeInfinityCount; + } else { + atomicAdd(&smemFinal.histo.data[binIdx], 1); + } } }; @@ -198,6 +204,12 @@ __device__ bool processHistogramStep( distributeToBins(logit, idx); } } + if constexpr (useTopK512Optimization) { + if (negativeInfinityCount > 0) { + const auto bin = extractBinIdx(-INFINITY); + atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount); + } + } // Make sure the histogram is ready. __syncthreads(); @@ -254,6 +266,12 @@ __device__ bool processHistogramStep( // The threshold bin. thresholdBinIdx = smemThresholdBinIdx[0]; + // Resolve a -inf cutoff with exact tie emission, keeping sort padding out. + const bool finalBinFits = + smemFinalBinSize[0] <= kNumFinalItems && + !(useTopK512Optimization && step < 3 && + isPartialMatch(-INFINITY, logitPattern) && + thresholdBinIdx == extractBinIdx(-INFINITY)); auto processBins = [&](float logit, int idx) { if (isPartialMatch(logit, logitPattern)) { @@ -262,8 +280,7 @@ __device__ bool processHistogramStep( // 1. This is step 0 and the threshold bin is small enough (no step 1) // 2. This is step >= 1 (where pattern matching filters correctly) // This prevents duplicates when step 0 and step 1 both run. - bool shouldWriteDirectly = - (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1); + bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1); if (binIdx < thresholdBinIdx && shouldWriteDirectly) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -279,8 +296,7 @@ __device__ bool processHistogramStep( } if constexpr (step < 3) { // Only fill the final items for sorting if the threshold bin fits - if (binIdx == thresholdBinIdx && - smemFinalBinSize[0] <= kNumFinalItems) { + if (binIdx == thresholdBinIdx && finalBinFits) { int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1); smemFinal.items.logits[dstIdx] = logit; if constexpr (mergeBlocks) { @@ -325,20 +341,21 @@ __device__ bool processHistogramStep( __syncthreads(); // Check if we should continue to next step - return smemFinalBinSize[0] > kNumFinalItems; + return !finalBinFits; } // Follows half - 11 - 11 - 10 bit iterations // Keep the adaptive kernel's device instantiation separate from legacy calls. template + bool deviceLengthAware = false, bool useTopK512Optimization = false> static __device__ void topKPerRowJob(const int* indices, const float* logits, int rowStart, int rowEnd, int* outIndices, float* outLogits, int stride1, int topK) { static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks); + static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024); // The number of slots for the final pass. - static constexpr int kNumFinalItems = 2048; + static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 2048; // The number of elements per thread for the final sort. static constexpr int kNumFinalItemsPerThread = kNumFinalItems / kNumThreadsPerBlock; @@ -404,7 +421,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, rowIt += kNumThreadsPerBlock) { outIndices[rowIt] = -1; if constexpr (multipleBlocksPerRow) { - outLogits[rowIt] = -FLT_MAX; + outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX; } } @@ -422,7 +439,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 0: Process first 11 bits of half representation bool continueToNextStep = processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -431,7 +449,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 1: Process next 11 bits continueToNextStep = processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -441,7 +460,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 2: Process next 11 bits continueToNextStep = processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -450,7 +470,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, if (continueToNextStep) { // Step 3: Process last 10 bits processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -460,50 +481,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // The histogram did not proceed to the final 10 bits, therefore we need to // sort the final items The logits of the elements to be sorted in the final // pass. - if constexpr (useRadixSort) { - // Sorting with radix sort - float finalLogits[kNumFinalItemsPerThread]; - // The indices of the elements to be sorted in the final pass. - int finalIndices[kNumFinalItemsPerThread]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - finalLogits[ii] = -FLT_MAX; - } - - // Read the elements from SMEM. -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < smemFinalDstIdx[0]) { - finalLogits[ii] = smemFinal.items.logits[srcIdx]; - finalIndices[ii] = smemFinal.items.indices[srcIdx]; - } - } - // Make sure the shared memory has been read. - __syncthreads(); - - // Sort the elements. - FinalSort(smemFinal.finalSort) - .SortDescendingBlockedToStriped(finalLogits, finalIndices); - - // Copy the data back to the shared memory storage. - int baseIdx = smemFoundTopKValues[0]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - int dstIdx = baseIdx + srcIdx; - - if (dstIdx < topK) { - smemOutput[dstIdx] = finalIndices[ii]; - if constexpr (multipleBlocksPerRow) { - reinterpret_cast(smemOutput + topK)[dstIdx] = - finalLogits[ii]; - } - } - } - } else { + auto insertionSort = [&]() { // Sorting with insertion sort auto baseIdx = smemFoundTopKValues[0]; for (int i = threadIdx.x; i < smemFinalDstIdx[0]; @@ -525,6 +503,57 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, } } } + }; + if constexpr (useRadixSort) { + if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) { + insertionSort(); + } else { + // Sorting with radix sort + float finalLogits[kNumFinalItemsPerThread]; + // The indices of the elements to be sorted in the final pass. + int finalIndices[kNumFinalItemsPerThread]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX; + if constexpr (useTopK512Optimization) finalIndices[ii] = -1; + } + + // Read the elements from SMEM. +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + if (srcIdx < smemFinalDstIdx[0]) { + finalLogits[ii] = smemFinal.items.logits[srcIdx]; + finalIndices[ii] = smemFinal.items.indices[srcIdx]; + } + } + // Make sure the shared memory has been read. + __syncthreads(); + + // Sort the elements. + FinalSort(smemFinal.finalSort) + .SortDescendingBlockedToStriped(finalLogits, finalIndices); + + // Copy the data back to the shared memory storage. + int baseIdx = smemFoundTopKValues[0]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + int dstIdx = baseIdx + srcIdx; + + if (dstIdx < topK) { + smemOutput[dstIdx] = finalIndices[ii]; + if constexpr (multipleBlocksPerRow) { + reinterpret_cast(smemOutput + topK)[dstIdx] = + finalLogits[ii]; + } + } + } + } + } else { + insertionSort(); } __syncthreads(); } @@ -669,6 +698,70 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware( stride1, topK); } +#ifdef USE_ROCM +// The grid and workspace stay fixed while graph replays change device lengths. +__device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows, + int maxBlocks) { + int activeBlocks = 1; + if (numRows <= 32 && rowLength > 16 * 1024) { + activeBlocks = rowLength <= 128 * 1024 ? 8 + : rowLength <= 256 * 1024 ? 10 + : 16; + } else if (numRows <= 128 && rowLength > 64 * 1024) { + activeBlocks = rowLength <= 512 * 1024 ? 4 : 10; + } else if (numRows > 128 && numRows <= 256 && rowLength > 256 * 1024) { + activeBlocks = numRows <= 160 ? 3 : 2; + } + // Every block in a merged row must contain at least k real indices. + return min(min(activeBlocks, maxBlocks), max(1, rowLength / 512)); +} + +template +static __global__ +__launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware( + const float* logits, const int* seqLens, int* outIndices, + int* outIndicesAux, float* outLogitsAux, int stride0, int next_n, + int seqLensIs2D, int maxBlocks) { + #if defined(__gfx950__) + constexpr int topK = 512; + constexpr int kNumBins = 2048; + const int rowIdx = blockIdx.x; + const int rowEnd = + seqLensIs2D + ? max(0, seqLens[rowIdx]) + : max(0, seqLens[rowIdx / next_n] - next_n + rowIdx % next_n + 1); + const int activeBlocks = + gfx950TopK512ActiveBlocks(rowEnd, gridDim.x, maxBlocks); + outIndices += static_cast(rowIdx) * topK; + if constexpr (mergeBlocks) { + if (activeBlocks == 1) return; + const int64_t offset = static_cast(rowIdx) * maxBlocks * topK; + topKPerRowJob(outIndicesAux + offset, outLogitsAux + offset, 0, + activeBlocks * topK, outIndices, nullptr, 1, topK); + } else { + if (blockIdx.y >= activeBlocks) return; + logits += static_cast(rowIdx) * stride0; + if (activeBlocks == 1) { + topKPerRowJob(nullptr, logits, 0, rowEnd, outIndices, nullptr, 1, + topK); + return; + } + const int blockSize = rowEnd / activeBlocks; + const int rowStart = blockSize * blockIdx.y; + const int blockRowEnd = + blockIdx.y + 1 == activeBlocks ? rowEnd : rowStart + blockSize; + const int64_t offset = + (static_cast(rowIdx) * maxBlocks + blockIdx.y) * topK; + topKPerRowJob(nullptr, logits, rowStart, blockRowEnd, + outIndicesAux + offset, outLogitsAux + offset, 1, topK); + } + #endif +} +#endif + } // namespace vllm void apply_repetition_penalties_( @@ -731,6 +824,60 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, // the same seq_len and the kernel computes the per-row offset itself. int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0; +#ifdef USE_ROCM + const bool useGfx950DeviceLengthAwareTopK512 = + topK == 512 && stride1 == 1 && numRows > 0 && numRows <= 384 && + numColumns <= 1024 * 1024 && + std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0; + if (useGfx950DeviceLengthAwareTopK512) { + int multipleBlocksPerRowConfig = 1; + if (numRows <= 32 && numColumns > 16 * 1024) { + multipleBlocksPerRowConfig = 16; + } else if (numRows <= 128 && numColumns > 64 * 1024) { + multipleBlocksPerRowConfig = 10; + } else if (numRows <= 160 && numColumns > 256 * 1024) { + multipleBlocksPerRowConfig = 3; + } else if (numRows <= 256 && numColumns > 256 * 1024) { + multipleBlocksPerRowConfig = 2; + } + + constexpr int kNumThreadsPerSplitBlock = 1024; + if (multipleBlocksPerRowConfig == 1) { + vllm::topKPerRowDecode512DeviceLengthAware + <<>>(logits.const_data_ptr(), + seqLens.const_data_ptr(), + indices.mutable_data_ptr(), nullptr, nullptr, + static_cast(stride0), static_cast(next_n), + seqLensIs2D, multipleBlocksPerRowConfig); + return; + } + const auto outIndicesAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Int, std::nullopt, logits.device()); + const auto outLogitsAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Float, std::nullopt, logits.device()); + vllm::topKPerRowDecode512DeviceLengthAware + <<>>( + logits.const_data_ptr(), seqLens.const_data_ptr(), + indices.mutable_data_ptr(), + outIndicesAux.mutable_data_ptr(), + outLogitsAux.mutable_data_ptr(), static_cast(stride0), + static_cast(next_n), seqLensIs2D, multipleBlocksPerRowConfig); + constexpr int kNumThreadsPerBlockMerge = 1024; + vllm::topKPerRowDecode512DeviceLengthAware + <<>>( + nullptr, seqLens.const_data_ptr(), + indices.mutable_data_ptr(), + outIndicesAux.mutable_data_ptr(), + outLogitsAux.mutable_data_ptr(), 0, static_cast(next_n), + seqLensIs2D, multipleBlocksPerRowConfig); + return; + } +#endif + if (numColumns < kSortingAlgorithmThreshold) { // Use insertion sort vllm::topKPerRowDecode diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index c629ae15de..591362020e 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -22,11 +22,22 @@ def _has_device_capability(major: int) -> bool: return current_platform.is_cuda() and current_platform.has_device_capability(major) +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() + + # DeepSelect is compiled for sm_100a/sm_103a only. requires_sm100 = pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="DeepSelect requires SM100a/SM103a", ) +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), reason="This test exercises the gfx950 launch configuration" +) COOPERATIVE_TOPK_BACKEND = pytest.param( @@ -487,6 +498,150 @@ def test_top_k_per_row_decode_gfx950_long_c4a_1d_seq_lens() -> None: ) +def _assert_exact_topk( + logits: torch.Tensor, indices: torch.Tensor, row_ends: torch.Tensor +) -> None: + ends = row_ends.reshape(-1, 1).clamp_min(0) + valid = (indices >= 0) & (indices < ends) + assert torch.equal(valid.sum(1), ends.clamp_max(indices.shape[1]).flatten()) + assert torch.all(valid | (indices == -1)) + ordered = indices.sort(1).values + assert not ((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any() + + width = int(ends.max()) + visible = logits[:, :width].clone() + columns = torch.arange(width, device=logits.device) + visible.masked_fill_(columns[None, :] >= ends, float("-inf")) + expected = visible.topk(min(indices.shape[1], width), dim=1).values + selected = logits.gather(1, indices.clamp_min(0).long()) + selected.masked_fill_(~valid, float("-inf")) + selected = selected.sort(1, descending=True).values + assert torch.equal(selected[:, : expected.shape[1]], expected) + + +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_1d_seq_lens() -> None: + """Honor causal offsets and exact compressed lengths, including short rows.""" + next_n, width, top_k = 6, 16_384, 512 + lengths = torch.tensor( + [0, 1, 511, 512, 513, 1023, 1024, 16_000], + dtype=torch.int32, + device="cuda", + ) + offsets = torch.arange(next_n, dtype=torch.int32, device="cuda") + row_ends = (lengths[:, None] - next_n + offsets + 1).clamp_min(0).flatten() + rows = row_ends.numel() + storage = torch.full((rows, width + 8), 1e20, dtype=torch.float32, device="cuda") + logits = storage[:, 1 : width + 1] + logits[:, :16_000] = torch.arange(16_000, dtype=torch.float32, device="cuda") + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") + _run_topk_backend( + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n + ) + _assert_exact_topk(logits, indices, row_ends) + + +@pytest.mark.parametrize( + ("rows", "row_length"), + [ + pytest.param(128, 65_536, id="single-before-64k"), + pytest.param(128, 65_537, id="four-splits-after-64k"), + pytest.param(129, 262_144, id="single-before-three-split-length"), + pytest.param(129, 262_145, id="three-splits-after-256k"), + pytest.param(160, 262_145, id="three-splits-upper-row-bound"), + pytest.param(161, 262_144, id="single-before-two-split-length"), + pytest.param(161, 262_145, id="two-splits-after-256k"), + pytest.param(256, 262_145, id="two-splits-upper-row-bound"), + pytest.param(257, 262_145, id="large-row-single-block"), + pytest.param(384, 262_145, id="optimized-row-limit"), + pytest.param(385, 262_145, id="native-fallback"), + ], +) +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_split_policy_boundaries( + rows: int, row_length: int +) -> None: + width, top_k = 262_145, 512 + row = torch.arange(width, dtype=torch.float32, device="cuda") + logits = row.expand(rows, -1) + lengths = torch.full((rows,), row_length, dtype=torch.int32, device="cuda") + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") + + _run_topk_backend("top_k_per_row_decode", logits, lengths, indices, top_k, width) + + expected = torch.arange( + row_length - top_k, row_length, dtype=torch.int32, device="cuda" + ).expand_as(indices) + assert torch.equal(indices.sort(1).values, expected) + + +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None: + """Preserve visible top-k scores when masks and lengths change on replay.""" + rows, next_n, width, top_k = 24, 6, 262_145, 512 + logits = torch.full((rows, width), 1e20, dtype=torch.float32, device="cuda") + lengths = torch.full( + (rows // next_n, next_n), 70_000, dtype=torch.int32, device="cuda" + ) + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") + logits[:, :70_000] = float("-inf") + bin_sizes = [128, 129, 1024, 1025] + for row in range(rows): + pattern = row % 9 + if pattern < 4: + count = [0, 1, 511, 513][pattern] + logits[row, :count] = torch.arange( + count, dtype=torch.float32, device="cuda" + ) + elif pattern == 4: + logits[row, :511] = 0 + logits[row, 511] = torch.finfo(torch.float32).min + else: + count = bin_sizes[pattern - 5] + higher = max(0, top_k - count // 2) + logits[row, :higher] = 2 + logits[row, higher : higher + count] = 1 + logits[row, higher] = 1 + 2**-23 + + def run() -> None: + _run_topk_backend( + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n + ) + + run() + _assert_exact_topk(logits, indices, lengths) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + indices.fill_(-777) + graph.replay() + _assert_exact_topk(logits, indices, lengths) + + bounds = [ + 0, + 511, + 512, + 513, + 16_384, + 16_385, + 131_072, + 131_073, + 262_144, + 262_145, + ] + row_bounds = torch.tensor(bounds, dtype=torch.int32, device="cuda") + repeats = (rows + len(bounds) - 1) // len(bounds) + lengths.flatten().copy_(row_bounds.repeat(repeats)[:rows]) + logits[:] = torch.arange(width, dtype=torch.float32, device="cuda") + indices.fill_(-777) + graph.replay() + _assert_exact_topk(logits, indices, lengths) + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") @torch.inference_mode() def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: @@ -620,6 +775,31 @@ def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: compress_ratio=4, num_rows=257, max_valid_seq_len=250_000, + num_columns=524_288, + on_gfx950=True, + ), + decode_kernel, + ), + ( + dict( + is_prefill=False, + compress_ratio=2, + num_rows=385, + max_valid_seq_len=250_000, + num_columns=524_288, + topk_tokens=512, + on_gfx950=True, + ), + decode_kernel, + ), + ( + dict( + is_prefill=False, + compress_ratio=2, + num_rows=384, + max_valid_seq_len=50_000, + num_columns=1_048_577, + topk_tokens=512, on_gfx950=True, ), decode_kernel, @@ -638,6 +818,15 @@ def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: ) ineligible_cases = [ dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True), + dict( + is_prefill=False, + compress_ratio=2, + num_rows=384, + max_valid_seq_len=1_048_576, + num_columns=1_048_576, + topk_tokens=512, + on_gfx950=True, + ), dict( is_prefill=False, compress_ratio=4, diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 44b6017a05..0e342cb0e0 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -58,6 +58,7 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024 _GFX950_C4A_NATIVE_MAX_ROWS = 256 +_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024 # Conservative perf gate, not a correctness bound: OPUS is correct for any query # count, but Triton stays faster below this measured crossover. _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 @@ -69,6 +70,8 @@ def _get_aiter_top_k_kernel( compress_ratio: int, num_rows: int, max_valid_seq_len: int | None = None, + num_columns: int | None = None, + topk_tokens: int = 1024, on_gfx950: bool = _ON_GFX950, ) -> Callable[..., None] | None: if compress_ratio <= 1 or not on_gfx950: @@ -76,6 +79,13 @@ def _get_aiter_top_k_kernel( if not is_prefill: assert max_valid_seq_len is not None + if ( + topk_tokens == 512 + and 0 < num_rows <= 384 + and num_columns is not None + and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS + ): + return None # AITER v0.1.19 decode is one-block only. This measured gfx950 # FP32/k=1024 compressed-row boundary is independent of the native # split-count boundary in sampler.cu. @@ -1261,6 +1271,8 @@ def rocm_aiter_sparse_attn_indexer( compress_ratio=compress_ratio, num_rows=num_rows, max_valid_seq_len=max_compressed_seq_len, + num_columns=logits.shape[1], + topk_tokens=topk_tokens, ) if aiter_topk_kernel is not None: _launch_aiter_top_k_per_row_decode( From 2fcc5247882849b675088b3146b27c006e87d87f Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 15 Sep 2026 14:41:54 -0500 Subject: [PATCH 112/139] [Bugfix] Avoid nested score cancellation handlers for /v1/score alias (#57024) Signed-off-by: Taneem Ibrahim --- vllm/entrypoints/pooling/scoring/api_router.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/entrypoints/pooling/scoring/api_router.py b/vllm/entrypoints/pooling/scoring/api_router.py index 83ed600fdd..8e99e9a1fe 100644 --- a/vllm/entrypoints/pooling/scoring/api_router.py +++ b/vllm/entrypoints/pooling/scoring/api_router.py @@ -57,8 +57,6 @@ async def create_score(request: ScoreRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation -@load_aware_call async def create_score_v1(request: ScoreRequest, raw_request: Request): logger.warning( "To indicate that Score API is not part of standard OpenAI API, we " From df42d112ee88dd4a9b64efbad55621af6a66a44b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 15 Sep 2026 21:44:26 +0200 Subject: [PATCH 113/139] [Config] Infer HiSparse attention config from HiSparseConnector (#57041) Signed-off-by: NickLucche --- tests/test_config.py | 102 +++++++++++++++++++++++++++++++++++++++ vllm/config/attention.py | 4 +- vllm/config/vllm.py | 9 +++- 3 files changed, 113 insertions(+), 2 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 4d39b56d26..f75592f089 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -337,6 +337,108 @@ def test_hisparse_rejects_non_cuda(monkeypatch): VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig())) +@pytest.mark.parametrize( + "kv_transfer_config", + [ + KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ), + KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "connectors": [ + { + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, + }, + { + "kv_connector": "HiSparseConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"host_pool_gib": 128}, + }, + ] + }, + ), + KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_consumer", + kv_connector_extra_config={ + "connectors": [ + {"kv_connector": "NixlConnector", "kv_role": "kv_consumer"}, + { + "kv_connector": "HiSparseConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"host_pool_gib": 128}, + }, + ] + }, + ), + ], + ids=["standalone", "multi-connector", "pd-decode"], +) +def test_hisparse_connector_implies_attention_config(monkeypatch, kv_transfer_config): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + config = VllmConfig( + kv_transfer_config=kv_transfer_config, + # Skip the HMA auto-detect block: it imports the connector class, + # which pulls the platform's compiled attention extensions. + scheduler_config=SchedulerConfig( + max_model_len=2048, + is_encoder_decoder=False, + disable_hybrid_kv_cache_manager=False, + ), + ) + assert isinstance(config.attention_config.hisparse_config, HiSparseConfig) + + +def test_hisparse_connector_preserves_explicit_attention_config(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + config = VllmConfig( + attention_config=AttentionConfig( + hisparse_config=HiSparseConfig(device_buffer_size=512) + ), + kv_transfer_config=KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ), + scheduler_config=SchedulerConfig( + max_model_len=2048, + is_encoder_decoder=False, + disable_hybrid_kv_cache_manager=False, + ), + ) + assert config.attention_config.hisparse_config.device_buffer_size == 512 + + +def test_hisparse_connector_without_cuda_still_rejected(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: False) + with pytest.raises(ValueError, match="requires NVIDIA CUDA"): + VllmConfig( + kv_transfer_config=KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ) + ) + + +def test_no_hisparse_connector_keeps_attention_config_unset(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + config = VllmConfig( + kv_transfer_config=KVTransferConfig( + kv_connector="OffloadingConnector", kv_role="kv_both" + ) + ) + assert config.attention_config.hisparse_config is None + + def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled MRV1 paths and off breakable cudagraphs by default.""" diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 05da048bdb..129fbcbb08 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -88,7 +88,9 @@ class AttentionConfig: hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental - Model Runner V2-only HiSparse sparse-MLA decode hot-buffering.""" + Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred + with defaults when HiSparseConnector is configured (directly or via + MultiConnector); set it explicitly only to tune its fields.""" use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ad8f077148..8e0846982b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -27,7 +27,7 @@ from vllm.triton_utils import HAS_TRITON from vllm.utils import random_uuid from vllm.utils.hashing import safe_hash -from .attention import AttentionConfig +from .attention import AttentionConfig, HiSparseConfig from .cache import CacheConfig from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode from .device import DeviceConfig @@ -1630,6 +1630,13 @@ class VllmConfig: self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() + if ( + self.attention_config.hisparse_config is None + and self.kv_transfer_config is not None + and self.kv_transfer_config.has_connector("HiSparseConnector") + ): + self.attention_config.hisparse_config = HiSparseConfig() + if self.attention_config.hisparse_config is not None: if not current_platform.is_cuda(): raise ValueError("HiSparse currently requires NVIDIA CUDA.") From 18f8aa04656de0a37eb5c7e9efc8ebb595ac73cf Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 15 Sep 2026 16:17:53 -0400 Subject: [PATCH 114/139] [Refactor] Remove dead kernel code (#56845) Signed-off-by: yewentao256 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../layers/mamba/ops/replayssm_config.py | 15 - .../models/glm5next/amd/ops/kpool_compress.py | 44 - vllm/models/inkling/amd/ops/lamport.py | 766 ------------------ vllm/models/inkling/amd/ops/mm_towers.py | 190 ----- .../cute_dsl/latent_moe_tail/primitives.py | 31 - .../v1/attention/ops/triton_fp8_mqa_logits.py | 262 ------ .../triton_turboquant_decode_v2.py | 572 ------------- .../triton_turboquant_unified_attention.py | 15 - 8 files changed, 1895 deletions(-) delete mode 100644 vllm/models/inkling/amd/ops/lamport.py delete mode 100644 vllm/models/inkling/amd/ops/mm_towers.py delete mode 100644 vllm/v1/attention/ops/triton_fp8_mqa_logits.py diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py index aff4fc1c3f..3cee44ba50 100644 --- a/vllm/model_executor/layers/mamba/ops/replayssm_config.py +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -8,7 +8,6 @@ auto-detected (Blackwell vs not) so call sites need not thread it through. """ import functools -from contextlib import contextmanager from vllm.platforms import current_platform from vllm.triton_utils import triton @@ -26,20 +25,6 @@ def _is_blackwell() -> bool: _overrides: dict[str, tuple] = {} -@contextmanager -def override_replayssm_config(kernel: str, config: tuple): - """Pin ``kernel``'s launch config for the duration of the context.""" - prev = _overrides.get(kernel) - _overrides[kernel] = config - try: - yield - finally: - if prev is None: - _overrides.pop(kernel, None) - else: - _overrides[kernel] = prev - - def _dstate_tile(dstate: int, tile: int) -> int: return max(16, min(tile, triton.next_power_of_2(dstate))) diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 64e6fd46d7..09e7a40aa5 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -65,50 +65,6 @@ def _hadamard128(x): return x * 0.08838834764831845 # 1/sqrt(128) -# Map pool ids to physical cache slots. - - -def compute_pooled_write_locs( - page_table_64: torch.Tensor, - pool_ids: torch.Tensor, - pool_size: int, -) -> torch.Tensor: - """Map logical pooled-K ids to physical flat cache slots. - - ``pool_size`` consecutive tokens share one pool slot that lives at the - *first* token page of each page-group. ``page_table_64`` maps token pages - to physical block ids; we gather the block id of each pool's page-group - and add the in-block pool offset. - """ - assert page_table_64.ndim == 1 - pool_ids = pool_ids.to(torch.int64) - block_size = 64 - pool_page_group = torch.div(pool_ids, block_size, rounding_mode="floor") - token_page_row = pool_page_group * pool_size - packed_page = page_table_64.index_select(0, token_page_row.to(torch.int64)) - return packed_page.to(torch.int64) * block_size + torch.remainder( - pool_ids, block_size - ) - - -def build_pooled_page_table( - page_table: torch.Tensor, - pool_size: int, -) -> torch.Tensor: - """Build a pool-granular page table by taking every ``pool_size``-th - token-page column (one pool maps to ``pool_size`` token pages). - - Uses gather (not strided slicing) so the result is always a fresh - row-major tensor — some downstream kernels require stride(-1) == 1. - """ - block_size = page_table.shape[-1] - assert block_size % pool_size == 0, ( - f"pool_size ({pool_size}) must divide page columns ({block_size})" - ) - idx = torch.arange(0, block_size, pool_size, device=page_table.device) - return page_table[..., idx].contiguous() - - # Fused pool compression and cache write. diff --git a/vllm/models/inkling/amd/ops/lamport.py b/vllm/models/inkling/amd/ops/lamport.py deleted file mode 100644 index 9705a9a3df..0000000000 --- a/vllm/models/inkling/amd/ops/lamport.py +++ /dev/null @@ -1,766 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm. - -The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``. - -Liveness --------- -Large grids are deliberately split at the two communication dependencies: - - 1. ``_publish_input_kernel`` only publishes rank partials. - 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard. - 3. ``_sconv_publish_kernel`` only computes and publishes the local result. - 4. ``_gather_norm_kernel`` waits, gathers, and normalizes. - -Without PDL, CUDA stream order completes a producer before its consumer. With -PDL, every producer CTA posts all peer stores before triggering its dependent, -and the consumer executes ``gdc_wait`` before polling. A consumer therefore -waits only for stores from a producer that is already running (or complete) on -another GPU. It never waits for another block in its own grid. Consequently -spinning consumers cannot occupy resources needed by any producer, and the -wait-for graph has no cycle. The proof is independent of grid size, block -dispatch order, and occupancy. - -For one token, the first three phases use eight independent channel slices to -expose enough CTA parallelism for decode latency. Each slice has exclusive -ownership of its cache and Lamport columns; the gather/RMSNorm phase retains -one CTA per token so no cross-CTA reduction or completion counter is needed. - -Immediate buffer reuse (including replay of a captured CUDA graph) is also -safe. Rank R cannot republish input for call n+1 until its gather for n has -finished; that gather waited for owner O's output, which O publishes only -after consuming R's input for n. Likewise, R cannot republish output for n+1 -until its reduction for n+1 has observed destination D's input; D publishes -that input only after its gather consumed R's output for n. Thus every prior -read happens-before a same-slot rewrite. Three generations reduce incidental -coupling for ordinary launches, but correctness does not depend on rotation. - -The payload itself is the Lamport flag. A 32-bit store publishes two bf16s -atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real -negative zeroes are changed to positive zero before publication. Consumers -use volatile 32-bit loads and restore the sentinel after consuming a slot. -""" - -from __future__ import annotations - -import os - -import torch - -from vllm.distributed import get_tp_group -from vllm.distributed.parallel_state import in_the_same_node_as -from vllm.logger import init_logger -from vllm.triton_utils import tl, triton - -logger = init_logger(__name__) - -_MAX_TOKENS = 16384 -_EMPTY_PAIR = tl.constexpr(0x80008000) - - -@triton.jit -def _pack_bf16_pairs(values): - """Pack bf16 values into atomic u32 pairs and reserve negative zero.""" - lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2])) - lo = lo.to(tl.uint16, bitcast=True) - hi = hi.to(tl.uint16, bitcast=True) - lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32) - hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32) - return lo | (hi << 16) - - -@triton.jit -def _unpack_bf16_pairs(values): - lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) - hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) - return tl.interleave(lo, hi) - - -@triton.jit -def _wait_pairs(ptr, offsets, mask): - values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) - while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0: - values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) - return values - - -@triton.jit -def _publish_input_kernel( - stage_ptr, - peer_ptrs, - peer_offset_u32, - stride_stage_t, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Publish this rank's full partial row into every shard owner's slots.""" - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - elem = tl.arange(0, CS_P2) - elem_mask = elem < CSS - ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) - # Address generation is independent of the preceding stream kernel. The - # acquire remains before stage/peer loads, which may consume its writes. - if USE_PDL: - tl.extra.cuda.gdc_wait() - - for owner in tl.static_range(WORLD): - values = tl.load( - stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem, - mask=elem_mask, - other=0.0, - ) - packed = _pack_bf16_pairs(values) - base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32)) - dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair - tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _reduce_insert_kernel( - input_peer_ptrs, - input_peer_offset_u32, - cache_ptr, - slot_ptr, - stride_cache_block, - stride_cache_head, - stride_cache_token, - stride_cache_dim, - block_size, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - CACHE_OFFSET: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Consume all partials for this rank and insert the reduced cache row.""" - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - source = tl.arange(0, WORLD) - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - offsets = ( - (token * WORLD + source)[:, None] * (CS // 2) - + split * (CSS // 2) - + pair[None, :] - ) - mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :] - input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) - input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32)) - input_u32 += input_peer_offset_u32 - # Slot metadata and the cache destination do not depend on input publish. - slot = tl.load(slot_ptr + token) - valid = slot >= 0 - channel = tl.arange(0, CS_P2) - channel_mask = channel < CSS - global_channel = split * CSS + channel - head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) - dim = CACHE_OFFSET + global_channel % HEAD_SIZE - safe_slot = tl.maximum(slot, 0).to(tl.int64) - dst = ( - cache_ptr - + (safe_slot // block_size) * stride_cache_block - + head * stride_cache_head - + (safe_slot % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - if USE_PDL: - tl.extra.cuda.gdc_wait() - packed = _wait_pairs(input_u32, offsets, mask) - - lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) - hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) - reduced = tl.interleave( - tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16), - tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16), - ) - tl.store( - input_u32 + offsets, - tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32), - mask=mask, - ) - - tl.store(dst, reduced, mask=valid & channel_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _sconv_publish_kernel( - peer_ptrs, - peer_offset_u32, - residual_ptr, - weight_ptr, - cache_ptr, - position_ptr, - sequence_ptr, - slot_ptr, - block_table_ptr, - stride_residual_t, - stride_cache_block, - stride_cache_head, - stride_cache_token, - stride_cache_dim, - stride_block_table_r, - max_blocks, - block_size, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - CACHE_OFFSET: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - WINDOW: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Compute this rank's shard, then publish it to every rank.""" - tl.static_assert(WINDOW == 4) - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - channel = tl.arange(0, CS_P2) - channel_mask = channel < CSS - global_channel = split * CSS + channel - head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) - dim = CACHE_OFFSET + global_channel % HEAD_SIZE - slot = tl.load(slot_ptr + token) - valid = slot >= 0 - position = tl.load(position_ptr + token) - sequence = tl.load(sequence_ptr + token) - safe_slot = tl.maximum(slot, 0).to(tl.int64) - own_ptr = ( - cache_ptr - + (safe_slot // block_size) * stride_cache_block - + head * stride_cache_head - + (safe_slot % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - # These loads were made visible before reduce/insert could signal us and - # are independent of its cache store. Hoisting them gives PDL useful work - # to overlap while retaining the acquire before every cache read. - residual = tl.load( - residual_ptr + token * stride_residual_t + RANK * CS + global_channel, - mask=channel_mask, - other=0.0, - ) - weight0 = tl.load( - weight_ptr + global_channel * WINDOW, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight1 = tl.load( - weight_ptr + global_channel * WINDOW + 1, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight2 = tl.load( - weight_ptr + global_channel * WINDOW + 2, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight3 = tl.load( - weight_ptr + global_channel * WINDOW + 3, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) - if USE_PDL: - tl.extra.cuda.gdc_wait() - current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0) - - conv = tl.zeros([CS_P2], tl.float32) - for tap_idx in tl.static_range(WINDOW): - source_position = position - (WINDOW - 1) + tap_idx - take = valid & (source_position >= 0) - if tap_idx == WINDOW - 1: - value = tl.where(take, current.to(tl.float32), 0.0) - else: - safe_position = tl.maximum(source_position, 0) - logical_block = tl.minimum(safe_position // block_size, max_blocks - 1) - physical_block = tl.load( - block_table_ptr + sequence * stride_block_table_r + logical_block, - mask=take, - other=0, - ).to(tl.int64) - source_ptr = ( - cache_ptr - + physical_block * stride_cache_block - + head * stride_cache_head - + (safe_position % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0) - value = tl.where(take, cached.to(tl.float32), 0.0) - if tap_idx == 0: - weight = weight0 - elif tap_idx == 1: - weight = weight1 - elif tap_idx == 2: - weight = weight2 - else: - weight = weight3 - conv += value * weight - - # Preserve both bf16 rounding points of the original sublayer. - short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16) - output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to( - tl.bfloat16 - ) - - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - packed = _pack_bf16_pairs(output) - row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair - - for destination in tl.static_range(WORLD): - base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32)) - tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _gather_norm_kernel( - output_peer_ptrs, - output_peer_offset_u32, - norm_weight_ptr, - normed_ptr, - residual_out_ptr, - eps, - stride_output_t, - C: tl.constexpr, - C_P2: tl.constexpr, - RANK: tl.constexpr, - HAS_NORM: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Consume a complete row; one CTA owns all outputs for one token.""" - token = tl.program_id(0).to(tl.int64) - pair = tl.arange(0, C_P2 // 2) - pair_mask = pair < C // 2 - output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) - output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32)) - output_u32 += output_peer_offset_u32 - offsets = token * (C // 2) + pair - channel = tl.arange(0, C_P2) - channel_mask = channel < C - # Gamma is independent of the preceding sconv publication. Keep the - # acquire immediately before polling the Lamport output slots. - if HAS_NORM: - weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0) - if USE_PDL: - tl.extra.cuda.gdc_wait() - packed = _wait_pairs(output_u32, offsets, pair_mask) - row = _unpack_bf16_pairs(packed) - tl.store( - residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask - ) - if HAS_NORM: - row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0) - inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps) - tl.store( - normed_ptr + token * stride_output_t + channel, - (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16), - mask=channel_mask, - ) - tl.store( - output_u32 + offsets, - tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32), - mask=pair_mask, - ) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _validate_lamport_init_kernel( - input_peer_ptrs, - output_peer_ptrs, - bad_ptr, - num_pairs, - RANK: tl.constexpr, - BLOCK: tl.constexpr, -): - """Validate both complete local allocations through their fabric pointers.""" - offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) - mask = offsets < num_pairs - ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) - ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) - local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32)) - local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32)) - value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR) - value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR) - bad = tl.max( - tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0) - ) - if bad != 0: - tl.atomic_max(bad_ptr, 1) - - -class LamportRSConv: - """Persistent symmetric buffers for one TP group.""" - - def _initialize_lamport_buffers(self) -> None: - """Arm every slot and collectively verify the exact sentinel bits. - - Do not replace this with a zero-fill: Lamport readiness distinguishes - the bf16 bit pattern for -0.0 (0x8000) from every published payload. - The validation is deliberately collective so one rank cannot enter the - first polling kernel while another rank still has an unarmed buffer. - """ - if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous(): - raise RuntimeError("Lamport symmetric buffers must be contiguous") - if self.buf_in.numel() % 2 or self.buf_out.numel() % 2: - raise RuntimeError("Lamport buffers must contain whole uint32 pairs") - - # Fill through int16 so each bf16 lane receives the exact -0.0 bits. - # This covers all three generations and all max-token slots, including - # slots that a smaller first invocation does not touch. - self.buf_in.view(torch.int16).fill_(-0x8000) - self.buf_out.view(torch.int16).fill_(-0x8000) - torch.accelerator.synchronize(self.device) - self.tp.barrier() - - # Scan the complete local allocations. Since every rank performs the - # scan and participates in MAX, success proves every symmetric backing - # allocation was armed before any rank is allowed to use generation 0. - bad = torch.logical_or( - self.buf_in.view(torch.int16).ne(-0x8000).any(), - self.buf_out.view(torch.int16).ne(-0x8000).any(), - ).to(dtype=torch.float32) - bad = self.tp.all_reduce(bad) - if int(bad.item()) != 0: - raise RuntimeError("Lamport sentinel initialization failed on a TP rank") - self.tp.barrier() - - def _initialize_mnnvl_buffers(self) -> None: - """Initialize and validate FlashInfer fabric-mapped Lamport storage.""" - self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16) - self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16) - torch.accelerator.synchronize(self.device) - self.tp.barrier() - - bad = torch.zeros((), dtype=torch.int32, device=self.device) - num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2 - _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)]( - self.input_peer_ptrs, - self.output_peer_ptrs, - bad, - num_pairs, - RANK=self.rank, - BLOCK=256, - num_warps=4, - ) - bad = self.tp.all_reduce(bad.to(torch.float32)) - if int(bad.item()) != 0: - raise RuntimeError("MNNVL Lamport sentinel initialization failed") - self.tp.barrier() - - def __init__( - self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS - ) -> None: - import torch.distributed._symmetric_memory as symm_mem - - tp = get_tp_group() - self.tp = tp - self.group = tp.device_group - self.world_size = tp.world_size - self.rank = tp.rank_in_group - self.device = torch.device(tp.device) - if self.world_size not in (2, 4, 8): - raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}") - if hidden_size % (2 * self.world_size) != 0: - raise ValueError("hidden size must produce an even shard on every rank") - if window_size != 4: - raise ValueError(f"short-conv window size must be 4, got {window_size}") - if max_tokens < 1 or max_tokens > _MAX_TOKENS: - raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]") - - is_cross_node = not all(in_the_same_node_as(tp.cpu_group)) - self.hidden_size = hidden_size - self.window_size = window_size - self.max_tokens = max_tokens - self.shard_size = hidden_size // self.world_size - # Three generations follow FlashInfer's Lamport layout. A generation - # is reused only after two intervening collective calls. - self.num_buffers = 3 - self.input_generation_bytes = max_tokens * hidden_size * 2 - self.output_generation_bytes = max_tokens * hidden_size * 2 - self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9 - if is_cross_node: - try: - from flashinfer.comm.mnnvl import ( - McastGPUBuffer, - TorchDistBackend, - is_mnnvl_fabric_supported, - ) - except ImportError as error: - raise RuntimeError( - "cross-node TP requires FlashInfer MNNVL support" - ) from error - - local_supported = int( - is_mnnvl_fabric_supported(torch.accelerator.current_device_index()) - ) - unsupported = torch.tensor( - 1 - local_supported, dtype=torch.float32, device=self.device - ) - unsupported = tp.all_reduce(unsupported) - if int(unsupported.item()) != 0: - raise RuntimeError("cross-node TP is supported only on MNNVL fabric") - - comm_backend = TorchDistBackend(self.group) - allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2 - self._mnnvl_input_handle = McastGPUBuffer( - allocation_bytes, - self.world_size, - self.rank, - self.device, - comm_backend, - ) - self._mnnvl_output_handle = McastGPUBuffer( - allocation_bytes, - self.world_size, - self.rank, - self.device, - comm_backend, - ) - self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev() - self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev() - self._initialize_mnnvl_buffers() - logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers") - else: - self.buf_in = symm_mem.empty( - self.num_buffers, - max_tokens, - self.world_size, - self.shard_size, - dtype=torch.bfloat16, - device=self.device, - ) - self.buf_out = symm_mem.empty( - self.num_buffers, - max_tokens, - hidden_size, - dtype=torch.bfloat16, - device=self.device, - ) - group_name = self.group.group_name - input_handle = symm_mem.rendezvous(self.buf_in, group_name) - output_handle = symm_mem.rendezvous(self.buf_out, group_name) - self.input_peer_ptrs = input_handle.buffer_ptrs_dev - self.output_peer_ptrs = output_handle.buffer_ptrs_dev - self._input_handle = input_handle - self._output_handle = output_handle - self._initialize_lamport_buffers() - self.generation = 0 - - def usable(self, num_tokens: int) -> bool: - return 0 < num_tokens <= self.max_tokens - - def rs_sconv_ag_add_norm( - self, - input_tensor: torch.Tensor, - residual: torch.Tensor, - conv_weight: torch.Tensor, - norm_weight: torch.Tensor | None, - eps: float, - cache: torch.Tensor, - positions: torch.Tensor, - block_table: torch.Tensor, - seq_idx: torch.Tensor, - slot_mapping: torch.Tensor, - off_s: int, - ws: int, - block_size: int, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``.""" - tokens, hidden_size = residual.shape - if not self.usable(tokens): - raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]") - if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16: - raise ValueError("residual must be bf16 [T, 6144]") - if ( - input_tensor.shape != residual.shape - or input_tensor.dtype != torch.bfloat16 - or input_tensor.stride(1) != 1 - ): - raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]") - shard_size = hidden_size // self.world_size - if conv_weight.shape != (shard_size, self.window_size): - raise ValueError( - f"conv_weight must have shape [{shard_size}, {self.window_size}]" - ) - if ( - conv_weight.dtype != torch.bfloat16 - or conv_weight.stride(0) != self.window_size - ): - raise ValueError("conv_weight must be contiguous bf16") - if norm_weight is not None and ( - norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16 - ): - raise ValueError("norm_weight must be bf16 [6144] or None") - if cache.dtype != torch.bfloat16 or cache.ndim != 4: - raise ValueError("cache must be a 4-D bf16 tensor") - if shard_size % ws != 0 or cache.shape[1] != shard_size // ws: - raise ValueError("cache head layout is inconsistent with ws") - if off_s < 0 or off_s + ws > cache.shape[3]: - raise ValueError("cache channel offset is out of bounds") - - index = self.generation - input_offset = index * self.input_generation_bytes // 4 - output_offset = index * self.output_generation_bytes // 4 - normed = torch.empty_like(residual) if norm_weight is not None else None - residual_out = torch.empty_like(residual) - phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1 - phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits) - phase_grid = (tokens, phase_splits) - # Wide CTAs win before the grid saturates; smaller CTAs reduce register - # pressure once high-throughput batches provide enough parallelism. - if 128 <= tokens <= 2048: - phase_warps = 16 - elif tokens > 2048: - phase_warps = 8 - else: - phase_warps = 4 - gather_warps = 4 if tokens >= 256 else 8 - _publish_input_kernel[phase_grid]( - input_tensor, - self.input_peer_ptrs, - input_offset, - input_tensor.stride(0), - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - RANK=self.rank, - WORLD=self.world_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _reduce_insert_kernel[phase_grid]( - self.input_peer_ptrs, - input_offset, - cache, - slot_mapping, - cache.stride(0), - cache.stride(1), - cache.stride(2), - cache.stride(3), - block_size, - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - HEAD_SIZE=ws, - CACHE_OFFSET=off_s, - RANK=self.rank, - WORLD=self.world_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _sconv_publish_kernel[phase_grid]( - self.output_peer_ptrs, - output_offset, - residual, - conv_weight, - cache, - positions, - seq_idx, - slot_mapping, - block_table, - residual.stride(0), - cache.stride(0), - cache.stride(1), - cache.stride(2), - cache.stride(3), - block_table.stride(0), - block_table.shape[1], - block_size, - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - HEAD_SIZE=ws, - CACHE_OFFSET=off_s, - RANK=self.rank, - WORLD=self.world_size, - WINDOW=self.window_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _gather_norm_kernel[(tokens,)]( - self.output_peer_ptrs, - output_offset, - norm_weight if norm_weight is not None else residual, - normed if normed is not None else residual_out, - residual_out, - eps, - residual.stride(0), - C=hidden_size, - C_P2=triton.next_power_of_2(hidden_size), - RANK=self.rank, - HAS_NORM=norm_weight is not None, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=gather_warps, - ) - self.generation = (index + 1) % self.num_buffers - return normed, residual_out - - -_STATE: LamportRSConv | None = None -_STATE_FAILED = False - - -def initialize_lamport_rs_conv( - hidden_size: int, window_size: int, max_num_batched_tokens: int -) -> None: - """Collectively initialize the TP-group state during model construction.""" - global _STATE, _STATE_FAILED - if _STATE is not None: - if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size: - raise RuntimeError("all Lamport users must share hidden and window sizes") - return - if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0": - return - try: - max_tokens = min(_MAX_TOKENS, max_num_batched_tokens) - _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens) - except Exception: - _STATE_FAILED = True - logger.exception("fused collective unavailable; use the NCCL fallback") - - -def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None: - """Return the state initialized with the model, or ``None`` for fallback.""" - if _STATE is not None and ( - _STATE.hidden_size != hidden_size or _STATE.window_size != window_size - ): - raise RuntimeError("all Lamport users must share hidden and window sizes") - return _STATE diff --git a/vllm/models/inkling/amd/ops/mm_towers.py b/vllm/models/inkling/amd/ops/mm_towers.py deleted file mode 100644 index 021cf5274b..0000000000 --- a/vllm/models/inkling/amd/ops/mm_towers.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused CUDA kernels for the Inkling vision/audio towers. - -Both kernels keep the reference paths' fp32 accumulation and per-op bf16 -rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently -bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order -(real-checkpoint-weight cosine vs reference > 0.9999998). -""" - -from __future__ import annotations - -import torch - -from vllm.triton_utils import tl, tldevice, triton - -from .norm import _get_num_warps_from_block_size - - -@triton.jit -def _dmel_embed_sum_norm_kernel( - idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB)) - w_ptr, # [NB * VOCAB, D] bf16 embedding table - norm_w_ptr, # [D] (unused if HAS_NORM=False) - out_ptr, # [T, D] bf16 - eps, - D, - stride_idx_t, - NB: tl.constexpr, - VOCAB: tl.constexpr, - D_P2: tl.constexpr, - HAS_NORM: tl.constexpr, -): - t = tl.program_id(0).to(tl.int64) - offs = tl.arange(0, D_P2) - mask = offs < D - # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)), - # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()). - acc = tl.zeros([D_P2], dtype=tl.float32) - for b in tl.static_range(NB): - v = tl.load(idx_ptr + t * stride_idx_t + b) - row = (b * VOCAB + v).to(tl.int64) - acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) - h = acc.to(tl.bfloat16) - if HAS_NORM: - # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded - # bf16 multiply with the bf16 weight. - x32 = h.to(tl.float32) - var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D - xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16) - w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0) - h = xn * w - tl.store(out_ptr + t * D + offs, h, mask=mask) - - -def dmel_embed_sum_norm( - dmel_idx: torch.Tensor, # [T, NB] int32 - weight: torch.Tensor, # [NB * VOCAB, D] bf16 - norm_weight: torch.Tensor | None, - eps: float, -) -> torch.Tensor: - """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no - [T, NB, D] intermediate).""" - T, nb = dmel_idx.shape - D = weight.shape[1] - assert weight.shape[0] % nb == 0 - vocab = weight.shape[0] // nb - out = torch.empty((T, D), dtype=weight.dtype, device=weight.device) - if T == 0: - return out - d_p2 = triton.next_power_of_2(D) - _dmel_embed_sum_norm_kernel[(T,)]( - dmel_idx, - weight, - norm_weight if norm_weight is not None else weight, - out, - eps, - D, - dmel_idx.stride(0), - NB=nb, - VOCAB=vocab, - D_P2=d_p2, - HAS_NORM=norm_weight is not None, - # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4% - # at large T (the 80-row gather chain is latency- not lane-bound). - num_warps=8, - ) - return out - - -@triton.jit -def _rmsnorm_gelu_kernel( - x_ptr, # [R, D] bf16 - w_ptr, # [D] - out_ptr, # [R, D] bf16 (or the folded layout when FOLD) - eps, - R, - D, - D_P2: tl.constexpr, - BLOCK_M: tl.constexpr, # rows per block (>1 for small D) - HAS_GELU: tl.constexpr, - FOLD: tl.constexpr, - # fold geometry: input rows index [N, T, H, W]; the store scatters each - # row to (out_row, slot) of fold_timespace_to_depth's output layout. - FT: tl.constexpr, - FH: tl.constexpr, - FW: tl.constexpr, - TF: tl.constexpr, - HF: tl.constexpr, -): - pid = tl.program_id(0).to(tl.int64) - rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) - rmask = rows < R - offs = tl.arange(0, D_P2) - mask = rmask[:, None] & (offs < D)[None, :] - x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to( - tl.float32 - ) - var = tl.sum(x32 * x32, axis=1) / D - xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16) - w = tl.load(w_ptr + offs, mask=offs < D, other=0.0) - h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm - if HAS_GELU: - # Exact (erf) GELU on the bf16-rounded norm output, fp32 math, - # matching F.gelu's opmath on a bf16 tensor. - g32 = h.to(tl.float32) - h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16) - if FOLD: - # Store directly into the next layer's folded layout (a pure - # permutation — replaces the separate fold copy pass). - t = (rows // (FH * FW)) % FT - hh = (rows // FW) % FH - ww = rows % FW - n = rows // (FT * FH * FW) - slot = ((t % TF) * HF + hh % HF) * HF + ww % HF - out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * ( - FW // HF - ) + ww // HF - base = (out_row * (TF * HF * HF) + slot) * D - tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask) - else: - tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask) - - -def rmsnorm_gelu( - x: torch.Tensor, # [..., D] bf16 contiguous - weight: torch.Tensor, - eps: float, - gelu: bool = True, - fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold -) -> torch.Tensor: - """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block - when D is small. With ``fold``, x must be [N, T, H, W, D] and the output - comes back as ``fold_timespace_to_depth(result, *fold)``.""" - D = x.shape[-1] - flat = x.reshape(-1, D) - assert flat.stride(1) == 1 and flat.stride(0) == D - R = flat.shape[0] - if fold is None: - out = torch.empty_like(flat) - ft = fh = fw = tf = hf = 1 - out_shape = x.shape - else: - tf, hf = fold - N, ft, fh, fw, _ = x.shape - out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D) - out = torch.empty(out_shape, dtype=x.dtype, device=x.device) - if R == 0: - return out.reshape(out_shape) - d_p2 = triton.next_power_of_2(D) - block_m = max(1, 4096 // d_p2) - _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)]( - flat, - weight, - out, - eps, - R, - D, - D_P2=d_p2, - BLOCK_M=block_m, - HAS_GELU=gelu, - FOLD=fold is not None, - FT=ft, - FH=fh, - FW=fw, - TF=tf, - HF=hf, - num_warps=_get_num_warps_from_block_size(d_p2 * block_m), - ) - return out.reshape(out_shape) diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py index 82782fcb96..7baec0b42b 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -587,34 +587,3 @@ def warp_sum_specialized( mask_and_clamp=31, ) return value - - -@cute.jit -def block_sum_specialized( - value: Float32, - warp_sums: cute.Tensor, - tidx: Int32, - warps: cutlass.Constexpr[int], - last_warp_lanes: cutlass.Constexpr[int], - last_warp_mask: cutlass.Constexpr[int], -) -> Float32: - """Upstream-equivalent FP32 block reduction.""" - - lane = cute.arch.lane_idx() - warp_idx = cute.arch.warp_idx() - value = warp_sum_specialized( - value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask - ) - if lane == 0: - warp_sums[warp_idx] = value - cute.arch.barrier() - - block_sum = Float32(0.0) - if warp_idx == 0: - if lane < Int32(warps): - block_sum = warp_sums[lane] - block_sum = cute.arch.warp_reduction_sum(block_sum) - if lane == 0: - warp_sums[0] = block_sum - cute.arch.barrier() - return warp_sums[0] diff --git a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py b/vllm/v1/attention/ops/triton_fp8_mqa_logits.py deleted file mode 100644 index 619d0ec50a..0000000000 --- a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py +++ /dev/null @@ -1,262 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Temporary gfx942 fallback for AITER's fp8_mqa_logits kernel. - -This module vendors AITER's Triton fp8_mqa_logits kernel with the gfx942 -tile-size workaround from ROCm/aiter#3257. It is used only while vLLM's -pinned AITER version lacks that fix. - -TODO: Remove this vendored copy once vLLM pins an AITER version that includes -ROCm/aiter#3257 bugfix for gfx942. -""" - -import torch - -from vllm.triton_utils import tl, triton - -# gfx942 (MI300X) has 64 KiB of LDS per CU. We accept the default -# (BLOCK_KV=128, num_stages=2) tile only when *both* of these hold: -# -# 1. Occupancy gate. With waves_per_eu=2 and num_warps=4 we target two -# workgroups co-resident on a CU -> per-WG LDS budget = 32 KiB. Triton -# keeps Q in registers (loop-invariant) and the fp32 scores accumulator -# in VGPRs (heavy VALU), so only the double-buffered KV tile is -# expected to live in LDS. A 0.9 safety factor leaves headroom for any -# LDS overhead the compiler may add. -# -# 2. Hardware ceiling. Defensive upper bound that also counts Q and -# scores against the 64 KiB CU limit, in case a Triton version (older -# or future) decides to spill them to LDS. False positives here only -# shrink the tile; false negatives are JIT-aborts, so we lean -# conservative. -_GFX942_CU_LDS_BYTES = 64 * 1024 -_GFX942_PER_WG_LDS_BUDGET_BYTES = _GFX942_CU_LDS_BYTES * 9 // 20 # ~28.8 KiB - - -def _gfx942_default_tile_fits_lds(num_heads: int, head_size: int) -> bool: - """Return True iff (BLOCK_KV=128, num_stages=2) fits in MI300X LDS.""" - BLOCK_KV = 128 - NUM_STAGES = 2 - kv_bytes = head_size * BLOCK_KV * NUM_STAGES - scores_bytes = num_heads * BLOCK_KV * 4 - q_bytes = num_heads * head_size - fits_occupancy = kv_bytes < _GFX942_PER_WG_LDS_BUDGET_BYTES - fits_hardware = q_bytes + kv_bytes + scores_bytes <= _GFX942_CU_LDS_BYTES - return fits_occupancy and fits_hardware - - -@triton.jit -def _fp8_mqa_logits_kernel( - Q_ptr, # fp8e4m3 [seq_len, H, D] - KV_ptr, # fp8e4m3 [seq_len_kv, D] - kv_scales_ptr, # fp32 [seq_len_kv] - weights_ptr, # fp32 [seq_len, H] - cu_start_ptr, # int32 [seq_len] - cu_end_ptr, # int32 [seq_len] - logits_ptr, # fp32 [seq_len, seq_len_kv] - seq_len, - seq_len_kv, - NUM_HEADS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - # strides - stride_q_s: tl.int64, - stride_q_h: tl.constexpr, - stride_q_d: tl.constexpr, - stride_kv_s: tl.int64, - stride_kv_d: tl.constexpr, - stride_w_s: tl.int64, - stride_w_h: tl.constexpr, - stride_logits_s: tl.int64, - stride_logits_k: tl.int64, - # block sizes - BLOCK_KV: tl.constexpr, -): - row_id = tl.program_id(0) - # go from larger to smaller in terms of work - # to reduce the tail effect - row_id = tl.num_programs(0) - row_id - 1 - tl.assume(row_id >= 0) - tl.assume(stride_q_s > 0) - tl.assume(stride_q_h > 0) - tl.assume(stride_q_d > 0) - tl.assume(stride_kv_s > 0) - tl.assume(stride_kv_d > 0) - tl.assume(stride_w_s > 0) - tl.assume(stride_w_h > 0) - - logits_row_ptrs = logits_ptr + row_id * stride_logits_s - - h_inds = tl.arange(0, NUM_HEADS)[:, None] - d_inds = tl.arange(0, HEAD_SIZE) - - # load Q[BLOCK_Q, NUM_HEADS, HEAD_SIZE] - q_ptrs = ( - Q_ptr + row_id * stride_q_s + h_inds * stride_q_h + d_inds[None, :] * stride_q_d - ) - - q_block = tl.load(q_ptrs, cache_modifier=".cg") - w_ptrs = weights_ptr + row_id * stride_w_s + h_inds * stride_w_h - w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32) - - # Load start/end for each row in this block - start_ind = tl.load(cu_start_ptr + row_id) - end_ind = tl.load(cu_end_ptr + row_id) - - start_ind = tl.maximum(start_ind, 0) - end_ind = tl.minimum(end_ind, seq_len_kv) - shifted_end = end_ind - start_ind - shifted_unmasked_end = shifted_end // BLOCK_KV * BLOCK_KV - - kv_col_offsets = tl.arange(0, BLOCK_KV) + start_ind - kv_ptrs = ( - KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d - ) - - kv_scales_ptrs = kv_scales_ptr + kv_col_offsets - - logits_ptrs = logits_row_ptrs + kv_col_offsets * stride_logits_k - - # Loop over KV tiles - for _ in tl.range(0, shifted_unmasked_end, BLOCK_KV): - kv_block = tl.load(kv_ptrs) - kv_scales = tl.load(kv_scales_ptrs) - - # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV] - scores = tl.dot(q_block, kv_block, input_precision="ieee") - # Multiply by kv_scales (broadcast along rows) - scores = scores * kv_scales[None, :] - # ReLU - scores = tl.maximum(scores, 0.0) - scores = scores * w_block - # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ] - scores = tl.sum(scores, axis=0) - tl.store(logits_ptrs, scores) - - kv_ptrs += BLOCK_KV * stride_kv_s - kv_scales_ptrs += BLOCK_KV - logits_ptrs += BLOCK_KV * stride_logits_k - kv_col_offsets += BLOCK_KV - - # masked load - kv_col_mask = kv_col_offsets < end_ind - kv_block = tl.load(kv_ptrs, mask=kv_col_mask[None, :], other=0.0) - kv_scales = tl.load(kv_scales_ptrs, mask=kv_col_mask, other=0.0) - - # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV] - scores = tl.dot(q_block, kv_block, input_precision="ieee") - # Multiply by kv_scales (broadcast along rows) - scores = scores * kv_scales[None, :] - # ReLU - scores = tl.maximum(scores, 0.0) - scores = scores * w_block - # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ] - scores = tl.sum(scores, axis=0) - # masked store - in_window = (kv_col_offsets >= start_ind) & (kv_col_offsets < end_ind) - tl.store(logits_ptrs, scores, mask=in_window) - - -def fp8_mqa_logits_gfx942( - q: torch.Tensor, - k_fp8: torch.Tensor, - kv_scales: torch.Tensor, - weights: torch.Tensor, - cu_starts: torch.Tensor, - cu_ends: torch.Tensor, -) -> torch.Tensor: - """Compute FP8 MQA logits on MI300X (gfx942) using the vendored kernel. - - Drop-in replacement for ``aiter.ops.triton.attention.fp8_mqa_logits. - fp8_mqa_logits`` on MI300X. Selects ``(BLOCK_KV, num_stages)`` based on - whether the default tile fits within the 64 KiB LDS budget of a gfx942 - CU (see module docstring). - - Args: - q: Query tensor of shape ``[M, H, D]``, FP8 dtype. - k_fp8: Key tensor of shape ``[N, D]``, FP8 dtype. - kv_scales: K scales of shape ``[N]`` (or ``[N, 1]`` -- viewed as - ``[N]``), float32. - weights: Per-head weights of shape ``[M, H]``, float32. - cu_starts: Start indices (inclusive) of shape ``[M]``, int32. - cu_ends: End indices (exclusive) of shape ``[M]``, int32. - - Returns: - Logits of shape ``[M, N]``, float32 -- positions outside - ``[cu_starts[i], cu_ends[i])`` for row ``i`` are pre-filled with - ``-inf`` so the caller can run a top-k without masking. - """ - seq_len, num_heads, head_size = q.shape - seq_len_kv = k_fp8.shape[0] - assert num_heads & (num_heads - 1) == 0, ( - f"num_heads must be a power of two (got {num_heads})" - ) - assert head_size & (head_size - 1) == 0, ( - f"head_size must be a power of two (got {head_size})" - ) - - # The kernel walks ``kv_scales`` as a 1-D contiguous array of size N - # (it indexes by ``kv_scales_ptr + kv_col_offsets``). The vLLM caller - # passes a ``[N, 4]`` uint8 view-cast-to-float32 which lands as - # ``[N, 1]`` contiguous -- byte-identical to ``[N]`` -- but flatten - # explicitly to keep the kernel's pointer arithmetic intent clear. - kv_scales_1d = kv_scales.reshape(-1) - - # Initialise with -inf so positions outside [cu_starts, cu_ends) read - # as ``-inf`` after the masked store path -- this matches AITER's - # ``fp8_mqa_logits`` semantics and is what the top-k consumer expects. - logits = torch.full( - (seq_len, seq_len_kv), - fill_value=-float("inf"), - dtype=torch.float32, - device=q.device, - ) - - if _gfx942_default_tile_fits_lds(num_heads, head_size): - block_kv = 128 - num_stages = 2 - else: - # DSv4 sparse indexer (NUM_HEADS=64, HEAD_SIZE=128) lands here: - # default tile spills past gfx942's 64 KiB LDS budget. (64, 1) - # needs ~33 KiB and clears the per-WG budget with margin. - block_kv = 64 - num_stages = 1 - - # heuristic for MFMA instruction shape, identical to AITER's choice - matrix_instr_nonkdim = 32 - if seq_len <= 1024: - matrix_instr_nonkdim = 16 - - stride_q_s, stride_q_h, stride_q_d = q.stride() - stride_kv_s, stride_kv_d = k_fp8.stride() - stride_w_s, stride_w_h = weights.stride() - stride_logits_s, stride_logits_k = logits.stride() - - _fp8_mqa_logits_kernel[(seq_len,)]( - Q_ptr=q, - KV_ptr=k_fp8, - kv_scales_ptr=kv_scales_1d, - weights_ptr=weights, - cu_start_ptr=cu_starts, - cu_end_ptr=cu_ends, - logits_ptr=logits, - seq_len=seq_len, - seq_len_kv=seq_len_kv, - NUM_HEADS=num_heads, - HEAD_SIZE=head_size, - stride_q_s=stride_q_s, - stride_q_h=stride_q_h, - stride_q_d=stride_q_d, - stride_kv_s=stride_kv_s, - stride_kv_d=stride_kv_d, - stride_w_s=stride_w_s, - stride_w_h=stride_w_h, - stride_logits_s=stride_logits_s, - stride_logits_k=stride_logits_k, - BLOCK_KV=block_kv, - num_warps=4, - num_stages=num_stages, - waves_per_eu=2, - matrix_instr_nonkdim=matrix_instr_nonkdim, - ) - - return logits diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py index 321f42bcb1..68ce0e0de2 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py @@ -21,15 +21,11 @@ Stage 2 is reused unchanged from triton_decode_attention.py. """ import math -from typing import Any import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.v1.attention.ops.triton_decode_attention import _fwd_kernel_stage2 - -from .triton_turboquant_decode import _use_fp8_e4b15 # ROCm prefers num_stages=1 in attention-like kernels to reduce shared-memory # pressure (mirrors the pattern used in triton_decode_attention.py and @@ -66,387 +62,6 @@ def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor: return lut.reshape(N * N, 2).contiguous() -# --------------------------------------------------------------------------- -# Stage 1 v2: Grouped-Q + pair LUT + tl.dot + exp2 -# --------------------------------------------------------------------------- - - -@triton.jit -def _tq_decode_stage1_v2( - Q_rot_ptr, - KV_cache_ptr, - KV_cache_u16_ptr, # uint16 view — Opt#3 SoA loads - Block_table_ptr, - Seq_lens_ptr, - Centroids_ptr, - Pair_lut_ptr, - Mid_o_ptr, - stride_qb, - stride_qh, - stride_cache_block, # bytes per block (bs*H*slot_aligned) - stride_bt_b, - stride_mid_b, - stride_mid_h, - stride_mid_s, - NUM_Q_HEADS: tl.constexpr, - NUM_KV_HEADS: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - PADDED_SLOT: tl.constexpr, - MAX_NUM_BLOCKS: tl.constexpr, - NUM_KV_SPLITS: tl.constexpr, - KV_GROUP_SIZE: tl.constexpr, - MSE_BITS: tl.constexpr, - MSE_BYTES: tl.constexpr, - VQB: tl.constexpr, - VAL_DATA_BYTES: tl.constexpr, - # Opt#3 SoA layout - KEY_DATA_BYTES: tl.constexpr, - META_REGION_OFFSET: tl.constexpr, - NUM_SOA_FIELDS: tl.constexpr, - SOA_K_NORM: tl.constexpr, - SOA_V_SCALE: tl.constexpr, - SOA_V_ZERO: tl.constexpr, - N_CENTROIDS: tl.constexpr, - ATTN_SCALE: tl.constexpr, - BLOCK_D: tl.constexpr, - TILE_SIZE: tl.constexpr, - BLOCK_M: tl.constexpr, - KEY_FP8: tl.constexpr, - NORM_CORRECTION: tl.constexpr = 0, - FP8_E4B15: tl.constexpr = 0, - USE_PAIR_LUT: tl.constexpr = 0, - USE_BF16_DOT: tl.constexpr = 0, -): - # Grid: (B, Hk, NUM_KV_SPLITS) — one program per KV head group - bid = tl.program_id(0) - kv_hid = tl.program_id(1) - sid = tl.program_id(2) - - seq_len = tl.load(Seq_lens_ptr + bid) - split_len = tl.cdiv(seq_len, NUM_KV_SPLITS) - split_start = split_len * sid - split_end = tl.minimum(split_start + split_len, seq_len) - - if split_start >= split_end: - return - - # Dimension and tile offsets - d_offs = tl.arange(0, BLOCK_D) - d_mask = d_offs < HEAD_DIM - m_offs = tl.arange(0, BLOCK_M) - t_offs = tl.arange(0, TILE_SIZE) - - # Q head indices for this KV group: [kv_hid * GS, kv_hid * GS + GS) - q_head_offs = kv_hid * KV_GROUP_SIZE + m_offs - q_mask = m_offs < KV_GROUP_SIZE - - # Load Q tile: [BLOCK_M, BLOCK_D] float32 - q_addrs = bid * stride_qb + q_head_offs[:, None] * stride_qh + d_offs[None, :] - Q = tl.load( - Q_rot_ptr + q_addrs, - mask=q_mask[:, None] & d_mask[None, :], - other=0.0, - ) - - # Pre-warm centroids in L1 - if not KEY_FP8: - _c_warm = tl.load(Centroids_ptr + tl.arange(0, N_CENTROIDS)) # noqa: F841 - - # Loop-invariant bit-extraction constants for MSE keys - if not KEY_FP8: - if MSE_BITS == 4: - pass # 4-bit byte extraction is computed inline in the tile loop - else: - # Generic bit extraction - mse_bit_off = d_offs * MSE_BITS - mse_byte_idx = mse_bit_off // 8 - mse_bit_shift = mse_bit_off % 8 - mse_mask_val = (1 << MSE_BITS) - 1 - - # Loop-invariant bit-extraction for 3-bit values - if VQB == 3: - val_bit_off = d_offs * 3 - val_byte_idx = val_bit_off // 8 - val_bit_shift = val_bit_off % 8 - - # exp2 pre-scaling - RCP_LN2: tl.constexpr = 1.4426950408889634 - LN2: tl.constexpr = 0.6931471805599453 - QK_SCALE = ATTN_SCALE * RCP_LN2 - - # Online softmax accumulators: [BLOCK_M] for max, [BLOCK_M] for L - M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) - L = tl.zeros([BLOCK_M], dtype=tl.float32) - - acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) - - bt_base = bid * stride_bt_b - - # ================================================================ - # TILED LOOP: TILE_SIZE KV tokens per iteration - # ================================================================ - for start_n in range(split_start, split_end, TILE_SIZE): - kv_offs = start_n + t_offs - kv_mask_1d = kv_offs < split_end - - page_idx = kv_offs // BLOCK_SIZE - page_off = kv_offs % BLOCK_SIZE - block_nums = tl.load( - Block_table_ptr + bt_base + page_idx, - mask=kv_mask_1d, - other=0, - ).to(tl.int64) - - # Opt#3 SoA addressing: data region then per-block SoA metadata. - slot_within_block = page_off.to(tl.int64) - block_base = block_nums * stride_cache_block - DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES - data_bases = ( - block_base - + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT) - + tl.cast(kv_hid, tl.int64) * DATA_BYTES_PER_SLOT - ) - head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast( - kv_hid, tl.int64 - ) * (NUM_SOA_FIELDS * BLOCK_SIZE) - knorm_u16_addrs = ( - head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block - ) - vscale_u16_addrs = ( - head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block - ) - vzero_u16_addrs = ( - head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block - ) - - # ============================================================ - # KEY DEQUANT → K_T: [BLOCK_D, TILE_SIZE] for tl.dot(Q, K_T) - # ============================================================ - if KEY_FP8: - k_addrs = data_bases[:, None] + d_offs[None, :] - k_raw = tl.load( - KV_cache_ptr + k_addrs, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ) - if FP8_E4B15: - K_f = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) - else: - K_f = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) - # K_f: [TILE_SIZE, BLOCK_D] → transpose for tl.dot - K_T = tl.trans(K_f) - # S: [BLOCK_M, TILE_SIZE] (bf16 dot on ROCm for better dynamic range) - if USE_BF16_DOT: - S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16)) - else: - S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16)) - else: - # --- MSE dequantization --- - if MSE_BITS == 4 and USE_PAIR_LUT: - # OPT#1 (load-halving): FLUTE §3.2 pair LUT, but with each - # packed byte loaded EXACTLY ONCE instead of twice (the - # previous code built byte_addrs from d_offs // 2 which has - # the pattern [0,0,1,1,2,2,...], duplicating every byte - # load). We load the packed bytes as a [TILE_SIZE, HALF_D] - # tile, decode both nibbles, then issue a single 3-D gather - # against pair_lut which returns (T[lo], T[hi]) in slots - # 0/1. Reshape [TILE_SIZE, HALF_D, 2] -> [TILE_SIZE, BLOCK_D] - # interleaves the pairs into the expected dim order - # (d=2k -> T[lo_k], d=2k+1 -> T[hi_k]). - HALF_D: tl.constexpr = BLOCK_D // 2 - half_offs = tl.arange(0, HALF_D) - # Each byte k covers dims (2k, 2k+1); it is in-range iff - # 2k < HEAD_DIM. HEAD_DIM is always even in our layouts, - # so this is equivalent to k < HEAD_DIM // 2. - byte_mask = (half_offs * 2) < HEAD_DIM - - byte_addrs = data_bases[:, None] + half_offs[None, :] - byte_raw = tl.load( - KV_cache_ptr + byte_addrs, - mask=kv_mask_1d[:, None] & byte_mask[None, :], - other=0, - ).to(tl.int32) - - lo_idx = byte_raw & 0xF # [TILE_SIZE, HALF_D] - hi_idx = (byte_raw >> 4) & 0xF # [TILE_SIZE, HALF_D] - pair_key = lo_idx * N_CENTROIDS + hi_idx # [TILE_SIZE, HALF_D] - - # Single 3-D gather: pair_lut[pair_key*2 + {0,1}] - # -> [TILE_SIZE, HALF_D, 2] = (T[lo_k], T[hi_k]) per element. - pair_slot = tl.arange(0, 2) - c_pair = tl.load( - Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :], - mask=(kv_mask_1d[:, None, None] & byte_mask[None, :, None]), - other=0.0, - ) - # Row-major reshape interleaves pairs: - # out[t, 2k] = c_pair[t, k, 0] = T[lo_k] - # out[t, 2k+1] = c_pair[t, k, 1] = T[hi_k] - c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D]) - elif MSE_BITS == 4: - # 4-bit without pair LUT: exploit byte alignment - half_idx = d_offs // 2 - nibble_shift = (d_offs % 2) * 4 - mse_addrs = data_bases[:, None] + half_idx[None, :] - mse_raw = tl.load( - KV_cache_ptr + mse_addrs, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF - c_vals = tl.load( - Centroids_ptr + mse_idx, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0.0, - ) - else: - # Generic 3-bit path: two byte loads → 16-bit → shift/mask - mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :] - mse_raw0 = tl.load( - KV_cache_ptr + mse_addrs0, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - mse_raw1 = tl.load( - KV_cache_ptr + mse_addrs0 + 1, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - raw16 = mse_raw0 | (mse_raw1 << 8) - mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val - c_vals = tl.load( - Centroids_ptr + mse_idx, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0.0, - ) - - # Opt#1: norm-correction is pre-folded into the stored per-token - # scalar at store time, so no per-tile sum+sqrt+divide here. The - # stored vec_norm carries ||k||/||c_vec||; multiply by c_vals - # below gives the original c_vec/||c_vec|| * ||k||. - - # Opt#3: K-norms from per-block SoA region. Single u16 load per - # token; contiguous across tokens within a block. - norm_u16 = tl.load( - KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask_1d, other=0 - ) - vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32) - - # Reconstruct K: K_recon[t, d] = norm[t] * centroid[t, d] - K_recon = c_vals * vec_norms[:, None] - # K_T: [BLOCK_D, TILE_SIZE] - K_T = tl.trans(K_recon) - - # S: [BLOCK_M, TILE_SIZE] via tensor core MMA - if USE_BF16_DOT: - S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16)) - else: - S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16)) - - # Mask out-of-range positions - S = tl.where(kv_mask_1d[None, :], S, float("-inf")) - - # ============================================================ - # ONLINE SOFTMAX with exp2 - # ============================================================ - m_j = tl.maximum(M, tl.max(S, axis=1)) - m_j = tl.where(m_j > float("-inf"), m_j, 0.0) - P = tl.math.exp2(S - m_j[:, None]) - l_j = tl.sum(P, axis=1) - alpha = tl.math.exp2(M - m_j) - - acc = acc * alpha[:, None] - L = L * alpha + l_j - M = m_j - - # ============================================================ - # VALUE DEQUANT → V: [TILE_SIZE, BLOCK_D] for tl.dot(P, V) - # Opt#3: V data lives at data_base + KEY_DATA_BYTES; scale/zero - # move to the SoA metadata region (one u16 load each per token). - # ============================================================ - val_bases = data_bases + KEY_DATA_BYTES - - if VQB == 3: - val_addrs0 = val_bases[:, None] + val_byte_idx[None, :] - val_raw0 = tl.load( - KV_cache_ptr + val_addrs0, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - val_raw1 = tl.load( - KV_cache_ptr + val_addrs0 + 1, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - raw16 = val_raw0 | (val_raw1 << 8) - v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32) - else: # VQB == 4 - # OPT#3 (load-halving, value path): same pattern as OPT#1 for - # keys -- load each byte exactly once as [TILE_SIZE, HALF_D], - # decode both nibbles, then tl.interleave to [TILE_SIZE, BLOCK_D] - # with pattern [v_lo0, v_hi0, v_lo1, v_hi1, ...]. - V_HALF_D: tl.constexpr = BLOCK_D // 2 - v_half_offs = tl.arange(0, V_HALF_D) - v_byte_mask = (v_half_offs * 2) < HEAD_DIM - - val_addrs = val_bases[:, None] + v_half_offs[None, :] - val_byte = tl.load( - KV_cache_ptr + val_addrs, - mask=kv_mask_1d[:, None] & v_byte_mask[None, :], - other=0, - ).to(tl.int32) - v_lo = (val_byte & 0xF).to(tl.float32) # [T, HALF_D] - v_hi = ((val_byte >> 4) & 0xF).to(tl.float32) # [T, HALF_D] - v_idx = tl.interleave(v_lo, v_hi) # [T, BLOCK_D] - - scale_u16 = tl.load( - KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask_1d, other=0 - ) - zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask_1d, other=0) - v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32) - v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32) - V = v_idx * v_scales[:, None] + v_zeros[:, None] - - # P·V accumulation via tensor core MMA - # P: [BLOCK_M, TILE_SIZE], V: [TILE_SIZE, BLOCK_D] - if USE_BF16_DOT: - acc += tl.dot(P.to(tl.bfloat16), V.to(tl.bfloat16)) - else: - acc += tl.dot(P.to(tl.float16), V.to(tl.float16)) - - # ================================================================ - # EPILOGUE: Store per-Q-head partial results for stage2 - # ================================================================ - # acc: [BLOCK_M, BLOCK_D], L: [BLOCK_M], M: [BLOCK_M] - # Stage2 expects per-Q-head: mid_o[B, Hq, split, D+1] - safe_L = tl.where(L > 0.0, L, 1.0) - one_over_L = 1.0 / safe_L - acc = acc * one_over_L[:, None] - - # Convert M from log2 space back to natural log for stage2 compatibility - lse = M * LN2 + tl.log(safe_L) - - # Store each Q head's result using 2D scatter - # q_head_offs: [BLOCK_M] = kv_hid * KV_GROUP_SIZE + m_offs - out_addrs = ( - bid * stride_mid_b - + q_head_offs[:, None] * stride_mid_h - + sid * stride_mid_s - + d_offs[None, :] - ) - tl.store( - Mid_o_ptr + out_addrs, - acc, - mask=q_mask[:, None] & d_mask[None, :], - ) - # Store LSE per Q head - lse_addrs = ( - bid * stride_mid_b + q_head_offs * stride_mid_h + sid * stride_mid_s + HEAD_DIM - ) - tl.store(Mid_o_ptr + lse_addrs, lse, mask=q_mask) - - # --------------------------------------------------------------------------- # Launcher v2 # --------------------------------------------------------------------------- @@ -480,190 +95,3 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor: shows up on a profile, cache by a hash-of-values fingerprint instead. """ return build_pair_lut(centroids) - - -def triton_turboquant_decode_attention_v2( - query: torch.Tensor, - kv_cache: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - Pi: torch.Tensor, - centroids: torch.Tensor, - scale: float, - mse_bits: int, - key_packed_size: int, - value_quant_bits: int, - value_packed_size: int, - max_seq_len: int = 0, # unused; kept for backward compatibility - key_fp8: bool = False, - norm_correction: bool = False, - PiT: torch.Tensor | None = None, - # Pre-allocated buffers (optional, required for CUDA-graph stability). - mid_o_buf: torch.Tensor | None = None, - output_buf: torch.Tensor | None = None, - lse_buf: torch.Tensor | None = None, - buf_holder: Any = None, - # Fixed split count — MUST be a compile-time constant across iterations - # for CUDA-graph capture/replay to work. Mirrors the v1 launcher contract. - max_num_kv_splits: int = 32, -) -> torch.Tensor: - """Launch optimized TQ decode attention (v2 stage1 + shared stage2). - - Follows the same buffer-reuse + fixed-grid contract as the v1 launcher - (triton_turboquant_decode_attention) so the backend can capture a CUDA - graph across both versions. - """ - B, Hq, D = query.shape - Hk = kv_cache.shape[2] - block_size = kv_cache.shape[1] - padded_slot = kv_cache.shape[3] - max_num_blocks = block_table.shape[1] - n_centroids = centroids.shape[0] - kv_group_size = Hq // Hk - device = query.device - del max_seq_len # no longer used: splits is fixed via max_num_kv_splits - - cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size) - - # Opt#3 SoA layout constants (match store-side computation). - key_data_bytes = D if key_fp8 else cfg["mse_bytes"] - data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"] - meta_region_offset = block_size * Hk * data_bytes_per_slot - num_soa_fields = 2 if key_fp8 else 3 - soa_k_norm = 0 - soa_v_scale = 0 if key_fp8 else 1 - soa_v_zero = 1 if key_fp8 else 2 - kv_cache_u16 = kv_cache.view(torch.uint16) - - # Compute q_rot = q @ Pi.T - if key_fp8: - q_rot = query.float().contiguous() - else: - q_float = query.float() - if PiT is None: - PiT = Pi.T.contiguous() - q_rot = (q_float @ PiT).contiguous() - - # BLOCK_M: pad KV_GROUP_SIZE to power of 2, minimum 16 for tensor cores - BLOCK_M = max(16, triton.next_power_of_2(kv_group_size)) - - # TILE_SIZE (BLOCK_KV): tokens per inner-loop iteration - TILE_SIZE = 16 - - # Fixed split count — must be constant across calls for cudagraph replay. - NUM_KV_SPLITS = max_num_kv_splits - - # --- mid_o buffer reuse (same pattern as v1 launcher) --- - if ( - mid_o_buf is not None - and mid_o_buf.shape[0] >= B - and mid_o_buf.shape[2] >= NUM_KV_SPLITS - ): - mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :] - else: - mid_o = torch.empty( - B, - Hq, - NUM_KV_SPLITS, - D + 1, - dtype=torch.float32, - device=device, - ) - if buf_holder is not None: - buf_holder._tq_mid_o_buf = mid_o - - fp8_e4b15 = _use_fp8_e4b15(device.index or 0) - - # Build pair LUT for 4-bit MSE (FLUTE §3.2) - use_pair_lut = mse_bits == 4 and not key_fp8 - pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids - - # Platform-dependent pipelining depth: ROCm prefers num_stages=1. - stage1_num_stages = 1 if _is_hip else 2 - - # --- Stage 1: v2 kernel --- - # Grid over KV heads (not Q heads) — each program handles all Q heads - grid = (B, Hk, NUM_KV_SPLITS) - _tq_decode_stage1_v2[grid]( - q_rot, - kv_cache, - kv_cache_u16, - block_table, - seq_lens, - centroids, - pair_lut, - mid_o, - q_rot.stride(0), - q_rot.stride(1), - kv_cache.stride(0), - block_table.stride(0), - mid_o.stride(0), - mid_o.stride(1), - mid_o.stride(2), - NUM_Q_HEADS=Hq, - NUM_KV_HEADS=Hk, - HEAD_DIM=D, - BLOCK_SIZE=block_size, - PADDED_SLOT=padded_slot, - MAX_NUM_BLOCKS=max_num_blocks, - NUM_KV_SPLITS=NUM_KV_SPLITS, - KV_GROUP_SIZE=kv_group_size, - MSE_BITS=mse_bits, - MSE_BYTES=cfg["mse_bytes"], - VQB=value_quant_bits, - VAL_DATA_BYTES=cfg["val_data_bytes"], - KEY_DATA_BYTES=key_data_bytes, - META_REGION_OFFSET=meta_region_offset, - NUM_SOA_FIELDS=num_soa_fields, - SOA_K_NORM=soa_k_norm, - SOA_V_SCALE=soa_v_scale, - SOA_V_ZERO=soa_v_zero, - N_CENTROIDS=n_centroids, - ATTN_SCALE=scale, - BLOCK_D=cfg["BLOCK_D"], - TILE_SIZE=TILE_SIZE, - BLOCK_M=BLOCK_M, - KEY_FP8=1 if key_fp8 else 0, - NORM_CORRECTION=1 if norm_correction else 0, - FP8_E4B15=fp8_e4b15, - USE_PAIR_LUT=1 if use_pair_lut else 0, - USE_BF16_DOT=1 if _is_hip else 0, - num_warps=4, - num_stages=stage1_num_stages, - ) - - # --- output / lse buffer reuse (same pattern as v1 launcher) --- - if output_buf is not None and output_buf.shape[0] >= B: - output = output_buf[:B, :Hq, :D] - else: - output = torch.empty(B, Hq, D, dtype=torch.float32, device=device) - if buf_holder is not None: - buf_holder._tq_output_buf = output - if lse_buf is not None and lse_buf.shape[0] >= B: - lse = lse_buf[:B, :Hq] - else: - lse = torch.empty(B, Hq, dtype=torch.float32, device=device) - if buf_holder is not None: - buf_holder._tq_lse_buf = lse - - # --- Stage 2: reduce across KV splits (unchanged) --- - grid2 = (B, Hq) - _fwd_kernel_stage2[grid2]( - mid_o, - output, - lse, - seq_lens, - mid_o.stride(0), - mid_o.stride(1), - mid_o.stride(2), - output.stride(0), - output.stride(1), - lse.stride(0), - NUM_KV_SPLITS=NUM_KV_SPLITS, - BLOCK_DV=cfg["BLOCK_D"], - Lv=D, - num_warps=4, - num_stages=2, - ) - - return output.to(query.dtype) diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py index 03ef36e5b6..d77592010a 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py @@ -876,21 +876,6 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor: return build_pair_lut(centroids) -# PiT in Q's dtype cached per (PiT-id, dtype) to avoid a host cast on every -# launcher call. Keyed by the fp32 PiT's data_ptr so ownership/lifetime is -# tied to the caller's tensor. -_pit_qdtype_cache: dict = {} - - -def _get_pit_in_query_dtype(PiT: torch.Tensor, qdtype: torch.dtype) -> torch.Tensor: - key = (PiT.data_ptr(), qdtype) - cached = _pit_qdtype_cache.get(key) - if cached is None: - cached = PiT.to(qdtype).contiguous() - _pit_qdtype_cache[key] = cached - return cached - - def triton_turboquant_unified_attention( query: torch.Tensor, # [num_tokens, Hq, D] - fp16/bf16 kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 From 38ca7a899cd3a9aa6f1e948b1b63f4fae7a685ce Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Tue, 15 Sep 2026 16:27:11 -0400 Subject: [PATCH 115/139] [Bugfix][HiSparse] Preserve per-layer offsets in KV cache bindings (#57027) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .buildkite/test_areas/lm_eval.yaml | 18 +++++++++++++++++ .../gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml | 20 +++++++++++++++++++ tests/evals/gsm8k/configs/models-hisparse.txt | 1 + vllm/v1/hisparse/binding.py | 15 ++++++++------ 4 files changed, 48 insertions(+), 6 deletions(-) create mode 100644 tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml create mode 100644 tests/evals/gsm8k/configs/models-hisparse.txt diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index d9afefc82e..ec45f7e69e 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -116,6 +116,24 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt +- label: ":nvidia: (B200) LM Eval HiSparse Nightly" + key: lm-eval-hisparse-nightly-4xb200 + timeout_in_minutes: 60 + device: b200-k8s + num_devices: 4 + optional: true + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/v1/hisparse/ + - vllm/distributed/kv_transfer/kv_connector/v1/hisparse/ + - vllm/v1/attention/backends/mla/ + - vllm/model_executor/layers/attention/sparse_mla_attention.py + - vllm/v1/worker/gpu/ + - csrc/libtorch_stable/hisparse_kernels.cu + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt + - label: ":nvidia: (B200) LM Eval PCP" key: lm-eval-pcp-4xb200 timeout_in_minutes: 360 diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml new file mode 100644 index 0000000000..0098719259 --- /dev/null +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml @@ -0,0 +1,20 @@ +model_name: "nvidia/GLM-5.2-NVFP4" +accuracy_threshold: 0.90 +tolerance: 0.0 +num_questions: 1319 +num_fewshot: 5 +max_tokens: 256 +max_concurrency: 32 +request_timeout_seconds: 1800 +startup_max_wait_seconds: 1200 +server_args: >- + --tensor-parallel-size 4 + --max-model-len 8192 + --max-num-seqs 32 + --max-num-batched-tokens 4096 + --enforce-eager + --kv-cache-dtype bfloat16 + --attention-config '{"hisparse_config":{}}' + --kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":32}}' +env: + VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/evals/gsm8k/configs/models-hisparse.txt b/tests/evals/gsm8k/configs/models-hisparse.txt new file mode 100644 index 0000000000..53ff181f4f --- /dev/null +++ b/tests/evals/gsm8k/configs/models-hisparse.txt @@ -0,0 +1 @@ +GLM-5.2-NVFP4-HiSparse.yaml diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index c23abc9cb4..716211fbb3 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -205,9 +205,12 @@ def bind_hisparse_kv_caches( """Bind existing cache storage and block tables; return the bound handles.""" assert host_pool.registered is not None tensor_configs = { - name: tensor_config + name: ( + tensor_config, + tensor_config.offset + layer_index * tensor_config.layer_stride, + ) for tensor_config in kv_cache_config.kv_cache_tensors - for name in tensor_config.layers + for layer_index, name in enumerate(tensor_config.layers) } resident_source_index = 0 for group_id, group in enumerate(kv_cache_config.kv_cache_groups): @@ -216,12 +219,12 @@ def bind_hisparse_kv_caches( for cache_name in group.layer_names: assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX) layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)] - tensor_config = tensor_configs[cache_name] + tensor_config, byte_offset = tensor_configs[cache_name] assert not tensor_config.host_resident cache_handle = _get_hisparse_cache(forward_context, layer_name) cache_handle.bind_cache( kv_caches[cache_name], - byte_offset=tensor_config.offset, + byte_offset=byte_offset, block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, @@ -249,11 +252,11 @@ def bind_hisparse_kv_caches( raise RuntimeError("HiSparse hot tensors must share one GPU backing.") layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)] cache_handle = _get_hisparse_cache(forward_context, layer_name) - tensor_config = tensor_configs[cache_name] + tensor_config, byte_offset = tensor_configs[cache_name] assert not tensor_config.host_resident cache_handle.runtime.bind_hot_cache( raw_tensor, - byte_offset=tensor_config.offset, + byte_offset=byte_offset, block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, From b6ce7143548a937b10764eaddfd03e459d0dab33 Mon Sep 17 00:00:00 2001 From: zwang86 <46699021+zwang86@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:35:54 -0700 Subject: [PATCH 116/139] [Bugfix] Format kernel-import errors eagerly so warning_once does not retain them (#54098) Signed-off-by: Zeyu Wang Co-authored-by: Zeyu Wang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/platforms/cpu.py | 12 ++++++++---- vllm/platforms/cuda.py | 4 +++- vllm/platforms/interface.py | 2 +- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index ecd2dc5fe6..89b68a401e 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -605,26 +605,30 @@ class CpuPlatform(Platform): try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once( + "Failed to import from vllm._C: %s", repr(e) + ) else: try: import vllm._C_AVX512 # noqa: F401 except ImportError as e: if ignored_msg not in e.msg: logger.warning_once( - "Failed to import from vllm._C_AVX512: %r", e + "Failed to import from vllm._C_AVX512: %s", repr(e) ) else: try: import vllm._C_AVX2 # noqa: F401 except ImportError as e: if ignored_msg not in e.msg: - logger.warning_once("Failed to import from vllm._C_AVX2: %r", e) + logger.warning_once( + "Failed to import from vllm._C_AVX2: %s", repr(e) + ) else: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %s", repr(e)) @classmethod def pack_kv_cache( diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 98bddbde29..30eed82ccb 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -239,7 +239,9 @@ class CudaPlatformBase(Platform): try: import vllm._C_stable_libtorch # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C_stable_libtorch: %r", e) + logger.warning_once( + "Failed to import from vllm._C_stable_libtorch: %s", repr(e) + ) with contextlib.suppress(ImportError): import vllm._moe_C_stable_libtorch # noqa: F401 with contextlib.suppress(ImportError): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7ad0b42da7..bd5d981912 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -365,7 +365,7 @@ class Platform: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %s", repr(e)) with contextlib.suppress(ImportError): import vllm._moe_C_stable_libtorch # noqa: F401 From f2aad6aa70749ccc4604d31d6f8a04c6728c3361 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 15 Sep 2026 13:36:51 -0700 Subject: [PATCH 117/139] [MRV2] Buffer util simplifications (#56888) Signed-off-by: Nick Hill --- tests/utils_/test_torch_utils.py | 25 +++++++++++++ .../v1/attention/test_indexer_dcp_localize.py | 2 +- tests/v1/worker/test_gpu_pcp_manager.py | 6 ++-- .../layers/attention/mla_attention.py | 35 +++++++++---------- vllm/model_executor/models/diffusion_gemma.py | 6 ++-- vllm/models/dots3_note/nvidia/attention.py | 22 ++++++++---- .../flash_linear_attention/ops/index.py | 5 +-- vllm/utils/torch_utils.py | 17 +++++++-- vllm/v1/attention/backends/flashinfer.py | 17 ++------- .../attention/backends/mla/flashmla_sparse.py | 9 +++-- vllm/v1/attention/backends/mla/indexer.py | 7 ++-- vllm/v1/sample/rejection_sampler.py | 5 ++- vllm/v1/spec_decode/llm_base_proposer.py | 1 - vllm/v1/utils.py | 7 ++-- vllm/v1/worker/cpu/shm.py | 9 +++-- vllm/v1/worker/gpu/buffer_utils.py | 19 ---------- vllm/v1/worker/gpu/model_runner.py | 13 +++---- vllm/v1/worker/gpu/pcp_manager.py | 12 +++---- vllm/v1/worker/gpu/pp_utils.py | 6 ++-- .../gpu/spec_decode/adaptive_verification.py | 9 ++--- vllm/v1/worker/gpu/structured_outputs.py | 5 ++- 21 files changed, 125 insertions(+), 112 deletions(-) diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index 6901a162a4..79dd3e320a 100644 --- a/tests/utils_/test_torch_utils.py +++ b/tests/utils_/test_torch_utils.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest import torch from vllm.utils.torch_utils import ( OMP_NUM_THREADS_SET_BY_VLLM, + async_tensor_h2d, available_cpu_count, common_broadcastable_dtype, current_stream, @@ -192,3 +194,26 @@ def test_runtime_threads_override_vllm_set_omp_num_threads( torch.set_num_threads(3) set_torch_threads_for_runtime() assert torch.get_num_threads() == 1 + + +@pytest.mark.parametrize( + "device", ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"] +) +def test_async_tensor_h2d_staging(device): + """Staging must preserve values while normalizing dtype and layout.""" + # Non-contiguous source with a dtype conversion: would previously sync. + src = torch.arange(12, dtype=torch.int64).reshape(3, 4).T + result = async_tensor_h2d(src, device=device, dtype=torch.int32) + assert result.dtype == torch.int32 + assert result.is_contiguous() + assert torch.equal(result.cpu(), src.to(torch.int32)) + + # numpy and list sources still work. + assert torch.equal( + async_tensor_h2d(np.arange(4, dtype=np.int64), device=device), + torch.arange(4, device=device), + ) + assert torch.equal( + async_tensor_h2d([1, 2, 3], device=device, dtype=torch.int32), + torch.tensor([1, 2, 3], dtype=torch.int32, device=device), + ) diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 812bce39d8..8bc409adb7 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -996,7 +996,7 @@ def test_pcp_plan_deinterleave_restores_global_order( monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False) monkeypatch.setattr( - "vllm.v1.attention.backends.mla.indexer.async_copy_to_gpu", + "vllm.v1.attention.backends.mla.indexer.async_tensor_h2d", lambda x, device: x.to(device), ) scheduled = np.array(req_lens, dtype=np.int64) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index d3ec562aa7..2bd7be2cfa 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -17,7 +17,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_c from vllm.v1.worker.gpu.pcp_manager import PCPManager -def _copy_to_cpu(value, out=None, device=None): +def _copy_to_cpu(value, device=None, dtype=None, out=None): tensor = torch.from_numpy(value) if isinstance(value, np.ndarray) else value if out is not None: return out.copy_(tensor) @@ -83,7 +83,7 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch): device=torch.device("cpu"), dcp_world_size=1, ) - monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu) + monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) segments_by_rank, per_rank_num_tokens = manager._build_batch_layout( num_scheduled_tokens=np.ones(3, dtype=np.int32), @@ -158,7 +158,7 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): device=torch.device("cpu"), dcp_world_size=1, ) - monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu) + monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) with pytest.raises(ValueError, match="smaller than the largest rank-local batch"): manager._build_batch_layout( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 874c43f8cd..1e673bd2da 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -269,11 +269,11 @@ from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, _resolve_layer_name, + async_tensor_h2d, direct_register_custom_op, get_dtype_size, is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, - np_to_pinned_tensor, ) from vllm.v1.attention.backend import ( AttentionBackend, @@ -1980,11 +1980,6 @@ def plan_mla_context_chunks( return plans -def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor: - """Pinned int32 CPU tensor backing one concatenated per-chunk field.""" - return np_to_pinned_tensor(np.asarray(values, dtype=np.int32)) - - def align_mla_chunked_context_workspace_size( vllm_config: VllmConfig, workspace_size: int, @@ -2161,24 +2156,28 @@ def build_mla_chunked_context_metadata( token_offset += num_tokens local_token_offset += num_local_tokens - seq_lens_cpu = _flat_int32(seq_lens_flat) - starts_and_context_lens = _flat_int32(starts_flat + context_lens).to( - device, non_blocking=True + seq_lens_cpu = torch.tensor( + seq_lens_flat, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY + ) + starts_and_context_lens = async_tensor_h2d( + starts_flat + context_lens, device=device, dtype=torch.int32 ) starts = starts_and_context_lens[: len(starts_flat)] context_lens_gpu = starts_and_context_lens[len(starts_flat) :] - cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True) - cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True) - token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to( - device, non_blocking=True + cu_seq_lens = async_tensor_h2d(cu_seq_lens_flat, device=device, dtype=torch.int32) + cu_seqlens_q = async_tensor_h2d(cu_seqlens_q_flat, device=device, dtype=torch.int32) + token_to_seq = async_tensor_h2d( + np.concatenate(token_to_seq_parts), device=device, dtype=torch.int32 ) if use_dcp: - padded_local_cu_seq_lens = _flat_int32(padded_local_cu_seq_lens_flat).to( - device, non_blocking=True + padded_local_cu_seq_lens = async_tensor_h2d( + padded_local_cu_seq_lens_flat, device=device, dtype=torch.int32 + ) + padded_local_token_to_seq = async_tensor_h2d( + np.concatenate(padded_local_token_to_seq_parts), + device=device, + dtype=torch.int32, ) - padded_local_token_to_seq = _flat_int32( - np.concatenate(padded_local_token_to_seq_parts) - ).to(device, non_blocking=True) chunks: list[MLACommonPrefillMetadata.ContextChunk] = [] for index, (plan, layout) in enumerate(zip(plans, layouts)): diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 603d68410d..c90d78bd6d 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -53,7 +53,7 @@ from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.outputs import LogprobsTensors from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.attn_utils import build_attn_metadata -from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores @@ -1159,7 +1159,7 @@ class DiffusionSampler: return # Move the slot indices across once, up front: indexing a device # tensor with a numpy array copies them over synchronously each time. - ps_gpu = async_copy_to_gpu( + ps_gpu = async_tensor_h2d( ps.astype(np.int64), device=states.is_encoder_phase.device ) states.init_canvas(ps_gpu) @@ -1266,7 +1266,7 @@ class DiffusionSampler: # was truncated near max_model_len, in which case the scheduler gave us # fewer than CL logits for that request. valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0] - valid_canvas_len = async_copy_to_gpu( + valid_canvas_len = async_tensor_h2d( valid_canvas_len_np.astype(np.int64), device=device ) diff --git a/vllm/models/dots3_note/nvidia/attention.py b/vllm/models/dots3_note/nvidia/attention.py index cbd9ad7f4d..27928b6e6f 100644 --- a/vllm/models/dots3_note/nvidia/attention.py +++ b/vllm/models/dots3_note/nvidia/attention.py @@ -19,7 +19,11 @@ from vllm.model_executor.layers.attention.mla_attention import ( ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import ( + PIN_MEMORY, + async_tensor_h2d, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import ( AttentionCGSupport, AttentionLayer, @@ -231,8 +235,12 @@ def _build_sliding_window_metadata( query_lens = query_lens_cpu[req_start:req_end] kv_lens = kv_lens_cpu[req_start:req_end] num_reqs = req_end - req_start - cu_seq_lens_q_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32) - cu_seq_lens_k_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32) + cu_seq_lens_q_cpu = torch.zeros( + num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY + ) + cu_seq_lens_k_cpu = torch.zeros( + num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY + ) torch.cumsum(query_lens, 0, out=cu_seq_lens_q_cpu[1:]) torch.cumsum(kv_lens, 0, out=cu_seq_lens_k_cpu[1:]) token_to_seq_cpu = torch.repeat_interleave( @@ -248,8 +256,8 @@ def _build_sliding_window_metadata( query_end=query_end, cu_seq_lens_q=cu_seq_lens_q_cpu.to(device, non_blocking=True), cu_seq_lens_k=cu_seq_lens_k_cpu.to(device, non_blocking=True), - starts=starts_cpu[req_start:req_end].to(device, non_blocking=True), - token_to_seq=token_to_seq_cpu.to(device, non_blocking=True), + starts=async_tensor_h2d(starts_cpu[req_start:req_end], device), + token_to_seq=async_tensor_h2d(token_to_seq_cpu, device), num_kv_tokens=num_kv_tokens, max_seq_len_q=int(query_lens.max().item()), max_seq_len_k=int(kv_lens.max().item()), @@ -396,8 +404,8 @@ class Dots3NoteMLAMetadataBuilder(TritonMLAMetadataBuilder): seq_lens = common_attn_metadata.seq_lens[ reqs_start + req_start : reqs_start + req_end ] - query_lens = query_lens_cpu[req_start:req_end].to( - self.device, non_blocking=True + query_lens = async_tensor_h2d( + query_lens_cpu[req_start:req_end], self.device ) kv_lens = torch.minimum(seq_lens, query_lens + self.sliding_window - 1) chunk.cu_seq_lens_k.zero_() diff --git a/vllm/third_party/flash_linear_attention/ops/index.py b/vllm/third_party/flash_linear_attention/ops/index.py index cc6f54d8af..880be7747d 100644 --- a/vllm/third_party/flash_linear_attention/ops/index.py +++ b/vllm/third_party/flash_linear_attention/ops/index.py @@ -10,6 +10,7 @@ import torch from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.torch_utils import async_tensor_h2d from vllm.triton_utils import triton from .utils import tensor_cache @@ -27,8 +28,8 @@ def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Te chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() indices = torch.cat([torch.arange(n) for n in chunk_counts]) chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1) - return chunk_indices.to( - device=cu_seqlens.device, dtype=cu_seqlens.dtype, non_blocking=True + return async_tensor_h2d( + chunk_indices, device=cu_seqlens.device, dtype=cu_seqlens.dtype ) diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 6f7dc14346..0b330edf00 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -705,17 +705,30 @@ def create_kv_caches_with_random( def async_tensor_h2d( data: list | np.ndarray | torch.Tensor, - device: str | torch.device, + device: str | torch.device | None = None, dtype: torch.dtype | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: """Copy list/numpy array/tensor async from host to device.""" + if dtype is None and out is not None: + dtype = out.dtype if isinstance(data, np.ndarray): data = torch.from_numpy(data) if isinstance(data, torch.Tensor): - t = data.pin_memory() if PIN_MEMORY else data + t = data + if PIN_MEMORY and not t.is_pinned(): + # Stage in pinned, contiguous buffer to ensure fully async copy. + t = torch.empty( + t.shape, dtype=dtype or t.dtype, device="cpu", pin_memory=True + ).copy_(t) else: t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu") assert t.is_cpu + + if out is not None: + assert out.dtype == dtype + return out.copy_(t, non_blocking=True) + assert device is not None, "must provide destination tensor or device" return t.to(device=device, dtype=dtype, non_blocking=True) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 0da68e3e3d..7ab2b2018e 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -1264,15 +1264,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): """ # write self.paged_kv_indptr_cpu inplace (0-index is always 0) np.cumsum( - num_blocks_np, - dtype=np.int32, - out=self.paged_kv_indptr.np[1 : num_reqs + 1], + num_blocks_np, dtype=np.int32, out=self.paged_kv_indptr.np[1 : num_reqs + 1] ) - paged_kv_indptr = self.paged_kv_indptr.gpu[: num_reqs + 1] - paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_reqs + 1] - if PIN_MEMORY: - paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory() - paged_kv_indptr.copy_(paged_kv_indptr_cpu, non_blocking=True) + paged_kv_indptr = self.paged_kv_indptr.copy_to_gpu(num_reqs + 1) # write self.paged_kv_indices inplace num_actual_pages = self.paged_kv_indptr.np[num_reqs] @@ -1292,12 +1286,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): page_size, paged_kv_last_page_len_np, ) - paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs] - if PIN_MEMORY: - paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory() - self.paged_kv_last_page_len.gpu[:num_reqs].copy_( - paged_kv_last_page_len_cpu, non_blocking=True - ) + self.paged_kv_last_page_len.copy_to_gpu(num_reqs) return paged_kv_indices def build( diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 9c440776ae..3a9ca2ffdb 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -20,7 +20,11 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import ( from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import current_stream, is_quantized_kv_cache +from vllm.utils.torch_utils import ( + async_tensor_h2d, + current_stream, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -48,7 +52,6 @@ from vllm.v1.attention.ops.flashmla import ( get_mla_metadata, ) from vllm.v1.kv_cache_interface import AttentionSpec -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: @@ -519,7 +522,7 @@ class FlashMLASparseMetadataBuilder( ) workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32)) max_prefill_buffer_size //= self.dcp_world_size - entry_rows = async_copy_to_gpu(row_bounds[:-1], device=self.device) + entry_rows = async_tensor_h2d(row_bounds[:-1], device=self.device) prefill_block_table = prefill_block_table.index_select(0, entry_rows) prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows) num_entries = len(workspace_rows) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 28bd450da1..837dbca817 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -27,7 +27,7 @@ from vllm.utils.deep_gemm import ( ) from vllm.utils.math_utils import round_down from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import PIN_MEMORY +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -47,7 +47,6 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, ) -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu logger = init_logger(__name__) @@ -338,14 +337,14 @@ def build_pcp_global_chunk_plan( + local ) - cu = async_copy_to_gpu(cu, device=device) + cu = async_tensor_h2d(cu, device=device) return PCPGlobalChunkPlan( row_start_cu=cu[0], global_cu=cu[1], padded_local_cu=cu[2], padded_local_total=padded_total, total=total, - deinterleave_idx=async_copy_to_gpu(idx, device=device), + deinterleave_idx=async_tensor_h2d(idx, device=device), ) diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 5aa8345473..85b623bd84 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -19,6 +19,7 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import ( triton_kernel_dispatcher_with_warmup, ) from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor from vllm.v1.sample.metadata import SamplingMetadata @@ -327,9 +328,7 @@ class RejectionSampler(nn.Module): num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu") original_indices = torch.arange(num_requests, device="cpu") repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens) - repeat_indices = repeat_indices_cpu.to( - device=logits.device, non_blocking=True - ) + repeat_indices = async_tensor_h2d(repeat_indices_cpu, logits.device) logits = self.apply_penalties( logits, sampling_metadata, metadata, repeat_indices, output_token_ids ) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index c0e456521c..1bd1bdc45b 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -255,7 +255,6 @@ class SpecDecodeBaseProposer: self.backup_next_token_ids = CpuGpuBuffer( self.max_batch_size, dtype=torch.int32, - pin_memory=PIN_MEMORY, device=device, with_numpy=True, ) diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index da17110d80..facafd2860 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -137,9 +137,10 @@ class CpuGpuBuffer: self.np = self.cpu.numpy() def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: - if n is None: - return self.gpu.copy_(self.cpu, non_blocking=True) - return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True) + cpu, gpu = self.cpu, self.gpu + if n is not None: + cpu, gpu = cpu[:n], gpu[:n] + return gpu.copy_(cpu.pin_memory() if PIN_MEMORY else cpu, non_blocking=True) def copy_to_cpu(self, n: int | None = None) -> torch.Tensor: """NOTE: Because this method is non-blocking, explicit synchronization diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index f92c6c3bf2..b615f4b5b5 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/vllm/v1/worker/cpu/shm.py @@ -76,14 +76,17 @@ import vllm.utils.torch_utils as torch_utils def async_tensor_h2d( data: list | np.ndarray | torch.Tensor, - device: str | torch.device, + device: str | torch.device | None = None, dtype: torch.dtype | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: if isinstance(data, np.ndarray): data = torch.from_numpy(data) - if isinstance(data, torch.Tensor): + if not isinstance(data, torch.Tensor): + data = torch.tensor(data, dtype=dtype, device="cpu") + elif out is None: return data.to(dtype=dtype) - return torch.tensor(data, dtype=dtype, device="cpu") + return data if out is None else out.copy_(data) torch_utils.async_tensor_h2d = async_tensor_h2d diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 0066f80907..0731072e59 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -10,7 +10,6 @@ from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import ( - PIN_MEMORY, async_tensor_h2d, get_accelerator_view_from_cpu_tensor, ) @@ -27,24 +26,6 @@ def set_default_max_concurrency(n: int) -> None: _DEFAULT_MAX_CONCURRENCY = max(2, n) -def async_copy_to_gpu( - x: torch.Tensor | np.ndarray, - out: torch.Tensor | None = None, - device: torch.device | None = None, -) -> torch.Tensor: - if isinstance(x, np.ndarray): - x = torch.from_numpy(x) - assert x.is_cpu - - if out is None: - assert device is not None - out = torch.empty_like(x, device=device) - - # pin_memory() is no-op if the memory is already pinned. - pinned = x.pin_memory() if PIN_MEMORY else x - return out.copy_(pinned, non_blocking=True) - - class UvaBuffer: def __init__(self, size: int | Sequence[int], dtype: torch.dtype): if not is_uva_available(): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 25395c3330..b66d552cb4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -64,7 +64,7 @@ from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -99,10 +99,7 @@ from vllm.v1.worker.gpu.attn_utils import ( init_kv_cache, ) from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.buffer_utils import ( - async_copy_to_gpu, - set_default_max_concurrency, -) +from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, @@ -1289,7 +1286,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): req_ids = batch_req_state.req_ids num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens idx_mapping_np = batch_req_state.idx_mapping_np - idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device) num_reqs = len(req_ids) # Get the number of draft tokens for each request. @@ -1324,7 +1321,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) cu_num_logits_np[0] = 0 np.cumsum(num_logits, out=cu_num_logits_np[1:]) - cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) + cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device) adaptive_verification = ( self.adaptive_verification if num_draft_tokens_per_req is not None else None @@ -1352,7 +1349,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Some attention backends like FA3 require query_start_loc to be non-decreasing. query_start_loc_np[num_reqs + 1 :] = num_tokens query_start_loc = self.input_buffers.query_start_loc - async_copy_to_gpu(query_start_loc_np, out=query_start_loc) + async_tensor_h2d(query_start_loc_np, out=query_start_loc) if adaptive_verification is not None: cu_num_logits, query_start_loc, total_num_draft_tokens = ( adaptive_verification.reallocate_drafts(req_ids, idx_mapping) diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index adfbf2f9f0..376e0f2510 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -10,9 +10,9 @@ import torch from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import ( InputBatch, InputBuffers, @@ -358,13 +358,13 @@ class PCPManager: dtype=np.int64, ) - self._hidden_restore_idx = async_copy_to_gpu( + self._hidden_restore_idx = async_tensor_h2d( hidden_restore_idx, device=self.device ) - self._padded_gather_idx = async_copy_to_gpu( + self._padded_gather_idx = async_tensor_h2d( padded_gather_idx, device=self.device ) - self._gathered_kv_write_mask = async_copy_to_gpu( + self._gathered_kv_write_mask = async_tensor_h2d( gathered_kv_write_mask, device=self.device ) return segments_by_rank, per_rank_num_tokens @@ -536,12 +536,12 @@ class PCPManager: local_query_start_loc_out = local_query_start_loc_np[1 : num_local_reqs + 1] np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out) local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens - async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc) + async_tensor_h2d(local_query_start_loc_np, out=input_buffers.query_start_loc) local_query_start_loc = input_buffers.query_start_loc[ : num_reqs_after_padding + 1 ] - local_to_global_req_idx = async_copy_to_gpu( + local_to_global_req_idx = async_tensor_h2d( local_to_global_req_idx_np, device=self.device ) seq_lens = input_buffers.seq_lens[:num_reqs_after_padding] diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 8c65a23ef7..ab1b991d1a 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -11,7 +11,7 @@ import torch from vllm.distributed.parallel_state import get_pp_group from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.worker.gpu.input_batch import InputBatch @@ -138,7 +138,7 @@ class PPHandler: return None # Filter excluded request indices. idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np) - idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device) self.main_stream.wait_event(slot.event) if slot.draft_tokens is not None and draft_tokens_to_update is not None: @@ -148,7 +148,7 @@ class PPHandler: keep = ~exclude_mask keep_t = torch.as_tensor(keep, device=self.device) draft_tokens = draft_tokens[keep_t] - draft_idx_mapping = async_copy_to_gpu( + draft_idx_mapping = async_tensor_h2d( slot.idx_mapping_np[keep], device=self.device ) draft_tokens_to_update[draft_idx_mapping] = draft_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index af58b8993b..245f304ac4 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -14,6 +14,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream @@ -21,7 +22,6 @@ from vllm.v1.worker.gpu.attn_utils import ( get_attn_cg_support, get_query_lens_mismatch_unsupported_backend, ) -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu logger = init_logger(__name__) _PROFILE_REPLAYS = 5 @@ -413,7 +413,7 @@ class AdaptiveVerificationManager: if draft_budget == 0: capacities.zero_() else: - async_copy_to_gpu(scheduled_drafts, out=capacities) + async_tensor_h2d(scheduled_drafts, out=capacities) if draft_budget < int(scheduled_drafts.sum()): _assign_draft_token_budget_compiled( self._confidence_probs, @@ -424,10 +424,7 @@ class AdaptiveVerificationManager: ) num_non_draft_tokens_gpu = self._num_non_draft_tokens[:num_reqs] - async_copy_to_gpu( - num_non_draft_tokens, - out=num_non_draft_tokens_gpu, - ) + async_tensor_h2d(num_non_draft_tokens, out=num_non_draft_tokens_gpu) self._cu_num_logits[:1].zero_() torch.cumsum( capacities + self.num_bonus_tokens, diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 8437b0aed9..057d6f7154 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -5,8 +5,7 @@ import torch from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import PIN_MEMORY -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.worker.gpu.input_batch import InputBatch @@ -68,7 +67,7 @@ class StructuredOutputsWorker: # Asynchronously copy the bitmask to GPU. with torch.cuda.stream(self.copy_stream): - bitmask = async_copy_to_gpu( + bitmask = async_tensor_h2d( grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]] ) From 12575123059c2142e39c9af4e217a4792c3c8176 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Wed, 16 Sep 2026 05:07:14 +0800 Subject: [PATCH 118/139] [DSA] Wire DeepGEMM sparse MQA logits into the DeepSeek V4.1 indexer (#56254) Signed-off-by: Jared Wen Co-authored-by: Claude Fable 5.1 --- .../test_fused_indexer_q_rope_quant.py | 52 ++ .../layers/test_mla_short_prefill_indexer.py | 544 ++++++++++++++++++ vllm/config/attention.py | 9 + .../attention/dsa/sparse_mqa_logits.py | 469 +++++++++++++++ .../layers/sparse_mqa_indexer.py | 209 +++++++ .../deepseek_v4/common/ops/fused_indexer_q.py | 24 +- .../nvidia/ops/fused_indexer_q_cutedsl.py | 17 +- vllm/models/deepseek_v41/attention.py | 87 ++- vllm/utils/deep_gemm.py | 183 ++++++ vllm/v1/attention/backends/mla/indexer.py | 7 +- .../attention/backends/mla/sparse_indexer.py | 235 ++++++++ vllm/v1/core/kv_cache_utils.py | 9 + vllm/v1/kv_cache_interface.py | 11 +- 13 files changed, 1830 insertions(+), 26 deletions(-) create mode 100644 vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py create mode 100644 vllm/model_executor/layers/sparse_mqa_indexer.py create mode 100644 vllm/v1/attention/backends/mla/sparse_indexer.py diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 8770882ddf..7464c90769 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -560,3 +560,55 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat torch.testing.assert_close( k_scale[emitted], expected_scale[emitted], rtol=0, atol=0 ) + + +@pytest.mark.parametrize("use_cutedsl", [False, True]) +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="MXFP4 indexer cache requires an SM100-family GPU", +) +@torch.inference_mode() +def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl): + """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's + sparse MQA-logits kernels take) instead of fp32; the values are the fp32 + result rounded once, so the scoring kernel needs no cast.""" + if use_cutedsl and not has_cutedsl(): + pytest.skip("cutedsl (cutlass) not installed") + + device = "cuda" + torch.manual_seed(0) + num_tokens, n_head = 257, 32 + q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device) + positions = torch.randint( + 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device + ) + cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device) + weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device) + softmax_scale, head_scale = HEAD_DIM**-0.5, 0.125 # pow2 head_scale, see above + + _, weights_ref = _reference( + positions, q, cos_sin_cache, weights, softmax_scale, head_scale, n_head, True + ) + cutedsl_patch = ( + mock.patch( + "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl", + return_value=False, + ) + if not use_cutedsl + else contextlib.nullcontext() + ) + with cutedsl_patch: + _, weights_fused = fused_indexer_q_rope_quant( + positions, + q.clone(), + cos_sin_cache, + weights, + softmax_scale, + head_scale, + use_fp4=True, + weights_out_dtype=torch.bfloat16, + ) + assert weights_fused.dtype == torch.bfloat16 + assert torch.equal(weights_fused, weights_ref.to(torch.bfloat16)) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index f8e7a89cca..41a06d0546 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -450,3 +450,547 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding( keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block) reference = torch.where(keep, 3.0, -torch.inf) torch.testing.assert_close(logits, reference, rtol=0, atol=0) + + +# --- DeepGEMM sparse-MQA indexer path (V4.1 two-level candidate filtering) --- + + +def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int): + """Production-faithful candidates: top min(avail, K) blocks per row — + full coverage for lens <= K*cbk, a fixed-size subset beyond that.""" + avail = ((lens.cpu() + cbk - 1) // cbk).tolist() + rows = [] + for a in avail: + n = min(a, num_candidates) + picks = torch.randperm(a)[:n].sort().values[:num_candidates].tolist() + rows.append(picks + [-1] * (num_candidates - len(picks))) + return torch.tensor(rows, dtype=torch.int32) + + +def _candidate_token_mask( + candidates: torch.Tensor, cbk: int, total: int, ks: torch.Tensor | None = None +): + """[rows, total] bool mask of the candidate-covered token positions. + + ``ks`` converts request-local candidate blocks to absolute positions + (prefill's packed workspace); pass None when already request-local. + """ + rows = candidates.shape[0] + device = candidates.device + # Sentinel block keeps -1 padding from scattering onto real columns. + width = max(total, (candidates.shape[1] + 1) * cbk) + if ks is not None: + width += int(ks.max()) + cbk + in_cand = torch.zeros(rows, width, dtype=torch.bool, device=device) + block_id = torch.where(candidates >= 0, candidates.long(), candidates.shape[1]) + tok = block_id.unsqueeze(2) * cbk + torch.arange(cbk, device=device) + if ks is not None: + tok = tok + ks.unsqueeze(1).unsqueeze(2) + keep = (candidates >= 0).unsqueeze(2).expand(-1, -1, cbk) + in_cand.scatter_(1, tok.flatten(1), keep.flatten(1)) + return in_cand[:, :total] + + +def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor: + """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid + count are -1 (the indexer's "no token" sentinel).""" + values, indices = masked.topk(topk, dim=1) + return torch.where(values > float("-inf"), indices.int(), -1) + + +def _assert_same_selection( + logits: torch.Tensor, + got: torch.Tensor, + expected: torch.Tensor, + row_starts: torch.Tensor | None = None, +) -> None: + """Both top-k index sets must pick the same multiset of logit values per + row (-1 padding must match too). Indices themselves may differ: the + top-k kernels break ties among equal values in nondeterministic order.""" + for r in range(logits.shape[0]): + base = int(row_starts[r]) if row_starts is not None else 0 + g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0] + assert g.numel() == e.numel(), f"row {r}: {g.numel()} vs {e.numel()} selected" + vg = logits[r][(base + g).long()].sort().values + ve = logits[r][(base + e).long()].sort().values + assert torch.equal(vg, ve), f"row {r}: selected values differ" + + +def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict: + """Caller-owned buffers the sparse path writes into (the metadata builder + owns these in production).""" + i32 = dict(dtype=torch.int32, device="cuda") + return dict( + sparse_indices=torch.zeros(rows, num_sparse, **i32), + end=torch.zeros(rows, **i32), + col_indices=torch.zeros(rows, topk, **i32), + ) + + +def _expand_candidates_reference( + candidates: torch.Tensor, ks: torch.Tensor, ke: torch.Tensor, cbk: int, sbk: int +) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand + each candidate block into ``cbk // sbk`` sparse blocks, keep those inside + ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by + repeating the last valid block, and count the valid sparse columns.""" + ratio = cbk // sbk + rows, num_candidates = candidates.shape + num_sparse = num_candidates * ratio + big = torch.iinfo(torch.int32).max + offsets = torch.arange(ratio, device=candidates.device, dtype=torch.int32) + blocks = (candidates.unsqueeze(-1) * ratio + offsets).flatten(1) + max_block = ((ke - ks).clamp(min=0) + sbk - 1) // sbk + valid = (blocks >= 0) & (blocks < max_block.unsqueeze(1)) + sorted_blocks, _ = torch.where(valid, blocks, big).sort(dim=1) + n_valid = (sorted_blocks != big).sum(dim=1, dtype=torch.int32) + base = (ks // sbk).unsqueeze(1) + kernel_blocks = torch.where(sorted_blocks != big, sorted_blocks, 0) + base + last_slot = (n_valid - 1).clamp(min=0).unsqueeze(1).long() + last_block = kernel_blocks.gather(1, last_slot) + pad = torch.arange(num_sparse, device=candidates.device) >= n_valid.unsqueeze(1) + sparse_indices = torch.where(pad, last_block, kernel_blocks) + last_start = last_block.squeeze(1) * sbk + ks % sbk + last_fill = (ke - last_start).clamp(min=1, max=sbk) + end = torch.where(n_valid > 0, (n_valid - 1) * sbk + last_fill, 0) + return sparse_indices.int(), end.int() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") +def test_candidate_blocks_to_sparse_indices_math(): + """Candidate->sparse-block expansion: ks anchoring, range filtering, + repeat-last padding, and the per-row valid column count (``end``).""" + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + + # Rows cover: exactly one block, an unaligned ks with a partial tail + # block, an empty row, duplicates / -1 padding / out-of-range garbage + # (warmup runs touch uninitialized buffers), and int32-overflow junk. + ks = torch.tensor([64, 13, 100, 8, 0], dtype=torch.int32, device="cuda") + ke = torch.tensor([72, 30, 100, 40, 1000], dtype=torch.int32, device="cuda") + candidates = torch.tensor( + [ + [0, -1, -1, -1], + [2, 0, -1, -1], + [3, -1, -1, -1], + [1, 1, 7, 0], + [2**30, -(2**30), 0, -1], + ], + dtype=torch.int32, + device="cuda", + ) + indices, end = candidate_blocks_to_sparse_indices( + candidates, ks, ke, candidate_block_size=8, sparse_block_kv=8 + ) + assert indices.tolist() == [ + [8, 8, 8, 8], + [1, 3, 3, 3], + [12, 12, 12, 12], + [1, 2, 2, 2], + [0, 0, 0, 0], + ] + # end: row 1's last block covers [29, 37) but ke = 30 -> one valid token; + # row 3 keeps the duplicated block (production candidates are unique). + assert end.tolist() == [8, 9, 0, 24, 8] + + # Wider candidate blocks expand into multiple sparse blocks. + idx2, end2 = candidate_blocks_to_sparse_indices( + torch.tensor([[1, 0]], dtype=torch.int32, device="cuda"), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([48], dtype=torch.int32, device="cuda"), + candidate_block_size=16, + sparse_block_kv=8, + ) + assert idx2.tolist() == [[0, 1, 2, 3]] + assert end2.tolist() == [32] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") +@pytest.mark.parametrize("cbk,sbk", [(8, 8), (16, 8), (16, 16)]) +def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int): + """The Triton expansion equals the PyTorch reference on production-shaped + random inputs, including strided (column-sliced) candidate rows and + strided output buffers.""" + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + + torch.manual_seed(0) + rows, num_candidates = 41, 2048 + ratio = cbk // sbk + lens = torch.randint(1, 40000, (rows,), dtype=torch.int32) + lens[0], lens[1] = 1, 3 + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke = ks.cuda(), (ks + lens).cuda() + candidates = _make_candidates(lens, num_candidates, cbk).cuda() + # Row 2 gets out-of-range garbage past the padding, as a CUDA-graph + # warm-up on uninitialized buffers would. + candidates[2, -3:] = torch.tensor([2**30, -(2**29), 10**6], dtype=torch.int32) + + ref_indices, ref_end = _expand_candidates_reference(candidates, ks, ke, cbk, sbk) + + cand_buf = torch.full((rows, num_candidates + 64), -7, dtype=torch.int32).cuda() + cand_buf[:, :num_candidates] = candidates + si_buf = torch.zeros(rows, num_candidates * ratio + 32, dtype=torch.int32).cuda() + end_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda") + ks_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda") + ks_buf[::2] = ks + got_indices, got_end = candidate_blocks_to_sparse_indices( + cand_buf[:, :num_candidates], + ks_buf[::2], + ke, + cbk, + sbk, + out=(si_buf[:, : num_candidates * ratio], end_buf[::2]), + ) + assert torch.equal(got_indices, ref_indices) + assert torch.equal(got_end, ref_end) + assert torch.equal( + cand_buf[:, num_candidates:], torch.full_like(cand_buf[:, num_candidates:], -7) + ) + assert not si_buf[:, num_candidates * ratio :].any() + assert not end_buf[1::2].any() + + +def _skip_unless_sm100_sparse_kernels(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + pytest.skip("SM100 required") + deep_gemm = pytest.importorskip("deep_gemm") + if not hasattr(deep_gemm, "fp8_fp4_sparse_mqa_logits"): + pytest.skip("DeepGEMM >= 2.8 required for sparse MQA logits") + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + has_deep_select, + ) + + if not has_deep_select(): + pytest.skip("DeepSelect extension (vllm._deepselect_C) required") + return deep_gemm + + +def _quant_fp4(x: torch.Tensor): + from deep_gemm.utils import per_token_cast_to_fp4 + + return per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): + """The sparse-MQA candidate path selects exactly the same logit values as + ``torch.topk`` over dense bf16 logits masked to the candidates (the two + DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA + graph capture, and a later layer reusing the step's expansion/schedule + selects the same values.""" + deep_gemm = _skip_unless_sm100_sparse_kernels() + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + sparse_mqa_logits_prefill_chunk, + ) + + torch.manual_seed(0) + rows, num_heads, head_dim = 37, 32, 128 + cbk, topk, num_candidates = 8, 512, 2048 # V4.1-Flash production values + lens = torch.randint(1, 30000, (rows,), dtype=torch.int32) + lens[0], lens[1], lens[2] = 1, 5, 2048 # single token, sub-block, long + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda() + total = int(lens.sum()) + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + k_fp, k_sf = _quant_fp4( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) + ) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(lens, num_candidates, cbk).cuda() + + # Reference: dense bf16 logits masked to the candidates, then torch.topk. + dense = deep_gemm.fp8_fp4_mqa_logits( + ( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + ), + (k_fp.view(torch.int8), k_sf.view(total)), + weights, + ks, + ke, + clean_logits=False, + logits_dtype=torch.bfloat16, + ) + cols = torch.arange(total, device="cuda") + in_range = (cols - ks.unsqueeze(1) >= 0) & ( + cols - ks.unsqueeze(1) < lens.unsqueeze(1) + ) + masked = dense.masked_fill( + ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)), + float("-inf"), + ).float() + ref_abs = _reference_topk(masked, topk) # packed-workspace positions + ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1) + + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + scratch = _sparse_scratch(rows, num_candidates, topk) + + def run_sparse(kernel_metadata=None): + return sparse_mqa_logits_prefill_chunk( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + k_fp.view(torch.int8), + k_sf.view(total), + weights, + ks, + ke, + candidates, + cbk, + cbk, + topk, + out, + kernel_metadata=kernel_metadata, + **scratch, + ) + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run_sparse() # warm up JIT + the DeepGEMM workspace pre-capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_sparse() + out.fill_(-1) + graph.replay() + else: + kernel_metadata = run_sparse() + _assert_same_selection(dense, out, ref, row_starts=ks) + out.fill_(-1) + run_sparse(kernel_metadata) + + _assert_same_selection(dense, out, ref, row_starts=ks) + + +def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): + """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' + fp4 data first, then their packed UE8M0 scales.""" + device = req_lens.device + pages_per = (req_lens + page_kv - 1) // page_kv + num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) + kv_cache = torch.zeros( + num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device + ) + block_table = torch.zeros( + len(req_lens), max_pages, dtype=torch.int32, device=device + ) + page_pool = torch.randperm(num_pages, device=device) + offset = 0 + for b in range(len(req_lens)): + n = int(pages_per[b]) + block_table[b, :n] = page_pool[offset : offset + n].to(torch.int32) + offset += n + k_fp, k_sf = _quant_fp4( + torch.randn(int(req_lens[b]), head_dim, device=device, dtype=torch.bfloat16) + ) + rows68 = torch.zeros(n * page_kv, 68, dtype=torch.uint8, device=device) + rows68[: int(req_lens[b]), : head_dim // 2] = k_fp.view(torch.uint8) + rows68[: int(req_lens[b]), head_dim // 2 :] = k_sf.view(torch.uint8).view(-1, 4) + by_page = rows68.view(n, page_kv, 68) + split = torch.cat( + [by_page[..., :64].reshape(n, -1), by_page[..., 64:].reshape(n, -1)], dim=1 + ).view(n, page_kv, 1, head_dim // 2 + 4) + kv_cache[block_table[b, :n].long()] = split + return kv_cache, block_table + + +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("capture", [False, True]) +def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool): + """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks + from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to + one query row each), CUDA graph capture/replay and schedule reuse by a + later layer.""" + deep_gemm = _skip_unless_sm100_sparse_kernels() + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + sparse_mqa_logits_paged_decode, + ) + + torch.manual_seed(0) + batch, num_heads, head_dim, page_kv = 9, 32, 128, 128 + cbk, topk, num_candidates = 8, 512, 2048 + rows = batch * next_n + req_lens = torch.randint(1, 30000, (batch,), dtype=torch.int32, device="cuda") + req_lens[0], req_lens[1] = 1, 30000 # tiny row + a filtered (subset) row + # Row (b, j) sees L_b - next_n + j + 1 tokens (native MTP layout). + row_lens = ( + (req_lens.unsqueeze(1) - next_n + torch.arange(1, next_n + 1, device="cuda")) + .clamp(min=0) + .to(torch.int32) + ) + kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim) + assert kv_cache.stride(0) % 512 == 0 # the paged sparse kernel requires it + + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + q = q_fp.view(torch.int8).view(batch, next_n, num_heads, head_dim // 2) + q_scale = q_sf.view(batch, next_n, num_heads) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda() + + def reference_topk(): + context_lens_2d = row_lens.view(batch, next_n) + sched = deep_gemm.get_paged_mqa_logits_metadata( + context_lens_2d, page_kv, deep_gemm.get_num_sms(), None + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + (q, q_scale), + kv_cache, + weights, + context_lens_2d, + block_table, + sched, + int(req_lens.max()), + False, + torch.bfloat16, + None, + ) + lens_flat = row_lens.flatten() + cols = torch.arange(logits.shape[1], device="cuda") + masked = logits.masked_fill( + ~( + (cols.unsqueeze(0) < lens_flat.unsqueeze(1)) + & _candidate_token_mask(candidates, cbk, logits.shape[1]) + ), + float("-inf"), + ).float() + return _reference_topk(masked, topk), logits + + expected, dense_logits = reference_topk() + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + + # The metadata builder flattens every query to one row: per-row context + # lengths, an expanded block table and a row -> request map. + context_lens = row_lens.flatten().contiguous() + row_block_table = block_table.repeat_interleave(next_n, dim=0) + row_indices = torch.arange(batch, dtype=torch.int32, device="cuda") + row_indices = row_indices.repeat_interleave(next_n) + row_ks = torch.zeros(rows, dtype=torch.int32, device="cuda") + scratch = _sparse_scratch(rows, num_candidates, topk) + + def run_sparse(o: torch.Tensor, kernel_metadata=None): + return sparse_mqa_logits_paged_decode( + q.view(rows, 1, num_heads, head_dim // 2), + q_scale.view(rows, 1, num_heads), + kv_cache, + weights, + context_lens, + row_block_table, + row_indices, + candidates, + cbk, + cbk, + topk, + o, + row_ks=row_ks, + kernel_metadata=kernel_metadata, + **scratch, + ) + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run_sparse(out) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_sparse(out) + out.fill_(-1) + graph.replay() + else: + kernel_metadata = run_sparse(out) + _assert_same_selection(dense_logits, out, expected) + out.fill_(-1) + run_sparse(out, kernel_metadata) + + _assert_same_selection(dense_logits, out, expected) + + +def test_sparse_mqa_logits_accuracy_vs_torch(): + """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized + MXFP4 inputs.""" + _skip_unless_sm100_sparse_kernels() + from deep_gemm.utils import cast_back_from_fp4 + + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + from vllm.utils.deep_gemm import ( + fp8_fp4_sparse_mqa_logits, + get_sparse_mqa_logits_metadata, + ) + + torch.manual_seed(0) + rows, num_heads, head_dim = 17, 32, 128 + sbk, num_candidates = 8, 2048 + lens = torch.randint(1, 1500, (rows,), dtype=torch.int32) + lens[0] = 1 + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda() + total = int(lens.sum()) + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + k_fp, k_sf = _quant_fp4( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) + ) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(lens, num_candidates, sbk).cuda() + + sparse_idx, end = candidate_blocks_to_sparse_indices(candidates, ks, ke, sbk, sbk) + metadata = get_sparse_mqa_logits_metadata( + ks, ke, total, sparse_idx, torch.int8, sbk + ) + logits = fp8_fp4_sparse_mqa_logits( + ( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + ), + (k_fp.view(torch.int8), k_sf.view(total)), + weights, + metadata, + sparse_idx.shape[1], + sbk, + ) + + # PyTorch fp32 reference on the dequantized inputs. + q_deq = cast_back_from_fp4(q_fp, q_sf, gran_k=32, use_packed_ue8m0=True) + k_deq = cast_back_from_fp4(k_fp, k_sf, gran_k=32, use_packed_ue8m0=True) + dots = torch.einsum( + "rhd,td->rht", + q_deq.view(rows, num_heads, head_dim).float(), + k_deq.float(), + ) + ref = (dots.relu() * weights.float().unsqueeze(-1)).sum(1) # [rows, total] + + # Compare on valid sparse columns (absolute packed-workspace positions). + width = sparse_idx.shape[1] * sbk + cols = torch.arange(width, device="cuda") + valid = cols.unsqueeze(0) < end.unsqueeze(1).long() + pos = ( + sparse_idx.repeat_interleave(sbk, dim=1).long() * sbk + + (ks % sbk).unsqueeze(1) + + (cols % sbk).unsqueeze(0) + ).clamp(min=0, max=total - 1) + abs_err = (logits.float() - ref.gather(1, pos)).abs()[valid] + row_scale = ref.abs().amax(dim=1) + norm_err = ((logits.float() - ref.gather(1, pos)).abs() / row_scale.unsqueeze(1))[ + valid + ] + print( + f"sparse vs torch fp32: max_abs={abs_err.max().item():.4f}; " + f"rowmax-normalized p50={norm_err.median().item():.5f} " + f"p99={norm_err.quantile(0.99).item():.5f} max={norm_err.max().item():.5f}" + ) + # bf16-level numerics: far below the value gaps that change top-k selection. + assert norm_err.quantile(0.99).item() < 1e-2 + assert norm_err.max().item() < 2e-2 diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 129fbcbb08..3646677728 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -86,6 +86,15 @@ class AttentionConfig: indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel support in the backend.""" + indexer_sparse_logits: bool = False + """DeepSeek V4.1 two-level indexer: score only the candidate blocks with + DeepGEMM's sparse MQA-logits kernels instead of computing dense logits over + the whole context and masking them. Requires `indexer_kv_dtype="mxfp4"`, + an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension + (the top-k runs on the kernels' bf16 logits). The sparse path costs + O(candidate blocks) per query regardless of context length, so it pays off + for long contexts (roughly 32K tokens and beyond) and is slower below.""" + hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py new file mode 100644 index 0000000000..839c581583 --- /dev/null +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM sparse MQA-logits path for the DeepSeek V4.1 two-level indexer. + +The candidate-source indexer publishes the top candidate blocks (in units of +``candidate_block_size`` compressed positions). Instead of scoring every KV +position and masking the dense logits down to those blocks, this path scores +only the candidate tokens with DeepGEMM's sparse MQA logits kernels +(``fp8_fp4_(paged_)sparse_mqa_logits``, DeepGEMM >= 2.8, SM100 only), runs +DeepSelect's top-k directly on the bf16 logits they produce, and remaps the +selected columns back to request-local positions. + +Only the MXFP4 indexer cache is supported: the kernels require UE8M0-packed +(granularity-32) Q/KV scales, which is exactly the MXFP4 cache layout. The +FP8 indexer cache stores one fp32 scale per quant block instead. + +The functions here take caller-owned output buffers (``sparse_indices``, +``end``, ``col_indices``) so the hot path allocates nothing; the indexer +metadata builder owns them. +""" + +import importlib.util + +import torch + +from vllm.model_executor.layers.indexer_topk import ( + deep_select_topk, + get_deep_select_stride_requirement, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import ( + fp8_fp4_paged_sparse_mqa_logits, + fp8_fp4_sparse_mqa_logits, + get_paged_sparse_mqa_logits_metadata, + get_sparse_mqa_logits_metadata, +) + +# Sparse-block sizes (in tokens) supported by the DeepGEMM kernels. +SPARSE_BLOCK_KV_CHOICES = (16, 8) + + +def has_deep_select() -> bool: + """Whether the DeepSelect top-k extension is built and this GPU runs it.""" + return ( + importlib.util.find_spec("vllm._deepselect_C") is not None + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ) + + +def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None: + """Raise unless DeepSelect accepts the sparse logits layout. + + The sparse kernels return contiguous bf16 ``[rows, num_sparse_cols]`` + logits and the top-k writes int32 ``[rows, topk_tokens]`` indices; + DeepSelect needs both row strides aligned (1024B / 32B on SM100). + """ + in_align, out_align = get_deep_select_stride_requirement() + logits_row_bytes = num_sparse_cols * torch.bfloat16.itemsize + if num_sparse_cols >= 2**23 or logits_row_bytes % in_align != 0: + raise ValueError( + f"DeepSelect needs the sparse logits row ({num_sparse_cols} bf16 " + f"columns, {logits_row_bytes} B) to be {in_align}B-aligned." + ) + if topk_tokens > 4096 or (topk_tokens * torch.int32.itemsize) % out_align != 0: + raise ValueError( + f"DeepSelect needs topk <= 4096 with a {out_align}B-aligned int32 " + f"output row, got topk={topk_tokens}." + ) + + +_INT32_BIG = tl.constexpr(2**31 - 1) + + +@triton.jit +def _expand_candidates_kernel( + cand_ptr, + cand_stride, + ks_ptr, + ks_stride, + ke_ptr, + ke_stride, + si_ptr, + si_stride, + end_ptr, + end_stride, + K: tl.constexpr, + S: tl.constexpr, + RATIO: tl.constexpr, + SBK: tl.constexpr, +): + """One program per row: expand + filter + sort candidate blocks. + + Each row's valid sparse blocks come out sorted ascending (no dedup pass: + production candidates are unique, and the kernel tolerates repeats), + padded by repeating the last valid block — DeepGEMM's padding convention. + ``end`` counts valid sparse-token columns; valid columns form a prefix + because block ids are sorted. + """ + row = tl.program_id(0) + ks = tl.load(ks_ptr + row * ks_stride) + ke = tl.load(ke_ptr + row * ke_stride) + cols = tl.arange(0, K) + cand = tl.load(cand_ptr + row * cand_stride + cols) # K is pow2, no mask + if RATIO == 2: + cand = tl.interleave(cand * RATIO, cand * RATIO + 1) + + # Keep only blocks inside the row's KV range [0, ke - ks); the lower + # bound is free since candidates are non-negative. + kv_len = tl.maximum(ke - ks, 0) + max_block = (kv_len + SBK - 1) // SBK + valid = (cand >= 0) & (cand < max_block) + vals = tl.where(valid, cand, _INT32_BIG) + vals = tl.sort(vals) + + # Anchor at the row's ks rounded down to a sparse-block boundary + # (identity for the paged path, where ks == 0). + base = ks // SBK + si = tl.where(vals != _INT32_BIG, vals, 0) + base + n_valid = tl.sum(valid.to(tl.int32)) + last_rel = tl.max(tl.where(valid, cand, -1)) + pad_val = tl.where(n_valid > 0, last_rel + base, base) + tl.store( + si_ptr + row * si_stride + tl.arange(0, S), + tl.where(tl.arange(0, S) < n_valid, si, pad_val), + ) + + # Valid sparse-token column count; only the last block can be partial. + # Absolute start of the last block: (base + last_rel) * SBK + ks % SBK, + # which is exactly last_rel * SBK + ks. + last_start = last_rel * SBK + ks + fill = tl.minimum(tl.maximum(ke - last_start, 1), SBK) + end = tl.where(n_valid > 0, (n_valid - 1) * SBK + fill, 0) + tl.store(end_ptr + row * end_stride, end) + + +@triton.jit +def _sparse_topk_remap_kernel( + col_ptr, + col_stride, + si_ptr, + si_stride, + ks_ptr, + ks_stride, + out_ptr, + out_stride, + k, + SBK: tl.constexpr, + K_POW2: tl.constexpr, +): + """Remap sparse top-k columns to request-local positions. + + Column ``j * SBK + o`` scores the token at + ``(si[row, j] - ks // SBK) * SBK + o``; -1 columns stay -1. + """ + row = tl.program_id(0) + ks_base = tl.load(ks_ptr + row * ks_stride) // SBK + cols = tl.arange(0, K_POW2) + c = tl.load(col_ptr + row * col_stride + cols, mask=cols < k, other=-1) + valid = c >= 0 + cc = tl.where(valid, c, 0) + blocks = tl.load(si_ptr + row * si_stride + cc // SBK) + pos = (blocks - ks_base) * SBK + cc % SBK + tl.store(out_ptr + row * out_stride + cols, tl.where(valid, pos, -1), mask=cols < k) + + +def pick_sparse_block_kv(candidate_block_size: int) -> int: + """Largest sparse-block size (tokens) dividing the candidate block size.""" + for sbk in SPARSE_BLOCK_KV_CHOICES: + if candidate_block_size % sbk == 0: + return sbk + raise ValueError( + f"candidate_block_size={candidate_block_size} is not a multiple of " + f"any supported sparse_block_kv {SPARSE_BLOCK_KV_CHOICES}." + ) + + +def candidate_blocks_to_sparse_indices( + candidate_blocks: torch.Tensor, + row_ks: torch.Tensor, + row_ke: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + out: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Expand request-local candidate blocks into DeepGEMM sparse block ids. + + Args: + candidate_blocks: [rows, K] int32 request-local candidate block ids + (-1 padded), in units of ``candidate_block_size`` positions. K + must be a power of two (the in-kernel sort); rows may be strided. + row_ks/row_ke: [rows] int32 per-row K range; bounds are in the same + (packed-workspace) coordinates the sparse kernel iterates over. + Pass zeros for the paged path, whose blocks are context-relative. + candidate_block_size: Positions per candidate block. + sparse_block_kv: Positions per sparse block (8 or 16). + out: Optional ``(sparse_indices, end)`` buffers to write into, + shaped [rows, K * ratio] and [rows], both int32. + + Returns: + sparse_indices: [rows, K * ratio] int32 block ids for DeepGEMM, + sorted ascending per row with the valid prefix first (unique in + production, where candidates are unique) and the remaining slots + repeating the last valid block (the kernel's padding convention). + With the unaligned-ks variant, block ``i`` starts at + ``i * sparse_block_kv + ks % sparse_block_kv``. + end: [rows] int32 count of valid sparse-token columns; valid columns + always form a prefix because the block ids are sorted. + """ + rows, num_candidates = candidate_blocks.shape + ratio = candidate_block_size // sparse_block_kv + assert ratio * sparse_block_kv == candidate_block_size and ratio in (1, 2), ( + candidate_block_size, + sparse_block_kv, + ) + assert num_candidates & (num_candidates - 1) == 0, num_candidates + num_sparse = num_candidates * ratio + if out is None: + device = candidate_blocks.device + out = ( + torch.empty(rows, num_sparse, dtype=torch.int32, device=device), + torch.empty(rows, dtype=torch.int32, device=device), + ) + sparse_indices, end = out + assert sparse_indices.shape == (rows, num_sparse) and end.shape == (rows,) + assert candidate_blocks.stride(1) == 1 and sparse_indices.stride(1) == 1 + for t in (candidate_blocks, row_ks, row_ke, sparse_indices, end): + assert t.dtype == torch.int32, t.dtype + + _expand_candidates_kernel[(rows,)]( + candidate_blocks, + candidate_blocks.stride(0), + row_ks, + row_ks.stride(0), + row_ke, + row_ke.stride(0), + sparse_indices, + sparse_indices.stride(0), + end, + end.stride(0), + K=num_candidates, + S=num_sparse, + RATIO=ratio, + SBK=sparse_block_kv, + num_warps=16, + ) + return sparse_indices, end + + +def sparse_topk_remap( + logits: torch.Tensor, + sparse_indices: torch.Tensor, + end: torch.Tensor, + row_ks: torch.Tensor, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + col_indices: torch.Tensor, +) -> None: + """Row top-k over sparse logits, remapped to request-local positions. + + Column ``j * sparse_block_kv + o`` of ``logits`` scores the token at + request-local position + ``(sparse_indices[row, j] - ks // sparse_block_kv) * sparse_block_kv + o``. + + DeepSelect consumes the bf16 logits as they come out of the sparse + kernels; the valid columns form a prefix (block ids are sorted), so the + per-row ``end`` bounds drive it directly. Slots beyond a row's valid + count come back as -1. + + Args: + col_indices: [rows, topk_tokens] int32 scratch for the sparse-column + top-k result (row stride 32B-aligned for DeepSelect). + """ + rows, width = logits.shape + assert logits.dtype == torch.bfloat16 and width >= topk_tokens + assert end.dtype == torch.int32 and end.shape == (rows,) + cols = col_indices[:rows] + deep_select_topk(logits, topk_tokens, end=end, output_idx=cols) + _sparse_topk_remap_kernel[(rows,)]( + cols, + cols.stride(0), + sparse_indices, + sparse_indices.stride(0), + row_ks, + row_ks.stride(0), + topk_indices, + topk_indices.stride(0), + topk_tokens, + SBK=sparse_block_kv, + K_POW2=triton.next_power_of_2(topk_tokens), + ) + + +def sparse_mqa_logits_prefill_chunk( + q: torch.Tensor, + q_scale: torch.Tensor, + k_quant: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + sparse_indices: torch.Tensor, + end: torch.Tensor, + col_indices: torch.Tensor, + kernel_metadata: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse-MQA logits + top-k for one prefill chunk (packed KV workspace). + + Args: + q: [rows, H, D] packed Q values (MXFP4 viewed as int8). + q_scale: [rows, H] int32 packed UE8M0 Q scales. + k_quant: [total_kv, D] packed K workspace (MXFP4 viewed as int8). + k_scale: [total_kv] int32 packed UE8M0 K scales. + weights: [rows, H] bf16 per-head weights; the sparse kernels take + bf16 and do not fold the Q scale in. + cu_seqlen_ks/cu_seqlen_ke: [rows] int32 per-token K bounds in the + packed workspace. + candidate_blocks: [rows, K] int32 request-local candidate block ids. + topk_indices: [rows, topk_tokens] output buffer. + sparse_indices/end/col_indices: Caller-owned scratch, see + `candidate_blocks_to_sparse_indices` and `sparse_topk_remap`. + kernel_metadata: DeepGEMM schedule from a previous call with the same + candidates and bounds (i.e. another indexer layer in the same + step). When given, the candidate expansion is skipped and + ``sparse_indices``/``end`` are assumed to be up to date. + + Returns: + The DeepGEMM schedule metadata used, for reuse by later layers. + """ + assert weights.dtype == torch.bfloat16, weights.dtype + if kernel_metadata is None: + candidate_blocks_to_sparse_indices( + candidate_blocks, + cu_seqlen_ks, + cu_seqlen_ke, + candidate_block_size, + sparse_block_kv, + out=(sparse_indices, end), + ) + kernel_metadata = get_sparse_mqa_logits_metadata( + cu_seqlen_ks, + cu_seqlen_ke, + k_quant.shape[0], + sparse_indices, + q.dtype, + sparse_block_kv, + ) + logits = fp8_fp4_sparse_mqa_logits( + (q, q_scale), + (k_quant, k_scale), + weights, + kernel_metadata, + sparse_indices.shape[1], + sparse_block_kv, + ) + sparse_topk_remap( + logits, + sparse_indices, + end, + cu_seqlen_ks, + sparse_block_kv, + topk_tokens, + topk_indices, + col_indices=col_indices, + ) + return kernel_metadata + + +def sparse_mqa_logits_paged_decode( + q: torch.Tensor, + q_scale: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + row_indices: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + row_ks: torch.Tensor, + sparse_indices: torch.Tensor, + end: torch.Tensor, + col_indices: torch.Tensor, + kernel_metadata: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse-MQA logits + top-k for paged decode rows. + + Every query is one row (spec-decode rows are flattened by the metadata + builder), so ``block_table`` is already per row. + + Args: + q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8). + q_scale: [rows, 1, H] int32 packed UE8M0 Q scales. + kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache + view; the page stride must be 512B-aligned. + weights: [rows, H] bf16 per-head weights. + context_lens: [rows] int32 per-row (compressed) context lengths. + block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``. + row_indices: [rows] int32 row -> request map (pairing only affects + scheduling). + candidate_blocks: [rows, K] int32 candidate block ids. + topk_indices: [rows, topk_tokens] output buffer. + row_ks: [rows] int32 zeros (paged blocks are context-relative). + sparse_indices/end/col_indices: Caller-owned scratch. + kernel_metadata: See `sparse_mqa_logits_prefill_chunk`. + + Returns: + The DeepGEMM schedule metadata used, for reuse by later layers. + """ + rows = q.shape[0] + page_kv = kv_cache.shape[1] + assert weights.dtype == torch.bfloat16, weights.dtype + assert kv_cache.stride(0) % 512 == 0, ( + "paged cache page stride must be 512B-aligned: " + f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}" + ) + assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_kv) + assert block_table.shape[0] == rows and block_table.stride(-1) == 1 + assert context_lens.shape == (rows,) and context_lens.is_contiguous() + + if kernel_metadata is None: + candidate_blocks_to_sparse_indices( + candidate_blocks, + row_ks, + context_lens, + candidate_block_size, + sparse_block_kv, + out=(sparse_indices, end), + ) + kernel_metadata = get_paged_sparse_mqa_logits_metadata( + context_lens, + block_table, + row_indices, + page_kv, + sparse_indices, + q.dtype, + sparse_block_kv, + ) + logits = fp8_fp4_paged_sparse_mqa_logits( + (q, q_scale), + kv_cache, + weights, + kernel_metadata, + sparse_indices.shape[1], + sparse_block_kv, + ) + sparse_topk_remap( + logits, + sparse_indices, + end, + row_ks, + sparse_block_kv, + topk_tokens, + topk_indices, + col_indices=col_indices, + ) + return kernel_metadata diff --git a/vllm/model_executor/layers/sparse_mqa_indexer.py b/vllm/model_executor/layers/sparse_mqa_indexer.py new file mode 100644 index 0000000000..78db1b3df5 --- /dev/null +++ b/vllm/model_executor/layers/sparse_mqa_indexer.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse Attention Indexer that scores only the candidate blocks. + +DeepSeek V4.1 two-level selection: the candidate-source indexer publishes +the top candidate blocks and later indexers pick their top-k inside them. +`SparseAttnIndexer` does that by computing dense logits over the whole +context and masking; this layer instead calls DeepGEMM's sparse MQA-logits +kernels on the candidate blocks only, so the work is O(candidate blocks) +instead of O(context). It requires the `DeepseekV41SparseIndexerBackend` +metadata (see ``AttentionConfig.indexer_sparse_logits``). +""" + +import torch +from torch import nn + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm.forward_context import get_forward_context +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + has_deep_select, + sparse_mqa_logits_paged_decode, + sparse_mqa_logits_prefill_chunk, +) +from vllm.model_executor.layers.sparse_attn_indexer import ( + _gather_workspace_shapes, + kv_cache_as_quant_view, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerPrefillChunkMetadata, +) +from vllm.v1.attention.backends.mla.sparse_indexer import ( + DeepseekV41SparseIndexerMetadata, +) +from vllm.v1.worker.workspace import current_workspace_manager + + +def _prefill_k_workspaces( + total_seq_lens: int, head_dim: int +) -> tuple[torch.Tensor, torch.Tensor]: + """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with + the dense indexer layers of the same model.""" + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True + ) + k_quant, k_scale = current_workspace_manager().get_simultaneous( + values_spec, scales_spec + ) + return k_quant, k_scale + + +def _gather_prefill_chunk_k( + kv_cache: torch.Tensor, + k_quant_full: torch.Tensor, + k_scale_full: torch.Tensor, + chunk: DeepseekV32IndexerPrefillChunkMetadata, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather one prefill chunk's paged K into the packed workspace.""" + assert chunk.local_cu_seq_lens is not None + k_quant = k_quant_full[: chunk.max_local_total_seq_lens] + k_scale = k_scale_full[: chunk.max_local_total_seq_lens] + if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0: + ops.cp_gather_indexer_k_quant_cache( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.local_cu_seq_lens, + ) + return k_quant, k_scale + + +class SparseMQAIndexer(nn.Module): + """Candidate-consuming indexer on DeepGEMM's sparse MQA-logits kernels. + + Only valid for indexer layers that read candidate blocks with the MXFP4 + indexer cache on SM100. The K cache is written by the model before this + runs; ``forward`` takes the same arguments as `SparseAttnIndexer` so the + attention layer can call either. + """ + + weights_dtype = torch.bfloat16 + """Per-head weights dtype the sparse kernels take. The fused Q RoPE-quant + kernel writes it directly so no cast runs per step.""" + + def __init__( + self, + k_cache, + topk_tokens: int, + head_dim: int, + max_total_seq_len: int, + topk_indices_buffer: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + ): + super().__init__() + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and has_deep_gemm_sparse_mqa() + ): + raise ValueError( + "SparseMQAIndexer requires an SM100-class GPU and DeepGEMM >= 2.8." + ) + if not has_deep_select(): + raise ValueError( + "SparseMQAIndexer requires the DeepSelect top-k extension " + "(vllm._deepselect_C)." + ) + self.k_cache = k_cache + self.topk_tokens = topk_tokens + self.head_dim = head_dim + self.max_total_seq_len = max_total_seq_len + self.topk_indices_buffer = topk_indices_buffer + self.candidate_blocks = candidate_blocks + self.candidate_block_size = candidate_block_size + + def _reserve_workspaces(self, device: torch.device) -> None: + """Profiling run: claim the K-gather workspace and the peak sparse + logits allocation so the memory estimate covers them.""" + _prefill_k_workspaces(self.max_total_seq_len, self.head_dim) + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + torch.empty(max_logits_bytes, dtype=torch.uint8, device=device) + + def forward( + self, + hidden_states: torch.Tensor, + q_quant: tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ) -> torch.Tensor: + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + self._reserve_workspaces(hidden_states.device) + return self.topk_indices_buffer + + metadata = attn_metadata[self.k_cache.prefix] + assert isinstance(metadata, DeepseekV41SparseIndexerMetadata), ( + "SparseMQAIndexer needs DeepseekV41SparseIndexerBackend metadata" + ) + assert k is None, "the model writes the indexer K cache" + q_values, q_scale = q_quant + kv_cache = self.k_cache.kv_cache + topk_tokens = self.topk_tokens + topk_indices_buffer = self.topk_indices_buffer + topk_indices_buffer[: hidden_states.shape[0]] = -1 + sparse_block_kv = metadata.sparse_block_kv + + if metadata.num_prefills > 0: + prefill = metadata.prefill + sparse_chunks = metadata.sparse_prefill + assert prefill is not None and sparse_chunks is not None + k_quant_full, k_scale_full = _prefill_k_workspaces( + self.max_total_seq_len, self.head_dim + ) + for chunk, rows in zip(prefill.chunks, sparse_chunks): + k_quant, k_scale = _gather_prefill_chunk_k( + kv_cache, k_quant_full, k_scale_full, chunk + ) + if chunk.local_total_seq_lens == 0: + continue # buffer already holds -1 + start, end = chunk.token_start, chunk.token_end + rows.kernel_metadata = sparse_mqa_logits_prefill_chunk( + q_values[start:end].view(torch.int8), + q_scale[start:end], + k_quant.view(torch.int8), + k_scale.view(torch.int32).squeeze(-1), + weights[start:end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + self.candidate_blocks[start:end], + self.candidate_block_size, + sparse_block_kv, + topk_tokens, + topk_indices_buffer[start:end, :topk_tokens], + sparse_indices=rows.sparse_indices, + end=rows.end, + col_indices=rows.col_indices, + kernel_metadata=rows.kernel_metadata, + ) + + if metadata.num_decodes > 0: + decode_rows = metadata.sparse_decode + assert decode_rows is not None and decode_rows.block_table is not None + assert decode_rows.row_indices is not None + num_rows = decode_rows.row_ke.shape[0] + if num_rows > 0: + decode_rows.kernel_metadata = sparse_mqa_logits_paged_decode( + q_values[:num_rows].view(torch.int8).unsqueeze(1), + q_scale[:num_rows].unsqueeze(1), + kv_cache_as_quant_view(kv_cache, self.head_dim, use_fp4_cache=True), + weights[:num_rows], + decode_rows.row_ke, + decode_rows.block_table, + decode_rows.row_indices, + self.candidate_blocks[:num_rows], + self.candidate_block_size, + sparse_block_kv, + topk_tokens, + topk_indices_buffer[:num_rows, :topk_tokens], + row_ks=decode_rows.row_ks, + sparse_indices=decode_rows.sparse_indices, + end=decode_rows.end, + col_indices=decode_rows.col_indices, + kernel_metadata=decode_rows.kernel_metadata, + ) + return topk_indices_buffer diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index cf6e7235f9..669186edbb 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -309,12 +309,22 @@ class FusedIndexerQRopeQuantTritonKernel( ) +def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]: + """Weights dtypes the model's indexer layers ask for: fp32 for the dense + scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer + (`SparseMQAIndexer`) is enabled.""" + if vllm_config.attention_config.indexer_sparse_logits: + return (torch.float32, torch.bfloat16) + return (torch.float32,) + + class FusedIndexerQRopeMxFp4TritonKernel( VllmTritonJitKernel["FusedIndexerQRopeMxFp4TritonKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: dtype: torch.dtype + weights_out_dtype: torch.dtype num_heads: int index_q_half_rot_dim: int index_q_head_dim: int @@ -434,12 +444,14 @@ class FusedIndexerQRopeMxFp4TritonKernel( self, *, dtype: torch.dtype, + weights_out_dtype: torch.dtype, num_heads: int, head_dim: int, rope_dim: int, ) -> CompileKey: return self.CompileKey( dtype=dtype, + weights_out_dtype=weights_out_dtype, num_heads=num_heads, index_q_half_rot_dim=rope_dim // 2, index_q_head_dim=head_dim, @@ -456,6 +468,7 @@ class FusedIndexerQRopeMxFp4TritonKernel( return self._trace_dispatch(self.dispatch)( dtype=vllm_config.model_config.dtype, + weights_out_dtype=_indexer_weights_out_dtypes(vllm_config), num_heads=num_heads, head_dim=head_dim, rope_dim=rope_dim, @@ -491,7 +504,7 @@ class FusedIndexerQRopeMxFp4TritonKernel( shape=(1, compile_key.num_heads, scale_head_dim), ), index_weights_out=TritonWarmupTensor( - torch.float32, + compile_key.weights_out_dtype, shape=(1, compile_key.num_heads), ), ) @@ -540,12 +553,17 @@ def fused_indexer_q_rope_quant( index_weights_softmax_scale: float, index_weights_head_scale: float, use_fp4: bool = False, + weights_out_dtype: torch.dtype = torch.float32, ) -> tuple[ torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor, ]: """Fused RoPE + quantize Q for the sparse indexer. + ``weights_out_dtype`` is the dtype the downstream scoring kernel takes: + fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse + MQA-logits kernels (CUDA MXFP4 path only). + Weight-fold semantics (important — the two paths differ): FP8 path (use_fp4=False, default): @@ -577,7 +595,7 @@ def fused_indexer_q_rope_quant( num_index_q_heads = index_q.shape[1] index_q_head_dim = index_q.shape[2] - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + index_weights_out = torch.empty_like(index_weights, dtype=weights_out_dtype) if use_fp4: assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( @@ -613,6 +631,7 @@ def fused_indexer_q_rope_quant( weights_out=index_weights_out, ) elif current_platform.is_xpu(): + assert weights_out_dtype == torch.float32, weights_out_dtype torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4( index_q, positions, @@ -647,6 +666,7 @@ def fused_indexer_q_rope_quant( index_q_scale.view(torch.int32).squeeze(-1), ), index_weights_out + assert weights_out_dtype == torch.float32, weights_out_dtype fp8_dtype = current_platform.fp8_dtype() use_fnuz = fp8_dtype == torch.float8_e4m3fnuz fp8_max = 224.0 if use_fnuz else 448.0 diff --git a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py index 8052178a03..f67a77e254 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py @@ -165,6 +165,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] rope_dim: int num_heads: int cos_sin_dtype: type[cutlass.Numeric] + weights_out_dtype: type[cutlass.Numeric] coarsen: int @staticmethod @@ -177,6 +178,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] # process multiple heads at the same time to armotize RoPE load costs coarsen = compile_key.coarsen tb_size = IndexerQMxFp4Kernel.tb_size + weights_out_dtype = compile_key.weights_out_dtype assert num_heads % coarsen == 0 # later we will use 32B load = 16 BF16 elems @@ -287,7 +289,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] weight_head_id = global_tid % num_heads weights_out[weight_token_id, weight_head_id] = ( weights[weight_token_id, weight_head_id].to(Float32) * scale - ) + ).to(weights_out_dtype) @cute.jit def host_entrypoint( @@ -320,11 +322,13 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] self, *, cos_sin_dtype: type[cutlass.Numeric], + weights_out_dtype: type[cutlass.Numeric], **compile_key_fields: int, ) -> CompileKey: return self.CompileKey( **compile_key_fields, cos_sin_dtype=cos_sin_dtype, + weights_out_dtype=weights_out_dtype, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -340,11 +344,17 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0: return [] + # The DeepSeek V4.1 sparse-logits indexer takes bf16 weights; the + # dense scoring kernels take fp32. + weights_out_dtypes: tuple[type[cutlass.Numeric], ...] = (Float32,) + if vllm_config.attention_config.indexer_sparse_logits: + weights_out_dtypes = (Float32, BFloat16) return self._trace_dispatch(self.dispatch)( head_dim=head_dim, rope_dim=rope_dim, num_heads=num_heads, cos_sin_dtype=Float32, + weights_out_dtype=weights_out_dtypes, coarsen=(1, 4), ) @@ -380,7 +390,9 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] divisibility=4, ) weights_out = make_fake_tensor( - Float32, (num_tokens, compile_key.num_heads), divisibility=4 + compile_key.weights_out_dtype, + (num_tokens, compile_key.num_heads), + divisibility=4, ) return ( positions, @@ -413,6 +425,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] rope_dim=cos_sin_cache.shape[-1], num_heads=num_heads, cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype), + weights_out_dtype=torch_to_cute_dtype(weights_out.dtype), coarsen=1 if num_tokens < 512 else 4, ) launch_args = ( diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index fa26351681..134328dbe9 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -4,6 +4,7 @@ DeepseekV4 MLA Attention Layer """ +import math from abc import ABC, abstractmethod from collections.abc import Callable from functools import cached_property @@ -25,6 +26,7 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, @@ -63,6 +65,9 @@ from vllm.v1.attention.backends.mla.indexer import ( dsa_indexer_uses_fp4, get_max_prefill_buffer_size, ) +from vllm.v1.attention.backends.mla.sparse_indexer import ( + DeepseekV41SparseIndexerBackend, +) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache from vllm.v1.kv_cache_interface import ( KVCacheSpec, @@ -997,7 +1002,9 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): self.cache_config = cache_config self.dtype = dtype self.compress_ratio = compress_ratio - compilation_config = get_current_vllm_config().compilation_config + vllm_config = get_current_vllm_config() + self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits + compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self @@ -1010,6 +1017,9 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): # head_dim already carries the fp8 scale padding # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout. uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + page_alignment = ( + 576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512 + ) return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, @@ -1020,16 +1030,21 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): # (the block stride is the sum of every page in the group), so it # has to carry the same padding: the packed record's TMA stride, # else 512B for FlashInfer sparse (#44577). - alignment=( - 576 - if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() - else 512 + alignment=page_alignment, + # DeepGEMM's paged sparse MQA logits address pages by the block + # stride, which packs every layer's page in this layout, and need + # it 512B-aligned; the main KV record read from the same blocks + # needs its own TMA stride, so keep both. + block_stride_alignment=( + math.lcm(512, page_alignment) if self.sparse_logits else None ), ) def forward(self): ... def get_attn_backend(self) -> type[AttentionBackend]: + if self.sparse_logits: + return DeepseekV41SparseIndexerBackend return DeepseekV41IndexerBackend @@ -1129,21 +1144,52 @@ class DeepseekV4Indexer(nn.Module): self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) self.k_cache = k_cache - self.indexer_op = SparseAttnIndexer( - self.k_cache, - self.quant_block_size, - self.scale_fmt, - self.topk_tokens, - self.head_dim, - self.max_model_len, - self.max_total_seq_len, - self.topk_indices_buffer, - skip_k_cache_insert=True, - use_fp4_cache=self.use_fp4_kv, - compress_ratio=self.compress_ratio, - candidate_blocks=candidate_block_buffer, - candidate_block_size=candidate_block_size, - candidate_write=candidate_write, + # Candidate consumers can score only the candidate blocks (opt-in); + # the candidate source and pre-candidate indexers stay dense. + use_sparse_logits = ( + vllm_config.attention_config.indexer_sparse_logits + and candidate_block_buffer is not None + and not candidate_write + ) + self.indexer_op: SparseAttnIndexer | SparseMQAIndexer + if use_sparse_logits: + if not self.use_fp4_kv: + raise ValueError( + "attention_config.indexer_sparse_logits requires the MXFP4 " + "indexer cache (attention_config.indexer_kv_dtype='mxfp4')." + ) + assert candidate_block_buffer is not None + assert topk_indices_buffer is not None + self.indexer_op = SparseMQAIndexer( + self.k_cache, + self.topk_tokens, + self.head_dim, + self.max_total_seq_len, + topk_indices_buffer, + candidate_block_buffer, + candidate_block_size, + ) + else: + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + compress_ratio=self.compress_ratio, + candidate_blocks=candidate_block_buffer, + candidate_block_size=candidate_block_size, + candidate_write=candidate_write, + ) + # The fused Q kernel writes the per-head weights in the dtype the + # scoring kernels take, so no cast runs per step. + self.indexer_weights_dtype = ( + SparseMQAIndexer.weights_dtype if use_sparse_logits else torch.float32 ) def _produce_k( @@ -1231,6 +1277,7 @@ class DeepseekV4Indexer(nn.Module): self.softmax_scale, self.n_head**-0.5, use_fp4=self.use_fp4_kv, + weights_out_dtype=self.indexer_weights_dtype, ) if isinstance(q_quant, tuple): q, q_scale = q_quant diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 5d05c29a84..91e315afb9 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -158,6 +158,10 @@ _grouped_fp4_impl: Callable[..., Any] | None = None _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None +_fp8_fp4_sparse_mqa_logits_impl: Callable[..., Any] | None = None +_fp8_fp4_paged_sparse_mqa_logits_impl: Callable[..., Any] | None = None +_get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None +_get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None _mega_mhc_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None @@ -232,6 +236,9 @@ def _lazy_init() -> None: global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl global _get_paged_mqa_logits_metadata_impl + global _fp8_fp4_sparse_mqa_logits_impl, _fp8_fp4_paged_sparse_mqa_logits_impl + global _get_sparse_mqa_logits_metadata_impl + global _get_paged_sparse_mqa_logits_metadata_impl global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl @@ -295,6 +302,17 @@ def _lazy_init() -> None: _get_paged_mqa_logits_metadata_impl = getattr( _dg, "get_paged_mqa_logits_metadata", None ) + # Sparse-indexer kernels (DeepGEMM >= 2.8, SM100 only). + _fp8_fp4_sparse_mqa_logits_impl = getattr(_dg, "fp8_fp4_sparse_mqa_logits", None) + _fp8_fp4_paged_sparse_mqa_logits_impl = getattr( + _dg, "fp8_fp4_paged_sparse_mqa_logits", None + ) + _get_sparse_mqa_logits_metadata_impl = getattr( + _dg, "get_sparse_mqa_logits_metadata", None + ) + _get_paged_sparse_mqa_logits_metadata_impl = getattr( + _dg, "get_paged_sparse_mqa_logits_metadata", None + ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) _mega_mhc_impl = getattr(_dg, "mega_mhc", None) _get_mn_major_tma_aligned_tensor_impl = getattr( @@ -681,6 +699,171 @@ def fp8_fp4_paged_mqa_logits( ) +def has_deep_gemm_sparse_mqa() -> bool: + """Whether the installed DeepGEMM provides the sparse-indexer kernels + (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only).""" + _lazy_init() + return ( + _fp8_fp4_sparse_mqa_logits_impl is not None + and _fp8_fp4_paged_sparse_mqa_logits_impl is not None + and _get_sparse_mqa_logits_metadata_impl is not None + and _get_paged_sparse_mqa_logits_metadata_impl is not None + ) + + +def get_sparse_mqa_logits_metadata( + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + num_kv_tokens: int, + sparse_kv_block_indices: torch.Tensor, + qk_dtype: torch.dtype, + sparse_block_kv: int, +) -> torch.Tensor: + """Build scheduling metadata for `fp8_fp4_sparse_mqa_logits`. + + Always uses the unaligned-ks variant, which degenerates to the aligned + variant when ``cu_seqlen_ks % sparse_block_kv == 0``, so callers never + need a host-side sync to pick one. + + Args: + cu_seqlen_ks/cu_seqlen_ke: Per-row K range bounds in the packed KV + workspace, shape [num_q_tokens], dtype int32. + num_kv_tokens: Total KV tokens in the packed workspace. + sparse_kv_block_indices: Per-row candidate block ids, shape + [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's + valid prefix must be sorted ascending (repeats are tolerated but + wasteful); pad by repeating the last valid block. Block ``i`` + covers tokens + ``[i * sparse_block_kv + ks % sparse_block_kv, ...)``. + qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4, + ``torch.float8_e4m3fn`` for FP8). + sparse_block_kv: Tokens per sparse block, 8 or 16. + """ + _lazy_init() + if _get_sparse_mqa_logits_metadata_impl is None: + return _missing() + return _get_sparse_mqa_logits_metadata_impl( + cu_seqlen_ks, + cu_seqlen_ke, + num_kv_tokens, + sparse_kv_block_indices, + qk_dtype, + sparse_block_kv, + use_unaligned_ks=True, + ) + + +def get_paged_sparse_mqa_logits_metadata( + context_lens: torch.Tensor, + block_table: torch.Tensor, + indices: torch.Tensor, + page_kv: int, + sparse_kv_block_indices: torch.Tensor, + qk_dtype: torch.dtype, + sparse_block_kv: int, +) -> torch.Tensor: + """Build scheduling metadata for `fp8_fp4_paged_sparse_mqa_logits`. + + Args: + context_lens: Per-row context length, shape [num_q_tokens], int32. + block_table: Per-row page table, shape [num_q_tokens, max_pages], + int32; rows of queries paired with the same request must be + identical (expand it per query row before calling). + indices: Request index for each query row, shape [num_q_tokens]. + page_kv: Tokens per KV page; must be a multiple of sparse_block_kv. + sparse_kv_block_indices: Per-row logical (context-relative) candidate + block ids, [num_q_tokens, num_max_sparse_blocks], int32. + qk_dtype: dtype of the packed Q values. + sparse_block_kv: Tokens per sparse block, 8 or 16. + """ + _lazy_init() + if _get_paged_sparse_mqa_logits_metadata_impl is None: + return _missing() + return _get_paged_sparse_mqa_logits_metadata_impl( + context_lens, + block_table, + indices, + page_kv, + sparse_kv_block_indices, + qk_dtype, + sparse_block_kv, + ) + + +def fp8_fp4_sparse_mqa_logits( + q: tuple[torch.Tensor, torch.Tensor], + kv: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + metadata: torch.Tensor, + num_max_sparse_blocks: int, + sparse_block_kv: int, +) -> torch.Tensor: + """Compute MQA logits only at the candidate blocks (prefill path). + + Args: + q: ``(q_values, q_scale)``. q_values is [M, H, D] (packed FP4 viewed + as int8, or FP8); q_scale is the packed UE8M0 scale tensor [M, H] + int32 (mandatory, unlike the dense kernels). + kv: ``(kv_values, kv_scale)`` — the packed KV workspace with UE8M0 + int32 scales. + weights: [M, H] ``torch.bfloat16`` (the Q scale is NOT folded in). + metadata: From `get_sparse_mqa_logits_metadata`. + num_max_sparse_blocks: Width of each sparse-index row. + sparse_block_kv: Tokens per sparse block, 8 or 16. + + Returns: + bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv]; + column ``j * sparse_block_kv + o`` scores the token at + ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``. + """ + _lazy_init() + if _fp8_fp4_sparse_mqa_logits_impl is None: + return _missing() + return _fp8_fp4_sparse_mqa_logits_impl( + q, + kv, + weights, + metadata, + num_max_sparse_blocks, + sparse_block_kv, + use_unaligned_ks=True, + ) + + +def fp8_fp4_paged_sparse_mqa_logits( + q: tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + weights: torch.Tensor, + metadata: torch.Tensor, + num_max_sparse_blocks: int, + sparse_block_kv: int, +) -> torch.Tensor: + """Compute MQA logits only at the candidate blocks (paged decode path). + + Args: + q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D]. + kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8, + page stride 512B-aligned. + weights: [num_q_tokens, H] ``torch.bfloat16``. + metadata: From `get_paged_sparse_mqa_logits_metadata`. + + Returns: + bf16 logits of shape + [num_q_tokens, num_max_sparse_blocks * sparse_block_kv]. + """ + _lazy_init() + if _fp8_fp4_paged_sparse_mqa_logits_impl is None: + return _missing() + return _fp8_fp4_paged_sparse_mqa_logits_impl( + q, + kv_cache, + weights, + metadata, + num_max_sparse_blocks, + sparse_block_kv, + ) + + def tf32_hc_prenorm_gemm( x: torch.Tensor, fn: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 837dbca817..53206afd93 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1116,6 +1116,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): for request_slice, query_slice in chunk_specs ] + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: + """Per-request KV lengths the prefill chunker budgets logits with; + subclasses whose logits rows are wider than the context override.""" + return seq_lens_cpu + @staticmethod def _split_indexer_prefill_chunks( compressed_seq_lens_cpu: torch.Tensor, @@ -1278,7 +1283,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): ) else: chunk_specs = self._split_indexer_prefill_chunks( - compressed_seq_lens_cpu[num_decodes:], + self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]), prefill_query_lens_cpu, self.max_prefill_buffer_size, max_logits_bytes, diff --git a/vllm/v1/attention/backends/mla/sparse_indexer.py b/vllm/v1/attention/backends/mla/sparse_indexer.py new file mode 100644 index 0000000000..c0b445534d --- /dev/null +++ b/vllm/v1/attention/backends/mla/sparse_indexer.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4.1 indexer backend for the DeepGEMM sparse MQA-logits path. + +Candidate-consuming indexer layers score only the candidate blocks published +by the candidate-source layer (``fp8_fp4_(paged_)sparse_mqa_logits``) instead +of dense logits over the whole context. This backend extends the dense +indexer metadata with the per-row scratch that path needs, so the layer op +allocates nothing per step, and lets all consumer layers of a step share one +candidate expansion and DeepGEMM schedule (they all read the same candidate +blocks). Opt in with ``AttentionConfig.indexer_sparse_logits``. +""" + +from dataclasses import dataclass, fields + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + check_deep_select_layout, + has_deep_select, + pick_sparse_block_kv, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + DeepseekV32IndexerMetadataBuilder, + DeepseekV41IndexerBackend, + dsa_indexer_uses_fp4, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +logger = init_logger(__name__) + + +@dataclass +class SparseMQARowsMetadata: + """Per-row state for one sparse-logits call: the decode rows, or one + prefill chunk. Buffers are builder-owned views sized to the rows. + + ``kernel_metadata`` (the DeepGEMM schedule) is filled by the first + consumer indexer layer of the step together with ``sparse_indices`` / + ``end``; later consumer layers reuse all three. + """ + + row_ks: torch.Tensor + """[rows] int32 K-range start per row; zeros for paged decode.""" + row_ke: torch.Tensor + """[rows] int32 K-range end (prefill) or compressed context length.""" + sparse_indices: torch.Tensor + """[rows, num_sparse_blocks] int32 DeepGEMM sparse block ids.""" + end: torch.Tensor + """[rows] int32 valid sparse-column count per row.""" + col_indices: torch.Tensor + """[rows, topk] int32 scratch for the sparse-column top-k.""" + block_table: torch.Tensor | None = None + """Decode only: [rows, pages] per-row page table.""" + row_indices: torch.Tensor | None = None + """Decode only: [rows] row -> request index.""" + kernel_metadata: torch.Tensor | None = None + + +@dataclass +class DeepseekV41SparseIndexerMetadata(DeepseekV32IndexerMetadata): + sparse_block_kv: int = 0 + sparse_decode: SparseMQARowsMetadata | None = None + sparse_prefill: list[SparseMQARowsMetadata] | None = None + """Parallel to ``prefill.chunks``.""" + + +class DeepseekV41SparseIndexerBackend(DeepseekV41IndexerBackend): + @staticmethod + def get_name() -> str: + return "DEEPSEEK_V41_SPARSE_INDEXER" + + @staticmethod + def get_builder_cls() -> type["DeepseekV41SparseIndexerMetadataBuilder"]: + return DeepseekV41SparseIndexerMetadataBuilder + + +class DeepseekV41SparseIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder): + """Dense indexer metadata plus the sparse-logits row scratch. + + Every requirement of the sparse kernels is checked here, at engine start, + rather than falling back to the dense path per step. + """ + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + **kwargs, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs) + hf_config = vllm_config.model_config.hf_text_config + candidate_block_size = getattr(hf_config, "candidate_block_size", 0) + candidate_topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0) + topk_tokens = getattr(hf_config, "index_topk", 0) + prefix = "attention_config.indexer_sparse_logits" + if candidate_block_size <= 0 or candidate_topk_blocks <= 0: + raise ValueError( + f"{prefix} requires a two-level (candidate block) indexer model " + "such as DeepSeek-V4.1." + ) + if not dsa_indexer_uses_fp4(vllm_config): + raise ValueError( + f"{prefix} requires the MXFP4 indexer cache " + "(attention_config.indexer_kv_dtype='mxfp4'): the sparse kernels " + "take UE8M0-packed granularity-32 scales." + ) + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ): + raise ValueError(f"{prefix} requires an SM100-class GPU.") + if not has_deep_gemm_sparse_mqa(): + raise ValueError( + f"{prefix} requires DeepGEMM >= 2.8 " + "(fp8_fp4_sparse_mqa_logits not found)." + ) + if self.dcp_world_size > 1 or self.use_pcp: + raise NotImplementedError( + f"{prefix} is not supported with decode/prefill context parallel." + ) + if not self.supports_varlen: + raise ValueError( + f"{prefix} requires the varlen paged MQA logits decode layout." + ) + # The top-k runs on the kernels' bf16 logits with DeepSelect; there is + # no fp32 fallback, so the other top-k backends cannot be honored. + topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend + if not has_deep_select(): + raise ValueError( + f"{prefix} requires the DeepSelect top-k extension " + "(vllm._deepselect_C) for this GPU." + ) + if topk_backend not in ("auto", "deep_select"): + raise ValueError( + f"{prefix} uses DeepSelect for the top-k; " + f"kernel_config.sparse_indexer_topk_backend='{topk_backend}' " + "is not supported with it." + ) + self.sparse_block_kv = pick_sparse_block_kv(candidate_block_size) + if kv_cache_spec.block_size % self.sparse_block_kv != 0: + raise ValueError( + f"{prefix}: indexer cache pages ({kv_cache_spec.block_size} " + f"tokens) must be a multiple of {self.sparse_block_kv} tokens." + ) + self.num_sparse_blocks = candidate_topk_blocks * ( + candidate_block_size // self.sparse_block_kv + ) + num_sparse_cols = self.num_sparse_blocks * self.sparse_block_kv + check_deep_select_layout(num_sparse_cols, topk_tokens) + # Sparse logits are bf16 [rows, num_sparse_cols]; make the prefill + # chunker (which budgets fp32 [rows, seq_len]) size chunks as if every + # row were this wide. + self.min_split_seq_len = (num_sparse_cols * 2 + 3) // 4 + + # Rows are indexed by batch token (decode rows first, then prefill + # chunks by token range), like ``arange_buffer``. + max_rows = self.arange_buffer.shape[0] + int32 = dict(dtype=torch.int32, device=device) + self.sparse_indices_buffer = torch.zeros( + (max_rows, self.num_sparse_blocks), **int32 + ) + self.end_buffer = torch.zeros(max_rows, **int32) + self.col_indices_buffer = torch.zeros((max_rows, topk_tokens), **int32) + self.zero_row_ks = torch.zeros(max_rows, **int32) + logger.info_once( + "DeepSeek V4.1 indexer: candidate consumers score only the candidate " + "blocks with DeepGEMM sparse MQA logits (%d sparse blocks of %d " + "tokens per row).", + self.num_sparse_blocks, + self.sparse_block_kv, + ) + + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: + return seq_lens_cpu.clamp(min=self.min_split_seq_len) + + def _rows(self, start: int, end: int, row_ks, row_ke, **kwargs): + return SparseMQARowsMetadata( + row_ks=row_ks, + row_ke=row_ke, + sparse_indices=self.sparse_indices_buffer[start:end], + end=self.end_buffer[start:end], + col_indices=self.col_indices_buffer[start:end], + **kwargs, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> DeepseekV41SparseIndexerMetadata: + base = super().build(common_prefix_len, common_attn_metadata, fast_build) + metadata = DeepseekV41SparseIndexerMetadata( + **{f.name: getattr(base, f.name) for f in fields(base)}, + sparse_block_kv=self.sparse_block_kv, + ) + if base.decode is not None: + decode = base.decode + rows = decode.seq_lens.shape[0] + # The varlen layout flattens every query to one row, so the base + # decode tensors are already per row. + assert decode.seq_lens.shape[1] == 1 and not decode.requires_padding + assert decode.block_table.shape[0] == rows + metadata.sparse_decode = self._rows( + 0, + rows, + row_ks=self.zero_row_ks[:rows], + row_ke=decode.seq_lens.view(-1), + block_table=decode.block_table, + row_indices=( + decode.indices + if decode.indices is not None + else self.arange_buffer[:rows] + ), + ) + if base.prefill is not None: + metadata.sparse_prefill = [ + self._rows( + chunk.token_start, + chunk.token_end, + row_ks=chunk.cu_seqlen_ks, + row_ke=chunk.cu_seqlen_ke, + ) + for chunk in base.prefill.chunks + ] + return metadata diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 28478b69f5..1c6e102c16 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1608,6 +1608,15 @@ def _get_kv_cache_bytes_per_block( ] if hot_page_sizes: bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes)) + stride_alignments = [ + spec.block_stride_alignment + for group in kv_cache_groups + for layer_name in group.layer_names + if isinstance(spec := _get_per_layer_spec(group, layer_name), MLAAttentionSpec) + and spec.block_stride_alignment + ] + if stride_alignments: + bytes_per_block = round_up(bytes_per_block, math.lcm(*stride_alignments)) return bytes_per_block diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 8cd867eb9c..7e91300b2f 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -653,6 +653,12 @@ class MLAAttentionSpec(FullAttentionSpec): is_index_group_leader: bool = False storage_block_size: int | None = None """Token width used to view storage when it differs from the kernel block.""" + block_stride_alignment: int | None = None + """Required alignment, in bytes, of the distance between consecutive + blocks of this cache. In block-major layouts that distance is the whole + block (all layers' pages), so the allocator rounds the block up to it. + DeepGEMM's paged sparse MQA-logits kernels address pages as + ``base + page * stride`` and need it 512B-aligned.""" # Group capability enabled when any member flattens a non-causal query block # into decode rows. Runtime metadata still selects causal vs. non-causal mode. non_causal_multi_token_decode: bool = False @@ -674,6 +680,7 @@ class MLAAttentionSpec(FullAttentionSpec): cache_role_set = {spec.cache_role for spec in specs} index_group_leader_set = {spec.is_index_group_leader for spec in specs} storage_block_size_set = set(spec.storage_block_size for spec in specs) + block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 @@ -681,10 +688,11 @@ class MLAAttentionSpec(FullAttentionSpec): and len(cache_role_set) == 1 and len(index_group_leader_set) == 1 and len(storage_block_size_set) == 1 + and len(block_stride_alignment_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, cache role, " - "index-sharing role, and storage block size." + "index-sharing role, storage block size and block stride alignment." ) merged_spec = cls( block_size=specs[0].block_size, @@ -701,6 +709,7 @@ class MLAAttentionSpec(FullAttentionSpec): cache_role=cache_role_set.pop(), is_index_group_leader=index_group_leader_set.pop(), storage_block_size=storage_block_size_set.pop(), + block_stride_alignment=block_stride_alignment_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs ), From c6fa1f05d1bc627171c3a58c05f23a1c41202ac1 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Tue, 15 Sep 2026 17:15:59 -0400 Subject: [PATCH 119/139] [Perf][Kernel] Add sampled filtering for persistent top-k (#56346) Signed-off-by: mgoin Co-authored-by: OpenAI Codex --- .../kernels/benchmark_persistent_topk.py | 134 ++++++++++ csrc/libtorch_stable/persistent_topk.cuh | 86 +++--- csrc/libtorch_stable/sampled_topk.cuh | 253 ++++++++++++++++++ csrc/libtorch_stable/topk.cu | 18 +- tests/kernels/test_top_k_per_row.py | 100 +++++++ 5 files changed, 556 insertions(+), 35 deletions(-) create mode 100644 benchmarks/kernels/benchmark_persistent_topk.py create mode 100644 csrc/libtorch_stable/sampled_topk.cuh diff --git a/benchmarks/kernels/benchmark_persistent_topk.py b/benchmarks/kernels/benchmark_persistent_topk.py new file mode 100644 index 0000000000..a77a04cfd6 --- /dev/null +++ b/benchmarks/kernels/benchmark_persistent_topk.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark sparse-indexer decode top-k on variable-length FP32 rows. + +Run on a reserved GPU, for example: + chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py + +Use --capacity to emulate graphs captured at a larger maximum context length. +Run the same command on the base and candidate builds to compare JSON results. +""" + +import argparse +import itertools +import json +import statistics +from pathlib import Path + +import torch +from flashinfer.testing import bench_gpu_time_with_cupti + +import vllm._custom_ops # noqa: F401 +from vllm.platforms import current_platform + + +def benchmark(rows: int, width: int, k: int, args) -> dict: + capacity = args.capacity or width + if capacity < width: + raise ValueError("Capacity must be at least the valid-length upper bound") + torch.manual_seed(args.seed) + logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32) + if args.full_length: + lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32) + else: + lengths = torch.randint( + int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32 + ) + invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None] + logits.masked_fill_(invalid, float("nan")) + out = torch.empty((rows, k), device="cuda", dtype=torch.int32) + workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8) + cooperative = ( + args.backend == "auto" + and rows <= 64 + and capacity % 4 == 0 + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk + + def run(): + op(logits, lengths, out, workspace, k, capacity) + + run() + valid = torch.arange(k, device="cuda")[None] < lengths[:, None] + assert torch.all(out[~valid] == -1) + assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid) + sorted_indices = out.sort(dim=1).values + assert torch.all( + (sorted_indices[:, 1:] != sorted_indices[:, :-1]) + | (sorted_indices[:, 1:] == -1) + ) + selected = logits.gather(1, out.clamp_min(0).long()).masked_fill( + ~valid, -float("inf") + ) + reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values + torch.testing.assert_close( + selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0 + ) + trials = [ + statistics.median( + bench_gpu_time_with_cupti( + run, + use_cuda_graph=True, + cold_l2_cache=not args.warm, + dry_run_iters=5, + repeat_iters=args.repeat, + ) + ) + * 1000 + for _ in range(args.trials) + ] + us = statistics.median(trials) + return dict( + rows=rows, + width=width, + capacity=capacity, + k=k, + backend="cooperative" if cooperative else "persistent", + us=us, + trials_us=trials, + # Minimum useful traffic: one valid-score read and one index write. + effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024]) + parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840]) + parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048]) + parser.add_argument("--capacity", type=int) + parser.add_argument( + "--full-length", action="store_true", help="Use exactly --widths valid scores" + ) + parser.add_argument("--backend", choices=["auto", "persistent"], default="auto") + parser.add_argument( + "--warm", action="store_true", help="Keep L2 warm between replays" + ) + parser.add_argument("--repeat", type=int, default=60) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + metadata = dict( + gpu=str(torch.cuda.get_device_properties(0)), + torch=torch.__version__, + cuda=torch.version.cuda, + args=vars(args), + ) + print(json.dumps(metadata, default=str)) + results = [] + for rows, width, k in itertools.product(args.rows, args.widths, args.top_k): + result = benchmark(rows, width, k, args) + results.append(result) + print(json.dumps(result), flush=True) + if args.output: + args.output.write_text( + json.dumps(dict(metadata=metadata, results=results), indent=2, default=str) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 9ae0fd4bcb..bf5b46ab18 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -1030,41 +1030,34 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB -/*! - * \brief Filtered Top-K kernel for ragged sequences. - * - * \tparam DType Data type (float, half, nv_bfloat16) - * \tparam IdType Index type (int32_t) - * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) - */ -template -__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) - FilteredTopKUnifiedKernel(const DType* __restrict__ input, - IdType* __restrict__ output, - const IdType* __restrict__ lengths, - uint32_t num_rows, uint32_t top_k, - uint32_t max_len) { +template +struct FilteredTopKStorage { + alignas(128) int histogram[2][256 + 128]; + alignas(128) int counter; + alignas(128) int threshold_bin; + alignas(128) int num_input[2]; + alignas(128) int indices[MAX_K]; + int last_remain; +}; + +// With CheckOverflow, return false before using a truncated stash so the caller +// can retry with exact full-row selection. +template +__device__ __forceinline__ bool filtered_topk_row( + const DType* score, IdType* dst, int length, uint32_t top_k, + FilteredTopKStorage& storage) { constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; constexpr int RADIX = 256; constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; - - const uint32_t bid = blockIdx.x; const int tx = threadIdx.x; - if (bid >= num_rows) return; - - const int length = - (lengths != nullptr) ? lengths[bid] : static_cast(max_len); - const DType* score = input + bid * max_len; - IdType* dst = output + bid * top_k; - // Trivial case: length <= top_k if (length <= static_cast(top_k)) { for (int i = tx; i < static_cast(top_k); i += BLOCK_SIZE) { dst[i] = (i < length) ? static_cast(i) : static_cast(-1); } - return; + return true; } // Short path @@ -1077,16 +1070,14 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) hist4096::histogram_4096_topk(score, dst, length, _smem_reg); } - return; + return true; } - // Static shared memory - alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128]; - alignas(128) __shared__ int s_counter; - alignas(128) __shared__ int s_threshold_bin_id; - alignas(128) __shared__ int s_num_input[2]; - alignas(128) __shared__ int s_indices[MAX_K]; - + auto& s_histogram_buf = storage.histogram; + auto& s_counter = storage.counter; + auto& s_threshold_bin_id = storage.threshold_bin; + auto& s_num_input = storage.num_input; + auto& s_indices = storage.indices; auto& s_histogram = s_histogram_buf[0]; // Dynamic shared memory for input double buffer @@ -1213,10 +1204,13 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) // Stage 2: refine with 8bit radix passes #pragma unroll for (int round = 0; round < NUM_ROUNDS; ++round) { - __shared__ int s_last_remain; + auto& s_last_remain = storage.last_remain; const auto r_idx = round % 2; const auto _raw_num_input = s_num_input[r_idx]; + if constexpr (CheckOverflow) { + if (_raw_num_input > SMEM_INPUT_SIZE) return false; + } const auto num_input = (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; @@ -1284,6 +1278,30 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const int idx = s_indices[base]; dst[base] = static_cast(idx); } + return true; +} + +/*! + * \brief Filtered Top-K kernel for ragged sequences. + * + * \tparam DType Data type (float, half, nv_bfloat16) + * \tparam IdType Index type (int32_t) + * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) + */ +template +__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) + FilteredTopKUnifiedKernel(const DType* __restrict__ input, + IdType* __restrict__ output, + const IdType* __restrict__ lengths, + uint32_t num_rows, uint32_t top_k, + uint32_t max_len) { + const uint32_t bid = blockIdx.x; + if (bid >= num_rows) return; + const int length = lengths ? lengths[bid] : static_cast(max_len); + __shared__ FilteredTopKStorage storage; + filtered_topk_row( + input + bid * max_len, output + bid * top_k, length, top_k, storage); } // Helper to compute GCD for VEC_SIZE selection diff --git a/csrc/libtorch_stable/sampled_topk.cuh b/csrc/libtorch_stable/sampled_topk.cuh new file mode 100644 index 0000000000..867abbd635 --- /dev/null +++ b/csrc/libtorch_stable/sampled_topk.cuh @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// Sampled top-k for long FP32 sparse-indexer rows, inspired by DeepSelect: +// https://github.com/deepseek-ai/DeepSelect +// A sample estimates a cutoff, then survivors are compacted in shared memory +// for exact FP32 selection. Too few or too many candidates trigger an exact +// full-row fallback. + +#pragma once + +#include +#include +#include + +#include "persistent_topk.cuh" + +namespace vllm::sampled_topk { + +constexpr int kThreads = 1024; +// Conservative crossover bounds from the B300 batch/length sweep. +template +constexpr int kMinSampledLength = K == 512 ? 98304 : 65536; +// Target half the buffer to leave room for sampling error. +constexpr int kSample = 4096; +constexpr int kCapacity = 8192; + +struct Storage { + uint2 candidates[kCapacity]; + int histogram[2048]; + typename cub::BlockScan::TempStorage scan; + int count; + int remaining; + int bin; + uint32_t prefix; +}; + +__device__ __forceinline__ uint32_t ordered(float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +__device__ __forceinline__ uint32_t coarse(uint32_t key) { + const uint32_t bits = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key; + return topk_histogram_4096::extract_coarse_bin_N<11>(__uint_as_float(bits)); +} + +// Select the bin containing the rank and update the rank within that bin. +template +__device__ void select_bin(Storage& s) { + constexpr int kItems = (Bins + kThreads - 1) / kThreads; + int sums[kItems]; +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int bin = Bins - 1 - (threadIdx.x * kItems + j); + sums[j] = bin >= 0 ? s.histogram[bin] : 0; + } + cub::BlockScan(s.scan).InclusiveSum(sums, sums); + const int remaining = s.remaining; + __syncthreads(); +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int bin = Bins - 1 - (threadIdx.x * kItems + j); + if (bin >= 0 && sums[j] >= remaining && + sums[j] - s.histogram[bin] < remaining) { + s.bin = bin; + s.remaining = remaining - (sums[j] - s.histogram[bin]); + } + } + __syncthreads(); +} + +// Find a coarse bin, then refine its full FP32 keys one byte at a time. +// Read shared candidates normally, or the full row for exact fallback. +template +__device__ uint32_t threshold(Storage& s, const float* row, int count, + int rank) { + const int tid = threadIdx.x; + if (tid == 0) { + s.remaining = rank; + s.prefix = 0; + } + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0; + __syncthreads(); + for (int i = tid; i < count; i += kThreads) { + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]); + atomicAdd(&s.histogram[coarse(key)], 1); + } + __syncthreads(); + select_bin<2048>(s); + const int coarse_bin = s.bin; + uint32_t mask = 0; + for (int shift = 24; shift >= 0; shift -= 8) { + if (tid < 256) s.histogram[tid] = 0; + __syncthreads(); + const uint32_t prefix = s.prefix; + for (int i = tid; i < count; i += kThreads) { + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]); + if (coarse(key) == coarse_bin && (key & mask) == prefix) { + atomicAdd(&s.histogram[(key >> shift) & 255], 1); + } + } + __syncthreads(); + select_bin<256>(s); + if (tid == 0) s.prefix = prefix | (static_cast(s.bin) << shift); + mask |= 255u << shift; + __syncthreads(); + } + return s.prefix; +} + +__device__ __forceinline__ int reserve(bool hit, int* counter) { + const uint32_t mask = __ballot_sync(0xffffffffu, hit); + if (mask == 0) return 0; + const int lane = threadIdx.x & 31; + int base = 0; + if (lane == 0) base = atomicAdd(counter, __popc(mask)); + base = __shfl_sync(0xffffffffu, base, 0); + return base + __popc(mask & ((1u << lane) - 1)); +} + +template +__device__ void emit(Storage& s, const float* row, int32_t* dst, int count, + uint32_t cutoff) { + const int tid = threadIdx.x; + if (tid == 0) s.count = 0; + __syncthreads(); + for (int base = 0; base < count; base += kThreads) { + const int i = base + tid; + const uint2 pair = i < count ? (Buffered ? s.candidates[i] + : make_uint2(i, ordered(row[i]))) + : make_uint2(0, 0); + bool keep = i < count && pair.y > cutoff; + if (i < count && pair.y == cutoff) { + keep = atomicSub(&s.remaining, 1) > 0; + } + const int offset = reserve(keep, &s.count); + if (keep) dst[offset] = pair.x; + } +} + +template +__global__ void __launch_bounds__(kThreads) + sampled_topk_kernel(const float* __restrict__ input, + const int32_t* __restrict__ lengths, + int32_t* __restrict__ output, int64_t stride, + int max_length) { + extern __shared__ __align__(16) unsigned char smem[]; + auto& s = *reinterpret_cast(smem); + const int tid = threadIdx.x; + const int length = max(0, min(lengths[blockIdx.x], max_length)); + const float* row = input + blockIdx.x * stride; + int32_t* dst = output + blockIdx.x * K; + // Graph allocations can hold rows much shorter than the host length bound. + if (length < kMinSampledLength) { + __shared__ filtered_topk::FilteredTopKStorage fallback_storage; + bool selected; + if ((stride & 3) == 0 && (reinterpret_cast(row) & 15) == 0) { + selected = + filtered_topk::filtered_topk_row( + row, dst, length, K, fallback_storage); + } else { + selected = + filtered_topk::filtered_topk_row( + row, dst, length, K, fallback_storage); + } + if (!selected) { + const uint32_t exact_cutoff = threshold(s, row, length, K); + emit(s, row, dst, length, exact_cutoff); + } + return; + } + if (tid == 0) s.remaining = max(1, (kCapacity / 2) * kSample / length); + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0; + __syncthreads(); + // Sample contiguous warps spread across the row to keep loads coalesced. + for (int i = tid; i < kSample; i += kThreads) { + const int index = (i / 32) * (length / (kSample / 32)) + i % 32; + atomicAdd(&s.histogram[coarse(ordered(row[index]))], 1); + } + __syncthreads(); + select_bin<2048>(s); + const uint16_t half_key = s.bin << 5; + const uint16_t half_bits = + (half_key & 0x8000u) ? (half_key ^ 0x8000u) : ~half_key; + const uint32_t cutoff = ordered(__half2float(__ushort_as_half(half_bits))); + if (tid == 0) s.count = 0; + __syncthreads(); + constexpr int kItems = 16; + // Compact survivors with one reservation per warp per tile. Keep counting + // beyond capacity so overflow triggers fallback instead of using a subset. + for (int base = 0; base < length; base += kItems * kThreads) { + uint32_t keys[kItems]; + uint32_t hit_mask = 0; + if ((reinterpret_cast(row) & 15) == 0 && + base + kItems * kThreads <= length) { +#pragma unroll + for (int j = 0; j < kItems / 4; ++j) { + const float4 values = + reinterpret_cast(row + base)[tid + j * kThreads]; + keys[4 * j] = ordered(values.x); + keys[4 * j + 1] = ordered(values.y); + keys[4 * j + 2] = ordered(values.z); + keys[4 * j + 3] = ordered(values.w); + } + } else { +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4; + keys[j] = i < length ? ordered(row[i]) : 0; + } + } +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4; + if (i < length && keys[j] >= cutoff) hit_mask |= 1u << j; + } + const int lane = tid & 31; + const int hits = __popc(hit_mask); + int inclusive = hits; +#pragma unroll + for (int delta = 1; delta < 32; delta *= 2) { + const int previous = __shfl_up_sync(0xffffffffu, inclusive, delta); + if (lane >= delta) inclusive += previous; + } + const int total = __shfl_sync(0xffffffffu, inclusive, 31); + int warp_base = 0; + if (lane == 31 && total != 0) warp_base = atomicAdd(&s.count, total); + warp_base = __shfl_sync(0xffffffffu, warp_base, 31); + int offset = warp_base + inclusive - hits; +#pragma unroll + for (int j = 0; j < kItems; ++j) { + if (hit_mask & (1u << j)) { + if (offset < kCapacity) { + s.candidates[offset] = make_uint2( + base + tid * 4 + (j / 4) * kThreads * 4 + j % 4, keys[j]); + } + ++offset; + } + } + } + __syncthreads(); + const int count = s.count; + if (count >= K && count <= kCapacity) { + const uint32_t exact_cutoff = threshold(s, row, count, K); + emit(s, row, dst, count, exact_cutoff); + } else { + const uint32_t exact_cutoff = threshold(s, row, length, K); + emit(s, row, dst, length, exact_cutoff); + } +} + +} // namespace vllm::sampled_topk diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index b3e1a8623f..fb3d946b96 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -9,6 +9,7 @@ #ifndef USE_ROCM #include "persistent_topk.cuh" + #include "sampled_topk.cuh" #endif namespace { @@ -36,7 +37,22 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, max_smem_per_block = device_prop->sharedMemPerBlockOptin; } - if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { + // Allow static fallback storage in addition to the 128 KiB dynamic buffer. + if (num_rows > 64 && + max_seq_len >= vllm::sampled_topk::kMinSampledLength && + max_smem_per_block >= 144 * 1024) { + auto kernel = vllm::sampled_topk::sampled_topk_kernel; + constexpr size_t smem_size = + vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC; + cudaError_t status = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + STD_TORCH_CHECK(status == cudaSuccess, + "sampled_topk smem failed: ", cudaGetErrorString(status)); + kernel<<>>( + logits.const_data_ptr(), lengths.const_data_ptr(), + output.mutable_data_ptr(), stride, + static_cast(std::min(max_seq_len, logits.size(1)))); + } else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { cudaError_t status = vllm::FilteredTopKRaggedTransform( logits.const_data_ptr(), output.mutable_data_ptr(), diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index 591362020e..2895394706 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -932,6 +932,106 @@ def test_deepseek_workspace_topk( ) +@pytest.mark.skipif(not _has_device_capability(80), reason="Requires SM80+") +@pytest.mark.parametrize("rows", [65, 128]) +@pytest.mark.parametrize("top_k", [512, 1024, 2048]) +@pytest.mark.parametrize( + "distribution", ["random", "10LSBits", "ascending", "constant", "sampled_peaks"] +) +@torch.inference_mode() +def test_persistent_topk_sampled_graph( + rows: int, top_k: int, distribution: str +) -> None: + """Sampling and both fallbacks preserve exact values across graph replays.""" + if torch.cuda.get_device_properties(0).shared_memory_per_block_optin < 144 * 1024: + pytest.skip("Sampled top-k requires at least 144 KiB of shared memory") + set_random_seed(42) + min_sampled_length = 98304 if top_k == 512 else 65536 + width = min_sampled_length + 3 + lengths = torch.full((rows,), width, dtype=torch.int32, device="cuda") + logits = create_random_logits( + torch.zeros_like(lengths), + lengths, + torch.float32, + 42, + False, + "10LSBits" if distribution == "10LSBits" else "random", + ) + if distribution == "ascending": + logits.copy_(torch.arange(width, device="cuda", dtype=torch.float32)) + elif distribution == "constant": + logits.fill_(1.0) + elif distribution == "sampled_peaks": + # A biased sample leaves fewer than k survivors and must fall back. + logits.zero_() + sample = torch.arange(4096, device="cuda") + positions = (sample // 32) * (width // 128) + sample % 32 + logits[:, positions] = 1 + sample.float() / 4096 + original = logits.clone() + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") + workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda") + + def run() -> None: + torch.ops._C.persistent_topk( + logits, + lengths.view(-1, 4 if rows == 128 else 1), + indices, + workspace, + top_k, + width, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + positions = torch.arange(width, device="cuda") + slots = torch.arange(top_k, device="cuda") + for step in range(3): + bounds = torch.tensor( + [ + -1, + 0, + 1, + top_k - 1, + top_k, + 32769, + min_sampled_length - 1, + min_sampled_length, + width, + width + 17, + ], + dtype=torch.int32, + device="cuda", + ) + lengths.copy_( + bounds[(torch.arange(rows, device="cuda") + step) % bounds.numel()] + ) + logits.copy_(original if step != 1 else original.flip(1)) + logits.masked_fill_(positions[None] >= lengths[:, None], float("nan")) + indices.fill_(-2) + graph.replay() + valid = slots[None] < lengths.clamp(max=top_k)[:, None] + assert torch.all(indices[~valid] == -1) + assert torch.all(((indices >= 0) & (indices < lengths[:, None])) == valid) + ordered_indices = indices.sort(dim=1).values + assert torch.all( + (ordered_indices[:, 1:] != ordered_indices[:, :-1]) + | (ordered_indices[:, 1:] == -1) + ) + selected = logits.gather(1, indices.clamp_min(0).long()) + selected.masked_fill_(~valid, -float("inf")) + expected = logits.masked_fill( + positions[None] >= lengths[:, None], -float("inf") + ) + torch.testing.assert_close( + selected.sort(dim=1, descending=True).values, + expected.topk(top_k, dim=1).values, + atol=0, + rtol=0, + ) + + def run_large_context_topk_test( batch_size: int, seq_lens: list[int], From 0fefffc9346644a01dc132e544b881ad646ea81b Mon Sep 17 00:00:00 2001 From: qizixi <22851944+zixi-qi@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:43:56 -0700 Subject: [PATCH 120/139] [KVConnector][NIXL] Support attention-HMA layouts in pipeline-parallel push prefill (#50494) Signed-off-by: zixi-qi Co-authored-by: Claude Co-authored-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Co-authored-by: Codex --- docs/design/nixl_kv_push_connector.md | 44 ++ docs/features/nixl_connector_compatibility.md | 18 + .../unit/test_nixl_connector_hma.py | 1 + .../unit/test_nixl_desc_geometry.py | 32 +- .../unit/test_nixl_push_connector.py | 400 +++++++++++++++++- .../kv_connector/v1/nixl/base_worker.py | 256 +++++++++-- .../kv_connector/v1/nixl/metadata.py | 7 +- .../kv_connector/v1/nixl/push_worker.py | 4 +- 8 files changed, 713 insertions(+), 49 deletions(-) diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md index 6f66be5eb9..a6debd7dba 100644 --- a/docs/design/nixl_kv_push_connector.md +++ b/docs/design/nixl_kv_push_connector.md @@ -181,6 +181,50 @@ The completion notif sent from P to D after a WRITE is the existing is D's own request id, taken from the registration), so the D-side accounting code is unchanged. +## Pipeline parallelism and hybrid KV caches + +The pull connector requires the local and remote workers to expose a +congruent list of KV regions — region *i* here corresponds to region +*i* there. That assumption breaks under pipeline parallelism (PP) +combined with a hybrid (HMA) KV layout: + +* a PP-sharded prefiller (**P**) holds only a slice of the model's + layers while the `PP=1` decoder (**D**) holds them all, so region + counts differ; +* with HMA, several layer names are pooled into one region and the layer + that represents a pooled region can differ between P and D. + +`NixlPushConnector` handles this for PP-sharded producers by routing +**by layer-name (member) identity** instead of by region index. Each +worker advertises which layer names back each of its NIXL regions in the +handshake metadata (`NixlAgentMetadata.region_members`). A producer that +needs member routing derives its member-major layout once, when it +registers its KV caches, and every transfer it issues uses that order. +`add_remote_agent` then selects exactly the remote regions this stage +owns and reorders them to match, so both sides stay paired regardless of +how each remote rank happens to order its metadata. + +Addresses, block lengths, strides, and per-region capacities follow the same +member order. Descriptor offsets use each member's region capacity, so P and +D need not allocate the same number of blocks. Physical allocations are still +registered once, even when multiple members share them. + +Invariants enforced when the remote regions are aligned: + +* every locally owned member must be advertised exactly once by the + remote; a missing member fails the handshake rather than silently + leaving that layer's KV stale, and remote-only members (owned by other + PP stages) are ignored; +* a remote that omits member metadata while the local layout requires + member routing fails the handshake instead of falling back to + region-index routing; +* member order is a property of the local layout alone, so the same local + source descriptors serve every remote engine and TP rank; only the + remote descriptor list is rebuilt per rank. + +Decode-side PP is unsupported because completions are counted per +consumer rank. Mamba/SSM hybrids are unsupported under PP. + ## Scheduler-side responsibilities `NixlPushConnectorScheduler` extends the base scheduler with: diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md index 031fba5b2f..a41d5868b1 100644 --- a/docs/features/nixl_connector_compatibility.md +++ b/docs/features/nixl_connector_compatibility.md @@ -104,6 +104,24 @@ By default, a **compatibility hash** is checked during handshake. P and D instan - `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting. - Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA. +### Pipeline parallelism + +The default (pull) NixlConnector does **not** support +`pipeline-parallel-size > 1` together with a hybrid KV cache (HMA) +layout: region indices are not a stable identity across the +prefill/decode layer split, so the connector raises at startup. Use the +push connector (`NixlPushConnector`) for pipeline parallelism with +hybrid KV caches — it routes transfers by layer-name (member) identity, +so a PP-sharded prefiller can write into a `PP=1` decoder. See +[NIXL push-mode KV transfer](../design/nixl_kv_push_connector.md). + +Current push PP + HMA limitations: + +- Only the prefiller (producer) may be PP-sharded; decode-side PP is not supported. +- Hybrid SSM/Mamba layouts are not supported under PP. +- HMA requires the same block size on P and D. +- Attention-HMA member routing requires decode TP to be no greater than prefill TP. + ### Quantized KV cache [Quantized KV cache](quantization/quantized_kvcache.md) (e.g., FP8) requires both P and D instances to use the **same** `cache_dtype`. Mismatched cache dtypes will fail the compatibility hash check during handshake. diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 60eb2d6614..ae2ddece7f 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -918,6 +918,7 @@ def _make_mock_worker_for_desc_ids( worker = MagicMock(spec=NixlConnectorWorker) worker.num_regions = num_regions worker._uses_region_group_mapping = False + worker._transfer_layer_group_ids = () worker._has_mamba = has_mamba worker._group_spec_types = group_spec_types worker.block_len_per_layer = block_len_per_layer or [100] diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index bfec51da69..a4bfa9ab83 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -12,6 +12,7 @@ mid-decode (silent corruption of an unrelated request). """ from collections import defaultdict +from threading import Event, Lock from unittest.mock import MagicMock, patch import numpy as np @@ -112,6 +113,7 @@ def test_local_descriptors_follow_each_region_pool_capacity(): worker.block_len_per_layer = [16, 16] worker.block_stride_per_layer = [16, 16] worker.region_num_blocks = [2, 3] + worker._transfer_layer_region_indices = () descriptors = worker._build_fa_local([100, 1000], block_size_ratio=1) @@ -119,7 +121,8 @@ def test_local_descriptors_follow_each_region_pool_capacity(): @pytest.mark.cpu_test -def test_overlaid_transfer_groups_share_region_geometry(): +@pytest.mark.parametrize("push_pp", [False, True]) +def test_overlaid_transfer_groups_share_region_geometry(push_pp): """Groups overlaid on one allocation share its transfer region.""" import msgspec @@ -127,6 +130,9 @@ def test_overlaid_transfer_groups_share_region_geometry(): from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlAgentMetadata, ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, + ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -153,7 +159,14 @@ def test_overlaid_transfer_groups_share_region_geometry(): } groups = [KVCacheGroupSpec([layer_name], spec) for layer_name in caches] - worker = object.__new__(NixlConnectorWorker) + worker_cls = NixlPushConnectorWorker if push_pp else NixlConnectorWorker + worker = object.__new__(worker_cls) + if push_pp: + worker._push_writer_stop = Event() + worker._push_writer_wake = Event() + worker._push_writer_thread = MagicMock() + worker._sending_transfers_lock = Lock() + worker._sending_transfers = defaultdict(list) worker.tp_rank = 0 worker.world_size = 1 worker.transfer_tp_rank = 0 @@ -198,6 +211,9 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker._desc_is_dram_by_block_size = {} worker._desc_pos_by_block_size = {} worker._dram_src_handles_by_block_size = {} + worker._transfer_layer_names = () + worker._transfer_layer_region_indices = () + worker._transfer_layer_group_ids = () worker._region_is_mla = [] worker.block_len_per_layer = [] worker.block_stride_per_layer = [] @@ -205,7 +221,8 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker.use_host_buffer = False worker.host_xfer_buffers = {} worker.device_kv_caches = {} - worker.pp_size = 1 + worker.pp_size = 2 if push_pp else 1 + worker._is_hma_required = True worker.dcp_size = 1 worker.pcp_size = 1 worker.kv_buffer_device = "cuda" @@ -240,7 +257,13 @@ def test_overlaid_transfer_groups_share_region_geometry(): expected_addrs = [ backing.data_ptr() + block * block_stride for block in range(num_blocks) ] - assert worker.src_blocks_data[:, 0].tolist() == expected_addrs + num_desc_regions = 2 if push_pp else 1 + assert worker.src_blocks_data[:, 0].tolist() == expected_addrs * num_desc_regions + assert worker.num_descs == num_blocks * num_desc_regions + assert ( + worker.dst_region_num_blocks[worker.engine_id] + == [num_blocks] * num_desc_regions + ) metadata = msgspec.msgpack.decode( worker.xfer_handshake_metadata.agent_metadata_bytes, @@ -248,6 +271,7 @@ def test_overlaid_transfer_groups_share_region_geometry(): ) assert metadata.region_group_ids == [-1] assert metadata.region_num_blocks == [num_blocks] + assert metadata.region_members == ([["layer.0", "layer.1"]] if push_pp else []) assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]] diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 913ed0dcba..a908d32462 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -32,9 +32,13 @@ from unittest.mock import MagicMock, patch import msgspec import pytest +from vllm.distributed.kv_transfer.kv_connector.utils import TransferTopology from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorTransferResults, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( + NixlBaseConnectorWorker, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( PUSH_REG_NOTIF_PREFIX, NixlAgentMetadata, @@ -43,6 +47,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( NixlPushConnectorWorker, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( get_base_request_id, ) @@ -354,6 +359,7 @@ class _StubWriterWorker(NixlPushConnectorWorker): w._failed_recv_reqs = queue.Queue() w._recv_failures = set() w._recving_transfers = defaultdict(list) + w._is_hma_required = False w._reqs_to_process = set() w._reqs_to_send = {} w.consumer_notification_counts_by_req = defaultdict(int) @@ -361,12 +367,16 @@ class _StubWriterWorker(NixlPushConnectorWorker): w.pcp_rank = 0 w.pcp_dcp_sharded = False w.world_size = 1 + w.pp_size = 1 w.engine_id = "test-decode-engine" w._remote_agents = {} w._handshake_lock = threading.RLock() w._physical_blocks_per_logical_kv_block = 1 w._uses_region_group_mapping = False w.region_group_ids = [0] + w._transfer_layer_names = () + w._transfer_layer_region_indices = () + w._transfer_layer_group_ids = () w.dst_region_num_blocks = {} w.dst_region_group_ids = {} w.dst_uses_region_group_mapping = {} @@ -922,12 +932,20 @@ class TestPushWriterNotifs: # Both handles were still cleaned up. assert request_id not in w._sending_transfers - def test_completed_push_send_waits_for_consumer_notif(self): + @pytest.mark.parametrize(("pp_size", "is_hma"), [(1, False), (2, True)]) + def test_completed_push_send_waits_for_consumer_notif(self, pp_size, is_hma): """P-side sends are completed by consumer notifs / lease expiry, not by local WRITE-handle completion.""" w = self._pollable_worker() + w.pp_size = pp_size + w._is_hma_required = is_hma request_id = self._make_sending_req(w) - w.nixl_wrapper.check_xfer_state.return_value = "DONE" + w.nixl_wrapper.check_xfer_state.side_effect = ["DONE", "PROC", "DONE"] + + results = w.get_transfer_results() + assert request_id not in results.finished_sending + assert w._sending_transfers[request_id] == [102] + assert request_id in w._reqs_to_send results = w.get_transfer_results() @@ -940,6 +958,12 @@ class TestPushWriterNotifs: results = w.get_transfer_results() assert request_id in results.finished_sending assert request_id not in w._reqs_to_send + assert request_id not in w._reqs_to_process + assert w._evict_finished_inbox.get_nowait() == request_id + assert w.get_transfer_results().finished_sending == set() + assert w._evict_finished_inbox.empty() + assert w.nixl_wrapper.release_xfer_handle.call_count == 2 + w.xfer_stats.record_kv_expired_req.assert_not_called() # ----------------------------------------------------------------- # @@ -1576,3 +1600,375 @@ class TestPushPrefixCaching: local, remote = self._written_block_ids(w) assert local == [10, 11, 12] assert remote == [500, 501, 502] + + @pytest.mark.parametrize("layer_name_routing", [False, True]) + def test_different_region_groups_require_layer_routing(self, layer_name_routing): + """A layer mapped to a pooled remote region can still receive its blocks.""" + w, engine_id = self._worker_driving_xfer() + w.dst_region_group_ids[engine_id] = [-1] + if layer_name_routing: + w._transfer_layer_group_ids = (0,) + reg = _registration_data("req-groups", local_block_ids=([500, 501],)) + + if not layer_name_routing: + with pytest.raises(NotImplementedError, match="cache-group layouts"): + NixlPushConnectorWorker._do_start_push_kv( + w, "req-groups", ([10, 11],), reg + ) + w.nixl_wrapper.make_prepped_xfer.assert_not_called() + return + + NixlPushConnectorWorker._do_start_push_kv(w, "req-groups", ([10, 11],), reg) + + local, remote = self._written_block_ids(w) + assert local == [10, 11] + assert remote == [500, 501] + + +def _agent_metadata( + region_layers: list[list[str]], + base_addresses: list[int], + block_lens: list[int], + block_strides: list[int] | None = None, +) -> NixlAgentMetadata: + return NixlAgentMetadata( + engine_id="remote-engine", + agent_metadata=b"agent", + kv_caches_base_addr=base_addresses, + device_id=7, + num_blocks=2, + block_lens=block_lens, + block_strides=list(block_lens if block_strides is None else block_strides), + kv_cache_layout="LBHNC", + block_size=16, + ssm_sizes=(0, 0), + attn_backend_name="FLASH_ATTN", + physical_blocks_per_logical_kv_block=1, + region_members=region_layers, + ) + + +def _layer_routing_worker( + region_layers: list[list[str]], + group_by_layer: dict[str, int], + pp_size: int = 2, + is_hma: bool = True, +) -> _StubWriterWorker: + worker = _StubWriterWorker.fresh() + worker.pp_size = pp_size + worker.dcp_size = 1 + worker.block_size = 16 + worker._has_mamba = False + worker._is_hma_required = is_hma + worker.kv_cache_config = MagicMock(transfer_group_index_by_layer=group_by_layer) + worker.transfer_topo = MagicMock() + worker._set_region_layers(region_layers) + return worker + + +@pytest.mark.parametrize( + ("region_num_blocks", "expected"), + [(None, [1, 2, 15, 21, 22]), ([4, 7, 5], [1, 2, 9, 12, 13])], +) +def test_layer_group_ids_route_descriptor_blocks(region_num_blocks, expected): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + desc_ids = worker._compute_desc_ids( + block_ids=[[1, 2], [5]], + dst_num_blocks=10, + block_size_ratio=None, + physical_blocks_per_logical=1, + region_num_blocks=region_num_blocks, + ) + assert desc_ids.tolist() == expected + + +def test_layer_metadata_round_trip(): + metadata = _agent_metadata([["L0", "L1"]], [0x10000], [256]) + metadata.region_num_blocks = [4] + metadata.region_group_ids = [-1] + metadata.region_names = ["L0"] + metadata.region_mem_types = ["VRAM"] + + encoded = msgspec.msgpack.encode(metadata) + assert msgspec.msgpack.Decoder(NixlAgentMetadata).decode(encoded) == metadata + + +def test_layer_identity_gate_preserves_the_non_hma_path(): + assert _layer_routing_worker([["a"]], {"a": 0})._transfer_layer_region_indices == ( + 0, + ) + + # A base worker (pull) never routes by layer name. + pull = object.__new__(NixlBaseConnectorWorker) + pull._transfer_layer_names = () + pull._transfer_layer_region_indices = () + pull._transfer_layer_group_ids = () + pull._set_region_layers([["a"]]) + assert pull._transfer_layer_region_indices == () + + # A non-HMA local layout does not require layer-name routing. + assert ( + _layer_routing_worker( + [["a"]], {"a": 0}, is_hma=False + )._transfer_layer_region_indices + == () + ) + + # PP=1 has congruent local/remote regions and keeps the region route even + # when its allocator is hybrid. + assert ( + _layer_routing_worker( + [["a"]], {"a": 0}, pp_size=1 + )._transfer_layer_region_indices + == () + ) + + +def test_layer_alignment_fails_loud_when_remote_omits_layers(): + # Falling back to region-index routing here would silently transfer stale + # KV, so an unannounced peer must fail the handshake instead. + worker = _layer_routing_worker([["a"]], {"a": 0}) + with pytest.raises(AssertionError, match="no region_members"): + worker._align_remote_regions_by_layer(_agent_metadata([], [0xA000], [128])) + + +def test_layer_alignment_expands_pooled_regions(): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + assert worker._transfer_layer_names == ("a", "a.swa", "b") + assert worker._transfer_layer_region_indices == (0, 0, 1) + assert worker._transfer_layer_group_ids == (0, 1, 0) + + metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128]) + metadata.region_num_blocks = [4, 6] + metadata.region_group_ids = [-1, 0] + metadata.region_names = ["a", "b"] + metadata.region_mem_types = ["VRAM", "VRAM"] + worker._align_remote_regions_by_layer(metadata) + + assert metadata.kv_caches_base_addr == [0xA000, 0xA000, 0xB000] + assert metadata.block_lens == [128, 128, 128] + assert metadata.block_strides == [128, 128, 128] + assert metadata.region_members == [["a"], ["a.swa"], ["b"]] + assert metadata.region_num_blocks == [4, 4, 6] + assert metadata.region_group_ids == [-1, -1, 0] + assert metadata.region_names == ["a", "a.swa", "b"] + assert metadata.region_mem_types == ["VRAM", "VRAM", "VRAM"] + + +def test_layer_alignment_filters_and_reorders_a_pp_stage(): + worker = _layer_routing_worker([["l2"], ["l3"]], {"l2": 0, "l3": 1}) + metadata = _agent_metadata( + [["l2"], ["l0"], ["l3"], ["l1"]], + [0xC000, 0xA000, 0xD000, 0xB000], + [65536, 65536, 32768, 32768], + [131072, 131072, 65536, 65536], + ) + metadata.region_num_blocks = [4, 6, 8, 10] + metadata.region_group_ids = [0, 0, 1, 1] + metadata.region_names = ["l2", "l0", "l3", "l1"] + metadata.region_mem_types = ["VRAM"] * 4 + + worker._align_remote_regions_by_layer(metadata) + + assert worker._transfer_layer_region_indices == (0, 1) + assert worker._transfer_layer_group_ids == (0, 1) + assert metadata.kv_caches_base_addr == [0xC000, 0xD000] + assert metadata.block_lens == [65536, 32768] + assert metadata.block_strides == [131072, 65536] + assert metadata.region_num_blocks == [4, 8] + assert metadata.region_group_ids == [0, 1] + assert metadata.region_names == ["l2", "l3"] + assert metadata.region_mem_types == ["VRAM", "VRAM"] + + +@pytest.mark.parametrize( + ("local_num_blocks", "remote_num_blocks"), + [([4, 4], [4, 4, 4]), ([4, 6], [3, 5, 7])], +) +def test_layer_descriptors_pair_layers_across_asymmetric_pp_split( + local_num_blocks, remote_num_blocks +): + """A PP split can leave each stage a different mix of attention types. + + Stage 1 owns L3 (full) pooled with L4 (sliding) in region 0, and L5 (full) + alone in region 1. The consumer holds every layer and pools them + differently, so region indices cannot be paired positionally. + """ + worker = _layer_routing_worker([["L3", "L4"], ["L5"]], {"L3": 0, "L4": 1, "L5": 0}) + worker.block_len_per_layer = [128, 128] + worker.block_stride_per_layer = [256, 512] + worker._region_is_mla = [False, False] + worker.num_blocks = 4 + worker.region_num_blocks = local_num_blocks + worker.device_id = 0 + + consumer = _agent_metadata( + [["L0", "L1"], ["L3", "L2"], ["L5", "L4"]], + [0xA000, 0xB000, 0xC000], + [128, 128, 128], + [256, 512, 1024], + ) + consumer.num_blocks = 4 + consumer.region_num_blocks = remote_num_blocks + worker._align_remote_regions_by_layer(consumer) + + # L3 -> remote region 1, L4 -> 2, L5 -> 2. Pairing by index would have sent + # L5 to region 1 and L4 to region 2's sibling. + assert consumer.kv_caches_base_addr == [0xB000, 0xC000, 0xC000] + + plan = TPMapping(((0,), (0,)), (0,), {0: 0}, 0) + local_descs = worker._build_fa_local([0x1000, 0x2000], block_size_ratio=1) + remote_descs = worker._build_fa_remote(plan, consumer, block_size_ratio=1) + local_counts = [local_num_blocks[i] for i in worker._transfer_layer_region_indices] + local_ids = worker._compute_desc_ids( + [[1, 2], [3]], 4, None, 1, region_num_blocks=local_counts + ) + remote_ids = worker._compute_desc_ids( + [[0, 1], [2]], 4, None, 1, region_num_blocks=consumer.region_num_blocks + ) + assert len(local_descs) == sum(local_counts) + assert len(remote_descs) == sum(consumer.region_num_blocks) + + # Each pair addresses the same layer's blocks despite different pooling + # and strides. L3 and L5 use group 0; L4 uses group 1. + assert local_descs[local_ids].tolist() == [ + [0x1100, 128, 0], + [0x1200, 128, 0], + [0x1300, 128, 0], + [0x2200, 128, 0], + [0x2400, 128, 0], + ] + assert remote_descs[remote_ids].tolist() == [ + [0xB000, 128, 7], + [0xB200, 128, 7], + [0xC800, 128, 7], + [0xC000, 128, 7], + [0xC400, 128, 7], + ] + + +def test_layer_alignment_is_canonical_across_remote_orderings(): + local, groups = [["x"], ["y"]], {"x": 0, "y": 1} + rank0 = _agent_metadata([["x"], ["y"]], [0x1000, 0x2000], [64, 128], [256, 512]) + rank1 = _agent_metadata([["y"], ["x"]], [0x2000, 0x1000], [128, 64], [512, 256]) + + _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank0) + _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank1) + + assert rank0.kv_caches_base_addr == rank1.kv_caches_base_addr == [0x1000, 0x2000] + assert rank0.block_lens == rank1.block_lens == [64, 128] + assert rank0.block_strides == rank1.block_strides == [256, 512] + + +def test_layer_alignment_is_idempotent(): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128]) + metadata.region_num_blocks = [4, 6] + metadata.region_group_ids = [-1, 0] + metadata.region_names = ["a", "b"] + metadata.region_mem_types = ["VRAM", "VRAM"] + + worker._align_remote_regions_by_layer(metadata) + aligned = msgspec.msgpack.encode(metadata) + worker._align_remote_regions_by_layer(metadata) + + assert msgspec.msgpack.encode(metadata) == aligned + + +def test_layer_alignment_rejects_missing_local_layer(): + worker = _layer_routing_worker([["l0"], ["l1"]], {"l0": 0, "l1": 1}) + metadata = _agent_metadata([["l0"]], [0xA000], [128]) + + with pytest.raises(AssertionError, match="missing locally owned layers"): + worker._align_remote_regions_by_layer(metadata) + + +def test_layer_alignment_rejects_duplicate_remote_layer(): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["a"]], [0xA000, 0xB000], [128, 128]) + + with pytest.raises(AssertionError, match="in multiple regions"): + worker._align_remote_regions_by_layer(metadata) + + +def test_layer_alignment_rejects_inconsistent_remote_metadata(): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["b"]], [0xA000], [128]) + + with pytest.raises(AssertionError, match="lengths disagree"): + worker._align_remote_regions_by_layer(metadata) + + +@pytest.mark.parametrize( + "field", + ["region_num_blocks", "region_group_ids", "region_names", "region_mem_types"], +) +def test_layer_alignment_rejects_inconsistent_region_geometry(field): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["b"]], [0xA000, 0xB000], [128, 128]) + setattr(metadata, field, []) + original = msgspec.msgpack.encode(metadata) + + with pytest.raises(AssertionError, match="lengths disagree"): + worker._align_remote_regions_by_layer(metadata) + assert msgspec.msgpack.encode(metadata) == original + + +def test_set_region_layers_rejects_duplicate_local_layer(): + with pytest.raises(AssertionError, match="spans multiple NIXL regions"): + _layer_routing_worker([["a"], ["a"]], {"a": 0}) + + +def test_set_region_layers_rejects_layer_outside_any_kv_group(): + with pytest.raises(AssertionError, match="outside any local group"): + _layer_routing_worker([["a"], ["b"]], {"a": 0}) + + +@pytest.mark.parametrize( + ("local_block_size", "remote_block_size", "remote_tp_size", "error"), + [ + (32, 16, 1, "identical P/D block sizes"), + (16, 32, 1, "identical P/D block sizes"), + (16, 16, 2, "decode TP greater"), + ], +) +def test_layer_handshake_rejects_unsupported_geometry( + local_block_size: int, remote_block_size: int, remote_tp_size: int, error: str +): + """Reject unsupported peers without registering agents or transfer state.""" + metadata = _agent_metadata([["a"]], [0xA000], [128]) + metadata.block_size = remote_block_size + worker = _layer_routing_worker([["a"]], {"a": 0}) + worker.block_size = local_block_size + worker.block_len_per_layer = [128] + worker.use_mla = False + worker.nixl_wrapper = MagicMock() + worker.kv_caches_base_addr = defaultdict(dict) + worker.dst_num_blocks = {} + worker.tp_mappings = {} + worker.transfer_topo = TransferTopology( + tp_rank=0, + tp_size=1, + block_size=local_block_size, + engine_id=worker.engine_id, + is_mla=False, + is_mamba=False, + total_num_kv_heads=8, + attn_backends=[], + ) + + with pytest.raises(NotImplementedError, match=error): + worker.add_remote_agent(metadata, remote_tp_size=remote_tp_size) + worker.nixl_wrapper.add_remote_agent.assert_not_called() + worker.nixl_wrapper.prep_xfer_dlist.assert_not_called() + assert not worker.tp_mappings + assert not worker.dst_num_blocks + assert not worker.kv_caches_base_addr + with pytest.raises(KeyError): + worker.transfer_topo.get_engine_info(metadata.engine_id) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 6f8f1e55dc..2704de589f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -164,6 +164,10 @@ class NixlBaseConnectorWorker: # Overridden by NixlPushConnectorWorker. _TRANSFER_MODE: str = "pull" + # Layer-name routing is supported only by NixlPushConnector for HMA KV + # caches under PP. + _supports_pp_hma = False + def _compute_desc_ids( self, block_ids: BlockIds, @@ -191,28 +195,42 @@ class NixlBaseConnectorWorker: num_blocks = dst_num_blocks if block_size_ratio is not None: num_blocks = int(num_blocks * block_size_ratio) - num_regions = self.num_regions or len(self.block_len_per_layer) + use_layer_name_routing = bool(self._transfer_layer_group_ids) + num_desc_regions = ( + len(self._transfer_layer_group_ids) + if use_layer_name_routing + else self.num_regions or len(self.block_len_per_layer) + ) if region_num_blocks is None: - region_num_blocks = [num_blocks] * num_regions + region_num_blocks = [num_blocks] * num_desc_regions elif block_size_ratio is not None: region_num_blocks = [ int(count * block_size_ratio) for count in region_num_blocks ] - if region_num_blocks: - num_regions = len(region_num_blocks) - assert len(region_num_blocks) == num_regions + if region_num_blocks and not use_layer_name_routing: + num_desc_regions = len(region_num_blocks) + assert len(region_num_blocks) == num_desc_regions region_offsets = np.cumsum([0, *region_num_blocks[:-1]]) num_fa_descs = sum(region_num_blocks) + if use_layer_name_routing: + # Descriptors in layer order, using each layer's KV-cache group. + return np.concatenate( + [ + np.asarray(block_ids[g], dtype=np.int32) + region_offsets[k] + for k, g in enumerate(self._transfer_layer_group_ids) + ] + ) + # All-attention fast path. if num_ssm_regions == 0: if region_group_ids is None: region_group_ids = self.region_group_ids if uses_region_group_mapping is None: uses_region_group_mapping = self._uses_region_group_mapping - if not region_group_ids and num_regions == 1: + if not region_group_ids and num_desc_regions == 1: region_group_ids = [0] - assert len(region_group_ids) == num_regions + assert len(region_group_ids) == num_desc_regions if uses_region_group_mapping is None: uses_region_group_mapping = len(set(region_group_ids)) > 1 region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32) @@ -392,6 +410,103 @@ class NixlBaseConnectorWorker: """ return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx] + def _set_region_layers(self, region_layers: list[list[str]]) -> None: + layer_names = [layer_name for region in region_layers for layer_name in region] + assert len(layer_names) == len(set(layer_names)), ( + "A KV cache layer spans multiple NIXL regions" + ) + self.region_members: list[list[str]] = region_layers + if not self._requires_layer_name_routing(): + return + group_by_layer = self.kv_cache_config.transfer_group_index_by_layer + ungrouped = [name for name in layer_names if name not in group_by_layer] + assert not ungrouped, f"KV cache layers outside any local group: {ungrouped}" + self._transfer_layer_names = tuple(layer_names) + self._transfer_layer_region_indices = tuple( + region_idx + for region_idx, region in enumerate(region_layers) + for _ in region + ) + self._transfer_layer_group_ids = tuple( + group_by_layer[name] for name in layer_names + ) + + def _requires_layer_name_routing(self) -> bool: + """Whether PP push must match HMA layers by name, not region index.""" + return ( + self._supports_pp_hma + and self.pp_size > 1 + and self._is_hma_required + and not self._has_mamba + ) + + def _align_remote_regions_by_layer( + self, nixl_agent_meta: NixlAgentMetadata + ) -> None: + """Select remote regions for this PP stage in local layer order.""" + remote_region_layers = nixl_agent_meta.region_members + # Region-index routing here would silently transfer stale KV. + assert remote_region_layers, "Remote advertised no region_members" + assert ( + len(nixl_agent_meta.kv_caches_base_addr) + == len(nixl_agent_meta.block_lens) + == len(nixl_agent_meta.block_strides) + == len(remote_region_layers) + ), "Remote region metadata lengths disagree" + assert all( + values is None or len(values) == len(remote_region_layers) + for values in ( + nixl_agent_meta.region_num_blocks, + nixl_agent_meta.region_group_ids, + nixl_agent_meta.region_names, + nixl_agent_meta.region_mem_types, + ) + ), "Remote region metadata lengths disagree" + + remote_region_by_layer: dict[str, int] = {} + for region_idx, layer_names in enumerate(remote_region_layers): + for layer_name in layer_names: + assert layer_name not in remote_region_by_layer, ( + f"Remote advertised layer {layer_name!r} in multiple regions" + ) + remote_region_by_layer[layer_name] = region_idx + + missing = [ + name + for name in self._transfer_layer_names + if name not in remote_region_by_layer + ] + assert not missing, f"Remote is missing locally owned layers: {missing}" + + remote_regions = [ + remote_region_by_layer[name] for name in self._transfer_layer_names + ] + nixl_agent_meta.kv_caches_base_addr = [ + nixl_agent_meta.kv_caches_base_addr[i] for i in remote_regions + ] + nixl_agent_meta.block_lens = [ + nixl_agent_meta.block_lens[i] for i in remote_regions + ] + nixl_agent_meta.block_strides = [ + nixl_agent_meta.block_strides[i] for i in remote_regions + ] + if nixl_agent_meta.region_num_blocks is not None: + nixl_agent_meta.region_num_blocks = [ + nixl_agent_meta.region_num_blocks[i] for i in remote_regions + ] + if nixl_agent_meta.region_group_ids is not None: + nixl_agent_meta.region_group_ids = [ + nixl_agent_meta.region_group_ids[i] for i in remote_regions + ] + if nixl_agent_meta.region_mem_types is not None: + nixl_agent_meta.region_mem_types = [ + nixl_agent_meta.region_mem_types[i] for i in remote_regions + ] + if nixl_agent_meta.region_names is not None: + nixl_agent_meta.region_names = list(self._transfer_layer_names) + # One layer per region keeps a second alignment pass a no-op. + nixl_agent_meta.region_members = [[name] for name in self._transfer_layer_names] + def __init__( self, vllm_config: "VllmConfig", @@ -638,18 +753,26 @@ class NixlBaseConnectorWorker: # transfers into the matching sub-range of a PP=1 remote's regions. self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self._remote_region_offset = 0 - # PP push slices regions per layer (uniform count); HMA breaks that. - if self.pp_size > 1 and self._is_hma_required: - raise NotImplementedError( - "NixlPushConnector does not support pipeline_parallel_size > 1 " - "with hybrid KV cache layouts (HMA) yet." - ) - # Decode-side PP is unsupported (completions counted per consumer rank). - if vllm_config.kv_transfer_config.kv_role == "kv_consumer" and self.pp_size > 1: - raise NotImplementedError( - "NixlPushConnector consumer (decode) does not support " - "pipeline_parallel_size > 1." - ) + if self.pp_size > 1: + if self._is_hma_required and not self._supports_pp_hma: + raise NotImplementedError( + "NixlConnector (pull) does not support pipeline_parallel_size " + "> 1 with hybrid KV cache layouts (HMA); use NixlPushConnector " + "for PP + HMA." + ) + # PP push routes HMA (hybrid) attention layouts by layer name; + # Mamba/SSM hybrids are not yet supported under PP. + if self._has_mamba: + raise NotImplementedError( + "NixlPushConnector does not support pipeline_parallel_size > 1 " + "with Mamba/SSM hybrid KV cache layouts yet." + ) + # Decode-side PP is unsupported (completions counted per consumer rank). + if vllm_config.kv_transfer_config.kv_role == "kv_consumer": + raise NotImplementedError( + "NixlPushConnector consumer (decode) does not support " + "pipeline_parallel_size > 1." + ) # Keep heartbeat handshakes to a PP-sharded producer notif-only. self._hb_handshake_notif_only = False @@ -773,6 +896,13 @@ class NixlBaseConnectorWorker: # within a block (BLHNC/BHLNC), where stride > block_len. self.block_stride_per_layer = list[int]() + # Region layer membership is populated for HMA layouts. + self.region_members = [] + # Local layer order, set by ``_set_region_layers`` only when + # routing by layer name. Empty means region-index routing. + self._transfer_layer_names = tuple[str, ...]() + self._transfer_layer_region_indices = tuple[int, ...]() + self._transfer_layer_group_ids = tuple[int, ...]() # Per-engine TP mappings. Generated during handshake. self.tp_mappings: dict[EngineId, TPMapping] = {} @@ -1306,6 +1436,11 @@ class NixlBaseConnectorWorker: ): compressed_region_owners.setdefault(cache.data_ptr(), cache) + track_region_layers = ( + self._supports_pp_hma and self._is_hma_required and not self._has_mamba + ) + region_layers: list[list[str]] = [] + # K and V are packed into the content dim, so each attention layer is a # single NIXL region whose block transfers as one unit. Mamba layers instead # register separate conv/ssm sub-regions (see `_build_mamba_local`). @@ -1508,6 +1643,12 @@ class NixlBaseConnectorWorker: self._region_is_mla.append(is_mla_region) region_mem_types.append(mem_type) + if track_region_layers: + if region_index == len(region_layers): + region_layers.append([layer_name]) + else: + region_layers[region_index].append(layer_name) + if isinstance(layer_spec, MambaSpec): if layer_spec.tp_replicated: assert self._is_csa_linear @@ -1563,7 +1704,9 @@ class NixlBaseConnectorWorker: if self._has_mamba: self.region_num_blocks = [self.num_blocks] * self.num_regions - if self.pp_size > 1: + self._set_region_layers(region_layers) + + if self.pp_size > 1 and not self._is_hma_required: start_layer, end_layer = self.model_config.get_layers_start_end_indices( self.vllm_config.parallel_config ) @@ -1573,7 +1716,12 @@ class NixlBaseConnectorWorker: self._remote_region_offset = regions_per_layer * start_layer # Total local FA descriptors (boundary between FA and mamba descs). - self.num_descs = sum(self.region_num_blocks) + xfer_region_num_blocks = ( + [self.region_num_blocks[i] for i in self._transfer_layer_region_indices] + if self._transfer_layer_region_indices + else self.region_num_blocks + ) + self.num_descs = sum(xfer_region_num_blocks) self._mixed_mem_types = len(set(region_mem_types)) > 1 if self._mixed_mem_types: @@ -1597,7 +1745,7 @@ class NixlBaseConnectorWorker: self.device_kv_caches = kv_caches self.dst_num_blocks[self.engine_id] = self.num_blocks - self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks + self.dst_region_num_blocks[self.engine_id] = xfer_region_num_blocks self.dst_region_group_ids[self.engine_id] = self.region_group_ids self.dst_uses_region_group_mapping[self.engine_id] = ( self._uses_region_group_mapping @@ -1647,6 +1795,7 @@ class NixlBaseConnectorWorker: region_mem_types=self.region_mem_types, dcp_size=self.dcp_size, pcp_size=self.pcp_size, + region_members=self.region_members, ) # Wrap metadata in payload with hash for defensive decoding assert self.compat_hash is not None @@ -1805,12 +1954,16 @@ class NixlBaseConnectorWorker: base_addresses: list[int], block_size_ratio: int, ) -> np.ndarray: - """Build local FA descriptors for all layers as an Nx3 uint64 array.""" + """Build local FA descriptors as an Nx3 uint64 array.""" assert self.transfer_topo is not None assert base_addresses, "Local KV cache base addresses must not be empty." device_id = self.device_id + region_indices = self._transfer_layer_region_indices or range( + len(base_addresses) + ) parts: list[np.ndarray] = [] - for i, base_addr in enumerate(base_addresses): + for i in region_indices: + base_addr = base_addresses[i] block_len = self.block_len_per_layer[i] // block_size_ratio logical_blocks = np.repeat( np.arange(self.region_num_blocks[i], dtype=np.uint64), @@ -1834,7 +1987,7 @@ class NixlBaseConnectorWorker: nixl_agent_meta: NixlAgentMetadata, block_size_ratio: int, ) -> np.ndarray: - """Build remote FA descriptors for all layers as an Nx3 uint64 array.""" + """Build remote FA descriptors as an Nx3 uint64 array.""" assert self.transfer_topo is not None assert nixl_agent_meta.kv_caches_base_addr, ( "Remote KV cache base addresses must not be empty." @@ -1851,10 +2004,14 @@ class NixlBaseConnectorWorker: device_id = nixl_agent_meta.device_id block_strides = nixl_agent_meta.block_strides parts: list[np.ndarray] = [] + use_layer_name_routing = bool(self._transfer_layer_region_indices) for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - replicated = self._is_region_replicated(i) + local_region = ( + self._transfer_layer_region_indices[i] if use_layer_name_routing else i + ) + replicated = self._is_region_replicated(local_region) # Read our whole local region size from remote.. - local_block_len = self.block_len_per_layer[i] + local_block_len = self.block_len_per_layer[local_region] remote_kv_block_len = local_block_len // block_size_ratio if block_size_ratio > 1: # ..using remote kv_block_len as transfer unit @@ -2021,9 +2178,22 @@ class NixlBaseConnectorWorker: ) return self._remote_agents[engine_id][(0, remote_tp_rank)] + assert self.transfer_topo is not None + transfer_topo = self.transfer_topo # Number of physical regions registered locally (one per layer/tensor). num_local_regions = len(self.block_len_per_layer) - if ( + if self._transfer_layer_region_indices: + if self.block_size != nixl_agent_meta.block_size: + raise NotImplementedError( + "Attention-HMA push requires identical P/D block sizes." + ) + if remote_tp_size > transfer_topo.tp_size and not self.use_mla: + raise NotImplementedError( + "Attention-HMA push does not support decode TP greater than " + "prefill TP yet" + ) + self._align_remote_regions_by_layer(nixl_agent_meta) + elif ( self.pp_size > 1 and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions ): @@ -2038,6 +2208,10 @@ class NixlBaseConnectorWorker: ] nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end] nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end] + if nixl_agent_meta.region_num_blocks is not None: + nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ + start:end + ] if nixl_agent_meta.region_group_ids is not None: nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[ start:end @@ -2048,14 +2222,8 @@ class NixlBaseConnectorWorker: nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[ start:end ] - if nixl_agent_meta.region_num_blocks is not None: - nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ - start:end - ] ### Register remote engine in TransferTopology (idempotent). - assert self.transfer_topo is not None - transfer_topo = self.transfer_topo physical_blocks_per_logical = ( nixl_agent_meta.physical_blocks_per_logical_kv_block ) @@ -2349,6 +2517,9 @@ class NixlBaseConnectorWorker: # the per-rank KV head ratio rather than the raw tp_ratio, because GQA # replication caps per-rank heads at 1 when tp > total_kv_heads # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here. + local_regions = self._transfer_layer_region_indices or range( + len(self.block_len_per_layer) + ) if self._has_mamba and self.use_mla: # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are # kernel-granularity views of the mamba-unified page. The MLA @@ -2365,7 +2536,7 @@ class NixlBaseConnectorWorker: f"remote={nixl_agent_meta.block_lens}." ) elif not self._has_mamba: - assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), ( + assert len(local_regions) == len(nixl_agent_meta.block_lens), ( "Number of KV layers must match between prefill and decode" ) model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated( @@ -2374,8 +2545,9 @@ class NixlBaseConnectorWorker: total_kv_heads = self.transfer_topo.total_num_kv_heads local_heads = self.transfer_topo.local_physical_heads remote_heads = max(1, total_kv_heads // remote_tp_size) - for i, local_len in enumerate(self.block_len_per_layer): - replicated = model_replicated or self._is_region_replicated(i) + for i, region_idx in enumerate(local_regions): + local_len = self.block_len_per_layer[region_idx] + replicated = model_replicated or self._is_region_replicated(region_idx) remote_len = nixl_agent_meta.block_lens[i] if replicated: assert local_len // block_size_ratio == remote_len, ( @@ -2406,8 +2578,7 @@ class NixlBaseConnectorWorker: # TP workers that handhshake with same remote have same #blocks. assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks - # Same number of regions/~layers. - assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer) + assert len(nixl_agent_meta.kv_caches_base_addr) == len(local_regions) num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr) if nixl_agent_meta.region_num_blocks is not None: assert len(nixl_agent_meta.region_num_blocks) == num_remote_regions @@ -2423,7 +2594,12 @@ class NixlBaseConnectorWorker: if nixl_agent_meta.region_group_ids is not None: assert len(nixl_agent_meta.region_group_ids) == num_remote_regions if nixl_agent_meta.region_names is not None: - assert nixl_agent_meta.region_names == self.region_names, ( + expected_names = ( + list(self._transfer_layer_names) + if self._transfer_layer_names + else self.region_names + ) + assert nixl_agent_meta.region_names == expected_names, ( "NIXL transfer regions must name the same model caches in the " "same order" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index be32a6969d..c2b4d9dce7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Metadata dataclasses and helpers for the NIXL connector.""" -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from vllm.config import VllmConfig @@ -47,8 +47,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:" # 9: Add block_strides # 10: Add dense virtual transfer pages for compressed MLA caches # 11: Add per-region transfer geometry and memory types to NixlAgentMetadata +# 12: Add per-region member names for PP push # -NIXL_CONNECTOR_VERSION: int = 11 +NIXL_CONNECTOR_VERSION: int = 12 @dataclass @@ -71,6 +72,8 @@ class NixlAgentMetadata: region_mem_types: list[str] | None = None dcp_size: int = 1 pcp_size: int = 1 + # Layer names sharing each advertised region, in region order. + region_members: list[list[str]] = field(default_factory=list) @dataclass diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 8b34e7ffbb..01c956a525 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -83,6 +83,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): # Distinguishes push from pull in the NIXL compatibility hash. _TRANSFER_MODE: str = "push" + _supports_pp_hma = True + def __init__( self, vllm_config: "VllmConfig", @@ -530,7 +532,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): local_region_groups = self.region_group_ids remote_region_groups = self.dst_region_group_ids[engine_id] groups_differ = local_region_groups != remote_region_groups - if groups_differ: + if groups_differ and not self._transfer_layer_group_ids: raise NotImplementedError( "NixlPushConnector does not support different producer and " "consumer cache-group layouts" From 031f5810c1957b6ec2a7666f65a3fe68c1f7e534 Mon Sep 17 00:00:00 2001 From: Shang Wang Date: Tue, 15 Sep 2026 16:29:13 -0700 Subject: [PATCH 121/139] [Build][NVIDIA] Update public Rubin dependencies and MSA compatibility (#56545) Signed-off-by: Shang Wang Co-authored-by: Codex --- cmake/external_projects/fmha_sm100.cmake | 2 +- docker/Dockerfile | 2 +- requirements/rubin-prerelease.txt | 22 ++++++++++++++++------ 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index 442097dec7..ae5bfec7b4 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -17,7 +17,7 @@ else() FetchContent_Declare( fmha_sm100 GIT_REPOSITORY https://github.com/vllm-project/MSA.git - GIT_TAG 087c161814d4d9c735b46c21212a09e5f8eb92fa + GIT_TAG f355c37eb4e1413f21ee2ad8bbad25079e6bef9d GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/docker/Dockerfile b/docker/Dockerfile index b3e32777e7..d1a7afcbfb 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -641,7 +641,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ python3 use_existing_torch.py --prefix; \ fi \ && if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ - sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\])[[:space:]]*==/Id' \ + sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\]|quack-kernels)[[:space:]]*==/Id' \ requirements/cuda.txt; \ fi diff --git a/requirements/rubin-prerelease.txt b/requirements/rubin-prerelease.txt index 18a644f6b1..8b9f7d3155 100644 --- a/requirements/rubin-prerelease.txt +++ b/requirements/rubin-prerelease.txt @@ -2,15 +2,25 @@ # Installed when INSTALL_RUBIN_PRERELEASE=true (see docker/Dockerfile). # TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels # (vllm-project/vllm#49735). ---pre --extra-index-url https://download.pytorch.org/whl/nightly/cu134 +--extra-index-url https://flashinfer.ai/whl/cu134 +--extra-index-url https://flashinfer.ai/whl -torch==2.15.0.dev20260903+cu134 -torchvision==0.30.0.dev20260903+cu134 -torchaudio==2.11.0.dev20260903+cu134 +# Match each architecture's torch pin to the latest torchvision wheel. +torch==2.15.0.dev20260909+cu134; platform_machine == "x86_64" +torch==2.15.0.dev20260908+cu134; platform_machine == "aarch64" +torchvision==0.30.0.dev20260909+cu134 +torchaudio==2.11.0.dev20260911+cu134 -# docker/Dockerfile strips these two from requirements/cuda.txt whenever +# docker/Dockerfile strips these packages from requirements/cuda.txt whenever # INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to # still get installed. nvidia-cutlass-dsl[cu13]==4.8.0.dev0 -flashinfer-python==0.6.18 +flashinfer-python==0.6.18.post1 +quack-kernels==0.6.5 + +cuda-python==13.4.1 +cuda-bindings==13.4.1 +flashinfer-jit-cache==0.6.18.post1+cu134 +flashinfer-cubin==0.6.18.post1 +nixl==1.4.1 From 85c1f58d50b80a6668f1f8158f0d8ef7b1d9112a Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Tue, 15 Sep 2026 17:31:25 -0700 Subject: [PATCH 122/139] [BugFix][Model Runner V2][Spec Decode] Fix decode instance's multi-layer MTP kv caches during P/D (#55055) Signed-off-by: Giancarlo Delfin --- .../unit/test_nixl_push_connector.py | 2 +- tests/v1/kv_connector/unit/utils.py | 7 ++- vllm/config/vllm.py | 23 ++++++++ .../kv_connector/v1/nixl/base_scheduler.py | 52 +++++++++++++------ vllm/v1/core/kv_cache_utils.py | 7 +-- vllm/v1/core/sched/scheduler.py | 14 +---- vllm/v1/worker/gpu/model_runner.py | 10 +--- 7 files changed, 69 insertions(+), 46 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index a908d32462..7ce84bf685 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -138,7 +138,7 @@ class TestPushScheduler: with patch.object( sched, - "_truncate_mamba_request_for_prefill", + "_truncate_request_for_prefill", side_effect=AssertionError("must not truncate after cache lookup"), ): assert sched.get_num_new_matched_tokens(request, 0) == (0, False) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 1ee1e81776..98633c233b 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -521,6 +521,8 @@ def make_nixl_scheduler( Only sets the flags needed by the tests. When *heartbeat=True* the scheduler-side heartbeat bookkeeping fields are also initialised. """ + from types import SimpleNamespace + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( NixlConnectorScheduler, ) @@ -532,6 +534,7 @@ def make_nixl_scheduler( block_size=16, mamba_enabled=has_mamba, ) + sched.vllm_config = SimpleNamespace(num_prefill_lookahead_tokens=0) if heartbeat: sched._heartbeat_by_engine = {} @@ -598,9 +601,11 @@ def make_nixl_push_scheduler( mamba_enabled=has_mamba, ) - # vllm_config is consulted for parallel_config.tensor_parallel_size. + # vllm_config is consulted for parallel_config.tensor_parallel_size, and by + # `_prefill_backoff` on both the P and D prefill paths. vllm_config = MagicMock() vllm_config.parallel_config.tensor_parallel_size = 1 + vllm_config.num_prefill_lookahead_tokens = 0 sched.vllm_config = vllm_config # Push-specific state. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 8e0846982b..7167ca5a90 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -650,6 +650,29 @@ class VllmConfig: return self.num_speculative_tokens return 0 + @property + def num_prefill_lookahead_tokens(self) -> int: + """Prefill tokens past the computed range that the drafter reads. + + Mid-prefill the drafter consumes tokens the target model has not been + scheduled for yet, so every component that has to keep them available + must apply this margin: the scheduler, which never ends a chunk within + it and shifts encoder scheduling by it, and the KV cache manager, which + treats the trailing `this - 1` tokens as re-prefillable rather than + finalized. Consumers must read this property rather than re-deriving + their own per-method lookahead, so those components cannot drift apart. + """ + speculative_config = self.speculative_config + if speculative_config is None or not speculative_config.use_eagle(): + return 0 + if speculative_config.use_multi_module_mtp(): + # Each MTP module reads one token further ahead than the one before + # it, so the chain needs num_speculative_tokens of runway at a + # chunked-prefill boundary. + return self.num_speculative_tokens + # Eagle-family drafters read only the immediate next token. + return 1 + @property def uniform_decode_query_len(self) -> int: """Query length of every request in a uniform decode batch. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 9cfd100730..54184aa339 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -199,8 +199,8 @@ class NixlBaseConnectorScheduler: def on_new_request(self, request: "Request") -> None: """Track a request that may need heartbeats.""" params = request.kv_transfer_params - if params is not None and params.get("do_remote_decode") and self._has_mamba: - self._truncate_mamba_request_for_prefill(request) + if params is not None and params.get("do_remote_decode"): + self._truncate_request_for_prefill(request) # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are # effectively disabled for Bidirectional KV transfer. @@ -382,36 +382,56 @@ class NixlBaseConnectorScheduler: (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)], ts) ) + def _prefill_backoff(self) -> int: + """Trailing prompt tokens the prefiller must not compute; the decoder + recomputes them locally. + + Mamba needs h(N-1) so the decoder can derive h(N) itself. Multi-module + MTP needs to keep its whole lookahead window off of the prefiller, which + would otherwise embed the unverified drafts in the MTP layer's KV cache. The + decoder would never rebuild them, because the update is sized by the rejection + count, which is zero for the first decode. + """ + return max( + 1 if self._has_mamba else 0, + self.vllm_config.num_prefill_lookahead_tokens - 1, + ) + def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: - """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1).""" - if self._has_mamba and num_prompt_tokens > 1: - return num_prompt_tokens - 1 + """D-side only. The number of prompt tokens to load from the prefiller. + Stops short of the trailing ``_prefill_backoff()`` tokens that the decoder + will recompute locally.""" + backoff = self._prefill_backoff() + if backoff and num_prompt_tokens > backoff: + return num_prompt_tokens - backoff return num_prompt_tokens - def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: - """P-side only: drop the last prompt token so the prefiller computes - h(N-1) instead of h(N). The decoder recomputes the last token to - derive h(N) correctly. + def _truncate_request_for_prefill(self, request: "Request") -> None: + """P-side only: drop the trailing ``_prefill_backoff()`` prompt tokens + so the prefiller stops short of what the decoder recomputes locally. + For Mamba that is the single token needed to yield h(N-1); for + multi-module MTP it is the drafter's whole lookahead window. Guarded by ``_p_side_truncated`` to avoid repeated truncation if the request is preempted and rescheduled.""" + backoff = self._prefill_backoff() params = request.kv_transfer_params if ( - params is not None + backoff + and params is not None # Guard against repeated truncation after preemption/reschedule. and not params.get("_p_side_truncated") - and request.num_prompt_tokens > 1 + and request.num_prompt_tokens > backoff ): if request.prompt_token_ids is not None: - request.prompt_token_ids.pop() + del request.prompt_token_ids[-backoff:] elif request.prompt_embeds is not None: - request.prompt_embeds = request.prompt_embeds[:-1] + request.prompt_embeds = request.prompt_embeds[:-backoff] else: return - request._all_token_ids.pop() - request.num_prompt_tokens -= 1 + del request._all_token_ids[-backoff:] + request.num_prompt_tokens -= backoff request.max_tokens = 1 params["_p_side_truncated"] = True diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 1c6e102c16..c55c06abf1 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2678,12 +2678,7 @@ def get_kv_cache_configs( # When speculating with more than 1 speculative module (e.g. multi-layered MTP) # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = ( - vllm_config.speculative_config.num_speculative_tokens - 1 - if vllm_config.speculative_config is not None - and vllm_config.speculative_config.use_multi_module_mtp() - else 0 - ) + extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) for layer_name, layer_spec in merged_kv_cache_specs.items(): if isinstance(layer_spec, SlidingWindowSpec): merged_kv_cache_specs[layer_name] = replace( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index daf2b316d0..f1e11e9b1f 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -262,13 +262,7 @@ class Scheduler(SchedulerInterface): self.use_eagle_block_drop = False self.num_spec_tokens = vllm_config.num_speculative_tokens self.num_lookahead_tokens = vllm_config.num_lookahead_tokens - # Positions past the computed tokens that the drafter reads mid-prefill. - # Eagle-family drafters read 1 ahead, but multi-module MTP reads - # num_spec_tokens ahead at chunked-prefill boundaries. Determines the - # encoder scheduling shift, the deferred encoder free, the KV cache - # manager's re-prefillable window (this minus 1), and how many tokens to - # reserve between a chunk boundary and the prefill end. - self.num_prefill_lookahead = 0 + self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens self.dynamic_sd_lookup: list[int] | None = None if speculative_config is not None: if speculative_config.num_speculative_tokens_per_batch_size: @@ -278,12 +272,6 @@ class Scheduler(SchedulerInterface): vllm_num_speculative_tokens=self.num_spec_tokens, ) self.use_eagle = speculative_config.use_eagle() - if self.use_eagle: - self.num_prefill_lookahead = ( - self.num_spec_tokens - if speculative_config.use_multi_module_mtp() - else 1 - ) self.use_eagle_block_drop = speculative_config.use_eagle_block_drop() if self.use_eagle and not self.use_eagle_block_drop: logger.warning( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b66d552cb4..82cc697e73 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -296,15 +296,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.is_pooling_model = self.model_config.runner_type == "pooling" self.pooling_runner: PoolingRunner | None = None - # Multi-module MTP feeds its modules the next num_speculative_steps prefill - # tokens during chunked prefill. Other speculators only read the immediate - # next one. - num_prefill_lookahead = ( - self.num_speculative_steps - if self.speculative_config is not None - and self.speculative_config.use_multi_module_mtp() - else 1 - ) + num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens) self.step_timing = StepTimingCollector() From 1fd119def5a841ada882fa3f33919b96783f9d50 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 15 Sep 2026 19:31:51 -0500 Subject: [PATCH 123/139] [Pooling] Cap max-length padding for chunked embeddings (#56505) Signed-off-by: Taneem Ibrahim Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/entrypoints/pooling/base/protocol.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index e30195a7b1..c1366734f7 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -171,6 +171,8 @@ class FixedMaxLenTokenizeParamsMixin(PoolingTokenizeParamsMixin): class EmbeddingTokenizeParamsMixin(PoolingTokenizeParamsMixin): + padding: Literal["max_length", "do_not_pad"] | None + def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: default_max_total_tokens = model_config.max_model_len max_total_tokens: int | None = default_max_total_tokens @@ -184,7 +186,7 @@ class EmbeddingTokenizeParamsMixin(PoolingTokenizeParamsMixin): max_embed_len = pooler_config.max_embed_len or default_max_total_tokens max_output_tokens = default_max_total_tokens - max_embed_len - return self._build_pooling_tok_params( + tok_params = self._build_pooling_tok_params( model_config, add_special_tokens=self.add_special_tokens, max_total_tokens=max_total_tokens, @@ -192,6 +194,20 @@ class EmbeddingTokenizeParamsMixin(PoolingTokenizeParamsMixin): max_output_tokens_param="max_model_len - max_embed_len", ) + if ( + self.padding == "max_length" + and pooler_config is not None + and pooler_config.enable_chunked_processing + ): + max_padding_len = pooler_config.max_embed_len + if max_padding_len is None: + max_padding_len = default_max_total_tokens + return tok_params.with_kwargs( + pad_prompt_tokens=min(default_max_total_tokens, max_padding_len) + ) + + return tok_params + class CompletionRequestMixin(OpenAIBaseModel): # --8<-- [start:completion-params] From c8d1cf077a7878ea30040a33ac61262003a3953a Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:25:29 -0700 Subject: [PATCH 124/139] [ROCm] [Bugfix] Enable Load and Inference of GLM-5.3-Flash Quark MXFP4 Checkpoint (#56176) Signed-off-by: Colin Zeng Co-authored-by: Andreas Karatzas --- tests/quantization/test_quark.py | 145 ++++++++++++++++++ vllm/models/glm5next/__init__.py | 4 +- vllm/models/glm5next/common/__init__.py | 2 + .../glm5next/{nvidia => common}/attention.py | 0 .../models/glm5next/{nvidia => common}/kda.py | 0 .../glm5next/{nvidia => common}/model.py | 15 +- .../models/glm5next/{nvidia => common}/mtp.py | 2 +- .../glm5next/{nvidia => common}/multimodal.py | 0 8 files changed, 164 insertions(+), 4 deletions(-) create mode 100644 vllm/models/glm5next/common/__init__.py rename vllm/models/glm5next/{nvidia => common}/attention.py (100%) rename vllm/models/glm5next/{nvidia => common}/kda.py (100%) rename vllm/models/glm5next/{nvidia => common}/model.py (98%) rename vllm/models/glm5next/{nvidia => common}/mtp.py (99%) rename vllm/models/glm5next/{nvidia => common}/multimodal.py (100%) diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index ae236d71f8..9f00adf783 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -1668,3 +1668,148 @@ def test_suffix_match_at_module_boundary(prefix, ignored_layer, expected): ) is expected ) + + +_GLM5_MXFP4_WEIGHT = { + "dtype": "fp4", + "qscheme": "per_group", + "group_size": 32, + "scale_format": "e8m0", + "is_dynamic": False, +} +_GLM5_BLOCK_FP8_WEIGHT = { + "dtype": "fp8_e4m3", + "qscheme": "per_block", + "is_dynamic": False, + "block_size": [128, 128], + "symmetric": True, +} +_GLM5_BLOCK_FP8_INPUT = { + "dtype": "fp8_e4m3", + "qscheme": "per_group", + "is_dynamic": True, + "group_size": 128, + "symmetric": True, +} +_GLM5_GATE_UP = "language_model.model.layers.0.mlp.gate_up_proj" + + +def _glm5_mixed_precision_config() -> QuarkConfig: + """Global MXFP4 with the dense-MLP gate/up shards marked block-FP8.""" + fp8 = {"weight": _GLM5_BLOCK_FP8_WEIGHT, "input_tensors": _GLM5_BLOCK_FP8_INPUT} + return QuarkConfig( + { + "global_quant_config": { + "weight": _GLM5_MXFP4_WEIGHT, + "input_tensors": None, + }, + "layer_type_quant_config": {}, + "layer_quant_config": {"*mlp.gate_proj": fp8, "*mlp.up_proj": fp8}, + "exclude": [], + } + ) + + +class _GLM5RecordingParam: + """Stand-in for a BF16 target param that records what gets loaded.""" + + def __init__(self): + self.loaded_weight: torch.Tensor | None = None + self.loaded_shard: object = "unset" + + def weight_loader(self, param, loaded_weight, shard_id=None): + self.loaded_weight = loaded_weight + self.loaded_shard = shard_id + + +def _glm5_block_fp8(out_dim: int, in_dim: int, block: int = 128): + """Return a (fp8 weight, f32 per-block scale) pair.""" + weight = (torch.randn(out_dim, in_dim) * 0.1).to(torch.float8_e4m3fn) + scale = torch.rand(out_dim // block, in_dim // block, dtype=torch.float32) + 0.5 + return weight, scale + + +def test_glm5next_gate_up_proj_mapping_is_genuine_fusion(): + from vllm.models.glm5next.common.model import Glm5NextForConditionalGeneration + + mapping = Glm5NextForConditionalGeneration.packed_modules_mapping + assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"] + + +def test_glm5next_genuine_fusion_resolves_gate_up_to_block_fp8(): + # The fused module must expand to its real shard names so each resolves to + # the per-layer block-FP8 entry (rather than the global MXFP4 scheme). + config = _glm5_mixed_precision_config() + config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]} + _, _, scheme_cls = config.get_scheme_cls(LinearBase, _GLM5_GATE_UP) + assert scheme_cls is QuarkW8A8Fp8PerBlock + + +def test_glm5next_attn_loader_accepts_quark_weight_scale(): + # Regression for KeyError on '...kv_a_proj_with_mqa.weight_scale': the fused + # q_a/kv_a projection is kept BF16 and dequantized on load; the Quark scale + # name must be recognized and routed to the correct fused shard. + from vllm.models.glm5next.common.model import ( + _dequant_fp8_block, + _try_load_fp8_attn_proj, + ) + + prefix = "layers.3.self_attn" + target = _GLM5RecordingParam() + params_dict = {f"{prefix}.fused_qkv_a_proj.weight": target} + buf: dict = {} + loaded: set = set() + + weight, scale = _glm5_block_fp8(256, 256) + # fp8 weight arrives first -> buffered, nothing loaded yet. + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.kv_a_proj_with_mqa.weight", weight, buf, params_dict, loaded, 0 + ) + is True + ) + assert target.loaded_weight is None + # Quark-named scale completes the pair -> dequantize + load. + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.kv_a_proj_with_mqa.weight_scale", + scale, + buf, + params_dict, + loaded, + 0, + ) + is True + ) + assert target.loaded_weight is not None + assert target.loaded_weight.dtype == torch.bfloat16 + assert target.loaded_shard == 1 # kv_a is shard 1 of fused_qkv_a_proj + assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128)) + assert f"{prefix}.fused_qkv_a_proj.weight" in loaded + + +def test_glm5next_attn_loader_accepts_deepseek_weight_scale_inv(): + # DeepSeek scale name keeps working unchanged. + from vllm.models.glm5next.common.model import ( + _dequant_fp8_block, + _try_load_fp8_attn_proj, + ) + + prefix = "layers.7.self_attn" + target = _GLM5RecordingParam() + params_dict = {f"{prefix}.o_proj.weight": target} + buf: dict = {} + loaded: set = set() + + weight, scale = _glm5_block_fp8(128, 256) + _try_load_fp8_attn_proj( + f"{prefix}.o_proj.weight", weight, buf, params_dict, loaded, 0 + ) + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.o_proj.weight_scale_inv", scale, buf, params_dict, loaded, 0 + ) + is True + ) + assert target.loaded_shard is None # o_proj is a direct (non-fused) proj + assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128)) diff --git a/vllm/models/glm5next/__init__.py b/vllm/models/glm5next/__init__.py index a1146e2602..b38806260b 100644 --- a/vllm/models/glm5next/__init__.py +++ b/vllm/models/glm5next/__init__.py @@ -6,8 +6,8 @@ from vllm.platforms import current_platform if current_platform.is_xpu(): raise NotImplementedError("GLM-5.3-Flash does not currently support XPU.") -from .nvidia.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration -from .nvidia.mtp import Glm5NextMTP +from .common.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration +from .common.mtp import Glm5NextMTP __all__ = [ "Glm5NextForCausalLM", diff --git a/vllm/models/glm5next/common/__init__.py b/vllm/models/glm5next/common/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/glm5next/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/common/attention.py similarity index 100% rename from vllm/models/glm5next/nvidia/attention.py rename to vllm/models/glm5next/common/attention.py diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/common/kda.py similarity index 100% rename from vllm/models/glm5next/nvidia/kda.py rename to vllm/models/glm5next/common/kda.py diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/common/model.py similarity index 98% rename from vllm/models/glm5next/nvidia/model.py rename to vllm/models/glm5next/common/model.py index 21918bdf6b..e4f669835e 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/common/model.py @@ -1036,6 +1036,17 @@ class Glm5NextForConditionalGeneration( has_inner_state: ClassVar[Literal[True]] = True is_hybrid: ClassVar[Literal[True]] = True + # GLM-5.3-Flash stores the dense-MLP gate/up as separate tensors (like + # ``Glm4vMoeForConditionalGeneration``, ``glm4_moe`` and ``deepseek_v2``), + # so the fused ``gate_up_proj`` must expand to its real shard names for + # per-layer quant-scheme resolution. The identity ``gate_up_proj`` entry + # inherited from ``Glm4vForConditionalGeneration`` (pre-fused gate_up_proj) + # would otherwise route the module to ``global_quant_config`` and mismatch + # at load for mixed-precision Quark checkpoints. + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + } + # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration # (``model.visual.`` -> ``visual.``, ``model.language_model.`` -> # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``), @@ -1259,7 +1270,9 @@ def _try_load_fp8_attn_proj( return False suffix, (key, target_base, shard_id, is_kva) = matched is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn - is_scale = "weight_scale_inv" in name + # Need to accept both the DeepSeek-native ``weight_scale_inv`` and the Quark + # ``weight_scale`` names before feeding the shared block dequant below. + is_scale = "weight_scale_inv" in name or name.endswith(".weight_scale") if not is_weight and not is_scale: return False diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/common/mtp.py similarity index 99% rename from vllm/models/glm5next/nvidia/mtp.py rename to vllm/models/glm5next/common/mtp.py index 35b452f44f..33a5402b8f 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/common/mtp.py @@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import ( from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.glm5next.nvidia.ops.fused_eh_norm import fused_eh_norm from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -33,7 +34,6 @@ from .model import ( _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) -from .ops.fused_eh_norm import fused_eh_norm class Glm5NextMultiTokenPredictorLayer(nn.Module): diff --git a/vllm/models/glm5next/nvidia/multimodal.py b/vllm/models/glm5next/common/multimodal.py similarity index 100% rename from vllm/models/glm5next/nvidia/multimodal.py rename to vllm/models/glm5next/common/multimodal.py From a4d2d9d95edd2db39cc7dec0a3433c764ff2c578 Mon Sep 17 00:00:00 2001 From: ppalanga Date: Tue, 15 Sep 2026 18:38:02 -0700 Subject: [PATCH 125/139] [ROCm][Kimi-K3] Enable a4w4 flydsl kernels for KimiK3 (#53940) Signed-off-by: ppalanga Signed-off-by: Poovaiah Palangappa Co-authored-by: Cursor Co-authored-by: Poovaiah Palangappa Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> --- tests/kernels/moe/test_rocm_aiter_moe.py | 68 +++++++++++++++++++ vllm/_aiter_ops.py | 38 +++++++++-- vllm/envs.py | 20 ++++-- .../fused_moe/experts/rocm_aiter_moe.py | 12 ++-- .../layers/fused_moe/oracle/mxfp4.py | 19 +++--- 5 files changed, 127 insertions(+), 30 deletions(-) diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index 89244bed1b..0967f9499e 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -565,6 +565,74 @@ def test_aiter_moe_shared_experts_enablement_follows_env( assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected +def test_aiter_moe_situv2_syncs_aiter_a4w4_env( + monkeypatch: pytest.MonkeyPatch, +): + """SiTUv2 MoE should route to AITER a4w4 and clear legacy a8w4 overrides.""" + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.delenv("AITER_SITUV2_A8W4", raising=False) + mp.delenv("AITER_SITUV2_A4W4", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "1") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert os.environ.get("AITER_SITUV2_A4W4") == "1" + assert "AITER_SITUV2_A8W4" not in os.environ + + +def test_aiter_moe_situv2_legacy_a8w4_alias_enables_a4w4( + monkeypatch: pytest.MonkeyPatch, +): + """Deprecated VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 still enables SiTUv2 a4w4.""" + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.delenv("AITER_SITUV2_A8W4", raising=False) + mp.delenv("AITER_SITUV2_A4W4", raising=False) + mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "1") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert os.environ.get("AITER_SITUV2_A4W4") == "1" + assert "AITER_SITUV2_A8W4" not in os.environ + + +def test_aiter_moe_situv2_clears_aiter_env_when_disabled( + monkeypatch: pytest.MonkeyPatch, +): + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.setenv("AITER_SITUV2_A4W4", "1") + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert "AITER_SITUV2_A4W4" not in os.environ + + @pytest.mark.parametrize("moe_padding", [True, False]) def test_aiter_moe_padding_env_var( moe_padding: bool, diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 9bea4f8b3a..8f2b2a24ad 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1752,6 +1752,30 @@ def _mhc_delayed_pre_tail( _OPS_REGISTERED = False +def _sync_aiter_situv2_moe_env() -> None: + """Mirror the SiTUv2 MoE toggle into AITER's a4w4 dispatch env. + + AITER selects afp8 vs afp4 activation kernels via AITER_SITUV2_A8W4 / + AITER_SITUV2_A4W4 (see ROCm/aiter fused_moe.py, A8W4 checked first). + When VLLM_ROCM_USE_AITER_MOE_SITUV2 is enabled we route to a4w4 + (afp4_wfp4_fp4 kernels) and clear any legacy AITER_SITUV2_A8W4 override. + + Requires AITER with ROCm/aiter#4463 (first tagged in v0.1.20): a4w4 + dispatch plus kimik3_a4w4_{un,}tuned_fmoe.csv. Older AITER still runs + afp4 FlyDSL but falls back to heuristic configs, which is not the + tuned a4w4 path this recipe is meant to use. + """ + import os + + import vllm.envs as envs + + if envs.VLLM_ROCM_USE_AITER_MOE_SITUV2: + os.environ["AITER_SITUV2_A4W4"] = "1" + os.environ.pop("AITER_SITUV2_A8W4", None) + else: + os.environ.pop("AITER_SITUV2_A4W4", None) + + class rocm_aiter_ops: """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM. @@ -1774,7 +1798,7 @@ class rocm_aiter_ops: VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply. VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings. VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion. - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: Controls a8w4 SiTU fused MoE variant. + VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4). VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM. Note: @@ -1841,7 +1865,7 @@ class rocm_aiter_ops: # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - _MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 + _MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2 # TODO: Consolidate under _LINEAR_ENABLED _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM # Lazily probed: whether aiter.topk_softmax supports the @@ -1870,7 +1894,8 @@ class rocm_aiter_ops: cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - cls._MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 + cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2 + _sync_aiter_situv2_moe_env() cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM @staticmethod @@ -1983,10 +2008,10 @@ class rocm_aiter_ops: @classmethod @if_aiter_supported - def is_fused_moe_situv2_a8w4_enabled(cls) -> bool: - # _MOE_SITUV2_A8W4 is a variant of aiter fused moe, so aiter + def is_fused_moe_situv2_enabled(cls) -> bool: + # _MOE_SITUV2 is a variant of aiter fused moe, so aiter # fused moe must be enabled as well. - return cls.is_fused_moe_enabled() and cls._MOE_SITUV2_A8W4 + return cls.is_fused_moe_enabled() and cls._MOE_SITUV2 @classmethod @if_aiter_supported @@ -4029,3 +4054,4 @@ class rocm_aiter_ops: rocm_aiter_ops.register_ops_once() +_sync_aiter_situv2_moe_env() diff --git a/vllm/envs.py b/vllm/envs.py index 6158f86f89..78af9bcb72 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -137,7 +137,7 @@ if TYPE_CHECKING: VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False VLLM_ROCM_USE_AITER_MOE: bool = True VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False + VLLM_ROCM_USE_AITER_MOE_SITUV2: bool = False VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" @@ -1251,12 +1251,18 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER_MOE": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1") ), - # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up- - # interleaved flydsl kernels instead of the default a16w4 separated path. - # This is the only flag users need: vLLM picks the kernels by passing - # gate_mode to AITER and sets the AITER-side workaround env at init. - "VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4": lambda: ( - os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() in ("true", "1") + # Route K3 SiTU MXFP4 MoE through the FlyDSL SiTUv2 path (a4w4 fp4 + # activations, separated gate/up layout) instead of default a16w4. vLLM + # sets AITER_SITUV2_A4W4 at init when this flag is on and clears any + # legacy AITER_SITUV2_A8W4 override (AITER checks A8W4 first). + # VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 is a deprecated alias for existing + # recipes; it does not select a8w4 kernels. + # Needs AITER >= v0.1.20 (ROCm/aiter#4463) for the a4w4 dispatch flag + # and tuned kimik3_a4w4_*_fmoe.csv rows; otherwise FlyDSL uses heuristics. + "VLLM_ROCM_USE_AITER_MOE_SITUV2": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0").lower() in ("true", "1") + or os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() + in ("true", "1") ), # MoE sorting dispatch policy for AITER fused MoE kernels. # 0 = auto (default): single-pass for small batches, multi-pass diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index cb4225c92a..62df395f69 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -385,14 +385,10 @@ def rocm_aiter_fused_experts( gate_mode = "" if activation == MoEActivation.SITU: - # a8w4 (VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1) uses the gate/up- - # interleaved (_gui_) fp8 flydsl kernels; default a16w4 SiTU stays - # separated. - gate_mode = ( - GateMode.INTERLEAVE.value - if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled() - else GateMode.SEPARATED.value - ) + # SiTUv2 flydsl (VLLM_ROCM_USE_AITER_MOE_SITUV2=1) uses a4w4 + # fp4 activations with separated gate/up weights (AITER #4463); + # default a16w4 SiTU also stays separated. + gate_mode = GateMode.SEPARATED.value elif quant_config.use_mxfp4_w4a16: gate_mode = GateMode.INTERLEAVE.value elif activation_interleave is not None: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3b415acd81..1c38e4378b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1572,12 +1572,6 @@ def convert_weight_to_mxfp4_moe_kernel_format( if w2_bias is not None: w2_bias = w2_bias.data.to(torch.float32) - import os - - # TODO: Remove this once AITER is fixed - # Necessary for AITER side from crashing - os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" - if activation == MoEActivation.SITU: from aiter.utility.fp4_utils import e8m0_shuffle @@ -1585,9 +1579,8 @@ def convert_weight_to_mxfp4_moe_kernel_format( fp4_dtype = torch.float4_e2m1fn_x2 e8m0_dtype = torch.float8_e8m0fnu - # a8w4 uses gate/up-interleaved flydsl kernels; - # default a16w4 keeps the separated layout. - guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled() + # SiTUv2 flydsl uses separated gate/up layout (a4w4). + guinterleave = False w13 = rocm_aiter_ops.shuffle_weight_a16w4( w13_weight.data.view(fp4_dtype), 16, guinterleave ) @@ -1606,6 +1599,14 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2.is_shuffled = True return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias) + import os + + # Interleaved a16w4 only (DeepSeekV4 etc.). AITER uses this bound to + # pick bf16 vs fp8 activations when gate_mode is INTERLEAVE. SiTUv2 + # a4w4 is separated and selects q_dtype_a independently, so the bound + # is unused on that path. + os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" + from aiter.ops.shuffle import shuffle_scale as _shuf_s from aiter.ops.shuffle import shuffle_weight as _shuf_w From b0898a493737ed34864efb2dbebc36a73ea043c5 Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Wed, 16 Sep 2026 10:02:37 +0800 Subject: [PATCH 126/139] [Misc] Allow install empty package from pip (#41074) Signed-off-by: wangxiyuan --- setup.py | 60 ++++++++++++++++++++++++++++------------------------ vllm/envs.py | 5 +++++ 2 files changed, 37 insertions(+), 28 deletions(-) diff --git a/setup.py b/setup.py index 13692f3d20..4012000e63 100644 --- a/setup.py +++ b/setup.py @@ -1257,35 +1257,39 @@ def get_vllm_version() -> str: version = get_version(write_to="vllm/_version.py") sep = "+" if "+" not in version else "." # dev versions might contain + - if _no_device(): - if envs.VLLM_TARGET_DEVICE == "empty": - version += f"{sep}empty" - elif _is_cuda(): - if USE_PRECOMPILED_EXTENSIONS and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX: - version += f"{sep}precompiled" + if not envs.VLLM_SKIP_VERSION_SUFFIX: + if _no_device(): + if envs.VLLM_TARGET_DEVICE == "empty": + version += f"{sep}empty" + elif _is_cuda(): + if ( + USE_PRECOMPILED_EXTENSIONS + and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX + ): + version += f"{sep}precompiled" + else: + cuda_version = str(get_nvcc_cuda_version()) + if cuda_version != envs.VLLM_MAIN_CUDA_VERSION: + cuda_version_str = cuda_version.replace(".", "")[:3] + # skip this for source tarball, required for pypi + if "sdist" not in sys.argv: + version += f"{sep}cu{cuda_version_str}" + elif _is_hip(): + # Get the Rocm Version + rocm_version = get_rocm_version() or torch.version.hip + if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION: + version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" + elif _is_tpu(): + version += f"{sep}tpu" + elif _is_cpu(): + # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above), + # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. + if VLLM_TARGET_DEVICE == "cpu": + version += f"{sep}cpu" + elif _is_xpu(): + version += f"{sep}xpu" else: - cuda_version = str(get_nvcc_cuda_version()) - if cuda_version != envs.VLLM_MAIN_CUDA_VERSION: - cuda_version_str = cuda_version.replace(".", "")[:3] - # skip this for source tarball, required for pypi - if "sdist" not in sys.argv: - version += f"{sep}cu{cuda_version_str}" - elif _is_hip(): - # Get the Rocm Version - rocm_version = get_rocm_version() or torch.version.hip - if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION: - version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" - elif _is_tpu(): - version += f"{sep}tpu" - elif _is_cpu(): - # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above), - # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. - if VLLM_TARGET_DEVICE == "cpu": - version += f"{sep}cpu" - elif _is_xpu(): - version += f"{sep}xpu" - else: - raise RuntimeError("Unknown runtime environment") + raise RuntimeError("Unknown runtime environment") return version diff --git a/vllm/envs.py b/vllm/envs.py index 78af9bcb72..bc4724e2e4 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2176,6 +2176,10 @@ environment_variables: dict[str, Callable[[], Any]] = { # Each op additionally checks its own shape / dtype constraints and falls # back to the eager path when they do not hold. "VLLM_ENABLE_HPC_OPS": lambda: bool(int(os.getenv("VLLM_ENABLE_HPC_OPS", "0"))), + # Whether to skip version suffix when building package + "VLLM_SKIP_VERSION_SUFFIX": lambda: bool( + int(os.getenv("VLLM_SKIP_VERSION_SUFFIX", "0")) + ), } @@ -2344,6 +2348,7 @@ def compile_factors() -> dict[str, object]: "LOCAL_RANK", "CUDA_VISIBLE_DEVICES", "NO_COLOR", + "VLLM_SKIP_VERSION_SUFFIX", } from vllm.config.utils import normalize_value From 7d4bba54985589176a60cbf3ad4424be3ee84999 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Tue, 15 Sep 2026 22:48:18 -0400 Subject: [PATCH 127/139] [Doc] Show how to get Responses prompt token IDs (#55912) Signed-off-by: Francisco Javier Arceo Co-authored-by: OpenAI Codex --- docs/serving/online_serving/README.md | 7 +++- .../openai_compatible_server.md | 5 +++ docs/serving/online_serving/renderer.md | 35 +++++++++++++++++-- 3 files changed, 44 insertions(+), 3 deletions(-) diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 1f5d4bdb6a..a292a4a3a4 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -130,7 +130,12 @@ Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out` ### Renderer APIs -For further details on renderer APIs, please refer to [this page](renderer.md). +Renderer APIs preprocess completion, chat, and Responses requests without running +inference. They handle tokenization, model-specific prompt formatting, and +multimodal preprocessing, returning prompt token IDs, sampling parameters, and +any processed multimodal inputs for generation. + +See the [renderer guide](renderer.md) for setup instructions and examples. - [Completions Render API](renderer.md) (`/v1/completions/render`) - Render completion requests diff --git a/docs/serving/online_serving/openai_compatible_server.md b/docs/serving/online_serving/openai_compatible_server.md index 22e2e9a608..c50910040d 100644 --- a/docs/serving/online_serving/openai_compatible_server.md +++ b/docs/serving/online_serving/openai_compatible_server.md @@ -199,6 +199,11 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai Code example: [examples/tool_calling/openai_responses_client_with_tools.py](../../../examples/tool_calling/openai_responses_client_with_tools.py) +To get prompt token IDs before generation, use +[`/v1/responses/render`](renderer.md#get-responses-prompt-token-ids). It preprocesses +a self-contained Responses request without inference and returns the prompt +token IDs, so a separate `/tokenize` call is not needed. + #### Extra parameters The following extra parameters in the request object are supported: diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index 9be23142c4..9d494df22e 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -26,27 +26,58 @@ vllm serve --enable-scale-out - [Responses Render API](renderer.md) (`/v1/responses/render`) - Render a self-contained Responses request +## Get Responses prompt token IDs + +Use `/v1/responses/render` to get prompt token IDs before choosing a model +replica. Rendering applies prompt construction and preprocessing without running +inference. + The Responses render endpoint uses the same prompt construction as `/v1/responses` and returns one token-in `GenerateRequest`. It is stateless: inline history is supported, but `previous_response_id` is not. Callers must resolve stored response state and include the resulting history in the request before rendering. +Configure the renderer and generation workers with the same model, tokenizer, +chat template, and preprocessing options. Render the full request, including +instructions, history, tools, and any template or truncation options, so the +returned IDs reflect the prompt the model will receive. + For multimodal requests, the `GenerateRequest` contains the model-processed multimodal payload, which can be substantially larger than the source image or video. The caller must forward that payload unchanged to the generation service and provision transport limits and memory accordingly. +For example, start a standard inference server with scale-out endpoints enabled: + ```bash -curl http://localhost:8000/v1/responses/render \ +VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + vllm serve meta-llama/Llama-3.1-8B-Instruct +``` + +Send a Responses request and use `jq` to extract the `token_ids` field: + +```bash +curl --fail --silent --show-error http://localhost:8000/v1/responses/render \ -H "Content-Type: application/json" \ -d '{ "model": "meta-llama/Llama-3.1-8B-Instruct", "input": "Explain prefix caching in one sentence.", "max_output_tokens": 32 - }' + }' | jq '.token_ids' ``` +To count the rendered prompt tokens for this text request, replace the `jq` +filter with `'.token_ids | length'`. The endpoint returns the full +`GenerateRequest`; `jq` filters the response on the client. Keep the complete +response if you will forward it to `/inference/v1/generate`, including any +multimodal features. + +If the server has `--api-key` or `VLLM_API_KEY` configured, add +`-H "Authorization: Bearer "` to the request. See +[API key authentication limitations](../../usage/security.md#api-key-authentication-limitations) +for the existing authentication boundaries. + For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md). ## Multimodal Render Features From 0673de980055a4283fcc152fc1ecc837ff534f1d Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:51:17 +0800 Subject: [PATCH 128/139] [Rust Frontend][Metrics] Add per-request preemption histogram (#57033) Co-authored-by: Bugen Zhao Signed-off-by: reidliu41 --- rust/src/llm/src/request_metrics.rs | 6 ++++++ rust/src/llm/tests/generate.rs | 12 ++++++++++++ rust/src/metrics/src/request.rs | 15 +++++++++++++++ 3 files changed, 33 insertions(+) diff --git a/rust/src/llm/src/request_metrics.rs b/rust/src/llm/src/request_metrics.rs index 280fa0f10f..43802eaf4a 100644 --- a/rust/src/llm/src/request_metrics.rs +++ b/rust/src/llm/src/request_metrics.rs @@ -45,6 +45,7 @@ pub(crate) struct RequestMetricsTracker { last_token_ts: f64, first_token_latency: f64, num_generation_tokens: u32, + num_preemptions: u64, latest_num_cached_tokens: u32, } @@ -69,6 +70,7 @@ struct RequestMetricHandles { request_max_num_generation_tokens: HistogramMetric, request_params_max_tokens: HistogramMetric, request_params_n: HistogramMetric, + request_num_preemptions: HistogramMetric, request_prefill_kv_computed_tokens: HistogramMetric, time_to_first_token_seconds: HistogramMetric, inter_token_latency_seconds: HistogramMetric, @@ -104,6 +106,7 @@ impl RequestMetricsTracker { last_token_ts: 0.0, first_token_latency: 0.0, num_generation_tokens: 0, + num_preemptions: 0, latest_num_cached_tokens: 0, } } @@ -184,6 +187,7 @@ impl RequestMetricsTracker { self.handles.request_params_max_tokens.observe(max_tokens_param as f64); } self.handles.request_params_n.observe(self.n_param as f64); + self.handles.request_num_preemptions.observe(self.num_preemptions as f64); self.handles .request_prefill_kv_computed_tokens .observe(prefill_kv_computed_tokens as f64); @@ -223,6 +227,7 @@ impl RequestMetricsTracker { } } EngineCoreEventType::Preempted => { + self.num_preemptions += 1; self.handles.num_preemptions.inc(); } } @@ -276,6 +281,7 @@ fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetri .get_or_create_owned(&labels), request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels), request_params_n: metrics.request_params_n.get_or_create_owned(&labels), + request_num_preemptions: metrics.request_num_preemptions.get_or_create_owned(&labels), request_prefill_kv_computed_tokens: metrics .request_prefill_kv_computed_tokens .get_or_create_owned(&labels), diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 2ff3faa8d4..2d09aff889 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -790,6 +790,12 @@ async fn generate_records_request_metrics_in_prometheus_output() { assert!(rendered.contains(&format!( "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1" ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"4\"}} 1.0" + ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" + ))); assert!(rendered.contains(&format!( "vllm:time_to_first_token_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" ))); @@ -883,6 +889,12 @@ async fn dropping_stream_records_abort_terminal_request_metrics() { assert!(rendered.contains(&format!( "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"5\"}} 1" ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"5\"}} 0.0" + ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"5\"}} 1" + ))); llm.shutdown().await.unwrap(); } diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index 691c789936..5bfc32e895 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -21,6 +21,7 @@ const REQUEST_LATENCY_BUCKETS: [f64; 21] = [ 480.0, 960.0, 1920.0, 7680.0, ]; const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0]; +const REQUEST_NUM_PREEMPTIONS_BUCKETS: [f64; 7] = [1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 20.0]; const ITERATION_TOKENS_BUCKETS: [f64; 13] = [ 1.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, 16384.0, ]; @@ -69,6 +70,10 @@ fn request_params_n_histogram() -> Histogram { Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied()) } +fn request_num_preemptions_histogram() -> Histogram { + Histogram::new(REQUEST_NUM_PREEMPTIONS_BUCKETS.iter().copied()) +} + fn iteration_tokens_histogram() -> Histogram { Histogram::new(ITERATION_TOKENS_BUCKETS.iter().copied()) } @@ -110,6 +115,7 @@ pub struct RequestMetrics { pub request_max_num_generation_tokens: HistogramFamily, pub request_params_max_tokens: HistogramFamily, pub request_params_n: HistogramFamily, + pub request_num_preemptions: HistogramFamily, pub request_prefill_kv_computed_tokens: HistogramFamily, pub time_to_first_token_seconds: HistogramFamily, pub inter_token_latency_seconds: HistogramFamily, @@ -216,6 +222,14 @@ impl RequestMetrics { request_params_n.clone(), ); + let request_num_preemptions = + Family::new_with_constructor(request_num_preemptions_histogram as fn() -> Histogram); + registry.register( + "vllm:request_num_preemptions", + "Histogram of the number of times a request was preempted.", + request_num_preemptions.clone(), + ); + let request_prefill_kv_computed_tokens = Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram); registry.register( @@ -302,6 +316,7 @@ impl RequestMetrics { request_max_num_generation_tokens, request_params_max_tokens, request_params_n, + request_num_preemptions, request_prefill_kv_computed_tokens, time_to_first_token_seconds, inter_token_latency_seconds, From bfd713bf8b3869a7b65a8f3198095c3d73c88875 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 16 Sep 2026 10:52:55 +0800 Subject: [PATCH 129/139] [Rust Frontend] Support HF config overrides `--hf-overrides` (#56931) Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 1 + rust/src/chat/src/backend/hf.rs | 22 ++--- rust/src/chat/src/backend/mod.rs | 3 + rust/src/cmd/Cargo.toml | 1 + rust/src/cmd/src/cli.rs | 17 ++++ rust/src/cmd/src/cli/tests.rs | 85 +++++++++++++++++ rust/src/cmd/src/cli/unsupported.rs | 5 - rust/src/managed-engine/src/cli.rs | 5 + .../examples/external_engine_openai_qwen.rs | 1 + rust/src/server/src/config.rs | 3 + rust/src/server/src/lib.rs | 1 + rust/src/server/src/render.rs | 4 + rust/src/text/Cargo.toml | 2 +- rust/src/text/src/backend/hf/mod.rs | 2 + rust/src/text/src/backend/hf/model_files.rs | 72 ++++++++++++++ rust/src/text/src/backend/hf/overrides.rs | 93 +++++++++++++++++++ 16 files changed, 298 insertions(+), 19 deletions(-) create mode 100644 rust/src/text/src/backend/hf/overrides.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 1b7dae6568..54a1640926 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5813,6 +5813,7 @@ dependencies = [ "vllm-engine-core-client", "vllm-managed-engine", "vllm-server", + "vllm-text", "vllm-tracing", ] diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 37212a9650..5b5b8772c5 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -129,7 +129,8 @@ pub(super) async fn load_model_backends( model_id: &str, options: LoadModelBackendsOptions, ) -> Result { - let files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?; + let mut files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?; + files.apply_overrides(&options.hf_overrides)?; let text_backend = HfTextBackend::from_resolved_model_files( files.clone(), model_id.to_string(), @@ -166,13 +167,11 @@ fn resolve_multimodal_render_info( #[cfg(test)] mod tests { use std::collections::HashMap; - use std::path::PathBuf; use std::sync::Arc; use tempfile::tempdir; use thiserror_ext::AsReport as _; use vllm_text::Prompt; - use vllm_text::backend::hf::TokenizerSource; use vllm_text::tokenizer::DynTokenizer; use vllm_tokenizer::test_utils::TestTokenizer; @@ -206,16 +205,12 @@ mod tests { write_json(&config_path, config_json); write_json(&tokenizer_config_path, tokenizer_config_json); - vllm_text::backend::hf::ResolvedModelFiles { - tokenizer: TokenizerSource::HuggingFace(PathBuf::from("/tmp/unused-tokenizer.json")), - tokenizer_config_path: Some(tokenizer_config_path), - generation_config_path: None, - preprocessor_config_path: None, - video_preprocessor_config_path: None, - processor_config_path: None, - chat_template_path: None, - config_path: Some(config_path), - } + write_json(&root.join("tokenizer.json"), "{}"); + futures::executor::block_on(vllm_text::backend::hf::ResolvedModelFiles::new( + root.to_str().unwrap(), + None, + )) + .unwrap() } fn test_tokenizer() -> DynTokenizer { @@ -232,6 +227,7 @@ mod tests { "test-model".to_string(), LoadModelBackendsOptions { revision: None, + hf_overrides: Default::default(), generation_config: Default::default(), renderer, language_model_only: false, diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index 028c0c40c8..eed4e95cbe 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; use std::sync::Arc; use serde_json::Value; +use vllm_text::backend::hf::HfOverrides; use vllm_text::{DynTextBackend, GenerationConfigMode, TextBackend}; use crate::error::Result; @@ -63,6 +64,8 @@ pub type DynChatTextBackend = Arc; pub struct LoadModelBackendsOptions { /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). pub revision: Option, + /// JSON Merge Patch applied to the model config before loading any backend. + pub hf_overrides: HfOverrides, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Which chat renderer implementation to use. diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index bf92e05ee1..e61c48ee3c 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -33,6 +33,7 @@ vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-managed-engine.workspace = true vllm-server.workspace = true +vllm-text.workspace = true vllm-tracing.workspace = true [dev-dependencies] diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 70b28018a3..cda2e6c1a8 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -33,6 +33,7 @@ use vllm_server::{ DEFAULT_KEEP_ALIVE_TIMEOUT, HttpListenerMode, LoraModulePath, ParserSelection, RenderConfig, RendererSelection, }; +use vllm_text::backend::hf::HfOverrides; use crate::cli::ssl::SslArgs; use crate::cli::unsupported::UnsupportedArgs; @@ -113,6 +114,9 @@ pub struct RenderArgs { /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). #[arg(long)] revision: Option, + /// JSON Merge Patch (RFC 7396) for config.json; null removes a field. + #[arg(long, value_parser = parse_json::, default_value = "{}", value_name = "JSON")] + hf_overrides: HfOverrides, /// HTTP bind host. #[arg(long, default_value = "127.0.0.1")] host: String, @@ -161,6 +165,7 @@ impl RenderArgs { RenderConfig { model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, served_model_name: self.served_model_name, host: self.host, port: self.port, @@ -204,6 +209,12 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub revision: Option, + /// JSON Merge Patch (RFC 7396) for config.json; null removes a field. + /// Objects merge recursively and arrays/scalars replace existing values. + #[arg(long, value_parser = parse_json::, default_value = "{}", value_name = "JSON")] + #[serde(default)] + pub hf_overrides: HfOverrides, + /// The source of generation-config sampling defaults. `"auto"` loads the /// model's defaults, while `"vllm"` uses vLLM's neutral defaults. #[arg(long, default_value_t)] @@ -500,6 +511,7 @@ impl SharedRuntimeArgs { }, model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd }, @@ -556,6 +568,7 @@ impl SharedRuntimeArgs { coordinator_mode: CoordinatorMode::MaybeInProc, model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode, @@ -761,6 +774,9 @@ impl ServeArgs { let reasoning_parser = effective_engine_reasoning_parser(&self.runtime.reasoning_parser, &self.runtime.model); let profiler_config = self.runtime.profiler_config_json(); + let hf_overrides = (!self.runtime.hf_overrides.is_empty()).then(|| { + serde_json::to_string(&self.runtime.hf_overrides).expect("JSON object serializes") + }); self.managed_engine.clone().into_config( self.runtime.model.clone(), @@ -773,6 +789,7 @@ impl ServeArgs { self.runtime.shutdown_timeout, handshake_port, self.runtime.limit_mm_per_prompt_json(), + hf_overrides, ) } } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 411d256295..5cc6b19214 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -177,6 +177,9 @@ fn serve_args_forward_python_flags_with_separator() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, @@ -1013,6 +1016,9 @@ fn frontend_args_accept_json() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: None, @@ -1120,6 +1126,73 @@ fn frontend_args_json_ignores_engine_owned_max_model_len() { assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B"); } +#[test] +fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() { + let patch = r#"{"text_config":{"rope_parameters":{"factor":4}},"sliding_window":null}"#; + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--hf-overrides", + patch, + ]) + .unwrap(); + let Command::Serve(args) = cli.command else { + panic!("expected serve args") + }; + let expected: serde_json::Value = serde_json::from_str(patch).unwrap(); + assert_eq!( + serde_json::to_value(&args.runtime.hf_overrides).unwrap(), + expected + ); + let engine = args.to_managed_engine_config(1234); + let index = engine.python_args.iter().position(|arg| arg == "--hf-overrides").unwrap(); + assert_eq!( + serde_json::from_str::(&engine.python_args[index + 1]).unwrap(), + expected + ); + assert!(!engine.python_args.iter().any(|arg| arg == "--hf-config-path")); + let config = args.to_frontend_config("tcp://127.0.0.1:1234".to_string()); + assert_eq!(serde_json::to_value(config.hf_overrides).unwrap(), expected); + + let payload = serde_json::json!({"model_tag":"Qwen/Qwen3-0.6B", "hf_overrides":expected}); + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + &payload.to_string(), + ]) + .unwrap(); + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args") + }; + assert_eq!( + serde_json::to_value(args.into_config().hf_overrides).unwrap(), + expected + ); +} + +#[test] +fn hf_overrides_reject_non_object_inputs() { + for patch in ["null", "[]", "1", r#""callable""#] { + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--hf-overrides", + patch, + ]) + .unwrap_err(); + assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation); + } +} + #[test] fn frontend_args_json_accepts_supported_non_default_fields() { let cli = Cli::try_parse_from([ @@ -1613,6 +1686,9 @@ fn serve_args_accept_handshake_aliases() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, @@ -1763,6 +1839,9 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: BindTcp { @@ -1852,6 +1931,9 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: BindTcp { @@ -1963,6 +2045,9 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present }, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: InheritedFd { diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 29ea5d06b0..8e514c9534 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -235,11 +235,6 @@ pub struct EngineUnsupportedArgs { #[arg(long, default_missing_value = "true", num_args = 0..=1)] pub hf_token: Option, - /// If a dictionary, contains arguments to be forwarded to the Hugging Face - /// config. If a callable, it is called to update the HuggingFace config. - #[arg(long)] - pub hf_overrides: Option, - /// Overrides or sets generation config. e.g. `{"temperature": 0.5}`. If /// used with `--generation-config auto`, the override parameters will be /// merged with the default config from the model. If used with diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index f0cd7bf171..0a81ae598e 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -93,6 +93,7 @@ impl ManagedEngineArgs { shutdown_timeout: u64, handshake_port: u16, limit_mm_per_prompt: Option, + hf_overrides: Option, ) -> ManagedEngineConfig { let mut python_args = self.python_args; // Manually forward some args to the Python engine. @@ -136,6 +137,10 @@ impl ManagedEngineArgs { python_args.push("--limit-mm-per-prompt".to_string()); python_args.push(limit_mm_per_prompt); } + if let Some(hf_overrides) = hf_overrides { + python_args.push("--hf-overrides".to_string()); + python_args.push(hf_overrides); + } ManagedEngineConfig { python: self.python, diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index fc9914e3e2..ad3872843a 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -61,6 +61,7 @@ async fn main() -> Result<()> { coordinator_mode: CoordinatorMode::MaybeInProc, model: args.model, revision: None, + hf_overrides: Default::default(), generation_config: Default::default(), served_model_name: vec![], listener_mode: HttpListenerMode::BindTcp { diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 0f0deac86d..4c337adbef 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -16,6 +16,7 @@ use vllm_chat::{ ChatTemplateContentFormatOption, GenerationConfigMode, ParserSelection, RendererSelection, }; use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; +use vllm_text::backend::hf::HfOverrides; /// Default keep-alive idle timeout (seconds); also the head-read bound /// when keep-alive is disabled (`0`). @@ -214,6 +215,8 @@ pub struct Config { pub model: String, /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). pub revision: Option, + /// JSON Merge Patch applied to the model config before backend initialization. + pub hf_overrides: HfOverrides, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Model name(s) exposed to clients via the OpenAI API. When non-empty, diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 95ae359f35..49d3d57f32 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -99,6 +99,7 @@ async fn build_state(config: &Config) -> Result> { &config.model, LoadModelBackendsOptions { revision: config.revision.clone(), + hf_overrides: config.hf_overrides.clone(), generation_config: config.generation_config, renderer: config.renderer, language_model_only: config.language_model_only, diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs index fe7033fa8c..4fc55fe60b 100644 --- a/rust/src/server/src/render.rs +++ b/rust/src/server/src/render.rs @@ -13,6 +13,7 @@ use vllm_chat::{ ParserSelection, RendererSelection, load_model_backends, }; use vllm_text::TextRequestProcessor; +use vllm_text::backend::hf::HfOverrides; use crate::{ HttpListenerMode, TlsConfig, @@ -25,6 +26,7 @@ use crate::{ pub struct RenderConfig { pub model: String, pub revision: Option, + pub hf_overrides: HfOverrides, pub served_model_name: Vec, pub host: String, pub port: u16, @@ -69,6 +71,7 @@ async fn build_state(config: &RenderConfig) -> Result> { &config.model, LoadModelBackendsOptions { revision: config.revision.clone(), + hf_overrides: config.hf_overrides.clone(), generation_config: Default::default(), renderer: config.renderer, language_model_only: true, @@ -153,6 +156,7 @@ mod tests { RenderConfig { model: "test-model".to_string(), revision: None, + hf_overrides: Default::default(), served_model_name: Vec::new(), host: "127.0.0.1".to_string(), port: 8000, diff --git a/rust/src/text/Cargo.toml b/rust/src/text/Cargo.toml index 7bda7f976e..36152d7cc4 100644 --- a/rust/src/text/Cargo.toml +++ b/rust/src/text/Cargo.toml @@ -16,6 +16,7 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true serde_with.workspace = true +tempfile.workspace = true thiserror.workspace = true thiserror-ext.workspace = true tracing.workspace = true @@ -28,7 +29,6 @@ vllm-tokenizer.workspace = true expect-test.workspace = true futures.workspace = true serial_test.workspace = true -tempfile.workspace = true tokio.workspace = true vllm-llm = { workspace = true, features = ["test-util"] } vllm-tokenizer = { workspace = true, features = ["test-utils"] } diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index ae768767d9..2f32f2cb48 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -3,6 +3,7 @@ mod config; mod model_files; +mod overrides; use std::collections::BTreeSet; use std::sync::Arc; @@ -18,6 +19,7 @@ pub use self::config::{ load_tokenizer_config, }; pub use self::model_files::{ResolvedModelFiles, TokenizerSource}; +pub use self::overrides::HfOverrides; use crate::backend::{GenerationConfigMode, SamplingHints, TextBackend}; use crate::error::Result; diff --git a/rust/src/text/src/backend/hf/model_files.rs b/rust/src/text/src/backend/hf/model_files.rs index de802fb439..3ffd7b2bb7 100644 --- a/rust/src/text/src/backend/hf/model_files.rs +++ b/rust/src/text/src/backend/hf/model_files.rs @@ -2,11 +2,13 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::path::{Path, PathBuf}; +use std::sync::Arc; use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo}; use hf_hub::{Cache, Repo, RepoType}; use thiserror_ext::AsReport as _; +use super::HfOverrides; use super::config::{HfTokenizerConfig, load_tokenizer_config}; use crate::error::{Error, Result}; @@ -50,9 +52,27 @@ pub struct ResolvedModelFiles { pub processor_config_path: Option, pub chat_template_path: Option, pub config_path: Option, + /// Keep a patched config alive across clones and backend ownership. + config_temp_path: Option>, } impl ResolvedModelFiles { + /// Materialize a merge-patched config owned by these resolved model files. + pub fn apply_overrides(&mut self, overrides: &HfOverrides) -> Result<()> { + if overrides.is_empty() { + return Ok(()); + } + let path = materialize_config(self.config_path.as_deref(), overrides).map_err(|error| { + Error::Tokenizer(format!( + "failed to materialize HF overrides: {}", + error.as_report() + )) + })?; + self.config_path = Some(path.to_path_buf()); + self.config_temp_path = Some(Arc::new(path)); + Ok(()) + } + /// Resolve tokenizer/config files from a local model directory first when /// `model_id` points to one, otherwise consult the local HF cache and /// finally the Hub. @@ -84,6 +104,24 @@ impl ResolvedModelFiles { } } +fn materialize_config( + path: Option<&Path>, + overrides: &HfOverrides, +) -> anyhow::Result { + let mut config: serde_json::Value = match path { + Some(path) => serde_json::from_slice(&std::fs::read(path)?)?, + None => serde_json::json!({}), + }; + anyhow::ensure!(config.is_object(), "HF model config must be a JSON object"); + overrides.apply(&mut config); + let file = tempfile::Builder::new() + .prefix("vllm-rs-hf-config-") + .suffix(".json") + .tempfile()?; + serde_json::to_writer(file.as_file(), &config)?; + Ok(file.into_temp_path()) +} + fn resolve_local_model_files(model_dir: &Path) -> Result { let tokenizer_config_path = local_file_if_exists(model_dir, "tokenizer_config.json"); let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?; @@ -101,6 +139,7 @@ fn resolve_local_model_files(model_dir: &Path) -> Result { processor_config_path: local_file_if_exists(model_dir, "processor_config.json"), chat_template_path: discover_chat_template_in_dir(model_dir), config_path: local_file_if_exists(model_dir, "config.json"), + config_temp_path: None, }) } @@ -164,6 +203,7 @@ async fn resolve_remote_model_files( processor_config_path, chat_template_path, config_path, + config_temp_path: None, }) } @@ -204,6 +244,7 @@ fn resolve_cached_model_files(cache: &Cache, repo: &Repo) -> Result); + +impl HfOverrides { + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + /// Apply the patch to an in-memory config, preserving the source file. + pub fn apply(&self, config: &mut Value) { + merge_object(config, &self.0); + } +} + +/// Apply RFC 7396 object-member updates at every nesting level. +fn merge_object(target: &mut Value, patch: &Map) { + if !target.is_object() { + *target = Value::Object(Map::new()); + } + let target = target.as_object_mut().unwrap(); + for (key, value) in patch { + match value { + Value::Null => { + target.remove(key); + } + Value::Object(object) => { + merge_object(target.entry(key.clone()).or_insert(Value::Null), object) + } + _ => { + target.insert(key.clone(), value.clone()); + } + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::HfOverrides; + + #[test] + fn merge_patch_preserves_siblings_replaces_arrays_and_deletes_null_members() { + let mut config = json!({ + "rope_parameters": {"rope_type": "yarn", "factor": 2}, + "text_config": {"rope_parameters": {"rope_type": "yarn", "factor": 2}}, + "architectures": ["Old", "Other"], + "sliding_window": 4096 + }); + let patch: HfOverrides = serde_json::from_value(json!({ + "rope_parameters": {"factor": 4}, + "text_config": {"rope_parameters": {"factor": 4}}, + "architectures": ["New"], + "sliding_window": null, + "missing": null + })) + .unwrap(); + patch.apply(&mut config); + assert_eq!( + config, + json!({ + "rope_parameters": {"rope_type": "yarn", "factor": 4}, + "text_config": {"rope_parameters": {"rope_type": "yarn", "factor": 4}}, + "architectures": ["New"] + }) + ); + } + + #[test] + fn merge_patch_creates_objects_and_preserves_nulls_inside_replacement_arrays() { + let mut config = json!({"a": 1, "b": [1, 2]}); + let patch: HfOverrides = serde_json::from_value(json!({ + "a": {"nested": {"removed": null}}, + "b": {}, + "c": [null, {"literal": null}] + })) + .unwrap(); + patch.apply(&mut config); + assert_eq!( + config, + json!({"a": {"nested": {}}, "b": {}, "c": [null, {"literal": null}]}) + ); + } +} From 22bb158c21fcba630f267a12ebab4e86296ccc42 Mon Sep 17 00:00:00 2001 From: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:59:43 +0900 Subject: [PATCH 130/139] [Bugfix][ROCm] Fix MiniMax-M3 fused MXFP8 block scale (#53674) Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> --- tests/kernels/test_minimax_m3_amd_ops.py | 24 ++++++++++++++++++-- vllm/models/minimax_m3/amd/ops/swiglu_oai.py | 8 +++++-- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index ec8f1fd4ea..997c619809 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -42,6 +42,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( # noqa: from vllm.models.minimax_m3.amd.ops import ( # noqa: E402 gemma_fused_add_rmsnorm, gemma_rmsnorm, + swiglu_oai_quantize_mxfp8, swiglu_oai_split, ) from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps # noqa: E402 @@ -167,6 +168,25 @@ def test_swiglu_oai_split(m, inter, limit, dtype): assert _relerr(got, ref) < 5e-3 +@torch.inference_mode() +def test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale(): + # Keep a small value in two MX blocks next to their maxima. The scale must + # use E4M3's finite range (448), otherwise that value underflows to zero. + # Keep the third block empty to cover the reference's finite tiny clamp. + m, inter = 3, 96 + gate = torch.full((m, inter), 2.0, device=DEVICE, dtype=torch.float16) + up = torch.full((m, inter), 2**-10, device=DEVICE, dtype=torch.float16) + up[:, :64:32] = 1.0 + up[:, 64:] = 0.0 + gate_up = torch.cat((gate, up), dim=-1) + + got_q, got_s = swiglu_oai_quantize_mxfp8(gate_up, alpha=0.0, beta=0.0, limit=None) + ref_q, ref_s = _mxfp8_e4m3_quantize_torch(up, is_sf_swizzled_layout=False) + + assert torch.equal(got_s, ref_s) + assert torch.equal(got_q, ref_q) + + # --------------------------------------------------------------------------- # # Fused MXFP8 activation quant (Triton vs torch reference) # --------------------------------------------------------------------------- # @@ -179,8 +199,8 @@ def test_mxfp8_quant_triton_matches_torch(shape, dtype): xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False) xq_k, s_k = _mxfp8_e4m3_quantize_triton(x) assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32) - # E8M0 block exponents share the floor(log2(amax))+127 algorithm; allow at - # most a 1-step difference at exact powers of two. + # Both paths use the E4M3-aware scale calculation; allow a 1-step + # difference at exact powers of two due to floating-point rounding. assert (s_k.int() - s_t.int()).abs().max().item() <= 1 # Dequantized values agree to fp8 granularity. deq_t = dequant_mxfp8_to_bf16(xq_t, s_t) diff --git a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py index 836649b725..df7abc0760 100644 --- a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py +++ b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py @@ -83,6 +83,8 @@ def _swiglu_oai_quant_kernel( limit, HAS_LIMIT: tl.constexpr, BLOCK_M: tl.constexpr, + FP8_MAX: tl.constexpr, + TINY: tl.constexpr, ): """SwiGLU-OAI (split layout) fused with per-32-block MXFP8 (E4M3 + E8M0) quant. Each program handles ``[BLOCK_M, 32]`` of the ``[M, I]`` output (one @@ -109,8 +111,8 @@ def _swiglu_oai_quant_kernel( gate = tl.minimum(gate, limit) up = tl.minimum(tl.maximum(up, -limit), limit) act = gate * tl.sigmoid(alpha * gate) * (up + beta) # [BLOCK_M, 32] fp32 - amax = tl.maximum(tl.max(tl.abs(act), axis=1), 1e-30) # [BLOCK_M] - sb = tl.minimum(tl.maximum(tl.floor(tl.log2(amax)) + 127.0, 0.0), 254.0) + amax = tl.maximum(tl.max(tl.abs(act), axis=1), TINY) # [BLOCK_M] + sb = tl.minimum(tl.maximum(tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0, 0.0), 254.0) descale = tl.exp2(sb - 127.0) aq = (act / descale[:, None]).to(aq_ptr.dtype.element_ty) tl.store( @@ -173,6 +175,8 @@ def swiglu_oai_quantize_mxfp8( 0.0 if limit is None else float(limit), HAS_LIMIT=limit is not None, BLOCK_M=block_m, + FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max), + TINY=float(torch.finfo(torch.float32).tiny), num_warps=4, ) return aq, asc From fdfcbbcac79b0d8134a3a447f2c33001f5de07c1 Mon Sep 17 00:00:00 2001 From: JulienDarve <86800349+JulienDarve@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:02:25 -0700 Subject: [PATCH 131/139] [CI/Build][Rust Frontend] Retire Buf schema publishing (#57046) Co-authored-by: OpenAI Codex Co-authored-by: Bugen Zhao Signed-off-by: Julien Darve --- .github/workflows/buf.yml | 29 ----------------------------- rust/proto/README.md | 15 +++++---------- rust/proto/buf.md | 8 -------- rust/proto/buf.yaml | 1 - 4 files changed, 5 insertions(+), 48 deletions(-) delete mode 100644 rust/proto/buf.md diff --git a/.github/workflows/buf.yml b/.github/workflows/buf.yml index fda4add1d7..dd695d28d5 100644 --- a/.github/workflows/buf.yml +++ b/.github/workflows/buf.yml @@ -5,12 +5,6 @@ on: paths: - "rust/proto/**" - ".github/workflows/buf.yml" - push: - tags: - - "v*" - schedule: - - cron: "0 8 * * *" - workflow_dispatch: permissions: contents: read @@ -21,37 +15,14 @@ concurrency: jobs: proto: - if: github.repository_owner == 'vllm-project' runs-on: ubuntu-latest steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - with: - fetch-depth: 0 - uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0 with: version: 1.72.0 input: rust/proto - token: ${{ github.repository == 'vllm-project/vllm' && github.event_name != 'pull_request' && secrets.BUF_TOKEN || '' }} format: false breaking: false pr_comment: false push: false - - name: Publish nightly schema - if: github.repository == 'vllm-project/vllm' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') - run: >- - buf push rust/proto - --create - --create-default-label main - --create-visibility public - --label nightly - --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}" - - name: Publish release schema - if: github.repository == 'vllm-project/vllm' && github.event_name == 'push' - run: >- - buf push rust/proto - --create - --create-default-label main - --create-visibility public - --label main - --label "${GITHUB_REF_NAME}" - --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}" diff --git a/rust/proto/README.md b/rust/proto/README.md index 62b134280b..3181cdc6a0 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -2,14 +2,10 @@ This directory is the canonical source for vLLM's gRPC schema. -The schema is published to `buf.build/vllm-project/vllm`: - -- A daily workflow publishes the latest Git `main` schema to the `nightly` label. -- The workflow can be run manually to retry nightly publication. -- Tags matching `v*` update the Buf `main` label and publish the corresponding release label. -- Buf commits and generated SDK versions are immutable and can be pinned by consumers. - -Repository setup requires a `BUF_TOKEN` GitHub Actions secret with permission to create and push the public Buf module. Register the Prost and Tonic generated SDKs for the `main` and `nightly` labels once so subsequent pushes generate them automatically. +Schema updates are no longer published to the Buf Schema Registry. Rust consumers +should use `vllm-proto` from crates.io; consumers in other languages can generate +bindings from the `.proto` files in this directory. Buf still builds and lints +the schemas on pull requests, including on forks. ## Rust crate @@ -47,8 +43,7 @@ change PR. This check becomes available after the first manual publication. The `proto-crate.yml` workflow checks that the tag matches the crate version, verifies the package, and publishes it to crates.io. -Use `proto-v*` tags so they do not trigger Buf's `v*` release workflow. There is -no scheduled crate publication. Pull requests and manual workflow runs only +There is no scheduled crate publication. Pull requests and manual workflow runs only verify packaging, including on forks. Retry a failed release by rerunning its tag workflow; crates.io versions cannot be overwritten. diff --git a/rust/proto/buf.md b/rust/proto/buf.md deleted file mode 100644 index 065124feee..0000000000 --- a/rust/proto/buf.md +++ /dev/null @@ -1,8 +0,0 @@ -# vLLM gRPC API - -Protocol definitions for the optional gRPC server exposed by the vLLM Rust frontend. - -- `Inference` provides unary and streaming generation. -- `Control` provides discovery and request control. - -The `nightly` label is updated daily from vLLM's `main` branch. The `main` label tracks the latest vLLM release, and version labels match tags such as `v0.27.0`. Pin a Buf commit or generated SDK version for reproducible builds. diff --git a/rust/proto/buf.yaml b/rust/proto/buf.yaml index 64764f561e..d30dc3297d 100644 --- a/rust/proto/buf.yaml +++ b/rust/proto/buf.yaml @@ -1,7 +1,6 @@ version: v2 modules: - path: . - name: buf.build/vllm-project/vllm lint: use: - MINIMAL From 4ace5cd9d80a8d56bf935c4a058a1369f1b94da3 Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Tue, 15 Sep 2026 20:05:43 -0700 Subject: [PATCH 132/139] [XPU] fix_test_worker_memory_snapshot_for_xpu (#50097) Signed-off-by: mayuyuace Co-authored-by: Kunshang Ji --- .../v1/worker/test_worker_memory_snapshot.py | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/tests/v1/worker/test_worker_memory_snapshot.py b/tests/v1/worker/test_worker_memory_snapshot.py index fe8a5a21f8..6ed62cb771 100644 --- a/tests/v1/worker/test_worker_memory_snapshot.py +++ b/tests/v1/worker/test_worker_memory_snapshot.py @@ -12,6 +12,7 @@ import torch from vllm.config import set_current_vllm_config from vllm.engine.arg_utils import EngineArgs +from vllm.platforms import current_platform from vllm.utils.mem_utils import MemorySnapshot from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment @@ -66,8 +67,22 @@ def worker_process( model="facebook/opt-125m", tensor_parallel_size=2, load_format="dummy" ).create_engine_config() + # Select worker class and patch target based on platform + if current_platform.is_xpu(): + from vllm.v1.worker.xpu_worker import XPUWorker + + worker_cls = XPUWorker + init_patch_target = ( + "vllm.v1.worker.xpu_worker.init_worker_distributed_environment" + ) + else: + worker_cls = Worker + init_patch_target = ( + "vllm.v1.worker.gpu_worker.init_worker_distributed_environment" + ) + # Create worker - worker = Worker( + worker = worker_cls( vllm_config=vllm_config, local_rank=rank, rank=rank, @@ -81,7 +96,7 @@ def worker_process( # Apply minimal patches to track operation order init_patch = patch( - "vllm.v1.worker.gpu_worker.init_worker_distributed_environment", + init_patch_target, side_effect=make_operation_tracker( "init_distributed", original_init_worker ), @@ -93,7 +108,7 @@ def worker_process( ) all_reduce_patch = patch( "torch.distributed.all_reduce", - side_effect=make_operation_tracker("nccl_all_reduce", original_all_reduce), + side_effect=make_operation_tracker("dist_all_reduce", original_all_reduce), ) with ( @@ -182,13 +197,13 @@ def test_init_distributed_is_called_before_memory_snapshot(): # Raises ValueError if the operation is not found init_distributed = rank_ops.index("init_distributed") - nccl_all_reduce = rank_ops.index("nccl_all_reduce") + dist_all_reduce = rank_ops.index("dist_all_reduce") memory_snapshot = rank_ops.index("memory_snapshot") # Verify order: init_distributed should happen before memory_snapshot - assert init_distributed < nccl_all_reduce < memory_snapshot, ( + assert init_distributed < dist_all_reduce < memory_snapshot, ( f"Rank {rank}: init_distributed (index {init_distributed}) " - f"must happen before nccl_all_reduce (index {nccl_all_reduce}) " + f"must happen before dist_all_reduce (index {dist_all_reduce}) " f"and memory_snapshot (index {memory_snapshot})" ) From ccbaed5b91e09f528180db363af7ebb9d4fafe32 Mon Sep 17 00:00:00 2001 From: Marceli Fylcek Date: Wed, 16 Sep 2026 05:05:54 +0200 Subject: [PATCH 133/139] [XPU] Route to fused_qk_rmsnorm_rope_gate triton kernel (#56096) Signed-off-by: Marceli Fylcek Co-authored-by: Kunshang Ji --- vllm/model_executor/models/qwen3_next.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 495650eb19..7ea3ed1319 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -378,7 +378,7 @@ class Qwen3NextAttention(nn.Module): self.use_fused_qk_norm_rope_gate = ( self.attn_output_gate and getattr(self.rotary_emb, "is_neox_style", False) - and current_platform.is_cuda() + and (current_platform.is_cuda() or current_platform.is_xpu()) and supports_dtype and (text_only or supports_mrope) ) From c16bb6068f70878fb8a2f7c4d6cda95cd03a778b Mon Sep 17 00:00:00 2001 From: Islam <43927639+almersawi@users.noreply.github.com> Date: Wed, 16 Sep 2026 07:10:58 +0400 Subject: [PATCH 134/139] [Bugfix][Rust Frontend] Tolerate NaN-corrupted logprobs in engine-core (#51026) Co-authored-by: Bugen Zhao Signed-off-by: Islam Signed-off-by: Bugen Zhao --- .../src/protocol/logprobs.rs | 7 +- .../src/protocol/logprobs/tests.rs | 118 ++++++++++++++++++ .../src/routes/openai/utils/logprobs.rs | 47 +++++++ 3 files changed, 168 insertions(+), 4 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index d5bc40e061..646c106a04 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -25,10 +25,13 @@ use crate::protocol::tensor::{WireArrayData, WireNdArray}; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct TokenLogprob { pub token_id: u32, + /// Preserves the engine's value, including NaN and infinities. pub logprob: f32, /// The sampled/selected token uses its actual vocab rank. Remaining entries /// use 1-based top-k ranks matching the engine's returned candidate /// order. + /// A sampled/selected rank of 0 occurs when its logprob is NaN: the engine's + /// `(logprobs >= selected_logprob).sum(-1)` counts no matching values. pub rank: u32, } @@ -53,10 +56,6 @@ impl PositionLogprobs { logprobs.len() ); } - if sampled_rank == 0 { - bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs"); - } - let mut entries = Vec::with_capacity(token_ids.len()); for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() { let rank = if index == 0 { diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs index 2470460b52..f69ff22f1a 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -396,3 +396,121 @@ fn rejects_zero_column_logprobs_with_rows() { "new_logprobs: zero-column logprobs payload with 2 rows" ); } + +fn logprobs_value(ids: &[i64], probs: &[f32], ranks: &[i64]) -> Value { + let rows = ranks.len(); + let cols = ids.len() / rows; + Value::Array(vec![ + ndarray_value( + " Date: Wed, 16 Sep 2026 11:53:23 +0800 Subject: [PATCH 135/139] [XPU] Honor device_ids for worker device placement (#56015) Signed-off-by: Chaojun Zhang Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/test-intel.yaml | 2 +- vllm/v1/engine/core.py | 11 +++--- vllm/v1/worker/xpu_worker.py | 50 +++++++++++++++++++++------ 3 files changed, 44 insertions(+), 19 deletions(-) diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index a76331f4d5..ef2b4ba073 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -97,7 +97,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'cd tests && + 'cd tests && pip install dpctl && pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py && pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py && pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" && diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ffc9a5de7e..62b6ef880e 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -2452,13 +2452,10 @@ class EngineCoreActorMixin: def _set_visible_devices(self, vllm_config: VllmConfig, local_dp_rank: int): from vllm.platforms import current_platform - if current_platform.is_xpu(): - pass - else: - device_control_env_var = current_platform.device_control_env_var - self._set_assigned_physical_gpu_ids( - vllm_config, local_dp_rank, device_control_env_var - ) + device_control_env_var = current_platform.device_control_env_var + self._set_assigned_physical_gpu_ids( + vllm_config, local_dp_rank, device_control_env_var + ) def _set_assigned_physical_gpu_ids( self, diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py index 68d6bde3a0..5c5d125f4b 100644 --- a/vllm/v1/worker/xpu_worker.py +++ b/vllm/v1/worker/xpu_worker.py @@ -115,27 +115,55 @@ class XPUWorker(Worker): self.local_rank += dp_local_rank * replica_world_size - assert self.local_rank < visible_device_count, ( - f"DP adjusted local rank {self.local_rank} is out of bounds. " - ) - assert parallel_config.local_world_size <= visible_device_count, ( - f"local_world_size ({parallel_config.local_world_size}) must " - f"be less than or equal to the number of visible devices " - f"({visible_device_count})." - ) - device = self.device_config.device if ( isinstance(device, torch.device) and device.type == "xpu" and current_platform.is_xpu() ): - self.device = torch.device(f"xpu:{self.local_rank}") + assigned_physical_gpu_ids = self.parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + assert self.local_rank < len(assigned_physical_gpu_ids), ( + f"local_rank {self.local_rank} is out of bounds for " + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" + ) + # NOTE: local_world_size is derived from parallel_config.nnodes, + # which is only set for the "mp" multi-node backend. With the + # "ray"/"external_launcher" backends nnodes stays 1, so + # local_world_size collapses to the full world_size and this + # check wrongly fires on cross-node deployments. + # assigned_physical_gpu_ids is already per-node and the + # local_rank bound above fully validates the mapping for + # these backends, so skip the check for them. + if parallel_config.distributed_executor_backend not in ( + "ray", + "external_launcher", + ): + assert parallel_config.local_world_size <= len( + assigned_physical_gpu_ids + ), ( + f"local_world_size ({parallel_config.local_world_size})" + " exceeds assigned_physical_gpu_ids count " + f"({len(assigned_physical_gpu_ids)})" + ) + else: + assert self.local_rank < torch.accelerator.device_count(), ( + f"DP adjusted local rank {self.local_rank} is out of " + f"bounds for {torch.accelerator.device_count()} devices." + ) + + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id(self.local_rank) + ) + self.device = torch.device(f"xpu:{visible_device_index}") torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) torch.accelerator.empty_cache() self.init_gpu_memory = torch.xpu.get_device_properties( - self.local_rank + visible_device_index ).total_memory else: raise RuntimeError(f"Unsupported device type: {self.device_config.device}") From 75dc5882699be65eaee144cb7d0b06c1f86d3cec Mon Sep 17 00:00:00 2001 From: alberto Date: Wed, 16 Sep 2026 04:57:38 +0100 Subject: [PATCH 136/139] [KV Offload] Add per-request `max_load_tokens control` (#55885) Signed-off-by: Alberto Perdomo --- docs/features/kv_offloading_usage.md | 29 ++++++ .../offloading_connector/test_scheduler.py | 93 ++++++++++++++++++- .../kv_connector/v1/offloading/scheduler.py | 30 ++++++ 3 files changed, 151 insertions(+), 1 deletion(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index a5c8684b42..22a2350d80 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -22,6 +22,35 @@ flowchart LR CPU <--> SN["..."] ``` +## Per-request load control + +Individual requests can cap how many tokens are loaded from offloaded storage +by setting `max_load_tokens` in `kv_transfer_params`. The cap applies to tokens +beyond those already available in the GPU prefix cache. Set it to `0` to +disable external loading for the request: + +```json +{ + "kv_transfer_params": { + "max_load_tokens": 0 + } +} +``` + +GPU prefix-cache reuse remains enabled, and tokens beyond the aligned load cap +are recomputed. Omitting the field leaves loading uncapped. The value must be a +non-negative integer; invalid values are ignored. Positive caps are rounded +down to a boundary supported by the configured KV cache groups. Offloading +remains enabled unless it is controlled separately with `max_offload_tokens`. + +!!! warning + `max_load_tokens` is experimental and subject to change. + +`kv_load_tiers` continues to select secondary tiers. CPU is always included: +it can satisfy a resident hit directly and is the required staging tier for +secondary-to-GPU loads. Consequently, an empty tier list allows CPU hits but +does not query any secondary tier. + ## Terminology: Chunks The unit of operation is a **chunk** — a fixed-size piece of KV data covering a group of tokens. By default, a chunk maps to a single accelerator block. A configurable `blocks_per_chunk` parameter allows larger chunks, yielding larger I/Os to the host and secondary tiers. diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index e7a83cb7b5..968474dd74 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, call import pytest @@ -195,15 +196,17 @@ def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler: def _make_partial_tail_request( scheduler: OffloadingConnectorScheduler, + kv_transfer_params: dict[str, Any] | None = None, ) -> MagicMock: request = MagicMock() request.request_id = "req" - request.kv_transfer_params = None + request.kv_transfer_params = kv_transfer_params request.num_prompt_tokens = 30 request.num_tokens = 30 request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)] request.all_token_ids = list(range(30)) request.lora_request = None + request.skip_reading_prefix_cache = False request.is_finished.return_value = False scheduler.on_new_request(request) return request @@ -427,6 +430,53 @@ def test_lookup_cap_stops_at_authoritative_prefix_boundary(): assert scheduler._req_status["req"].partial_tail_boundary == 20 +def test_max_load_tokens_zero_skips_partial_tail_lookup(): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 0} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 16) == (0, False) + scheduler.manager.lookup.assert_not_called() + assert all( + state.num_hit_chunks == 1 for state in scheduler._req_status["req"].group_states + ) + + +def test_max_load_tokens_caps_tokens_beyond_gpu_prefix(): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 8} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 16) == (8, True) + assert scheduler._req_status["req"].partial_tail_boundary == 24 + + +def test_max_load_tokens_rounds_down_without_partial_tail(): + scheduler = _make_partial_tail_scheduler() + scheduler.config = scheduler.config._replace(supports_partial_tail=False) + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 20} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 0) == (16, True) + + +@pytest.mark.parametrize("value", ["8", 8.5, -1, True]) +def test_invalid_max_load_tokens_is_ignored(value): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": value} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 0) == (28, True) + + def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary(): """A recurrent group may hold unhashed blocks inside the computed prefix. @@ -2964,6 +3014,47 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): runner.manager.lookup.assert_not_called() +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_max_load_tokens_limits_external_load(request_runner, async_scheduling: bool): + """The load cap limits external reads without changing the store path.""" + block_size = 4 + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + blocks_per_chunk=blocks_per_chunk, + ) + + runner.new_request(token_ids=[0] * tokens_per_chunk) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0, 1, 2), + ) + + runner.scheduler.reset_prefix_cache() + runner.manager.lookup.reset_mock() + runner.new_request( + token_ids=[0] * tokens_per_chunk, + kv_transfer_params={"max_load_tokens": 0}, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_loaded=(), + expected_stored=(0, 1, 2), + ) + + runner.manager.lookup.assert_not_called() + + # --------------------------------------------------------------------------- # Eagle/MTP test class # --------------------------------------------------------------------------- diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 2fb3c5bb90..8cd5e90b6f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -354,6 +354,8 @@ class RequestOffloadState: req_context: ReqContext offloading_context: RequestOffloadingContext group_states: tuple[RequestGroupState, ...] = field(init=False) + # upper bound on tokens to load for this request; None means no cap + max_load_tokens: int | None = None # upper bound on tokens to offload for this request; None means no cap max_offload_tokens: int | None = None # number of hits in the GPU cache @@ -376,6 +378,20 @@ class RequestOffloadState: ) params = self.req.kv_transfer_params + # NOTE: This field is experimental and subject to change in the future. + raw = params.get("max_load_tokens") if params else None + if type(raw) is int and raw >= 0: + self.max_load_tokens = raw + logger.debug( + "Request %s: max_load_tokens set to %d", + self.req.request_id, + raw, + ) + elif raw is not None: + logger.warning( + "max_load_tokens must be a non-negative int, got %r; ignoring", raw + ) + # NOTE: This field is experimental and subject to change in the future. raw = params.get("max_offload_tokens") if params else None if type(raw) is int and raw >= 0: @@ -767,6 +783,11 @@ class OffloadingConnectorScheduler: max_hit_size_tokens = min( max_hit_size_tokens, num_computed_tokens + max_num_new_tokens ) + if req_status.max_load_tokens is not None: + max_hit_size_tokens = min( + max_hit_size_tokens, + num_computed_tokens + req_status.max_load_tokens, + ) if self._sliding_window_groups: # the last prompt token has to be recomputed to get the logprobs # for sliding window attention, we must reduce by 1 to make sure @@ -810,6 +831,10 @@ class OffloadingConnectorScheduler: max_hit_size_tokens = min( max_hit_size_tokens, len(offload_keys) * tokens_per_chunk ) + if req_status.max_load_tokens is not None: + max_hit_size_tokens = round_down( + max_hit_size_tokens, tokens_per_chunk + ) if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk: # We can only load less than a chunk, so skip. return 0 @@ -965,6 +990,11 @@ class OffloadingConnectorScheduler: max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1) if max_num_new_tokens is not None: max_boundary = min(max_boundary, local_tokens + max_num_new_tokens) + if req_status.max_load_tokens is not None: + max_boundary = min( + max_boundary, + local_tokens + req_status.max_load_tokens, + ) max_boundary = round_down(max_boundary, tokens_per_hash) if max_boundary <= complete_boundary: return complete_hit From a31ec3a68bbbedd4b5d59490c762bcf32abf1a17 Mon Sep 17 00:00:00 2001 From: danisereb Date: Wed, 16 Sep 2026 07:20:28 +0300 Subject: [PATCH 137/139] [LoRA][Nemotron] Add LoRA support for Nemotron VL models (for the language model only) (#56231) Signed-off-by: Daniel Serebrenik --- docs/models/supported_models.md | 7 +++++ .../model_executor/models/nano_nemotron_vl.py | 29 +++++++++++++++++-- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index d55348ad72..ddc50ad41c 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -602,6 +602,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ | | `MuseGlimmerForCausalLM`, `MuseGlimmerForConditionalGeneration` | Muse Glimmer | T + I+ + V+ | `meta-models/Muse-Glimmer-30B` | ✅︎ | ✅︎ | +| `NemotronH_Nano_Omni_Reasoning_V3`, `NemotronH_Nano_VL_V2` | Nemotron Nano 3 Omni, Nemotron Nano V2 VL | T + IE+ + V+ + A* | `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16` | ✅︎ | | | `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ | | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + IE+ | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ | | `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + IE+ + VE+ | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ | @@ -701,6 +702,12 @@ Some models are supported only via the [Transformers modeling backend](#transfor `meta-models/Muse-Glimmer-30B-assistant` checkpoint, which vLLM serves through its [DFlash](../features/speculative_decoding/README.md) path. +!!! note + Both Nemotron Nano architecture names map to the same vLLM implementation. + Audio inputs are only supported by the Omni variants; `NemotronH_Nano_VL_V2` + checkpoints such as `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` accept text, + image and video only. + !!! note The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now. For more details, please see: diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 2edf56871d..6f5e073bd6 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -23,6 +23,7 @@ from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.logger import init_logger +from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.layers.activation import ReLUSquaredActivation from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.model_loader.weight_utils import default_weight_loader @@ -30,6 +31,7 @@ from vllm.model_executor.models.interfaces import ( HasInnerState, IsHybrid, MultiModalEmbeddings, + SupportsLoRA, SupportsMultiModal, SupportsMultiModalPruning, ) @@ -897,11 +899,22 @@ class NanoNemotronVLDummyInputsBuilder( dummy_inputs=NanoNemotronVLDummyInputsBuilder, ) class NemotronH_Nano_VL_V2( - nn.Module, HasInnerState, IsHybrid, SupportsMultiModal, SupportsMultiModalPruning + nn.Module, + HasInnerState, + IsHybrid, + SupportsMultiModal, + SupportsMultiModalPruning, + SupportsLoRA, ): requires_sequential_video_encoding = True """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet""" + # LoRA covers the language model only + is_non_gated_moe = NemotronHForCausalLM.is_non_gated_moe + packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping + embedding_modules = NemotronHForCausalLM.embedding_modules + lora_skip_prefixes = NemotronHForCausalLM.lora_skip_prefixes + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "language_model.backbone": "language_model.model", @@ -1331,7 +1344,19 @@ class NemotronH_Nano_VL_V2( # Create final video embeddings, merging text embeddings for indicator # tokens with video embeddings - text_embeddings = self.get_language_model().embed_input_ids(repl_token_ids) + + # LoRA support - + # These replacement tokens are produced inside the encoder, so they are + # absent from the request-token batch that the LoRA adapter index + # mapping is built from, and there can be more + # of them than max_num_batched_tokens. + # Embed them with the base weights - + # the LoRA delta is undefined for tokens with no mapping entry. + embed_tokens = self.get_language_model().model.embed_tokens + if isinstance(embed_tokens, BaseLayerWithLoRA): + embed_tokens = embed_tokens.base_layer + text_embeddings = embed_tokens(repl_token_ids) + final_video_embeddings = _merge_multimodal_embeddings( inputs_embeds=text_embeddings, multimodal_embeddings=video_embeddings, From 6ecd97f1b7fb13991065370f6d6be3dffe9fe1db Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 16 Sep 2026 12:34:55 +0800 Subject: [PATCH 138/139] [Perf][DSpark] Add KV-only context insertion across V4.1 cache formats (#56441) Signed-off-by: 0z5a <0z5a@users.noreply.github.com> Signed-off-by: Yongye Zhu Co-authored-by: 0z5a <0z5a@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Yongye Zhu Co-authored-by: Claude Opus 5 (1M context) --- ...deepseek_v4_qnorm_rope_kv_insert_kernel.cu | 87 ++++++++++++- csrc/libtorch_stable/ops.h | 7 ++ csrc/libtorch_stable/torch_bindings.cpp | 8 ++ tests/kernels/test_compressor_kv_cache.py | 116 ++++++++++++++++++ vllm/models/deepseek_v41/nvidia/dspark.py | 77 +++--------- 5 files changed, 234 insertions(+), 61 deletions(-) diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 254b9783bc..344b82fbda 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -657,7 +657,12 @@ static void launchFusedDeepseekV4Templated( config.attrs = attrs; config.numAttrs = (sm_version >= 90) ? 1 : 0; - if (num_tokens_full < NUM_TOKEN_CUTOFF) { + // The reduced grid pays off by iterating a token's Q heads inside one CTA. + // With zero query heads there is nothing to iterate, so one warp per block + // would do all the work and the other seven idle; keep the warp-per-slot + // grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single + // insert, where it is worth ~2x. + if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) { cudaLaunchKernelEx(&config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, @@ -993,6 +998,86 @@ static void launchFullCacheKernel( // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper // ──────────────────────────────────────────────────────────────────────────── +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8) { + using torch::headeronly::ScalarType; + STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous() && + kv.dim() == 2 && kv.size(1) == 512 && + kv.scalar_type() == ScalarType::BFloat16, + "kv must be contiguous CUDA bfloat16 [tokens, 512]"); + auto const device = kv.get_device_index(); + STD_TORCH_CHECK(k_cache.device().is_cuda() && + k_cache.get_device_index() == device && + k_cache.dim() == 3 && cache_block_size > 0 && + k_cache.size(1) == cache_block_size && + k_cache.stride(2) == 1, + "k_cache must be a paged cache on the KV device"); + for (auto const* tensor : {&slot_mapping, &position_ids}) { + STD_TORCH_CHECK(tensor->device().is_cuda() && + tensor->get_device_index() == device && + tensor->scalar_type() == ScalarType::Long && + tensor->dim() == 1 && tensor->is_contiguous(), + "slots and positions must be contiguous int64 on the KV device"); + } + STD_TORCH_CHECK(position_ids.size(0) == kv.size(0) && + slot_mapping.size(0) <= kv.size(0), + "positions must match KV rows; slots may omit padded rows"); + STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() && + cos_sin_cache.get_device_index() == device && + cos_sin_cache.scalar_type() == ScalarType::Float && + cos_sin_cache.dim() == 2 && + cos_sin_cache.size(1) == 64 && cos_sin_cache.is_contiguous(), + "cos_sin_cache must be contiguous float32 [max_pos, 64]"); + bool const packed = k_cache.scalar_type() == ScalarType::Byte; + bool const fp8 = k_cache.scalar_type() == ScalarType::Float8_e4m3fn; + STD_TORCH_CHECK(packed || fp8 || + k_cache.scalar_type() == ScalarType::BFloat16, + "cache must be uint8, bfloat16 or float8_e4m3fn"); + STD_TORCH_CHECK(k_cache.size(2) == (packed ? (kv_mxfp8 ? 528 : 584) : 512), + "cache row width does not match its dtype"); + if (fp8) { + STD_TORCH_CHECK(fp8_scale.has_value() && fp8_scale->device().is_cuda() && + fp8_scale->get_device_index() == device && + fp8_scale->scalar_type() == ScalarType::Float && + fp8_scale->numel() == 1, + "fp8 cache requires one float32 KV scale on the KV device"); + } + int const num_tokens = static_cast(slot_mapping.size(0)); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard(device); + const cudaStream_t stream = get_current_cuda_stream(device); + using scalar_t = c10::BFloat16; + auto const* input = reinterpret_cast(kv.const_data_ptr()); + auto* cache = reinterpret_cast(k_cache.mutable_data_ptr()); + auto const* slots = slot_mapping.const_data_ptr(); + auto const* positions = position_ids.const_data_ptr(); + auto const* cos_sin = cos_sin_cache.const_data_ptr(); + // Zero query heads schedules only the existing KV branch, preserving its + // RoPE rounding and quantization without allocating any query tensors. + if (packed) { + vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated( + nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f, + num_tokens, num_tokens, 0, static_cast(cache_block_size), + static_cast(k_cache.stride(0)), kv_mxfp8, stream); + } else if (fp8) { + vllm::deepseek_v4_fused_ops::launchFullCacheKernel( + nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, + fp8_scale->const_data_ptr(), nullptr, 0.0f, num_tokens, num_tokens, + 0, static_cast(cache_block_size), k_cache.stride(0), + k_cache.stride(1), false, "fused_deepseek_v4_kv_rope_insert", stream); + } else { + vllm::deepseek_v4_fused_ops::launchFullCacheKernel( + nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, nullptr, + nullptr, 0.0f, num_tokens, num_tokens, 0, static_cast(cache_block_size), + k_cache.stride(0) * 2, k_cache.stride(1) * 2, false, + "fused_deepseek_v4_kv_rope_insert", stream); + } +} + torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16 torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only) diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index e19f49d603..bc08d90a48 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -259,6 +259,13 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, torch::stable::Tensor& position_ids, int64_t forced_token_heads_per_warp); +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8); + torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 54548dfd05..3918ab8498 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -425,6 +425,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "bool is_neox, Tensor position_ids, " "int forced_token_heads_per_warp=-1) -> ()"); + ops.def( + "fused_deepseek_v4_kv_rope_insert(" + "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " + "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " + "bool kv_mxfp8=False) -> " + "()"); ops.def( "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" "Tensor q_in, Tensor kv, Tensor! k_cache, " @@ -807,6 +813,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { // Positional encoding kernels (shared CUDA/ROCm) ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_kv_rope_insert", + TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); ops.impl( diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 7e092229fe..c1de94be92 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -42,6 +42,122 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize( + "cache_dtype,kv_mxfp8", + [ + (torch.uint8, False), + (torch.uint8, True), + (torch.bfloat16, False), + (torch.float8_e4m3fn, False), + ], +) +@pytest.mark.parametrize("num_tokens", [1, 17, 1023, 1024]) +@pytest.mark.parametrize("use_graph", [False, True]) +def test_dspark_context_kv_matches_query_insert( + cache_dtype, kv_mxfp8, num_tokens, use_graph +): + """KV-only insertion must preserve every cache byte, including graph replay.""" + from vllm.models.deepseek_v41.nvidia.dspark import _insert_context_kv + + torch.manual_seed(42) + block_size = 256 + num_blocks = math.ceil((num_tokens + 7) / block_size) + row_size = (528 if kv_mxfp8 else 584) if cache_dtype == torch.uint8 else 512 + page_stride = math.ceil(block_size * row_size / 576) * 576 + 576 + backing = torch.full((num_blocks, page_stride), 3, device="cuda").to(cache_dtype) + cache = backing.as_strided( + (num_blocks, block_size, row_size), (page_stride, row_size, 1) + ) + reference_backing = backing.clone() + reference = reference_backing.as_strided(cache.shape, cache.stride()) + kv = torch.randn(num_tokens + 3, 512, device="cuda", dtype=torch.bfloat16) + positions = torch.arange(num_tokens + 3, device="cuda") + 7 + angles = torch.randn(num_tokens + 16, 32, device="cuda") + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + slots = torch.randperm(num_blocks * block_size, device="cuda")[:num_tokens] + slots[::5] = -1 + scale = torch.tensor([0.7], device="cuda") + # Only what _insert_context_kv actually reads: the record width comes off + # the cache tensor, so no per-record width attribute is stubbed here. + attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(kv_cache=cache, block_size=block_size), + kv_mxfp8=kv_mxfp8, + rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin), + _flashinfer_fp8_kv_scale=scale, + ) + + def legacy_insert(): + q = torch.zeros(kv.shape[0], 8, 512, dtype=kv.dtype, device="cuda") + if cache_dtype == torch.uint8: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + reference.view(num_blocks, -1), + slots, + positions, + cos_sin, + 8, + 1e-20, + block_size, + True, + kv_mxfp8, + ) + elif cache_dtype == torch.bfloat16: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + q, + kv, + reference, + slots, + positions, + cos_sin, + 1e-20, + block_size, + ) + else: + q_fp8 = torch.empty_like(q, dtype=cache_dtype) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + q, + kv, + q_fp8, + reference, + slots, + positions, + cos_sin, + scale, + scale, + 1e-20, + block_size, + ) + + def insert(): + _insert_context_kv(attn, kv, positions, slots) + + insert() + graph = None + if use_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + insert() + for iteration in range(3): + kv.normal_() + positions.add_(1) + if iteration == 2: + slots.fill_(-1) + backing.view(torch.uint8).fill_(165) + reference_backing.copy_(backing) + before = kv.clone() + legacy_insert() + insert() if graph is None else graph.replay() + torch.testing.assert_close( + backing.view(torch.uint8), + reference_backing.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close(kv, before, rtol=0, atol=0) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="graph capture coverage" ) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index c19dca8005..214681f3a9 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -234,67 +234,24 @@ def _insert_context_kv( positions: torch.Tensor, slot_mapping: torch.Tensor, ) -> None: - """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. - - Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy - query and discard it, since context tokens have no query). Mirrors - ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``. - """ - swa_cache = attn.swa_cache_layer.kv_cache - block_size = attn.swa_cache_layer.block_size - cos_sin_cache = attn.rotary_emb.cos_sin_cache - cache_dtype = swa_cache.dtype - n_ctx = kv.shape[0] - dummy_q = torch.zeros( - (n_ctx, attn.n_local_heads, attn.head_dim), - dtype=kv.dtype, - device=kv.device, + """Insert normalized context KV without constructing an unused query.""" + # The bound cache already is [num_blocks, block_size, bytes_per_token], so + # the record width needs no per-layer attribute -- one whose name the + # attention layer is free to change (kv_bytes_per_token has to become + # swa_bytes_per_token once the sliding-window and compressed records can + # differ). block_size stays the layer's, so the op's size(1) check still + # cross-checks the two rather than restating the tensor. + cache = attn.swa_cache_layer.kv_cache + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + kv, + cache, + slot_mapping, + positions, + attn.rotary_emb.cos_sin_cache, + attn.swa_cache_layer.block_size, + attn._flashinfer_fp8_kv_scale if cache.dtype == torch.float8_e4m3fn else None, + attn.kv_mxfp8, ) - if cache_dtype == torch.uint8: - # fp8_ds_mla UE8M0 paged layout - swa_2d = swa_cache.view(swa_cache.shape[0], -1) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - dummy_q, - kv, - swa_2d, - slot_mapping, - positions, - cos_sin_cache, - attn.padded_heads, - attn.eps, - block_size, - True, # apply_q_norm (unused: the query is a discarded dummy) - attn.kv_mxfp8, - ) - elif cache_dtype == torch.bfloat16: - swa_3d = swa_cache.view(-1, block_size, attn.head_dim) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( - dummy_q, - kv, - swa_3d, - slot_mapping, - positions, - cos_sin_cache, - attn.eps, - block_size, - ) - else: # per-tensor fp8 (torch.float8_e4m3fn) - # TODO(ben): double-check if this is being dispatched correctly for FI backend - swa_3d = swa_cache.view(-1, block_size, attn.head_dim) - dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( - dummy_q, - kv, - dummy_q_fp8, - swa_3d, - slot_mapping, - positions, - cos_sin_cache, - attn._flashinfer_fp8_kv_scale, - attn._flashinfer_fp8_q_scale_inv, - attn.eps, - block_size, - ) class DSparkDeepseekV4ForCausalLM(nn.Module): From af1c01499b289be555c475669ba50a88e96d846e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 15 Sep 2026 21:37:08 -0700 Subject: [PATCH 139/139] [Bugfix] Make GPU sync checks safe under torch.compile (#56904) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: Sherlock Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../generation/test_voxtral_realtime.py | 3 ++- tests/utils_/test_gpu_sync_debug.py | 18 ++++++++++++++++++ vllm/utils/gpu_sync_debug.py | 14 ++++++++++++-- 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 1bd9f62303..2c13e07ff3 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -21,7 +21,7 @@ from vllm.utils.math_utils import cdiv from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.kv_cache_interface import SlidingWindowSpec -from ....utils import ROCM_ENGINE_KWARGS +from ....utils import ROCM_ENGINE_KWARGS, create_new_process_for_each_test MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602" AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn" @@ -192,6 +192,7 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypa "cudagraph_mode", [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], ) +@create_new_process_for_each_test(method="spawn") def test_voxtral_realtime_cudagraph( audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode ): diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py index dc8b122800..10b550c0a1 100644 --- a/tests/utils_/test_gpu_sync_debug.py +++ b/tests/utils_/test_gpu_sync_debug.py @@ -123,6 +123,24 @@ def test_suppressing_works_while_compiling(monkeypatch): with_gpu_sync_check(suppressed)() +@create_new_process_for_each_test() +def test_copy_checker_is_traceable_while_compiling(monkeypatch): + """The copy wrapper must bypass ContextVars while Dynamo is tracing.""" + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + gsd._install_copy_checkers() + + def convert_like(x, like): + return x.to(device=like.device, dtype=like.dtype) + + compiled = torch.compile(convert_like, fullgraph=True) + result = compiled( + torch.ones(4, device="cuda", dtype=torch.float64), + torch.ones(4, device="cuda", dtype=torch.float32), + ) + assert result.dtype == torch.float32 + + @create_new_process_for_each_test() def test_sync_debug_mode_restored_after_checked_call(monkeypatch): """The mode is armed only for the duration of a checked call. Leaving it diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py index 53fcdfe710..8d0e7ae187 100644 --- a/vllm/utils/gpu_sync_debug.py +++ b/vllm/utils/gpu_sync_debug.py @@ -225,6 +225,16 @@ _torch_to = torch.Tensor.to _torch_cuda = torch.Tensor.cuda _torch_copy_ = torch.Tensor.copy_ + +# Once ``Tensor.to`` is replaced below, Dynamo sees the saved C++ descriptor +# through this module global instead of as a recognized tensor method. Put the +# descriptor behind a weak-referenceable Python callable so ``allow_in_graph`` +# can admit it and the downstream compiler can lower the original operation. +@torch.compiler.allow_in_graph +def _torch_to_in_graph(self, *args, **kwargs): + return _torch_to(self, *args, **kwargs) + + # torch's own parser for `Tensor.to`'s overloaded signature, returning # (device, dtype, non_blocking, memory_format) (also used by torch._dynamo). _parse_to = getattr(torch._C._nn, "_parse_to", None) @@ -232,7 +242,7 @@ _parse_to = getattr(torch._C._nn, "_parse_to", None) def _active_check_mode() -> str | None: """The mode the calling thread is being checked in, if any.""" - if not _sync_check_enabled or _allow_depth.get() or torch.compiler.is_compiling(): + if not _sync_check_enabled or torch.compiler.is_compiling() or _allow_depth.get(): return None return _checking.get() @@ -296,7 +306,7 @@ def _checked_to(self, *args, **kwargs): else: if non_blocking and (reason := _cpu_copy_stall_reason(self, device)): _report_implicit_sync("to", "H2D", reason, mode) - return _torch_to(self, *args, **kwargs) + return _torch_to_in_graph(self, *args, **kwargs) def _checked_cuda(self, *args, **kwargs):