From 1adab90a7db51ca84a02fe8137d2fa07c80703b7 Mon Sep 17 00:00:00 2001 From: bofengl Date: Sun, 18 Oct 2026 01:57:17 -0700 Subject: [PATCH 001/406] [Perf][Kernel][DSv4.1] Copy-based Q zero-pad for large eager mega-attention steps With FlashMLA mega attention the Q half of fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert only zero-pads the chunk-interleaved wq_b output to the kernel's 64 heads. With 32 live heads (TP2 of a 64-head model) that is 32 KiB read and 64 KiB written per token, against about 1.8 KiB for the KV insert sharing the launch. On eager steps of at least 1024 tokens (the op's NUM_TOKEN_CUTOFF), such layers now run the op in its existing KV-only mode (q_head_padded=0) and build the padded Q with _pad_interleaved_q, a Triton copy over (token, 16-dim chunk) rows that copies the live slots and writes +0.0 into the padding slots. The kernel is registered for JIT warmup. Captured CUDA graphs, smaller steps, other head counts, other attention backends and other platforms keep the fused launch. The result is bit-identical for NaN-free Q: with norm and RoPE off the op's live-Q path round-trips each element through fp32 unchanged and pads with +0.0, and the KV-only mode runs the same per-token KV code. test_q_interleaved_pad_copy_matches_fused_insert compares cache bytes and padded Q bits with the fused launch at 17, 1024 and 2049 tokens, with and without DP-padding rows. The PCP stand-in layer in test_compressor_kv_cache.py declares the new _pads_q_by_copy attribute. Signed-off-by: bofengl Co-Authored-By: Claude Opus 5.5 --- tests/kernels/test_compressor_kv_cache.py | 1 + ..._fused_deepseek_v4_qnorm_rope_kv_insert.py | 56 +++++++++ vllm/models/deepseek_v41/attention.py | 106 +++++++++++++++++- .../nvidia/flash_mla_mega_attn.py | 4 +- 4 files changed, 164 insertions(+), 3 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 3774d14b0c87..5e0016fb3b47 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -239,6 +239,7 @@ def test_v41_pcp_swa_cache_matches_non_pcp_writer(kv_mxfp8, monkeypatch): _flashinfer_fp8_kv_scale=None, kv_mxfp8=kv_mxfp8, accepts_unnormed_unroped_query=False, + _pads_q_by_copy=False, n_local_heads=num_heads, padded_heads=num_heads, eps=1e-20, 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 a14c889caae1..b3b856283b76 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 @@ -715,6 +715,62 @@ def run(q_in, is_q_interleaved): torch.testing.assert_close(cache_inter, cache_major, rtol=0, atol=0) +@pytest.mark.parametrize("num_tokens", [17, 1024, 2049]) +@pytest.mark.parametrize("dp_pad", [0, 5]) +def test_q_interleaved_pad_copy_matches_fused_insert(num_tokens: int, dp_pad: int): + """A KV-only launch plus ``_pad_interleaved_q`` equals the fused launch. + + Same cache bytes and same padded Q bits; the copy writes into a NaN-filled + buffer, and the last ``dp_pad`` rows have no slot, as under DP padding. + """ + from vllm.models.deepseek_v41.attention import _pad_interleaved_q + + torch.manual_seed(8) + device = "cuda" + dtype = torch.bfloat16 + block_size, n_heads, q_head_padded = 64, 32, 64 + total = num_tokens + dp_pad + + q = torch.randn(total, n_heads, HEAD_DIM, dtype=dtype, device=device) + kv = torch.randn(total, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(total, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(4096, 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) + + def insert(pad_to: int): + cache = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + out = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + cache, + slot_mapping, + positions, + cos_sin_cache, + pad_to, + 1e-6, + block_size, + False, # apply_q_norm + True, # kv_mxfp8 + False, # apply_q_rope + True, # is_q_interleaved + ) + return out, cache + + want, want_cache = insert(q_head_padded) + _, cache = insert(0) + assert want_cache.any() + torch.testing.assert_close(cache, want_cache, rtol=0, atol=0) + + got = torch.full_like(want, float("nan")) + _pad_interleaved_q(q, got) + torch.testing.assert_close( + got.view(torch.int16), want.view(torch.int16), rtol=0, atol=0 + ) + + # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ───────────────────── diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 76300d57eb61..e89e3202fc66 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -60,6 +60,12 @@ 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.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) +from vllm.models.deepseek_v41.common.ops.fused_layout import Q_CHUNK 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 @@ -90,6 +96,75 @@ # shard per TP rank and all-gathers the result. _SHARD_FUSED_WQA_WKV_MIN_TOKENS = 2048 +# From this many tokens, eager steps build the padded interleaved Q with +# _pad_interleaved_q instead of in the fused KV-insert op. Kept equal to the +# op's NUM_TOKEN_CUTOFF, where it switches to its reduced-grid kernel. +_Q_PAD_COPY_MIN_TOKENS = 1024 + + +@triton.jit +def _pad_interleaved_q_kernel( + q_ptr, + out_ptr, + num_rows, + LIVE: tl.constexpr, + PAD: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + # One row per (token, 16-dim chunk): LIVE copied elements, then PAD zeros. + pid = tl.program_id(0) + rows = pid.to(tl.int64) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS).to(tl.int64) + row_mask = rows[:, None] < num_rows + cols = tl.arange(0, LIVE) + x = tl.load(q_ptr + rows[:, None] * LIVE + cols[None, :], mask=row_mask) + out_rows = out_ptr + rows[:, None] * (LIVE + PAD) + tl.store(out_rows + cols[None, :], x, mask=row_mask) + pad_cols = tl.arange(0, PAD) + tl.store( + out_rows + LIVE + pad_cols[None, :], + tl.zeros([BLOCK_ROWS, PAD], dtype=x.dtype), + mask=row_mask, + ) + + +def _pad_interleaved_q_warmup_inputs( + *, + live_heads: int, + padded_heads: int, + head_dim: int, + dtype: torch.dtype, +) -> dict[str, object]: + return dict( + q=TritonWarmupTensor(dtype, shape=(1, live_heads, head_dim)), + out=TritonWarmupTensor(dtype, shape=(1, padded_heads, head_dim)), + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_pad_interleaved_q_kernel, + warmup_inputs=_pad_interleaved_q_warmup_inputs, +) +def _pad_interleaved_q(q: torch.Tensor, out: torch.Tensor) -> DispatchSpec: + """Copy chunk-interleaved ``q`` [N, H, D] into ``out`` [N, HP, D]. + + Writes the fused KV-insert op's interleaved Q layout: in each 16-dim chunk, + the H live heads copied bit for bit, then +0.0 in the HP - H padding slots. + Both tensors are dense and share a dtype; ``16 * H`` and ``16 * (HP - H)`` + must be powers of two. + """ + num_tokens, live_heads, head_dim = q.shape + padded_heads = out.shape[1] + assert out.shape == (num_tokens, padded_heads, head_dim) and out.dtype == q.dtype + num_rows = num_tokens * (head_dim // Q_CHUNK) + block_rows = 4 + return (triton.cdiv(num_rows, block_rows),), dict( + num_rows=num_rows, + LIVE=live_heads * Q_CHUNK, + PAD=(padded_heads - live_heads) * Q_CHUNK, + BLOCK_ROWS=block_rows, + num_warps=4, + ) + def _replace_layer_index(prefix: str, layer_id: int) -> str: """Swap the layer index inside a ``...layers....`` prefix.""" @@ -618,6 +693,15 @@ def __init__( else: self.compressed_cache_prefix = None + # _pad_interleaved_q needs power-of-two chunk widths: 32 live heads of + # 64 (e.g. TP2) qualify, 16 or 8 (768 or 896 pad columns) do not. + self._pads_q_by_copy = ( + self.accepts_unnormed_unroped_query + and self.n_local_heads == 32 + and self.padded_heads == 64 + and self.max_num_batched_tokens >= _Q_PAD_COPY_MIN_TOKENS + ) + if vllm_config.kernel_config.enable_jit_warmup: from vllm.v1.attention.backends.mla.sparse_swa import ( _COMPUTE_PREFILL_METADATA_KERNEL, @@ -630,6 +714,14 @@ def __init__( block_size=self.swa_cache_layer.block_size, ) + if self._pads_q_by_copy: + _pad_interleaved_q.register_warmup( + live_heads=self.n_local_heads, + padded_heads=self.padded_heads, + head_dim=self.head_dim, + dtype=vllm_config.model_config.dtype, + ) + if self.compress_ratio > 1: from vllm.v1.attention.backends.mla.compressor_utils import ( _COMPRESSED_SLOT_MAPPING_KERNEL, @@ -1061,6 +1153,8 @@ def _fused_qnorm_rope_kv_insert( attention kernel norms and rotates Q itself and reads it in its own chunk-interleaved layout, so the Q half of the launch is a zero-pad to ``padded_heads`` -- and nothing at all once the shard is that wide. + On large eager steps of a ``_pads_q_by_copy`` layer, a KV-only launch + leaves that pad to ``_pad_interleaved_q``. """ if not isinstance(attn_metadata, dict): # Profile run: kernel doesn't fire; produce a padded tensor so @@ -1126,6 +1220,12 @@ def _fused_qnorm_rope_kv_insert( and self.n_local_heads == self.padded_heads else self.padded_heads ) + # Captured graphs keep the fused pad. + copy_q = ( + self._pads_q_by_copy + and q.shape[0] >= _Q_PAD_COPY_MIN_TOKENS + and not torch.cuda.is_current_stream_capturing() + ) q_padded = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( q, kv, @@ -1133,7 +1233,7 @@ def _fused_qnorm_rope_kv_insert( slot_mapping, positions, cos_sin_cache, - pad_to, + 0 if copy_q else pad_to, self.eps, swa_metadata.block_size, False, # apply_q_norm: qr is normed before wq_b @@ -1141,6 +1241,10 @@ def _fused_qnorm_rope_kv_insert( not self.accepts_unnormed_unroped_query, # apply_q_rope self.accepts_unnormed_unroped_query, # is_q_interleaved ) + if copy_q: + out = q.new_empty((q.shape[0], self.padded_heads, q.shape[2])) + _pad_interleaved_q(q, out) + return out return q if pad_to == 0 else q_padded assert not self.accepts_unnormed_unroped_query, ( diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py index 7a7f4f8c9016..0a757ded94ae 100644 --- a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py +++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py @@ -5,7 +5,7 @@ One kernel does Q RoPE, sparse attention, the inverse RoPE of the output and its FP8 cast, and writes straight into the buffer ``wo_a`` consumes. So the layer declares ``accepts_unnormed_unroped_query`` -- the kernel RoPEs Q itself, -leaving the fused KV insert only the zero-pad to the kernel's head count -- and +leaving Q only the zero-pad to the kernel's head count -- and its ``_alloc_attn_out`` / ``_o_proj`` pair speaks QuantizedActivation instead of bf16, since the output needs no rotation or quantization here. @@ -253,7 +253,7 @@ def forward_mqa( return assert isinstance(attn_metadata, dict) - # q was zero-padded to the kernel's head count by the fused KV insert. + # _fused_qnorm_rope_kv_insert zero-padded q to the kernel's head count. assert q.shape[1] == self.padded_heads flashmla_metadata = cast( DeepseekV4FlashMLAMetadata | None, From 274a5e19a18460ed41d47117398494b721e3df16 Mon Sep 17 00:00:00 2001 From: Qing Date: Tue, 27 Oct 2026 13:45:57 -0700 Subject: [PATCH 002/406] [Perf][DSv4.1] Capture the decoder replay graphs on the model graphs' stream (#63155) Signed-off-by: Qing Signed-off-by: qing.xiao Co-authored-by: vllm-agent Co-authored-by: qing.xiao --- ...test_deepseek_v41_decoder_replay_layers.py | 14 +++++++++++-- .../v1/cudagraph/test_breakable_cudagraph.py | 5 ++++- tests/v1/worker/test_gpu_model_runner_v2.py | 21 ++++++++++++++----- vllm/distributed/parallel_state.py | 6 ++++-- vllm/models/deepseek_v41/amd/model.py | 1 + .../deepseek_v41/decoder_replay_layers.py | 4 ++++ .../nvidia/decoder_replay_cudagraph.py | 18 +++++++++++++--- vllm/models/deepseek_v41/nvidia/model.py | 1 + .../models/deepseek_v41/nvidia/model_state.py | 3 ++- .../specialized_decode/runtime/model_state.py | 3 ++- vllm/v1/worker/gpu/cudagraph_utils.py | 8 ++++++- vllm/v1/worker/gpu/model_runner.py | 1 + vllm/v1/worker/gpu/model_states/interface.py | 4 +++- 13 files changed, 72 insertions(+), 17 deletions(-) diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index b7f5ca29ec2d..fafcebb610f8 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DecoderReplayLayers runs the replay layers on the forward's replay batch.""" +from contextlib import contextmanager from unittest.mock import MagicMock import pytest @@ -9,6 +10,7 @@ from vllm.config import CompilationConfig, CUDAGraphMode from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS, NATIVE_MEGA_MOE_BACKENDS +from vllm.distributed.parallel_state import GraphCaptureContext from vllm.forward_context import ( BatchDescriptor, ForwardContext, @@ -158,8 +160,16 @@ def test_replay_graph_matches_eager(monkeypatch, moe_backend, ids_dtype): """The replay graph of the next captured size pads and runs the rows. Input ids arrive as the model passes them: int32, cast to int64 for native MegaMoE only.""" monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock) - capture = torch.cuda.stream(torch.cuda.Stream()) - monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture) + + @contextmanager + def fake_graph_capture(device, stream=None): + if stream is None: + stream = torch.cuda.Stream(device=device) + stream.wait_stream(torch.cuda.current_stream(device)) + with torch.cuda.stream(stream): + yield GraphCaptureContext(stream) + + monkeypatch.setattr(cudagraph_utils, "graph_capture", fake_graph_capture) compilation = CompilationConfig(decoder_replay_cudagraph_capture_sizes=[4]) cfg = MagicMock(compilation_config=compilation, speculative_config=None) # Mocking it true would capture into the cuMem pool, which breaks on ROCm. diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index 4e19912f4fef..6a5970f26e25 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -6,6 +6,7 @@ import threading from contextlib import nullcontext +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -68,10 +69,11 @@ def forward_fn(cg_mode): return forward_fn + capture_stream = object() with ( patch( "vllm.v1.worker.gpu.cudagraph_utils.graph_capture", - return_value=nullcontext(), + return_value=nullcontext(SimpleNamespace(stream=capture_stream)), ), patch( "vllm.v1.worker.gpu.cudagraph_utils.is_global_first_rank", @@ -80,6 +82,7 @@ def forward_fn(cg_mode): ): manager.capture(create_forward_fn) + assert manager.capture_stream is capture_stream assert [warmup for warmup, _ in create_calls] == [True, False] assert [mode for _, mode, _ in forward_calls] == [ CUDAGraphMode.NONE, diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 752107c2f851..9aebddceda3f 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -283,23 +283,24 @@ def _make_capture_runner( cudagraph_manager's needs_capture decision.""" runner = GPUModelRunner.__new__(GPUModelRunner) runner.model_state = SimpleNamespace( - supports_mm_inputs=False, capture_inner_cudagraphs=lambda *args: None + supports_mm_inputs=False, capture_inner_cudagraphs=Mock() ) runner.cudagraph_manager = SimpleNamespace( needs_capture=lambda: captured, capture=lambda *args, **kwargs: None, + capture_stream=object(), ) runner.lora_config = None monkeypatch.setattr(runner, "maybe_setup_dummy_loras", _no_dummy_loras) runner.speculator = None runner.adaptive_verification = None runner.model = None - runner.input_buffers = None + runner.input_buffers = object() runner.pcp_manager = None runner.intermediate_tensors = None - runner.block_tables = None - runner.attn_groups = None - runner.kv_cache_config = None + runner.block_tables = object() + runner.attn_groups = object() + runner.kv_cache_config = object() runner.use_aux_hidden_state_outputs = False runner.kv_connector = model_runner_module.NO_OP_KV_CONNECTOR return runner @@ -325,6 +326,16 @@ def test_capture_model_locks_workspace_after_capture(monkeypatch): runner.capture_model() assert lock_calls == ["lock"] + assert runner.cudagraph_manager is not None + capture_inner_cudagraphs = runner.model_state.capture_inner_cudagraphs + assert isinstance(capture_inner_cudagraphs, Mock) + capture_inner_cudagraphs.assert_called_once_with( + runner.input_buffers, + runner.block_tables, + runner.attn_groups, + runner.kv_cache_config, + runner.cudagraph_manager.capture_stream, + ) def test_capture_model_skips_lock_when_nothing_captured(monkeypatch): diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 897518247f13..a9c9d0b87115 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1692,7 +1692,7 @@ def _aiter_registers_graph_buffers(group: GroupCoordinator) -> bool: @contextmanager -def graph_capture(device: torch.device): +def graph_capture(device: torch.device, stream: torch.cuda.Stream | None = None): """`graph_capture` is a context manager which should surround the code that is capturing the CUDA graph. Its main purpose is to ensure that some operations will be run after the graph is captured, before the graph @@ -1704,12 +1704,14 @@ def graph_capture(device: torch.device): the graph capture is running on a separate stream from the default stream, in order to explicitly distinguish the kernels to capture from other kernels possibly launched on background in the default stream. + A given ``stream`` is reused instead, so graphs captured on it share the + caching allocator's free blocks of the graph pool. """ from vllm.device_allocator.alloc_conf import non_expandable_allocations from vllm.platforms import current_platform tp, pp, dp = get_tp_group(), get_pp_group(), get_dp_group() - context = GraphCaptureContext(torch.cuda.Stream(device=device)) + context = GraphCaptureContext(stream or torch.cuda.Stream(device=device)) # Let AITER all-reduce register captured buffers instead of copying them. # Not with DP AITER collectives: they reduce workspaces allocated earlier. non_expandable = ( diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index e7c513047140..b9094c697f50 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -604,6 +604,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.layers[cut].write_kv, metadata_prefixes, self.layers[cut].attn.swa_cache_layer.prefix, + lambda: any(i < cut for i in self.aux_hidden_state_layers), ) logger.info_once( "Decoder SWA bounded replay: in eager prefill steps, layer %d " diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py index b109d8634a1a..a08cd8063597 100644 --- a/vllm/models/deepseek_v41/decoder_replay_layers.py +++ b/vllm/models/deepseek_v41/decoder_replay_layers.py @@ -46,6 +46,8 @@ class DecoderReplayLayers: every row. ``metadata_prefixes`` are the attention metadata keys the layers read, so the replay batch builds only those. ``first_swa_prefix`` keys the first replay layer's SWA metadata, built with the batch's window starts. + ``has_early_aux`` tells whether layers before the replay layers capture aux + hidden states, which the forward holds across the replay layers. """ def __init__( @@ -55,12 +57,14 @@ def __init__( write_batch_kv: Callable[..., None], metadata_prefixes: set[str], first_swa_prefix: str, + has_early_aux: Callable[[], bool] = lambda: False, ) -> None: self.window = window self.run_layers = run_layers self.write_batch_kv = write_batch_kv self.metadata_prefixes = metadata_prefixes self.first_swa_prefix = first_swa_prefix + self.has_early_aux = has_early_aux # Set by the model state every step; None runs the layers on the batch. self.replay_batch: ReplayBatch | None = None # Set by the model state with replay graphs: PIECEWISE model graphs of diff --git a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py index c985f973bb49..e8a8e2c33539 100644 --- a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py +++ b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py @@ -68,9 +68,21 @@ def __init__(self, vllm_config: VllmConfig, device, layers: DecoderReplayLayers) ] def capture_replay_graphs( - self, prepare: Callable[[BatchExecutionDescriptor], None] + self, + prepare: Callable[[BatchExecutionDescriptor], None], + stream: torch.cuda.Stream | None = None, ) -> None: - """Capture each size on the replay batch ``prepare`` sets on the layers.""" + """Capture each size on the replay batch ``prepare`` sets on the layers. + + On the model graphs' ``stream`` they share its free blocks of the graph + pool: a replay graph runs in a model graph's eager break, after it has + gathered its inputs into ``input_buffers`` and before the model graph + reads its outputs from the break's own buffers, so no model graph + tensor in the pool is live while it runs. Aux hidden states captured + before the replay layers would be, so with those it keeps its own. + """ + if self.layers.has_early_aux(): + stream = None def create_forward_fn(desc: BatchExecutionDescriptor, warmup: bool): prepare(desc) @@ -86,7 +98,7 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None: return forward_fn - self.capture(create_forward_fn, "Capturing decoder replay CUDA graphs") + self.capture(create_forward_fn, "Capturing decoder replay CUDA graphs", stream) self.layers.replay_batch = None def run(self, rows: torch.Tensor, inputs: tuple) -> tuple[torch.Tensor, ...]: diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index b71dadfe9755..fb694faf5ea7 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -889,6 +889,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.layers[cut].write_kv, metadata_prefixes, self.layers[cut].attn.swa_cache_layer.prefix, + lambda: any(i < cut for i in self.aux_hidden_state_layers), ) logger.info_once( "Decoder SWA bounded replay: in eager prefill steps, layer %d " diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index ed09ae3549ec..03a9109484d6 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -635,6 +635,7 @@ def capture_inner_cudagraphs( block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, + stream: torch.cuda.Stream | None = None, ) -> None: def prepare(desc: BatchExecutionDescriptor) -> None: # Capture reuses the runtime path on a dummy batch, so the graph reads @@ -654,7 +655,7 @@ def prepare(desc: BatchExecutionDescriptor) -> None: if self.replay_cudagraphs is not None: # profile_only captures use pre-KV-cache groups: drop the cache. self._replay_attn_groups = None - self.replay_cudagraphs.capture_replay_graphs(prepare) + self.replay_cudagraphs.capture_replay_graphs(prepare, stream) def _kept_input_batch( self, diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py index 2fc7ea5c8ff1..251acde01edc 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/runtime/model_state.py @@ -154,6 +154,7 @@ def capture_inner_cudagraphs( block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, + stream: torch.cuda.Stream | None = None, ) -> None: # Record the metadata build of every decode step the FULL decode graphs # serve, on dummy batches of that step's shape. @@ -173,5 +174,5 @@ def capture_inner_cudagraphs( finally: self._recording = None super().capture_inner_cudagraphs( - input_buffers, block_tables, attn_groups, kv_cache_config + input_buffers, block_tables, attn_groups, kv_cache_config, stream ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index d29fd3341bb5..33555fcc7aee 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -179,6 +179,9 @@ def __init__( self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {} self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None + # The stream of the last capture; graphs replayed inside these ones (in + # their eager breaks) can capture on it to share the pool's free blocks. + self.capture_stream: torch.cuda.Stream | None = None self._graphs_captured = False @@ -425,6 +428,7 @@ def capture( self, create_forward_fn: CreateForwardFn, progress_bar_desc: str = "Capturing CUDA graphs", + stream: torch.cuda.Stream | None = None, ) -> None: """Capture CUDA graphs. @@ -435,9 +439,11 @@ def capture( because attention backends may mutate or lazily initialize metadata during warmup. progress_bar_desc: Description shown on the capture progress bar. + stream: Capture on this stream instead of a new one. """ - with graph_capture(device=self.device), ExitStack() as stack: + with graph_capture(self.device, stream) as context, ExitStack() as stack: + self.capture_stream = context.stream 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) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 025fc08567d4..b8d34076b082 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1155,6 +1155,7 @@ def capture_model(self, *, profile_only: bool = False) -> int: self.block_tables, self.attn_groups, self.kv_cache_config, + self.cudagraph_manager.capture_stream, ) if self.speculator is not None: with use_workspace_lane(self._draft_workspace_lane): diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 3db3fc9e7ae6..98eb552dd304 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -139,8 +139,10 @@ def capture_inner_cudagraphs( block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, + stream: torch.cuda.Stream | None = None, ) -> None: - """Capture the CUDA graphs this state runs inside the model's forward.""" + """Capture the CUDA graphs this state runs inside the model's forward, + on ``stream``: the model graphs' capture stream.""" return None def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]: From bb7ea479c6e11d18d8c0eb1e195e1ffb27c87e01 Mon Sep 17 00:00:00 2001 From: seungrokj <144636725+seungrokj@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:08:27 -0700 Subject: [PATCH 003/406] [ROCm][Bugfix][Kimi-K3] Pass zero positions to the fused MLA decode q-prep kernel (#63409) Signed-off-by: seungrokj Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .../kimi_k3/test_amd_mla_direct_return.py | 71 +++++++++++++++++-- vllm/models/kimi_k3/amd/mla.py | 34 +++++++-- 2 files changed, 94 insertions(+), 11 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_mla_direct_return.py b/tests/models/kimi_k3/test_amd_mla_direct_return.py index 57df5bad3344..47c5e373bfac 100644 --- a/tests/models/kimi_k3/test_amd_mla_direct_return.py +++ b/tests/models/kimi_k3/test_amd_mla_direct_return.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace import pytest import torch @@ -91,9 +92,7 @@ def attention_output(path, output_shape): _fused_qk_prep=fused_decode, _qkv_gate_fused=False, _fold_qb_uk=False, - _fused_decode=lambda q, kv, pe, pos, shape, *args: attention_output( - "fused", shape - ), + _fused_decode=lambda q, kv, pe, shape, *args: attention_output("fused", shape), g_proj=None, o_proj=lambda value: (value, None), ) @@ -119,3 +118,67 @@ def attention_output(path, output_shape): assert output is projected assert paths == ["fused" if fused_decode else "regular"] + + +def test_fused_decode_feeds_zero_positions_to_kernel(monkeypatch) -> None: + """The fused kernel's identity RoPE must see position 0 for every slot. + + Not every AITER variant clamps positions into the one-row cos/sin cache, and + the kernel takes its token count from slot_mapping, padding included. + """ + from vllm.models.kimi_k3.amd.mla import KimiK3MultiHeadLatentAttentionWrapper + + calls = [] + cache_ops = ModuleType("aiter.ops.cache") + monkeypatch.setattr( + cache_ops, + "fused_qk_rope_concat_and_cache_mla", + lambda *args, **kwargs: calls.append(args), + raising=False, + ) + monkeypatch.setitem(sys.modules, "aiter", ModuleType("aiter")) + monkeypatch.setitem(sys.modules, "aiter.ops", ModuleType("aiter.ops")) + monkeypatch.setitem(sys.modules, "aiter.ops.cache", cache_ops) + + num_tokens, num_heads, nope, rope, lora, v_dim = 4, 2, 4, 2, 4, 4 + head_size = lora + rope + zero_positions = torch.zeros(16, dtype=torch.int64) + output = torch.randn(num_tokens, num_heads * v_dim, dtype=torch.bfloat16) + attn = SimpleNamespace( + is_aiter_triton_fp4_bmm_enabled=False, + W_UK_T=torch.randn(num_heads, nope, lora, dtype=torch.bfloat16), + kv_lora_rank=lora, + impl=SimpleNamespace(forward_mqa=lambda q, *args: (q, None)), + _v_up_proj=lambda attn_out, out: out.copy_(output), + ) + wrapper = SimpleNamespace( + mla_attn=attn, + qk_nope_head_dim=nope, + qk_rope_head_dim=rope, + _identity_rope=None, + _fused_q_heads=4, + _zero_positions=zero_positions, + ) + # A CUDA-graph padded batch: the last slot is a pad. + slot_mapping = torch.tensor([[5], [9], [2], [-1]]) + + result = KimiK3MultiHeadLatentAttentionWrapper._fused_decode( + wrapper, + torch.randn(num_tokens, num_heads, nope + rope, dtype=torch.bfloat16), + torch.randn(num_tokens, lora, dtype=torch.bfloat16), + torch.randn(num_tokens, 1, rope, dtype=torch.bfloat16), + tuple(output.shape), + None, + SimpleNamespace(_k_scale=torch.ones(1), _q_scale=torch.ones(1)), + torch.zeros(2, 4, head_size, dtype=torch.uint8), + slot_mapping, + ) + + assert len(calls) == 1 + kernel_slots, positions = calls[0][6], calls[0][9] + assert torch.equal(kernel_slots, slot_mapping.flatten()) + assert positions.shape == kernel_slots.shape + assert not positions.any() + # A view of the persistent buffer, so captured graphs allocate nothing. + assert positions.data_ptr() == zero_positions.data_ptr() + assert torch.equal(result, output) diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py index 73f7e506b76d..0c30e0d634b8 100644 --- a/vllm/models/kimi_k3/amd/mla.py +++ b/vllm/models/kimi_k3/amd/mla.py @@ -7,6 +7,7 @@ import torch from vllm._aiter_ops import rocm_aiter_ops +from vllm.config import get_current_vllm_config from vllm.model_executor.layers.attention.attention import get_attention_context from vllm.model_executor.layers.attention.mla_attention import split_kv_b_proj from vllm.model_executor.layers.fusion.quant_activation import ( @@ -46,6 +47,18 @@ def __init__(self, *args, **kwargs) -> None: self._fused_q_heads = AiterMLAHelper.get_actual_mla_num_heads( self.num_heads ) + # Positions fed to the fused kernel's identity RoPE; see _fused_decode. + # Allocated once here so every captured CUDA graph shares one zeroed + # buffer instead of capturing a memset per layer. + zero_positions = None + if self._fused_qk_prep: + vllm_config = get_current_vllm_config() + max_tokens = max( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.compilation_config.max_cudagraph_capture_size or 0, + ) + zero_positions = torch.zeros(max_tokens, dtype=torch.int64) + self.register_buffer("_zero_positions", zero_positions, persistent=False) # The fold reads q_b_proj.weight as plain bf16/fp16 rows, so a quantized # q_b_proj keeps the unfused path. self._fold_qb_uk = ( @@ -174,7 +187,6 @@ def _fused_decode( q: torch.Tensor, kv_c_normed: torch.Tensor, k_pe: torch.Tensor, - positions: torch.Tensor, output_shape: tuple[int, int], attn_metadata, layer, @@ -228,9 +240,9 @@ def _fused_decode( if self._identity_rope is None: # cos = 1, sin = 0 makes the kernel's RoPE the identity, which is - # what a NoPE model needs. One row suffices because the `_opt` - # kernel clamps `pos` into [0, cos_cache.size(0)); a full-length - # cache would cost ~268 MB at K3's 1M max_position_embeddings. + # what a NoPE model needs. A full-length cache would cost ~268 MB at + # K3's 1M max_position_embeddings, so it holds one row and the + # kernel is given all-zero positions below. half = self.qk_rope_head_dim // 2 opts = {"dtype": kv_c_normed.dtype, "device": kv_c_normed.device} self._identity_rope = ( @@ -253,6 +265,15 @@ def _fused_decode( device=ql_nope.device, ) q_out = q_out_padded[:, :num_heads] + # The kernel must not see the real positions: not every AITER variant + # clamps them into the one-row cos/sin cache. The head-grouped kernel + # (AITER v0.1.26) reads cos = sin = 0 past the cache at one head per + # block, zeroing q_pe and the cached k_pe, and faults at more heads per + # block. RoPE is the identity, so position 0 is exact for every token. + # The kernel takes its token count from slot_mapping. + slot_mapping = slot_mapping.flatten() + num_slots = slot_mapping.shape[0] + assert num_slots <= self._zero_positions.shape[0] fused_qk_rope_concat_and_cache_mla( ql_nope, q_pe, @@ -260,10 +281,10 @@ def _fused_decode( k_pe.squeeze(1), kv_cache.view(kv_cache.shape[0], -1, head_size), q_out, - slot_mapping.flatten(), + slot_mapping, layer._k_scale, layer._q_scale, - positions, + self._zero_positions[:num_slots], cos_cache, sin_cache, is_neox=True, @@ -383,7 +404,6 @@ def forward( q, kv_c_normed, k_pe, - positions, output_shape, attn_metadata, layer, From c6f6e6a9835c3e4df9c06bb6c564c7d866be9f22 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:08:59 -0400 Subject: [PATCH 004/406] [GLM5.3 Perf] Enable dspark fp8 kv cache for draft model, 5.38 GB GPU Memory saved per rank (#63841) Signed-off-by: yewentao256 Co-authored-by: vllm-agent --- vllm/model_executor/models/config.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 40d00345a9ca..e40171abcbaa 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -99,12 +99,10 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ): cache_config.cache_dtype = "fp8_e4m3" logger.info_once("Using fp8 kv-cache for GlmMoeDsaForCausalLM on SM10x") - # MTP shares the target's KV cache layout; a separate draft model - # must not inherit a default it was never validated with. spec_config = vllm_config.speculative_config if ( spec_config is not None - and spec_config.method != "mtp" + and spec_config.method not in ("mtp", "dspark") and spec_config.kv_cache_dtype is None ): spec_config.kv_cache_dtype = "auto" From c4f2bb4abbdbad9ed5b0a71acd0bd585056eeab8 Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Tue, 27 Oct 2026 17:09:32 -0400 Subject: [PATCH 005/406] [Bugfix][Spec Decode] Preserve grammar-invalid drafts for rejection sampling (#61557) Signed-off-by: Yifan Zong --- .buildkite/test_areas/spec_decode.yaml | 2 + tests/v1/core/test_scheduler.py | 104 +----------------- tests/v1/e2e/general/test_async_scheduling.py | 44 ++++---- .../general/test_async_scheduling_accuracy.py | 23 ++-- .../spec_decode/test_distribution_shift.py | 76 +++++++++++++ .../test_grammar_backends.py | 73 ++++++++++++ .../test_structured_output_manager.py | 98 +++++++++-------- vllm/v1/core/sched/scheduler.py | 7 +- vllm/v1/engine/core.py | 9 +- vllm/v1/structured_output/__init__.py | 77 ++++++------- vllm/v1/structured_output/backend_guidance.py | 37 +++++-- vllm/v1/structured_output/backend_types.py | 2 + vllm/v1/structured_output/backend_xgrammar.py | 12 +- 13 files changed, 315 insertions(+), 249 deletions(-) create mode 100644 tests/v1/e2e/spec_decode/test_distribution_shift.py create mode 100644 tests/v1/structured_output/test_grammar_backends.py diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 61b6677ae586..411e37aa2db6 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -456,6 +456,7 @@ steps: commands: # DFlash2 NVFP4 (27B) does not fit a 35GB MIG slice; runs on B200 below. - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" + - pytest -v -s v1/e2e/spec_decode/test_distribution_shift.py mirror: amd: label: ":amd: (MI355 DPX) Spec Decode AL DFlash Nightly" @@ -465,6 +466,7 @@ steps: timeout_in_minutes: 120 commands: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" + - pytest -v -s v1/e2e/spec_decode/test_distribution_shift.py depends_on: - image-build-amd source_file_dependencies: diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 5fa8e206dd4e..18bb469cd720 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -8312,93 +8312,6 @@ def test_encoder_input_skipped_when_connector_already_has_the_item(ec_role: str) assert not output.scheduled_encoder_inputs.get(req_id) -class _RecordingGrammar: - """Records what reaches ``validate_tokens`` and raises on a negative id. - - llguidance raises ``OverflowError`` on a negative token id, so this mirrors - the real failure: the tests below fail on the unpatched scheduler (the -1 - padding reaches ``validate_tokens``) and pass once it is stripped first. - """ - - def __init__(self) -> None: - self.seen: list[list[int]] = [] - - def validate_tokens(self, tokens: list[int]) -> list[int]: - if any(t < 0 for t in tokens): - raise OverflowError("llguidance rejects negative token ids") - self.seen.append(list(tokens)) - return tokens - - -def _decode_ready_request(scheduler): - """Add one request and advance it out of prefill into decode.""" - request = create_requests(num_requests=1, num_tokens=1)[0] - scheduler.add_request(request) - output = scheduler.schedule() - scheduler.update_from_output( - 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=[], - ), - ) - return request - - -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, reasoning_ended=True - ) - - scheduler.update_draft_token_ids( - DraftTokenIds([request.request_id], [[10, 11, -1, -1]]) - ) - - assert grammar.seen == [[10, 11]] - assert request.spec_token_ids == [10, 11] - - -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 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, reasoning_ended=True - ) - - scheduler_output = SimpleNamespace( - scheduled_spec_decode_tokens={request.request_id: [0, 0, 0, 0]} - ) - scheduler.update_draft_token_ids_in_output( - DraftTokenIds([request.request_id], [[10, 11, -1, -1]]), - scheduler_output, - ) - - # The grammar only saw the stripped prefix, never a -1. - assert grammar.seen == [[10, 11]] - # Two drafts were rejected (4 scheduled - 2 valid), padded back with -1. - assert scheduler_output.scheduled_spec_decode_tokens[request.request_id] == [ - 10, - 11, - -1, - -1, - ] - assert scheduler_output.num_invalid_spec_tokens == {request.request_id: 2} - - def test_diffusion_canvas_width_defaults_to_the_served_canvas(): def req(extra): return SimpleNamespace(sampling_params=SimpleNamespace(extra_args=extra)) @@ -8464,13 +8377,10 @@ def test_diffusion_scheduler_narrows_the_canvas_per_request(): assert narrow.num_output_placeholders == 4 -@pytest.mark.parametrize("structured", [False, True]) @pytest.mark.parametrize("async_scheduling", [False, True]) @pytest.mark.usefixtures("diffusion_model_runner") -def test_diffusion_scheduler_trims_full_width_worker_drafts( - structured, async_scheduling -): - """Padded worker drafts must be narrowed before scheduling or grammar validation.""" +def test_diffusion_scheduler_trims_full_width_worker_drafts(async_scheduling): + """Padded worker drafts must be narrowed before scheduling.""" scheduler = _diffusion_scheduler(async_scheduling=async_scheduling) wide = _diffusion_request("wide", {}) narrow = _diffusion_request("narrow", {"diffusion_canvas_length": 4}) @@ -8479,12 +8389,7 @@ def test_diffusion_scheduler_trims_full_width_worker_drafts( prefill = scheduler.schedule() _model_output(scheduler, prefill, [[], []]) - if structured: - for request in (wide, narrow): - request.structured_output_request = SimpleNamespace( - grammar=_RecordingGrammar(), reasoning_ended=True - ) - tokens = list(range(8)) if structured else [-1] * 8 + tokens = [-1] * 8 drafts = DraftTokenIds(["wide", "narrow"], [tokens.copy(), tokens.copy()]) if async_scheduling: output = scheduler.schedule() @@ -8498,9 +8403,6 @@ def test_diffusion_scheduler_trims_full_width_worker_drafts( "narrow": tokens[:4], } assert output.num_scheduled_tokens == {"wide": 8, "narrow": 4} - if structured: - assert wide.structured_output_request.grammar.seen == [tokens] - assert narrow.structured_output_request.grammar.seen == [tokens[:4]] @pytest.mark.parametrize( diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py index 28f9943b3b59..3378e0ac3557 100644 --- a/tests/v1/e2e/general/test_async_scheduling.py +++ b/tests/v1/e2e/general/test_async_scheduling.py @@ -241,15 +241,15 @@ def run_tests( print(f"BASELINE: config=[{baseline_config}], accept_rates={baseline_acceptances}") failure = None - for test_config, test_outputs, test_acceptance_rates in outputs[1:]: - for (base_outs, base_logprobs), base_acceptance_rate, ( + for test_config, test_outputs, test_acceptance_lengths in outputs[1:]: + for (base_outs, base_logprobs), base_acceptance_length, ( test_outs, test_logprobs, - ), test_acceptance_rate, params in zip( + ), test_acceptance_length, params in zip( baseline_tests, baseline_acceptances or repeat(None), test_outputs, - test_acceptance_rates or repeat(None), + test_acceptance_lengths or repeat(None), test_sampling_params, ): reason = None @@ -272,8 +272,8 @@ def run_tests( if reason is None: try: if ( - base_acceptance_rate is not None - and test_acceptance_rate is not None + base_acceptance_length is not None + and test_acceptance_length is not None ): if "spec_mml=None" in test_config: # Preemption causes more variance in acceptance rates @@ -285,14 +285,14 @@ def run_tests( else: tolerance = 0.05 assert ( - test_acceptance_rate > base_acceptance_rate - or test_acceptance_rate - == pytest.approx(base_acceptance_rate, rel=tolerance) + test_acceptance_length > base_acceptance_length + or test_acceptance_length + == pytest.approx(base_acceptance_length, rel=tolerance) ) else: # Currently the reported acceptance rate is expected to be # lower when we sometimes skip drafting altogether. - assert test_acceptance_rate > 0.1 + assert test_acceptance_length > 1.1 except AssertionError as e: reason = "accept ", e @@ -300,14 +300,14 @@ def run_tests( print( f"\033[32mPASSED\033[0m: " f"config=[{test_config}], params={params}" - f" accept_rate={test_acceptance_rate}" + f" accept_rate={test_acceptance_length}" ) else: reason_str, _ = reason print( f"\033[31mFAILED\033[0m({reason_str}): " f"config=[{test_config}], params={params}" - f" accept_rate={test_acceptance_rate}" + f" accept_rate={test_acceptance_length}" ) if failure is None: _, failure = reason @@ -366,7 +366,7 @@ def run_test( **cache_arg, ) as vllm_model: results = [] - acceptance_rates: list[float] | None = [] if spec_decoding else None + acceptance_lengths: list[float] | None = [] if spec_decoding else None for override_params in sampling_param_tests: metrics_before = vllm_model.llm.get_metrics() print(f"----------- RUNNING PARAMS: {override_params}") @@ -378,10 +378,12 @@ def run_test( ) ) metrics_after = vllm_model.llm.get_metrics() - if acceptance_rates is not None: - acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after) - acceptance_rates.append(acceptance_rate) - print(f"ACCEPTANCE RATE {acceptance_rate}") + if acceptance_lengths is not None: + acceptance_length = _get_acceptance_length( + metrics_before, metrics_after + ) + acceptance_lengths.append(acceptance_length) + print(f"ACCEPTANCE RATE {acceptance_length}") if test_preemption: preemptions = _get_count( @@ -404,7 +406,7 @@ def run_test( ) assert _all_logprobs_match(results[0][1], other_test_logprobs) - return test_config, results, acceptance_rates + return test_config, results, acceptance_lengths def _all_logprobs_match(req_a, req_b) -> bool: @@ -438,10 +440,10 @@ def _logprobs_match( ) -def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float: - draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens") +def _get_acceptance_length(before: list[Metric], after: list[Metric]) -> float: + draft = _get_count(before, after, "vllm:spec_decode_num_drafts") accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens") - return accept / draft if draft > 0 else 0.0 + return 1 + accept / draft if draft > 0 else 1.0 def _get_count(before: list[Metric], after: list[Metric], name: str) -> int: diff --git a/tests/v1/e2e/general/test_async_scheduling_accuracy.py b/tests/v1/e2e/general/test_async_scheduling_accuracy.py index 0fd208b1ea7c..1133a4dbf10b 100644 --- a/tests/v1/e2e/general/test_async_scheduling_accuracy.py +++ b/tests/v1/e2e/general/test_async_scheduling_accuracy.py @@ -292,7 +292,7 @@ def run_tests( f"{context}: acceptance={actual}, baseline={baseline}" ) else: - assert actual > 0.1, f"{context}: acceptance={actual}" + assert actual > 1.1, f"{context}: acceptance={actual}" def run_test( @@ -346,7 +346,7 @@ def run_test( **cache_arg, ) as vllm_model: results = [] - acceptance_rates: list[float] | None = [] if spec_decoding else None + acceptance_lengths: list[float] | None = [] if spec_decoding else None for override_params in sampling_param_tests: metrics_before = vllm_model.llm.get_metrics() print(f"----------- RUNNING PARAMS: {override_params}") @@ -357,10 +357,12 @@ def run_test( ) ) metrics_after = vllm_model.llm.get_metrics() - if acceptance_rates is not None: - acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after) - acceptance_rates.append(acceptance_rate) - print(f"ACCEPTANCE RATE {acceptance_rate}") + if acceptance_lengths is not None: + acceptance_length = _get_acceptance_length( + metrics_before, metrics_after + ) + acceptance_lengths.append(acceptance_length) + print(f"ACCEPTANCE RATE {acceptance_length}") if test_preemption: preemptions = _get_count( @@ -387,7 +389,7 @@ def run_test( f"{test_config}: structured sampling parameters had no effect" ) - return test_config, results, acceptance_rates + return test_config, results, acceptance_lengths def _result_signature(requests): @@ -402,12 +404,11 @@ def _result_signature(requests): ] -def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float: - draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens") +def _get_acceptance_length(before: list[Metric], after: list[Metric]) -> float: + draft = _get_count(before, after, "vllm:spec_decode_num_drafts") accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens") assert draft > 0, "speculative test did not draft any tokens" - assert 0 <= accept <= draft, f"invalid acceptance counters: {accept=}, {draft=}" - return accept / draft + return 1 + accept / draft def _get_count(before: list[Metric], after: list[Metric], name: str) -> int: diff --git a/tests/v1/e2e/spec_decode/test_distribution_shift.py b/tests/v1/e2e/spec_decode/test_distribution_shift.py new file mode 100644 index 000000000000..7cbe4827611b --- /dev/null +++ b/tests/v1/e2e/spec_decode/test_distribution_shift.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import Counter + +import pytest +import regex as re +from scipy.stats import chi2_contingency + +from tests.utils import single_gpu_only +from vllm import SamplingParams +from vllm.sampling_params import StructuredOutputsParams + +from .utils import compute_acceptance_rate + + +@pytest.mark.parametrize("rejection_sample_method", ["standard", "block"]) +@pytest.mark.parametrize("async_scheduling", [False, True], ids=["sync", "async"]) +@single_gpu_only +def test_structured_output_preserves_sampling_distribution( + rejection_sample_method, async_scheduling, vllm_runner +): + spec_config = { + "method": "dflash", + "model": "z-lab/Qwen3-4B-DFlash-b16", + "num_speculative_tokens": 3, + "draft_sample_method": "probabilistic", + "rejection_sample_method": rejection_sample_method, + } + regex = r'\{"city": "(Paris|Berlin|Madrid)"\}' + num_samples = 4000 + cities_counts = [] + for config in (None, spec_config): + with vllm_runner( + "Qwen/Qwen3-4B", + speculative_config=config, + async_scheduling=async_scheduling, + enforce_eager=True, + enable_prefix_caching=False, + gpu_memory_utilization=0.8, + disable_log_stats=False, + ) as runner: + params = [ + SamplingParams( + temperature=1.0, + max_tokens=32, + seed=i, + structured_outputs=StructuredOutputsParams(regex=regex), + ) + for i in range(num_samples) + ] + outputs = runner.get_llm().chat( + [[{"role": "user", "content": "Name a random European capital."}]] + * num_samples, + params, + use_tqdm=False, + chat_template_kwargs={"enable_thinking": False}, + ) + if config is not None: + assert compute_acceptance_rate(runner.get_llm().get_metrics()) > 0 + texts = [o.outputs[0].text for o in outputs] + assert all(re.fullmatch(regex, text) for text in texts), ( + "Grammar-invalid output" + ) + cities_counts.append(Counter(texts)) + assert_grammar_distribution_unchanged(*cities_counts) + + +def assert_grammar_distribution_unchanged(counts1: Counter[str], counts2: Counter[str]): + outcomes = sorted(counts1.keys() | counts2.keys()) + row1 = [counts1[outcome] for outcome in outcomes] + row2 = [counts2[outcome] for outcome in outcomes] + _, p, _, _ = chi2_contingency([row1, row2], correction=False) + assert p > 1e-3, ( + f"Grammar output distributions differ: {counts1=}, {counts2=}, {p=}" + ) diff --git a/tests/v1/structured_output/test_grammar_backends.py b/tests/v1/structured_output/test_grammar_backends.py new file mode 100644 index 000000000000..6ed0756b31ae --- /dev/null +++ b/tests/v1/structured_output/test_grammar_backends.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable + +import pytest +from transformers import AutoTokenizer + +from vllm.config import DeviceConfig, VllmConfig +from vllm.v1.structured_output.backend_guidance import GuidanceBackend +from vllm.v1.structured_output.backend_lm_format_enforcer import ( + LMFormatEnforcerBackend, +) +from vllm.v1.structured_output.backend_outlines import OutlinesBackend +from vllm.v1.structured_output.backend_types import ( + StructuredOutputBackend, + StructuredOutputOptions, +) +from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend + +TOKENIZER = "gpt2" +REGEX = "1a" +BACKENDS = ( + pytest.param("xgrammar", id="xgrammar"), + pytest.param("guidance", id="guidance"), + pytest.param("outlines", id="outlines"), + pytest.param( + "lm-format-enforcer", + marks=pytest.mark.xfail( + reason=( + "lm-format-enforcer 0.11.3 imports the removed " + "vllm.transformers_utils.tokenizer path" + ), + strict=True, + ), + id="lm-format-enforcer", + ), +) +BACKEND_CLASSES: dict[str, Callable[..., StructuredOutputBackend]] = { + "xgrammar": XgrammarBackend, + "guidance": GuidanceBackend, + "outlines": OutlinesBackend, + "lm-format-enforcer": LMFormatEnforcerBackend, +} + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(TOKENIZER) + + +def _single_token(tokenizer, text: str) -> int: + token_ids = tokenizer.encode(text) + assert len(token_ids) == 1, (text, token_ids) + return token_ids[0] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_accept_tokens_rejection_is_atomic(tokenizer, backend: str): + structured_output_backend = BACKEND_CLASSES[backend]( + VllmConfig(device_config=DeviceConfig(device="cpu")), + tokenizer, + len(tokenizer), + ) + grammar = structured_output_backend.compile_grammar( + StructuredOutputOptions.REGEX, REGEX + ) + + first = _single_token(tokenizer, "1") + invalid = _single_token(tokenizer, "z") + + assert not grammar.accept_tokens("request", [first, invalid]) + assert grammar.accept_tokens("request", [first]) diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py index 8d3b98be7a5b..c68f58f9e705 100644 --- a/tests/v1/structured_output/test_structured_output_manager.py +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -47,9 +47,9 @@ def tokenizer(): @dataclass(frozen=True) class FlowCase: - """One validate -> bitmask -> accept -> post-accept bitmask scenario.""" + """One bitmask -> accept -> post-accept bitmask scenario.""" - raw_drafts: tuple[str, ...] + raw_drafts: tuple[str | int, ...] expected_row_pattern: str expected_reasoning: bool | None expect_terminated: bool @@ -102,10 +102,14 @@ def _single_token(tokenizer, text: str) -> int: return token_ids[0] -def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]: - """Convert literal text (or EOS sentinel) to token IDs.""" +def _to_token_ids(tokenizer, texts: tuple[str | int, ...]) -> list[int]: + """Convert literal text, EOS, or speculative-padding tokens to IDs.""" return [ - tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t) + t + if isinstance(t, int) + else tokenizer.eos_token_id + if t == EOS + else _single_token(tokenizer, t) for t in texts ] @@ -182,6 +186,40 @@ def _build_harness( return manager, request +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + ("raw_drafts", "expected_length"), + [ + ((THINK_END, "1", "2", -1), 3), + ((THINK_END, "1", "z", "2"), 3), + (("1", "z", "2", -1), 3), + (("z", THINK_END, "1", "z", "2", -1), 4), + ], +) +def test_trim_spec_tokens_with_grammar( + tokenizer, + backend: StructuredOutputsBackend, + raw_drafts: tuple[str | int, ...], + expected_length: int, +): + manager, request = _build_harness( + tokenizer, + backend, + reasoning_ended=False, + reasoning_parser_kwargs={ + "marker": _single_token(tokenizer, THINK_END), + }, + structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX), + ) + + spec_tokens = _to_token_ids(tokenizer, raw_drafts) + + # The reasoning marker makes the manager find the grammar boundary from + # the real request state instead of hard-coding a boundary in the test. + trimmed = manager.trim_spec_tokens(request, spec_tokens) + assert trimmed == spec_tokens[:expected_length] + + def _row_pattern(bitmask, masked_ids=()) -> str: """U: unconstrained; R: all but `masked_ids` allowed; C: grammar rows.""" reasoning_row = np.full(bitmask.shape[1], -1, dtype=np.int32) @@ -210,23 +248,16 @@ def _run_real_flow( ) -> 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}, + scheduled_spec_decode_tokens={request.request_id: raw_drafts}, ) assert bitmask is not None assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1 row_pattern = _row_pattern(bitmask[: len(raw_drafts) + 1], masked_ids or ()) assert row_pattern == 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 @@ -238,8 +269,7 @@ def _run_real_flow( 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. + # StructuredOutputManager.accept_tokens() advances the grammar. request.append_output_token_ids(list(expected_validated)) assert manager.accept_tokens(request, list(expected_validated)) @@ -268,8 +298,9 @@ def _run_real_flow( pytest.param( FlowCase( prefix='{"a"', - raw_drafts=(":", ' "', "b", '"}'), - expected_row_pattern="CCCCC", + raw_drafts=(":", ' "', "b", '"}', -1, -1), + expected_validated=(":", ' "', "b", '"}'), + expected_row_pattern="CCCCCUU", expected_reasoning=True, expect_terminated=False, reasoning_ended=True, @@ -490,7 +521,6 @@ def test_reasoning_rows_mask_grammar_invalid_implicit_end( ) manager.reasoner_cls = MockEngineReasoner drafts = _to_token_ids(tokenizer, raw_drafts) - assert manager.validate_tokens(request, drafts) == drafts bitmask = manager.grammar_bitmask( requests={request.request_id: request}, @@ -562,19 +592,17 @@ def test_initial_constraint_activation( if reasoning_ended is not None else None }, + num_spec_tokens=1, ) - # "{" 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={}, + scheduled_spec_decode_tokens={request.request_id: [open_brace]}, ) assert bitmask is not None + assert _row_pattern(bitmask) == "CC" request.append_output_token_ids([open_brace]) assert manager.accept_tokens(request, [open_brace]) @@ -625,30 +653,6 @@ def test_regex_flow( ) -@pytest.mark.parametrize("backend", REGEX_BACKENDS) -def test_rejected_draft_keeps_later_rows_constrained( - tokenizer, backend: StructuredOutputsBackend -): - """A draft the grammar rejects during bitmask fill must not unconstrain - the remaining rows or the bonus row.""" - manager, request = _build_harness( - tokenizer, - backend, - use_reasoner=False, - structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX), - ) - drafts = _to_token_ids(tokenizer, ("1", "z", "2")) - 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 _row_pattern(bitmask[: len(drafts) + 1]) == "CCCC" - assert (bitmask[2] == bitmask[1]).all() - assert (bitmask[3] == bitmask[1]).all() - - def test_outlines_termination(tokenizer): """outlines_core never advances on EOS; termination must still be tracked exactly once per accepted EOS and undone by rollback.""" diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index cb8ee19ae3be..5d97ef2dd91e 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2850,7 +2850,7 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: continue # Add newly generated spec token ids to the request. - request.spec_token_ids = self.structured_output_manager.validate_tokens( + request.spec_token_ids = self.structured_output_manager.trim_spec_tokens( request, spec_token_ids ) @@ -2877,10 +2877,7 @@ def update_draft_token_ids_in_output( # Trim drafts to scheduled number of spec tokens # (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. - spec_token_ids = self.structured_output_manager.validate_tokens( - request, spec_token_ids - ) + # Preserve draft slots; rejection sampling handles grammar violations. # 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/engine/core.py b/vllm/v1/engine/core.py index c0510396315b..854388eb5bdc 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -807,14 +807,13 @@ def step_with_batch_queue( # in a field and do it immediately once step_with_batch_queue is # re-called. The latter slightly favors TTFT over TPOT/throughput. if deferred_scheduler_output: - # When draft tokens are used with structured output, validate them - # before computing the grammar bitmask for the deferred request. + # Update the deferred scheduler output before computing its grammar + # bitmask. if self.check_for_draft_tokens: draft_token_ids = self.model_executor.take_draft_token_ids() if draft_token_ids is not None: - # Update the draft token ids in the scheduler output to - # filter out the invalid spec tokens, which will be padded - # with -1 and skipped by the grammar bitmask computation. + # Update the draft token ids in the scheduler output while + # preserving the worker-provided draft slots. self.scheduler.update_draft_token_ids_in_output( draft_token_ids, deferred_scheduler_output ) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 69e34514728c..9911fa06add8 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -3,6 +3,7 @@ import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor +from itertools import chain from typing import TYPE_CHECKING, NamedTuple import numpy as np @@ -359,12 +360,15 @@ def _get_constraint_bounds( return ConstraintBounds(i + 1, i + 1) return ConstraintBounds(1, 1) - def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]: - """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`.""" + def trim_spec_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]: + """Trim structured-output padding and grammar-invalid suffixes.""" if not request.use_structured_output: + # Synchronous unstructured decoding uses -1 values as scheduler-side + # markers for speculative slots, so preserve their length. return spec_tokens spec_tokens = strip_speculative_padding(spec_tokens) + grammar_start, _ = self._get_constraint_bounds(request, spec_tokens) if grammar_start >= len(spec_tokens): return spec_tokens @@ -375,9 +379,11 @@ def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[in grammar = structured_req.grammar if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) + if grammar.is_terminated(): + return spec_tokens[:grammar_start] prefix = spec_tokens[:grammar_start] - validated = grammar.validate_tokens(spec_tokens[grammar_start:]) - return prefix + validated + len_validated = len(grammar.validate_tokens(spec_tokens[grammar_start:])) + return prefix + spec_tokens[grammar_start : grammar_start + len_validated + 1] def grammar_bitmask( self, @@ -449,7 +455,7 @@ def grammar_bitmask( else: # Fallback to serial filling of bitmasks for small-batch-size cases for grammar_idx, req_id in enumerate(structured_output_request_ids): - cumulative_index = grammar_idx * mask_stride + base_index = grammar_idx * mask_stride request = requests[req_id] structured_output_request = request.structured_output_request @@ -460,62 +466,43 @@ def grammar_bitmask( assert isinstance(grammar, StructuredOutputGrammar) req_tokens = scheduled_spec_decode_tokens.get(req_id, list()) + unpadded_tokens = strip_speculative_padding(req_tokens) grammar_start, constraint_start = self._get_constraint_bounds( - request, strip_speculative_padding(req_tokens) + request, unpadded_tokens ) # Rows before constraint_start are still reasoning. masked_ids = ( self._masked_reasoning_ends(request) if constraint_start else () ) + # Diffusion LLMs don't sample a bonus token after the + # scheduled positions, so skip its bitmask in that case. + tokens_to_mask = ( + chain(req_tokens, (-1,)) + if not (is_diffusion and req_tokens) + else req_tokens + ) + state_advancements = 0 - seen_padding = False - # Row filled from the last valid grammar state before a draft - # was rejected; later rows reuse it rather than unconstraining. - failed_index: int | None = None - bitmask = self._grammar_bitmask - for i, token in enumerate(req_tokens): - if failed_index is not None: - bitmask[cumulative_index].copy_(bitmask[failed_index]) - cumulative_index += 1 + padding_or_grammar_invalid = False + for i, token in enumerate(tokens_to_mask): + if padding_or_grammar_invalid: + self._fill_bitmasks(((grammar, base_index + i, False, ()),)) continue - apply_bitmask = not seen_padding and i >= constraint_start - reasoning = not seen_padding and i < constraint_start - row_masked = masked_ids if reasoning else () + + apply_bitmask = i >= constraint_start + row_masked = masked_ids if i < constraint_start else () self._fill_bitmasks( - ((grammar, cumulative_index, apply_bitmask, row_masked),) + ((grammar, base_index + i, apply_bitmask, row_masked),) ) if token == -1: - seen_padding = True - elif not seen_padding and i >= grammar_start: + padding_or_grammar_invalid = True + elif i >= grammar_start: if not grammar.is_terminated() and grammar.accept_tokens( req_id, [token] ): state_advancements += 1 else: - failed_index = cumulative_index - 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 (is_diffusion and req_tokens): - if failed_index is not None: - bitmask[cumulative_index].copy_(bitmask[failed_index]) - else: - bonus_apply = not seen_padding and constraint_start <= len( - req_tokens - ) - bonus_reasoning = not seen_padding and not bonus_apply - bonus_masked = masked_ids if bonus_reasoning else () - self._fill_bitmasks( - ((grammar, cumulative_index, bonus_apply, bonus_masked),) - ) - cumulative_index += 1 + padding_or_grammar_invalid = True if state_advancements > 0: grammar.rollback(state_advancements) diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 08aab576df1c..95c34e477df9 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -236,12 +236,15 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: Returns True if the parser was advanced successfully. Returns False if the parser failed to advance. """ - if self.ll_tokenizer.eos_token in tokens: - if self.ll_matcher.is_stopped() and not self.terminated: - self.rollback_lag = 1 - self.terminated = True + eos_token = self.ll_tokenizer.eos_token + eos_index = next( + (index for index, token in enumerate(tokens) if token == eos_token), + None, + ) + if eos_index is not None and eos_index != len(tokens) - 1: + return False - if self.ll_matcher.is_stopped(): + if self.terminated: return True # TODO - Add jump decoding support in the future: @@ -251,11 +254,31 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: # For conversion between the two, see # https://github.com/guidance-ai/llguidance/blob/main/docs/fast_forward.md - r = self.ll_matcher.consume_tokens(tokens) + num_tokens = self.ll_matcher.try_consume_tokens(tokens) + if num_tokens != len(tokens): + if ( + eos_index is not None + and num_tokens == eos_index + and self.ll_matcher.is_stopped() + and self.ll_matcher.is_accepting() + ): + # LLMatcher stops at an accepting grammar state and leaves the + # wrapper-level EOS token unconsumed. + self.rollback_lag = 1 + self.terminated = True + self.check_error() + return True + if num_tokens > 0: + self.ll_matcher.rollback(num_tokens) + self.check_error() + return False + + if eos_index is not None: + self.terminated = True self.check_error() - return r + return True def validate_tokens(self, tokens: list[int]) -> list[int]: """Checks if the list of tokens are accepted by the parser in sequence. diff --git a/vllm/v1/structured_output/backend_types.py b/vllm/v1/structured_output/backend_types.py index b3695efd009c..6291e4a32ad0 100644 --- a/vllm/v1/structured_output/backend_types.py +++ b/vllm/v1/structured_output/backend_types.py @@ -36,6 +36,8 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: """Determines whether the provided tokens are accepted for the given request. + A false result must leave the grammar state unchanged. + Args: request_id (str): The unique identifier for the request. tokens (list[int]): A list of token IDs to evaluate. diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index aae72873b2e6..3772d371ad01 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -270,14 +270,12 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: """ if self._is_terminated: return True - for token in tokens: + for num_accepted, token in enumerate(tokens): if not self.matcher.accept_token(token): - logger.error( - "Failed to advance FSM for request %s " - "for tokens %s. Please file an issue.", - request_id, - token, - ) + if num_accepted: + self.matcher.rollback(num_accepted) + self.num_processed_tokens -= num_accepted + self._is_terminated = self.matcher.is_terminated() return False self.num_processed_tokens += 1 self._is_terminated = self.matcher.is_terminated() From 70cb50dab240a03f7db9ab87a613b5b52f1339fe Mon Sep 17 00:00:00 2001 From: zhenwenqi2024 <155598497+zhenwenqi2024@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:18:40 +0800 Subject: [PATCH 006/406] [Feature][HiSparse] Add a platform backend for device operations (#62592) Signed-off-by: zhenwenqi2024 Signed-off-by: QiuChunshuo Signed-off-by: Matthew Bonanni Co-authored-by: Codex Co-authored-by: QiuChunshuo Co-authored-by: vllm-agent Co-authored-by: Matthew Bonanni --- docs/design/hisparse.md | 18 ++- tests/test_config.py | 54 +++++-- .../v1/attention/test_sparse_mla_backends.py | 35 ++++- tests/v1/core/test_kv_cache_planning.py | 10 +- tests/v1/core/test_kv_cache_utils.py | 9 +- .../unit/test_hisparse_connector.py | 9 +- tests/v1/worker/test_attn_utils.py | 26 +++- .../worker/test_kv_cache_allocation_scope.py | 4 + tests/v1/worker/test_utils.py | 105 ++++++++----- vllm/config/vllm.py | 4 +- .../kv_connector/v1/hisparse/worker.py | 68 ++++----- vllm/platforms/cuda.py | 4 + vllm/platforms/interface.py | 5 + vllm/v1/attention/backends/mla/index_group.py | 30 ++-- vllm/v1/hisparse/backend.py | 142 ++++++++++++++++++ vllm/v1/hisparse/binding.py | 13 +- vllm/v1/hisparse/runtime.py | 104 ++++++------- 17 files changed, 450 insertions(+), 190 deletions(-) create mode 100644 vllm/v1/hisparse/backend.py diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md index f5fab3c0f5d2..64aab6f52715 100644 --- a/docs/design/hisparse.md +++ b/docs/design/hisparse.md @@ -221,9 +221,16 @@ hot row? ──────── yes ──► existing hot physical row + upda choose GPU LRU victim ──► copy pinned host row ──► hot physical row ``` -ROCm is not currently supported because the fused HiSparse cache operations are -implemented only by CUDA kernels. A future platform-specific worker may provide -the same command, output, and cache-resolution boundaries. +Device operations are selected by `Platform.get_hisparse_backend_cls()`. The base +platform returns `None`; CUDA selects `CUDAHiSparseBackend` in `platforms/cuda.py`. +Hardware plugins can provide a `HiSparseBackend` for host registration, cache +operations, IPC events, capture queries, and transfers while reusing the common +runtime and connector. HiSparse validates availability and cache formats before +allocation. Streams, ordinary events, current-stream queries, and synchronization +use `torch.Stream`, `torch.Event`, and `torch.accelerator` directly. The backend +provides IPC event creation/import/export only for shared host pools. +Without HiSparse, sparse attention keeps its normal stream and event path. +ROCm does not currently provide a HiSparse backend. ## Main classes @@ -260,6 +267,5 @@ the same command, output, and cache-resolution boundaries. The command/result and attention-layer boundaries can be shared. The host allocator, copy implementation, hot layout, and replacement policy should stay platform-specific. NVIDIA uses the current accelerator LRU and fused host/hot -kernel. ROCm is not currently supported; AMD or other accelerator backends can -implement their own worker without forcing NVIDIA's policy into the shared -boundary. +kernel. Other accelerators can implement `HiSparseBackend` and return its class +path from their platform override without changing the shared control flow. diff --git a/tests/test_config.py b/tests/test_config.py index 1f71e1c966d2..a53843e2a93a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -622,7 +622,11 @@ def test_hisparse_requires_v2_model_runner(): def test_hisparse_rejects_decode_context_parallelism(monkeypatch): - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", + ) monkeypatch.setattr(current_platform, "device_count", lambda: 2) with pytest.raises(ValueError, match="decode context parallelism"): VllmConfig( @@ -635,7 +639,11 @@ def test_hisparse_rejects_decode_context_parallelism(monkeypatch): def test_hisparse_rejects_pipeline_parallelism(monkeypatch): - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", + ) monkeypatch.setattr(current_platform, "device_count", lambda: 2) with pytest.raises(ValueError, match="pipeline parallelism"): VllmConfig( @@ -645,7 +653,11 @@ def test_hisparse_rejects_pipeline_parallelism(monkeypatch): def test_hisparse_rejects_full_cudagraph_mode(monkeypatch): - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", + ) monkeypatch.setattr(current_platform, "support_static_graph_mode", lambda: True) monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) with pytest.raises(ValueError, match="does not support cudagraph_mode=FULL"): @@ -656,7 +668,11 @@ def test_hisparse_rejects_full_cudagraph_mode(monkeypatch): def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch): - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", + ) monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) with pytest.raises(ValueError, match="requires the hybrid KV cache manager"): VllmConfig( @@ -670,7 +686,11 @@ def test_hisparse_rejects_disabled_hybrid_kv_cache_manager(monkeypatch): def test_hisparse_rejects_disabled_full_isl_reservation(monkeypatch): - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", + ) with pytest.raises(ValueError, match="requires --scheduler-reserve-full-isl"): VllmConfig( attention_config=AttentionConfig(hisparse_config=HiSparseConfig()), @@ -682,9 +702,9 @@ def test_hisparse_rejects_disabled_full_isl_reservation(monkeypatch): ) -def test_hisparse_rejects_non_cuda(monkeypatch): - monkeypatch.setattr(current_platform, "is_cuda", lambda: False) - with pytest.raises(ValueError, match="requires NVIDIA CUDA"): +def test_hisparse_rejects_unsupported_platform(monkeypatch): + monkeypatch.setattr(current_platform, "get_hisparse_backend_cls", lambda: None) + with pytest.raises(ValueError, match="platform does not support HiSparse"): VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig())) @@ -732,7 +752,11 @@ def test_hisparse_rejects_non_cuda(monkeypatch): 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( + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", + ) monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) config = VllmConfig( kv_transfer_config=kv_transfer_config, @@ -748,7 +772,11 @@ def test_hisparse_connector_implies_attention_config(monkeypatch, kv_transfer_co def test_hisparse_connector_preserves_explicit_attention_config(monkeypatch): - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", + ) monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) config = VllmConfig( attention_config=AttentionConfig( @@ -768,9 +796,9 @@ def test_hisparse_connector_preserves_explicit_attention_config(monkeypatch): 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"): +def test_hisparse_connector_with_unsupported_platform_rejected(monkeypatch): + monkeypatch.setattr(current_platform, "get_hisparse_backend_cls", lambda: None) + with pytest.raises(ValueError, match="platform does not support HiSparse"): VllmConfig( kv_transfer_config=KVTransferConfig( kv_connector="HiSparseConnector", diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 565a517532be..78113b63a93a 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -114,6 +114,7 @@ ) from vllm.v1.attention.ops import flashmla from vllm.v1.hisparse import runtime as hisparse_runtime +from vllm.v1.hisparse.backend import get_hisparse_backend from vllm.v1.hisparse.runtime import ( HiSparseCacheHandle, HiSparseRuntime, @@ -2014,6 +2015,27 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int, reuse_buffers assert torch.all(result[row, num_valid:] == -1) +def test_hisparse_generic_stream_orders_ipc_event(): + """Generic copy streams must record IPC events and restore the compute stream.""" + from vllm.v1.hisparse.backend import CUDAHiSparseBackend + + backend = CUDAHiSparseBackend() + compute_stream = torch.accelerator.current_stream() + copy_stream = torch.Stream(device=torch.device(DEVICE_TYPE)) + values = torch.arange(8, dtype=torch.float32, device=DEVICE_TYPE) + input_ready = torch.Event() + input_ready.record(compute_stream) + copy_stream.wait_event(input_ready) + output_ready = backend.create_ipc_event() + with copy_stream: + values.add_(1) + output_ready.record(copy_stream) + assert torch.cuda.current_stream().stream_id == compute_stream.stream_id + compute_stream.wait_event(output_ready) + + torch.testing.assert_close(values.cpu(), torch.arange(8, dtype=torch.float32) + 1) + + def test_sparse_mla_index_group_converts_decode_indices_once(monkeypatch): device = torch.device(DEVICE_TYPE) logical = torch.full((2, 128), -1, dtype=torch.int32, device=device) @@ -2949,6 +2971,7 @@ def test_hisparse_row_dma_copies_discontiguous_spans_across_layers(): for _ in resident_caches ) worker = object.__new__(HiSparseConnectorWorker) + worker.backend = get_hisparse_backend() worker.is_host_writer = True worker.kernel_block_size = block_size worker.resident_caches = resident_caches @@ -2960,7 +2983,7 @@ def test_hisparse_row_dma_copies_discontiguous_spans_across_layers(): ) for _ in resident_caches ] - worker.dma_stream = torch.cuda.Stream(device=device) + worker.dma_stream = torch.Stream(device=device) worker.host_write_event = torch.Event() worker._dma_free_descriptors = [] worker._pending_dma_descriptors = deque() @@ -3612,8 +3635,10 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows(): .to(device) ) - staged = HiSparseRuntime.gather_prefill_cache( - None, host_cache, plan, resident_cache=resident_cache + runtime = _make_hisparse_runtime(row_width=row_width, block_size=block_size) + runtime.bind_source_cache(host_cache) + staged = runtime.gather_prefill_cache( + host_cache, plan, resident_cache=resident_cache ) staged_flat = staged.view(-1, row_width).cpu() @@ -4772,7 +4797,7 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): compute_stream = object() prefill_stream = MagicMock() gather = MagicMock() - monkeypatch.setattr(index_group_module, "current_stream", lambda: compute_stream) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: compute_stream) monkeypatch.setattr( index_group_module.ops, "cp_gather_and_upconvert_fp8_kv_cache", gather ) @@ -4783,7 +4808,7 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch): view=SimpleNamespace(cache=resident_cache, block_size=4), block_table=torch.tensor([[2, 3]], dtype=torch.int32), runtime=SimpleNamespace( - request_state_indices=torch.tensor([0], dtype=torch.int32) + request_state_indices=torch.tensor([0], dtype=torch.int32), ), ) index_group = object.__new__(HiSparseMLAIndexGroup) diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py index 5d2a192920e7..b1558be3d84d 100644 --- a/tests/v1/core/test_kv_cache_planning.py +++ b/tests/v1/core/test_kv_cache_planning.py @@ -12,7 +12,6 @@ import vllm.v1.core.kv_cache_planning as kv_cache_planning import vllm.v1.core.kv_cache_utils as kv_cache_utils -import vllm.v1.hisparse.runtime as hisparse_runtime_module from tests.v1.core.utils import ( get_kv_cache_configs, new_kv_cache_spec, @@ -27,6 +26,7 @@ VllmConfig, ) from vllm.config.attention import HiSparseConfig +from vllm.platforms import current_platform from vllm.utils.mem_constants import GiB_bytes from vllm.v1.core.kv_cache_planning import ( check_enough_kv_cache_memory, @@ -70,7 +70,9 @@ def test_hisparse_hma_uses_resolved_gpu_block_size( monkeypatch, gpu_block_size, shared_host_pool ): monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", ) specs = { "model.layers.0.self_attn": MLAAttentionSpec( @@ -190,7 +192,9 @@ def test_hisparse_pool_must_fit_max_model_len( plus the null block. Auto-fit (-1) must pick a length that passes the same check.""" monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", ) monkeypatch.setattr( kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 917ed57c0ae7..5d0d9cd5a900 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -13,7 +13,6 @@ import vllm.v1.core.kv_cache_planning as kv_cache_planning import vllm.v1.core.kv_cache_utils as kv_cache_utils -import vllm.v1.hisparse.runtime as hisparse_runtime_module from tests.v1.core.utils import ( new_kv_cache_spec, new_mamba_spec, @@ -35,6 +34,7 @@ MultiModalKwargsItem, PlaceholderRange, ) +from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager @@ -138,7 +138,9 @@ def test_hisparse_steady_state_concurrency_excludes_spilled_resident_pages(): def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok): """Every page needs a host block, so one max_model_len request must fit.""" monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + current_platform, + "get_hisparse_backend_cls", + lambda: "vllm.v1.hisparse.backend.CUDAHiSparseBackend", ) specs = { "model.layers.0.self_attn": MLAAttentionSpec( @@ -181,9 +183,6 @@ def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch): backing; KV cache layout resolution must skip them.""" from vllm.v1.attention.backends.utils import resolve_kv_cache_layout - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False) specs: dict[str, KVCacheSpec] = { "model.layers.0.self_attn": MLAAttentionSpec( diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index bdfbee28d4c0..5d6eaf7dd2f5 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import contextlib import ctypes from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -247,11 +246,8 @@ def swap_blocks_batch(src, dst, sizes): monkeypatch.setattr(torch, "Event", _FakeEvent) monkeypatch.setattr(torch, "Stream", lambda *args, **kwargs: MagicMock()) - monkeypatch.setattr(torch.cuda, "Stream", lambda *args, **kwargs: MagicMock()) - monkeypatch.setattr(torch.cuda, "stream", lambda _: contextlib.nullcontext()) compute_stream = MagicMock() - monkeypatch.setattr(worker_module, "current_stream", lambda: compute_stream) - monkeypatch.setattr(worker_module.ops, "swap_blocks_batch", swap_blocks_batch) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: compute_stream) monkeypatch.setattr( runtime_module, "get_forward_context", @@ -290,6 +286,7 @@ def swap_blocks_batch(src, dst, sizes): ), MagicMock(), ) + monkeypatch.setattr(worker.backend, "swap_blocks_batch", swap_blocks_batch) worker.initialize( handles, ["target", "draft"], @@ -412,7 +409,7 @@ def test_zero_token_step_holds_tp_ranks_after_host_write_wait( compute_stream.wait_event.side_effect = lambda event: calls.append(("wait", event)) tp_group = MagicMock() tp_group.barrier.side_effect = lambda: calls.append(("barrier", None)) - monkeypatch.setattr(worker_module, "current_stream", lambda: compute_stream) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: compute_stream) monkeypatch.setattr(worker_module, "get_tp_group", lambda: tp_group) worker = object.__new__(HiSparseConnectorWorker) diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 197929731f42..5b284f620626 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -26,6 +26,7 @@ from vllm.config.compilation import CompilationConfig, CUDAGraphMode from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.backend import CUDAHiSparseBackend from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -502,7 +503,10 @@ def test_profiling_cleanup_keeps_host_pool_for_real_kv_cache(monkeypatch): pool = attn_utils_module.get_hisparse_host_pool(vllm_config, kv_cache_config) backing = pool.allocate(8) runtime = SimpleNamespace( - _host_cache=backing, registered_host_pool=pinned, hot_backing=object() + _host_cache=backing, + registered_host_pool=pinned, + hot_backing=object(), + host_pool=pool, ) forward_context = { "layer": SimpleNamespace(hisparse_cache=SimpleNamespace(runtime=runtime)) @@ -511,6 +515,7 @@ def test_profiling_cleanup_keeps_host_pool_for_real_kv_cache(monkeypatch): attn_utils_module.release_hisparse_profiling_cache(forward_context) assert not hasattr(runtime, "_host_cache") + assert runtime.host_pool is None reused = attn_utils_module.get_hisparse_host_pool(vllm_config, kv_cache_config) assert reused is pool assert reused.allocate(8) is backing @@ -518,10 +523,21 @@ def test_profiling_cleanup_keeps_host_pool_for_real_kv_cache(monkeypatch): hisparse_runtime_module.drop_hisparse_host_pool(vllm_config) -@pytest.mark.parametrize("failure_phase", ["allocation", "binding", "buffers"]) +@pytest.mark.parametrize("failure_phase", ["allocation", "views", "binding", "buffers"]) @pytest.mark.parametrize("shared", [False, True]) def test_init_hisparse_rolls_back_host_pool(monkeypatch, failure_phase, shared): """Failed initialization must release shared or private registrations.""" + + class Backend(CUDAHiSparseBackend): + def prepare_kv_cache_views(self, *args): + if failure_phase == "views": + raise RuntimeError("initialization failed") + + backend = Backend() + monkeypatch.setattr(attn_utils_module, "get_hisparse_backend", lambda: backend) + monkeypatch.setattr( + hisparse_runtime_module, "get_hisparse_backend", lambda: backend + ) region = _FakeSharedHostRegion() registered = torch.empty(4096, dtype=torch.int8) cudart = MagicMock() @@ -544,6 +560,7 @@ def allocate(*args): else: args[-1].registered = registered args[-1].backing = registered[:8] + args[-1].pinned = [registered] if failure_phase == "allocation": raise RuntimeError("initialization failed") return {} @@ -571,10 +588,7 @@ def buffers(*args, **kwargs): SimpleNamespace(), ) assert region.cleanup_calls == int(shared) - if shared: - cudart.cudaHostUnregister.assert_not_called() - else: - cudart.cudaHostUnregister.assert_called_once_with(registered.data_ptr()) + cudart.cudaHostUnregister.assert_called_once_with(registered.data_ptr()) assert attn_utils_module.get_hisparse_host_pool(vllm_config, None) is not pools[0] hisparse_runtime_module.drop_hisparse_host_pool(vllm_config) diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index b5b288d9b5e4..3570d64d906e 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -11,6 +11,7 @@ import vllm.v1.hisparse.binding as hisparse_binding import vllm.v1.worker.gpu.attn_utils as attn_utils import vllm.v1.worker.gpu_model_runner as gpu_model_runner +from vllm.v1.hisparse.backend import CUDAHiSparseBackend from vllm.v1.worker.gpu_worker import Worker @@ -60,6 +61,9 @@ def bind_hisparse(**kwargs): ) ] + monkeypatch.setattr( + hisparse_binding, "get_hisparse_backend", CUDAHiSparseBackend + ) monkeypatch.setattr( hisparse_binding, "get_hisparse_host_pool", diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index ef162e4b0ccb..e004ed94d380 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -28,6 +28,7 @@ from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import MemorySnapshot from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.backend import CUDAHiSparseBackend, HiSparseBackend from vllm.v1.hisparse.runtime import HiSparseCacheHandle from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror from vllm.v1.worker.utils import ( @@ -38,8 +39,26 @@ ) +@pytest.mark.parametrize("packed_bytes", [False, True]) +def test_hisparse_resident_gather_preserves_stride_and_payload(packed_bytes): + # Include a gap between physical pages and select reordered/repeated rows. + backing = torch.arange(24, dtype=torch.float32) + source = backing.as_strided((2, 2, 4), (12, 4, 1)) + expected = torch.stack([backing[16:20], backing[0:4], backing[12:16], backing[0:4]]) + if packed_bytes: + source = source.view(torch.uint8) + destination = torch.empty((4, 4), dtype=torch.float32) + + HiSparseBackend().gather_resident_rows( + source, destination, torch.tensor([3, 0, 2, -1]) + ) + + torch.testing.assert_close(destination, expected, rtol=0, atol=0) + + def _make_hisparse_worker() -> HiSparseConnectorWorker: worker = object.__new__(HiSparseConnectorWorker) + worker.backend = CUDAHiSparseBackend() worker._slot_mapping_staging = None worker._row_mirror_num_rows = 0 worker._per_layer_mirrored = set() @@ -85,9 +104,7 @@ def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeyp ) worker.leader_runtimes = [SimpleNamespace(index_group=group)] # type: ignore[list-item] # runtime stand-in monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - monkeypatch.setattr( - hisparse_worker_module, "current_stream", lambda: compute_stream - ) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: compute_stream) assert worker.get_kv_connector_stats() is None compute_stream.wait_stream.assert_called_once_with(group.copy_stream) @@ -168,8 +185,7 @@ def test_hisparse_appends_reference_slots_within_a_mirror_phase(monkeypatch): worker._row_mirror_num_rows = 1 worker._submitted_mirror_layers = set() main_stream = MagicMock() - monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: main_stream) - monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: main_stream) handle.slot_mapping = torch.tensor([7, 8], dtype=torch.int64) worker._stage_row_mirror_mapping(2) @@ -223,7 +239,7 @@ def _hisparse_parallel_config(**overrides): def test_hisparse_shares_host_pool_only_for_local_tp(monkeypatch): monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + hisparse_runtime_module, "get_hisparse_backend", CUDAHiSparseBackend ) config = SimpleNamespace(parallel_config=_hisparse_parallel_config()) @@ -297,7 +313,7 @@ def pin_tensor(tensor): parallel_config=_hisparse_parallel_config(), ) - pools, private_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools( + pools, pinned_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools( config, [24, 40], num_blocks=4, @@ -307,7 +323,7 @@ def pin_tensor(tensor): assert region is not None assert isinstance(region, FakeSharedOffloadRegion) - assert private_pools == [] + assert [t.data_ptr() for t in pinned_pools] == [t.data_ptr() for t in pinned] assert region.kwargs == { "engine_id": "hisparse_instance_dp3", "num_chunks": 1, @@ -325,14 +341,18 @@ def pin_tensor(tensor): (tensor.data_ptr() - region.base_tensor.data_ptr(), tensor.nbytes) for tensor in pinned ] == [(0, page), (page, 3 * page)] - assert region.is_pinned + assert not region.is_pinned @pytest.mark.parametrize("fail_registration", [None, 1, 2]) def test_shared_host_pool_tracks_successful_registrations( monkeypatch, fail_registration ): - """Only successful registrations may be unregistered on allocation failure.""" + """Only successful registrations may be unpinned on allocation failure.""" + monkeypatch.setattr( + hisparse_runtime_module, "get_hisparse_backend", CUDAHiSparseBackend + ) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) page = mmap.PAGESIZE tp_group = MagicMock(rank_in_group=0, world_size=2) monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", lambda: tp_group) @@ -350,6 +370,7 @@ def test_shared_host_pool_tracks_successful_registrations( cudart.cudaHostRegister.side_effect = [ SimpleNamespace(value=int(fail_registration == i)) for i in (1, 2) ] + cudart.cudaHostUnregister.return_value = SimpleNamespace(value=0) monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) config = SimpleNamespace( instance_id="test", parallel_config=SimpleNamespace(data_parallel_index=0) @@ -359,16 +380,13 @@ def test_shared_host_pool_tracks_successful_registrations( if fail_registration else nullcontext() ) + result = None with context: - hisparse_runtime_module.allocate_hisparse_host_pools( + result = hisparse_runtime_module.allocate_hisparse_host_pools( config, [3 * page], 3, page, use_shared_host_pool=True ) successful = 2 if fail_registration is None else fail_registration - 1 - assert region.pinned_addresses == [ - backing.data_ptr() + i * page for i in range(successful) - ] - assert region.is_pinned == bool(successful) - assert region.cleanup.call_count == int(fail_registration is not None) + pinned_ptrs = [backing.data_ptr() + i * page for i in range(successful)] attempts = 2 if fail_registration is None else fail_registration assert ( cudart.cudaHostRegister.call_args_list @@ -377,6 +395,18 @@ def test_shared_host_pool_tracks_successful_registrations( call(backing.data_ptr() + page, 2 * page, 0), ][:attempts] ) + assert region.pinned_addresses == [] + assert not region.is_pinned + if fail_registration is None: + assert result is not None + assert [t.data_ptr() for t in result[1]] == pinned_ptrs + cudart.cudaHostUnregister.assert_not_called() + region.cleanup.assert_not_called() + else: + assert cudart.cudaHostUnregister.call_args_list == [ + call(ptr) for ptr in reversed(pinned_ptrs) + ] + region.cleanup.assert_called_once_with() def test_shared_host_pool_registers_layer_spans_in_one_backing(monkeypatch): @@ -600,6 +630,9 @@ def test_hisparse_tail_restore_preserves_imported_rows( """Restore externally pinned host rows on every rank under sync checking.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA is required for the device copy check") + monkeypatch.setattr( + hisparse_runtime_module, "get_hisparse_backend", CUDAHiSparseBackend + ) worker = _make_hisparse_worker() worker.is_host_writer = is_host_writer worker.kernel_block_size = 4 @@ -636,7 +669,7 @@ def test_hisparse_tail_restore_preserves_imported_rows( worker, "_record_transfer_completion", record_transfer_completion ) stream = torch.cuda.current_stream() if device == "cuda" else MagicMock() - monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: stream) transfer = SparseKVPageTransfer(7, 2, (1, 3), False, restore=True) restore_pages = worker._restore_pages @@ -703,14 +736,9 @@ def test_hisparse_dma_row_mirror_builds_descriptors(monkeypatch): current_stream = MagicMock() event = MagicMock() swap_blocks_batch = MagicMock() - monkeypatch.setattr( - hisparse_worker_module, "current_stream", lambda: current_stream - ) - monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: current_stream) monkeypatch.setattr(torch, "Event", lambda: event) - monkeypatch.setattr( - hisparse_worker_module.ops, "swap_blocks_batch", swap_blocks_batch - ) + monkeypatch.setattr(worker.backend, "swap_blocks_batch", swap_blocks_batch) worker._set_row_mirrors((SparseKVRowMirror((3,), 4, 1),)) worker._enqueue_row_dma(range(1)) @@ -802,9 +830,7 @@ def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch): worker.host_write_event = MagicMock() worker._forward_ready_event = MagicMock() current_stream = MagicMock() - monkeypatch.setattr( - hisparse_worker_module, "current_stream", lambda: current_stream - ) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: current_stream) worker.finish_forward() @@ -945,9 +971,7 @@ def test_hisparse_finish_forward_orders_next_forward_after_dma(monkeypatch): worker._dma_free_descriptors = [] worker._forward_ready_event = MagicMock() monkeypatch.setattr(worker, "_finish_mirror_phase", MagicMock()) - monkeypatch.setattr( - hisparse_worker_module, "current_stream", lambda: current_stream - ) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: current_stream) worker.finish_forward() @@ -1086,7 +1110,7 @@ def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer worker._per_layer_mirrored = set() worker._submitted_mirror_layers = set() stream = MagicMock() - monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: stream) worker.start_step( SimpleNamespace( @@ -1172,7 +1196,7 @@ def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch): enqueue_host_mirror = MagicMock(wraps=worker._enqueue_host_mirror) monkeypatch.setattr(worker, "_enqueue_host_mirror", enqueue_host_mirror) stream = MagicMock() - monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) + monkeypatch.setattr(torch.accelerator, "current_stream", lambda: stream) worker.start_step( SimpleNamespace( @@ -1195,6 +1219,9 @@ def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch): def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatch): """Followers must not allocate duplicate runtime state before profiling.""" + monkeypatch.setattr( + hisparse_runtime_module, "get_hisparse_backend", CUDAHiSparseBackend + ) config = SimpleNamespace( scheduler_config=SimpleNamespace( max_num_seqs=2, @@ -1209,7 +1236,7 @@ def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatc device_buffer_size=8, max_union_rows=8, ) - monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True) + monkeypatch.setattr(CUDAHiSparseBackend, "check_available", lambda self: None) monkeypatch.setattr( hisparse_runtime_module.ResolvedHiSparseConfig, "from_vllm_config", @@ -1230,8 +1257,16 @@ def create_stream(_device): hisparse_runtime_module, "_create_shared_topk_state", create_shared_state ) monkeypatch.setattr(hisparse_runtime_module, "_create_copy_stream", create_stream) - monkeypatch.setattr(index_group_module, "_create_side_stream", lambda _: object()) - monkeypatch.setattr(index_group_module, "_create_event", lambda: object()) + monkeypatch.setattr( + torch, + "Stream", + lambda device: object(), + ) + monkeypatch.setattr( + torch, + "Event", + lambda: object(), + ) index_group_builder = index_group_module.SparseMLAIndexGroupBuilder( torch.empty((2, 4), dtype=torch.int32) ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 078da543555d..69b4d64469a2 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1882,8 +1882,8 @@ def has_blocked_weights(): 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.") + if current_platform.get_hisparse_backend_cls() is None: + raise ValueError("The current platform does not support HiSparse.") if self.parallel_config.pipeline_parallel_size > 1: raise ValueError("HiSparse does not support pipeline parallelism.") if self.parallel_config.decode_context_parallel_size > 1: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index 6e3d3415c3ad..3639e8031bc0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -13,7 +13,6 @@ import numpy as np import torch -from vllm import _custom_ops as ops from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.stats import ( HiSparseKVConnectorStats, @@ -23,11 +22,12 @@ get_tp_group, ) from vllm.utils.gpu_sync_debug import gpu_sync_allowed -from vllm.utils.torch_utils import current_stream from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.hisparse.backend import get_hisparse_backend from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX from vllm.v1.hisparse.runtime import ( HiSparseCacheHandle, + HiSparseHostPool, drop_hisparse_host_pool, release_pinned_state, update_hisparse_residency, @@ -164,6 +164,7 @@ def _create_hisparse_host_events( is_host_writer: bool, device: torch.device, ) -> tuple[torch.Event, torch.Event]: + backend = get_hisparse_backend() if shared_host_region is None: return torch.Event(), torch.Event() @@ -171,18 +172,18 @@ def _create_hisparse_host_events( ipc_handles = None if is_host_writer: events = ( - torch.cuda.Event(interprocess=True), - torch.cuda.Event(interprocess=True), + backend.create_ipc_event(), + backend.create_ipc_event(), ) - stream = current_stream() + stream = torch.accelerator.current_stream() for event in events: event.record(stream) - ipc_handles = tuple(event.ipc_handle() for event in events) + ipc_handles = tuple(backend.export_event(event) for event in events) ipc_handles = get_tp_group().broadcast_object(ipc_handles, src=0) if events is None: events = ( - torch.cuda.Event.from_ipc_handle(device, ipc_handles[0]), - torch.cuda.Event.from_ipc_handle(device, ipc_handles[1]), + backend.import_event(device, ipc_handles[0]), + backend.import_event(device, ipc_handles[1]), ) return events @@ -191,6 +192,7 @@ class HiSparseConnectorWorker: """Own HiSparse host/hot state and execute its worker-side transfers.""" def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None: + self.backend = get_hisparse_backend() self.vllm_config = vllm_config self.kv_cache_config = kv_cache_config self._initialized = False @@ -211,26 +213,20 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: if not cache_handles: raise RuntimeError("HiSparse connector found no hot-cache handles.") hot_backings: dict[int, torch.Tensor] = {} - registered_host_pools: dict[int, torch.Tensor] = {} - shared_host_regions: dict[int, SharedOffloadRegion] = {} + host_pools: dict[int, HiSparseHostPool] = {} for cache in cache_handles: hot_backing = cache.runtime.hot_backing hot_backings[hot_backing.untyped_storage().data_ptr()] = hot_backing - registered_pool = cache.runtime.registered_host_pool - registered_host_pools[registered_pool.data_ptr()] = registered_pool - if (region := cache.runtime.shared_host_region) is not None: - shared_host_regions[id(region)] = region + assert cache.runtime.host_pool is not None + host_pools[id(cache.runtime.host_pool)] = cache.runtime.host_pool if len(hot_backings) != 1: raise RuntimeError("HiSparse hot tensors must share one GPU backing.") hot_backing = next(iter(hot_backings.values())) - if len(shared_host_regions) > 1: - raise RuntimeError("HiSparse caches must share one host region.") - shared_host_region = next(iter(shared_host_regions.values()), None) - pinned_host_pools = ( - [] - if shared_host_region is not None - else list(registered_host_pools.values()) - ) + if len(host_pools) != 1: + raise RuntimeError("HiSparse caches must share one host pool.") + host_pool = next(iter(host_pools.values())) + shared_host_region = host_pool.shared_region + pinned_host_pools = host_pool.pinned is_host_writer = _is_hisparse_host_writer(shared_host_region) resident = cache_handles[0].view @@ -280,9 +276,7 @@ def initialize( self.pinned_host_pools = pinned_host_pools self.shared_host_region = shared_host_region self.is_host_writer = is_host_writer - self.dma_stream = ( - torch.cuda.Stream(device=device) if self.is_host_writer else None - ) + self.dma_stream = torch.Stream(device=device) if self.is_host_writer else None self._slot_mapping_staging = None if self.is_host_writer: max_mirror_rows = ( @@ -371,7 +365,7 @@ def set_request_state_indices(self, indices: torch.Tensor) -> None: "HiSparse request-state mapping exceeds max_num_seqs: " f"{indices.numel()} > {self.request_state_indices.numel()}." ) - if torch.cuda.is_current_stream_capturing(): + if self.backend.is_current_stream_capturing(): return # Attention indexes persistent request state by input-batch row. Refresh # that indirection after every batch compaction or reorder. @@ -417,7 +411,7 @@ def start_step( previous_host_write_event = self.host_write_event self.host_write_event = self.host_write_events[self._next_host_write_event] self._next_host_write_event ^= 1 - current_stream().wait_event(previous_host_write_event) + torch.accelerator.current_stream().wait_event(previous_host_write_event) if self.shared_host_region is not None and not num_tokens: # Without a forward, no TP collective keeps rank 0 from recording # a shared event again before every rank has issued its wait. @@ -499,9 +493,9 @@ def _stage_row_mirror_mapping(self, num_tokens: int) -> None: "HiSparse row mapping exceeds staging capacity: " f"{end} > {state.slots.shape[0]}." ) - main_stream = current_stream() + main_stream = torch.accelerator.current_stream() state.stream.wait_stream(main_stream) - with torch.cuda.stream(state.stream): + with state.stream: state.slots[start:end].copy_( slots[:num_tokens], non_blocking=True, @@ -579,9 +573,9 @@ def get_kv_connector_stats(self) -> HiSparseKVConnectorStats | None: if ( self._metrics_calls % _METRICS_INTERVAL == 0 and not self._metrics_pending - and not torch.cuda.is_current_stream_capturing() + and not self.backend.is_current_stream_capturing() ): - compute_stream = current_stream() + compute_stream = torch.accelerator.current_stream() for runtime in self.leader_runtimes: group = runtime.index_group compute_stream.wait_stream(group.copy_stream) @@ -614,12 +608,12 @@ def _submit_dma_descriptors( stream = self.dma_stream assert stream is not None if ready_event is None: - stream.wait_stream(current_stream()) + stream.wait_stream(torch.accelerator.current_stream()) else: stream.wait_event(ready_event) completion_event = torch.Event() - with torch.cuda.stream(stream): - ops.swap_blocks_batch( + with stream: + self.backend.swap_blocks_batch( descriptors.src[:descriptor_count], descriptors.dst[:descriptor_count], descriptors.sizes[:descriptor_count], @@ -785,7 +779,9 @@ def _restore_pages(self, transfers: list[SparseKVPageTransfer]) -> None: source[host_start : host_start + self.kernel_block_size], non_blocking=True, ) - self._record_transfer_completion(transfers, stream=current_stream()) + self._record_transfer_completion( + transfers, stream=torch.accelerator.current_stream() + ) def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None: if self.cache_handles[0].runtime.eager_host_mirror: @@ -915,7 +911,7 @@ def _finish_mirror_phase(self, ready_event: torch.Event | None = None) -> None: self._clear_forward_mirror_state() def finish_forward(self) -> None: - compute_stream = current_stream() + compute_stream = torch.accelerator.current_stream() self._forward_ready_event.record() self._finish_mirror_phase(self._forward_ready_event) if self.is_host_writer: diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index c76587787132..87330e6d5b11 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -235,6 +235,10 @@ class CudaPlatformBase(Platform): "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", ] + @classmethod + def get_hisparse_backend_cls(cls) -> str | None: + return "vllm.v1.hisparse.backend.CUDAHiSparseBackend" + @classmethod def import_kernels(cls) -> None: """Import CUDA kernel extensions (_C_stable_libtorch, optional _qutlass_C).""" diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index de899f7a1bf2..a872805f7cf8 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -196,6 +196,11 @@ def supported_dtypes(self) -> list[torch.dtype]: def is_cuda(self) -> bool: return self._enum == PlatformEnum.CUDA + @classmethod + def get_hisparse_backend_cls(cls) -> str | None: + """Return the device implementation for the common HiSparse runtime.""" + return None + def is_rocm(self) -> bool: return self._enum == PlatformEnum.ROCM diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index ba787b38da55..67069860e85c 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -24,14 +24,6 @@ ) -def _create_side_stream(device: torch.device) -> torch.Stream: - return torch.Stream(device=device) - - -def _create_event() -> torch.Event: - return torch.Event() - - @dataclass class SparseMLAIndexGroup: """Layers that consume one sparse-indexer result.""" @@ -45,6 +37,9 @@ class SparseMLAIndexGroup: has_indexer: bool num_layers: int = 0 + def _current_stream(self) -> Any: + return current_stream() + def register_layer( self, vllm_config: VllmConfig | None = None, @@ -66,7 +61,7 @@ def set_logical_topk_ready(self, layer_index: int) -> None: # touches this event), so skip the record as well. if in_piecewise_cudagraph(): return - self.logical_topk_ready.record(current_stream()) + self.logical_topk_ready.record(self._current_stream()) def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None: pass @@ -136,7 +131,7 @@ def _convert_once( physical_topk_indices = self.physical_topk_indices[:num_tokens] valid_topk_counts = self.valid_topk_counts[:num_tokens] - compute_stream = current_stream() + compute_stream = self._current_stream() if layer_index == 0: if self.has_indexer and not in_piecewise_cudagraph(): self.side_stream.wait_event(self.logical_topk_ready) @@ -173,6 +168,9 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup): prefill_stream: torch.Stream | None = None prefill_ready_events: list[torch.Event] = field(default_factory=list) + def _current_stream(self) -> Any: + return torch.accelerator.current_stream() + def register_layer( self, vllm_config: VllmConfig | None = None, @@ -185,7 +183,7 @@ def register_layer( assert kv_cache_dtype is not None layer_index = super().register_layer() if self.prefill_stream is None: - self.prefill_stream = _create_side_stream(self.logical_topk_indices.device) + self.prefill_stream = torch.Stream(device=self.logical_topk_indices.device) if kv_cache_dtype == "fp8_ds_mla": row_width = FP8_DS_MLA_ROW_BYTES kv_dtype = torch.uint8 @@ -207,7 +205,7 @@ def register_layer( assert cache is not None self.caches.append(cache) cache.index_group_caches = self.caches - self.prefill_ready_events.append(_create_event()) + self.prefill_ready_events.append(torch.Event()) return layer_index def cache(self, layer_index: int) -> HiSparseCacheHandle: @@ -325,7 +323,7 @@ def gather_fp8_prefill( prefill_stream = self.prefill_stream assert prefill_stream is not None block_table = plan.block_table[request_start : request_start + batch_size] - prefill_stream.wait_stream(current_stream()) + prefill_stream.wait_stream(torch.accelerator.current_stream()) with prefill_stream: ops.cp_gather_and_upconvert_fp8_kv_cache( cache.view.cache, @@ -385,9 +383,9 @@ def register_layer( dtype=torch.int32, device=self.logical_topk_indices.device, ), - side_stream=_create_side_stream(self.logical_topk_indices.device), - logical_topk_ready=_create_event(), - physical_topk_ready=_create_event(), + side_stream=torch.Stream(device=self.logical_topk_indices.device), + logical_topk_ready=torch.Event(), + physical_topk_ready=torch.Event(), has_indexer=is_index_producing_layer, ) group = self.current_group diff --git a/vllm/v1/hisparse/backend.py b/vllm/v1/hisparse/backend.py new file mode 100644 index 000000000000..2ed099e32219 --- /dev/null +++ b/vllm/v1/hisparse/backend.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Device operations used by the common HiSparse runtime and connector.""" + +from __future__ import annotations + +from functools import cache +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.platforms import current_platform +from vllm.utils.import_utils import resolve_obj_by_qualname + +if TYPE_CHECKING: + from vllm.v1.hisparse.runtime import HiSparseRuntime + from vllm.v1.kv_cache_interface import KVCacheConfig + + +class HiSparseBackend: + supports_shared_host_pool = False + stats_dtype = torch.uint64 + + def check_available(self) -> None: + raise NotImplementedError + + def pin_tensor(self, tensor: torch.Tensor) -> None: + raise NotImplementedError + + def unpin_tensor(self, tensor: torch.Tensor) -> None: + raise NotImplementedError + + def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: + raise NotImplementedError + + def validate_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: + """Validate device cache formats before allocating or pinning storage.""" + + def prepare_kv_cache_views( + self, + kv_cache_config: KVCacheConfig, + kv_caches: dict[str, Any], + forward_context: dict[str, Any], + ) -> None: + """Adapt device views before runtime binding; retain allocation ownership. + + Called once per cache initialization, inside the common rollback scope. + Backends must not put per-step metadata or transfers in this hook. + """ + + def create_ipc_event(self) -> Any: + """Create an event that can be exported to another process.""" + raise NotImplementedError + + def export_event(self, event: Any) -> Any: + raise NotImplementedError + + def import_event(self, device: torch.device, handle: Any) -> Any: + raise NotImplementedError + + def gather_resident_rows( + self, source: torch.Tensor, destination: torch.Tensor, row_ids: torch.Tensor + ) -> None: + """Copy resident rows into staging without changing their payload format.""" + rows = row_ids.to(torch.long).clamp_min(0) + block_size = source.shape[1] + payload = source[rows // block_size, rows % block_size] + if payload.dtype != destination.dtype: + payload = payload.contiguous().view(destination.dtype) + destination.copy_(payload) + + def is_current_stream_capturing(self) -> bool: + raise NotImplementedError + + def swap_blocks_batch( + self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor + ) -> None: + """Enqueue device-to-host copies on the current stream.""" + raise NotImplementedError + + +class CUDAHiSparseBackend(HiSparseBackend): + supports_shared_host_pool = True + + def check_available(self) -> None: + names = ( + "hisparse_resolve_residency", + "hisparse_invalidate_written_slots", + "hisparse_gather_plan", + "hisparse_gather_compact", + ) + if not all(hasattr(torch.ops._C_cache_ops, name) for name in names): + raise RuntimeError( + "HiSparse requires its compiled _C_cache_ops CUDA kernels. " + "Rebuild vLLM from source with csrc/hisparse_kernels.cu." + ) + + def pin_tensor(self, tensor: torch.Tensor) -> None: + from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor + + pin_tensor(tensor) + + def unpin_tensor(self, tensor: torch.Tensor) -> None: + cudart = torch.cuda.cudart() + error = cudart.cudaHostUnregister(tensor.data_ptr()) + if error.value != 0: + cudart.cudaGetLastError() + raise RuntimeError(f"cudaHostUnregister failed (code={error.value}).") + + def create_cache_ops(self, runtime: HiSparseRuntime) -> Any: + return torch.ops._C_cache_ops + + def create_ipc_event(self) -> Any: + return torch.cuda.Event(interprocess=True) + + def export_event(self, event: Any) -> Any: + return event.ipc_handle() + + def import_event(self, device: torch.device, handle: Any) -> Any: + return torch.cuda.Event.from_ipc_handle(device, handle) + + def is_current_stream_capturing(self) -> bool: + return torch.cuda.is_current_stream_capturing() + + def swap_blocks_batch( + self, src: torch.Tensor, dst: torch.Tensor, sizes: torch.Tensor + ) -> None: + from vllm import _custom_ops as ops + + ops.swap_blocks_batch(src, dst, sizes) + + +def get_hisparse_backend() -> HiSparseBackend: + backend_cls = current_platform.get_hisparse_backend_cls() + if backend_cls is None: + raise RuntimeError("The current platform does not support HiSparse.") + return _get_hisparse_backend(backend_cls) + + +@cache +def _get_hisparse_backend(backend_cls: str) -> HiSparseBackend: + return resolve_obj_by_qualname(backend_cls)() diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 6a02dd571cc3..6d591cd66e40 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -9,6 +9,7 @@ import torch from vllm.config import VllmConfig +from vllm.v1.hisparse.backend import get_hisparse_backend from vllm.v1.hisparse.layout import ( HISPARSE_HOT_SUFFIX, HISPARSE_RESIDENT_SUFFIX, @@ -126,8 +127,10 @@ def init_hisparse_kv_cache( vllm_config: VllmConfig, forward_context: dict[str, Any], block_tables: "BlockTables", -) -> dict[str, torch.Tensor]: +) -> dict[str, Any]: """Allocate and bind HiSparse caches within the caller's allocation context.""" + backend = get_hisparse_backend() + backend.validate_kv_cache_config(kv_cache_config) host_pool = get_hisparse_host_pool(vllm_config, kv_cache_config) try: kv_caches = allocate_hisparse_kv_caches( @@ -137,6 +140,7 @@ def init_hisparse_kv_cache( kernel_block_sizes, host_pool, ) + backend.prepare_kv_cache_views(kv_cache_config, kv_caches, forward_context) cache_handles = bind_hisparse_kv_caches( forward_context=forward_context, kv_cache_config=kv_cache_config, @@ -150,10 +154,7 @@ def init_hisparse_kv_cache( ) return kv_caches except Exception: - if host_pool.shared_region is not None: - host_pool.shared_region.cleanup() - elif host_pool.registered is not None: - release_pinned_state([], [host_pool.registered]) + release_pinned_state([], host_pool.pinned, host_pool.shared_region) drop_hisparse_host_pool(vllm_config) raise @@ -277,7 +278,7 @@ def bind_hisparse_kv_caches( host_pool.registered.untyped_storage().data_ptr() ) cache_handle.draft_layer = forward_context[layer_name].is_draft_layer - cache_handle.runtime.shared_host_region = host_pool.shared_region + cache_handle.runtime.host_pool = host_pool cache_handle.runtime.bind_source_cache( source_cache, registered_host_pool=host_pool.registered, diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 85866ac0e9b7..e9d0b09f37ed 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NVIDIA HiSparse resident, host, and hot-cache data plane.""" +"""HiSparse resident, host, and hot-cache data plane.""" from __future__ import annotations @@ -22,11 +22,11 @@ from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import get_max_shared_memory_bytes -from vllm.utils.torch_utils import async_tensor_h2d, current_stream +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import max_decode_query_len +from vllm.v1.hisparse.backend import get_hisparse_backend from vllm.v1.hisparse.types import SparseKVResidencyUpdate from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion -from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor logger = init_logger(__name__) @@ -212,8 +212,12 @@ def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]: return registered[:size], registered +def pin_tensor(tensor: torch.Tensor) -> None: + get_hisparse_backend().pin_tensor(tensor) + + class HiSparseHostPool: - """Keep the single host backing and its exact CUDA registration range.""" + """Keep the single host backing and its exact registration range.""" def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None: self.vllm_config = vllm_config @@ -221,6 +225,7 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> N self.backing: torch.Tensor | None = None self.registered: torch.Tensor | None = None self.shared_region: SharedOffloadRegion | None = None + self.pinned: list[torch.Tensor] = [] def allocate(self, size: int) -> torch.Tensor: if self.backing is not None: @@ -239,7 +244,7 @@ def allocate(self, size: int) -> torch.Tensor: ] assert config.hisparse_host_num_blocks is not None assert config.hisparse_host_block_stride is not None - _, _, self.shared_region = allocate_hisparse_host_pools( + _, self.pinned, self.shared_region = allocate_hisparse_host_pools( self.vllm_config, tensor_sizes, config.hisparse_host_num_blocks, @@ -252,6 +257,7 @@ def allocate(self, size: int) -> torch.Tensor: else: check_hisparse_host_memory(size) self.backing, self.registered = allocate_pinned_host_pool(size) + self.pinned = [self.registered] return self.backing @@ -276,7 +282,7 @@ def use_shared_hisparse_host_pool(vllm_config: VllmConfig) -> bool: """Whether replicated MLA host KV can share one local mmap.""" parallel = vllm_config.parallel_config return ( - current_platform.is_cuda_alike() + get_hisparse_backend().supports_shared_host_pool and parallel.tensor_parallel_size > 1 and parallel.pipeline_parallel_size == 1 and parallel.prefill_context_parallel_size == 1 @@ -400,21 +406,21 @@ def allocate_hisparse_host_pools( creator_memory_check=check_hisparse_host_memory, populate_shard=(tp_group.rank_in_group, tp_group.world_size), ) + pinned: list[torch.Tensor] = [] try: for start, end in _hisparse_registration_ranges( tensor_sizes, num_blocks, host_block_stride ): tensor = region.base_tensor[start:end] pin_tensor(tensor) - region.pinned_addresses.append(tensor.data_ptr()) - region.is_pinned = True + pinned.append(tensor) pools = [ region.create_next_canonical_view(size).view(-1) for size in tensor_sizes ] except Exception: - region.cleanup() + release_pinned_state([], pinned, region) raise - return pools, [], region + return pools, pinned, region def release_pinned_state( @@ -424,31 +430,31 @@ def release_pinned_state( ) -> None: """Synchronize and release registered host KV pools.""" if pinned_host_pools or shared_host_region is not None: + backend = get_hisparse_backend() try: torch.accelerator.synchronize() except RuntimeError as e: logger.warning( - "HiSparse: CUDA context unusable at teardown (%s); leaving " + "HiSparse: device context unusable at teardown (%s); leaving " "%d host-pool tensors pinned for kernel exit reclaim.", e, - len(pinned_host_pools) + int(shared_host_region is not None), + len(pinned_host_pools), ) return - cudart = torch.cuda.cudart() release_start = time.perf_counter() freed_bytes = 0 while pinned_host_pools: tensor = pinned_host_pools[-1] - err = cudart.cudaHostUnregister(tensor.data_ptr()) - if err.value != 0: + try: + backend.unpin_tensor(tensor) + except RuntimeError as error: logger.warning( - "HiSparse: cudaHostUnregister failed (code=%d); leaving " + "HiSparse: host unregister failed (%s); leaving " "%d host-pool tensors pinned for kernel exit reclaim.", - err.value, + error, len(pinned_host_pools), ) - cudart.cudaGetLastError() break freed_bytes += tensor.nbytes pinned_host_pools.pop() @@ -460,6 +466,8 @@ def release_pinned_state( ) for runtime in runtimes: + runtime.cache_ops = None + runtime.host_pool = None del runtime._host_cache del runtime.registered_host_pool del runtime.hot_backing @@ -630,14 +638,11 @@ def build_hisparse_prefill_staging_plan( def _has_hisparse_ops() -> bool: - if not hasattr(torch.ops, "_C_cache_ops"): + try: + get_hisparse_backend().check_available() + except RuntimeError: return False - return ( - hasattr(torch.ops._C_cache_ops, "hisparse_resolve_residency") - and hasattr(torch.ops._C_cache_ops, "hisparse_invalidate_written_slots") - and hasattr(torch.ops._C_cache_ops, "hisparse_gather_plan") - and hasattr(torch.ops._C_cache_ops, "hisparse_gather_compact") - ) + return True class _SharedTopKState: @@ -687,7 +692,10 @@ def __init__( dtype=torch.int32, device=device, ) - self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device) + # NPU arange does not support int16 output, so generate int32 then cast. + self.lru_init = torch.arange( + region_stride, dtype=torch.int32, device=device + ).to(torch.int16) self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous() self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k) # Request ids the resolver reads, built on the copy stream: a @@ -699,9 +707,10 @@ def __init__( ) self.logical_topk_ready = logical_topk_ready self.followers: list[HiSparseRuntime] = [] - self.swap_stats = torch.zeros(2, dtype=torch.uint64, device=device) + stats_dtype = get_hisparse_backend().stats_dtype + self.swap_stats = torch.zeros(2, dtype=stats_dtype, device=device) self.swap_stats_host = torch.empty( - 2, dtype=torch.uint64, device="cpu", pin_memory=True + 2, dtype=stats_dtype, device="cpu", pin_memory=True ) self.stats_row_bytes = 0 @@ -727,13 +736,9 @@ def __init__( copy_stream: torch.Stream | None = None, logical_topk_ready: torch.Event | None = None, ) -> None: - if not _has_hisparse_ops(): - raise RuntimeError( - "HiSparse requires its compiled _C_cache_ops CUDA kernels " - "(host-resident decode has no Python fallback). Rebuild vLLM " - "from source so " - "csrc/hisparse_kernels.cu is included." - ) + self.backend = get_hisparse_backend() + self.backend.check_available() + self.cache_ops: Any = None self.max_num_reqs = max_num_reqs self.row_width = row_width self.kv_dtype = kv_dtype @@ -774,7 +779,7 @@ def __init__( self.eager_host_mirror = config.eager_host_mirror self.resident_source_index = -1 self.request_state_indices: torch.Tensor | None = None - self.shared_host_region: SharedOffloadRegion | None = None + self.host_pool: HiSparseHostPool | None = None @property def host_cache(self) -> torch.Tensor: @@ -833,6 +838,7 @@ def bind_source_cache( self.registered_host_pool = ( registered_host_pool if registered_host_pool is not None else kv_cache ) + self.cache_ops = self.backend.create_cache_ops(self) def gather_prefill_cache( self, @@ -858,14 +864,10 @@ def gather_prefill_cache( device=plan.block_table.device, ) if plan.gpu_row_ids is not None and resident_cache is not None: - gpu_rows = plan.gpu_row_ids[0].to(torch.long) - src = gpu_rows.clamp_min(0) - resident_block_size = resident_cache.shape[1] - rows = resident_cache[src // resident_block_size, src % resident_block_size] - if rows.dtype != staged.dtype: - rows = rows.contiguous().view(staged.dtype) - staged.view(-1, row_width).copy_(rows) - torch.ops._C_cache_ops.hisparse_gather_plan( + self.backend.gather_resident_rows( + resident_cache, staged.view(-1, row_width), plan.gpu_row_ids[0] + ) + self.cache_ops.hisparse_gather_plan( kv_cache.view(-1, row_width), staged, plan.row_ids, @@ -880,7 +882,7 @@ def gather_prefill_cache( def reset_hot_state(self) -> None: """Drop all hot-buffer bookkeeping (hits become misses).""" group = self.index_group - compute_stream = current_stream() + compute_stream = torch.accelerator.current_stream() compute_stream.wait_stream(group.copy_stream) group.device_global_indices.fill_(-1) group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots)) @@ -912,7 +914,7 @@ def invalidate_written_slots( if num_tokens == 0: return assert self.request_state_indices is not None - torch.ops._C_cache_ops.hisparse_invalidate_written_slots( + self.cache_ops.hisparse_invalidate_written_slots( self.index_group.device_global_indices, self.request_state_indices, req_id_per_token[:num_tokens], @@ -979,7 +981,7 @@ def _resolve_residency( # Padded rows are skipped by the kernel (request_state_indices) and must # come out as -1 so the attention kernel masks them. - torch.ops._C_cache_ops.hisparse_resolve_residency( + self.cache_ops.hisparse_resolve_residency( host_cache, hot.cache, self.hot_block_table, @@ -1012,7 +1014,7 @@ def _resolve_residency( def _swap_rows(self, shared_rows: slice) -> None: hot = self.hot shared = self.index_group.shared_topk - torch.ops._C_cache_ops.hisparse_gather_compact( + self.cache_ops.hisparse_gather_compact( self.host_cache, hot.cache, shared.swap_host_physical_rows[shared_rows], @@ -1035,7 +1037,7 @@ def _resolve_and_stage_group( num_valid_rows: torch.Tensor, ) -> None: group = self.index_group - compute_stream = current_stream() + compute_stream = torch.accelerator.current_stream() if group.logical_topk_ready is not None and not in_piecewise_cudagraph(): group.copy_stream.wait_event(group.logical_topk_ready) else: @@ -1072,7 +1074,7 @@ def _resolve_and_stage_group( def _stage_this_layer(self, shared_rows: slice) -> None: group = self.index_group - compute_stream = current_stream() + compute_stream = torch.accelerator.current_stream() group.copy_stream.wait_stream(compute_stream) with group.copy_stream: self._swap_rows(shared_rows) @@ -1084,7 +1086,7 @@ def _stage_this_layer(self, shared_rows: slice) -> None: def wait_for_staged_swap(self) -> None: if not self._swap_staged or self._layer_ready_event is None: raise RuntimeError("HiSparse layer swap was not prefetched.") - current_stream().wait_event(self._layer_ready_event) + torch.accelerator.current_stream().wait_event(self._layer_ready_event) self._swap_staged = False def swap_in( From 7bf13a3db1d0be2b1365e3ec5e5d5ae91e2db6fd Mon Sep 17 00:00:00 2001 From: John Calderon <81483067+johncalesp@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:33:08 -0400 Subject: [PATCH 007/406] [CI][Test] Adjust TRT-LLM and GDN numerical tolerances for Rubin (#64081) Signed-off-by: John Calderon --- tests/kernels/attention/test_flashinfer_trtllm_attention.py | 4 ++-- tests/kernels/fla/test_fused_gdn_post_conv.py | 6 +++++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/kernels/attention/test_flashinfer_trtllm_attention.py b/tests/kernels/attention/test_flashinfer_trtllm_attention.py index 074e53d014c9..957bd2fdc5eb 100644 --- a/tests/kernels/attention/test_flashinfer_trtllm_attention.py +++ b/tests/kernels/attention/test_flashinfer_trtllm_attention.py @@ -342,7 +342,7 @@ def test_flashinfer_trtllm_decode_with_baseline( if is_nvfp4_kv: rtol, atol = 1.0, 1.0 # nvfp4 has higher quantization error elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE: - rtol, atol = 7e-2, 9e-2 + rtol, atol = 7e-2, 1e-1 elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE: rtol, atol = 3e-2, 4e-2 elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype: @@ -637,7 +637,7 @@ def test_flashinfer_trtllm_prefill_with_baseline( elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE: rtol, atol = 4e-2, 6e-2 elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype: - rtol, atol = 2e-2, 3e-2 + rtol, atol = 2e-2, 6e-2 else: rtol, atol = 1e-2, 1e-2 diff --git a/tests/kernels/fla/test_fused_gdn_post_conv.py b/tests/kernels/fla/test_fused_gdn_post_conv.py index a29395edac9b..f6d8500c5a9c 100644 --- a/tests/kernels/fla/test_fused_gdn_post_conv.py +++ b/tests/kernels/fla/test_fused_gdn_post_conv.py @@ -343,6 +343,10 @@ def test_fused_gdn_decode_post_conv_mtp_head_ratios( if query_lengths[-1] == 0: state_indices[-1].zero_() + output_rtol = 5e-4 + if impl == "cuda" and state_dtype == torch.float32 and norm_dtype == torch.bfloat16: + output_rtol = 1e-3 + for step, accepted_tokens in enumerate((1, min(2, state_width), state_width)): num_accepted_tokens.fill_(accepted_tokens) if query_lengths[-1] == 0: @@ -392,7 +396,7 @@ def test_fused_gdn_decode_post_conv_mtp_head_ratios( output_error = (actual.float() - expected.float()).norm() output_relative_l2 = output_error / expected.float().norm().clamp_min(1e-20) - assert output_relative_l2 < 5e-4, ( + assert output_relative_l2 < output_rtol, ( f"MTP output relative L2 mismatch at step {step}: " f"{output_relative_l2.item():.6g}" ) From 4d95f222e68bd43acbdfdffad6d46eea70146dcc Mon Sep 17 00:00:00 2001 From: Sherif Waly Date: Tue, 27 Oct 2026 21:36:02 +0000 Subject: [PATCH 008/406] [Feature][Core] Enable best-effort NUMA binding by default (#61258) --- docs/configuration/optimization.md | 53 +- tests/engine/test_arg_utils.py | 5 + tests/utils_/test_numa_utils.py | 455 +++++++++++++++++- tests/utils_/test_system_utils.py | 37 +- .../distributed/test_dense_dp_world_size.py | 141 +++++- .../v1/simple_kv_offload/test_integration.py | 103 +++- vllm/config/parallel.py | 41 +- vllm/engine/arg_utils.py | 2 +- vllm/utils/numa_utils.py | 205 ++++++-- vllm/utils/system_utils.py | 9 +- 10 files changed, 964 insertions(+), 87 deletions(-) diff --git a/docs/configuration/optimization.md b/docs/configuration/optimization.md index 18122ea2c8bc..4ef9108db7cb 100644 --- a/docs/configuration/optimization.md +++ b/docs/configuration/optimization.md @@ -156,10 +156,47 @@ GPU. vLLM can pin each worker with `numactl` before the Python subprocess starts so the interpreter, imports, and early allocator state are created with the desired NUMA policy from the beginning. -Use `--numa-bind` to enable the feature. By default, vLLM auto-detects the -GPU-to-NUMA mapping and uses `--cpunodebind= --membind=` for each -worker. When you need a custom CPU policy, add `--numa-bind-cpus` and vLLM will -switch to `--physcpubind= --membind=`. +By default, vLLM binds automatically when the environment supports it: +CUDA-alike platform, multiple NUMA nodes, `numactl` installed, permission to +set NUMA memory policy, unrestricted CPU affinity, a default NUMA memory policy, +the `spawn` multiprocessing start method (the CLI default), and a local `mp` or +`uni` executor with multiprocessing DP, including supported prefill context +parallelism (PCP) launches. Ray, external launchers and `--device-ids` remapping +are excluded from automatic binding. +Otherwise vLLM starts unbound. Use +`--numa-bind` to make a missing `numactl` or undetectable GPU-to-NUMA topology a +startup error, or `--no-numa-bind` to disable binding. Passing +`--numa-bind-nodes` or `--numa-bind-cpus` explicitly enables binding even without +`--numa-bind`: either list requests strict `--membind`, not automatic soft memory +preference or CPU-only binding. Neither can be combined with `--no-numa-bind`. +Automatic binding preserves an existing non-default memory policy, such as +interleave, preferred-node, or bind. Explicit enablement may override it. + +Automatic binding uses `--cpunodebind= --preferred=` for each worker: +memory allocations prefer the GPU-local node but can fall back to other nodes. +EngineCore prefers its shard's nodes with `--preferred-many=` when the +shard spans multiple NUMA nodes, or `--preferred=` for a single-node shard. +Explicit `--numa-bind` or binding lists retain strict `--membind=`. +When you need a custom CPU policy, add `--numa-bind-cpus` and vLLM will switch +to `--physcpubind=`. + +Automatic eligibility and topology validation happen in each local launcher +before starting its first bound subprocess, using freshly discovered topology. +The mapping must cover all visible GPUs and local workers, and each mapped NUMA +node must have CPUs. The resolved decision and mapping are passed to spawned +children, so an already-bound EngineCore does not repeat detection under its +restricted CPU affinity. +PCP workers use the full TP × PCP × PP shard size for local DP GPU offsets; +EngineCore's CPU affinity covers all workers in its shard. +Multi-node multiprocessing launchers resolve their own host-local mapping; +NUMA node IDs are never interpreted as cluster-wide IDs. + +Before each bound launch, vLLM probes the selected policy with +`numactl true`. If CPU-and-memory binding is rejected, it retries with +CPU-only binding; if that is also rejected, the subprocess starts unbound. +Logs report the selected policy or fallback. If `--preferred-many` is unsupported, +automatic mode falls back to CPU-only binding, never strict memory binding. +These probes test policy acceptance, not future allocation capacity. These `--numa-bind*` options only apply to GPU execution processes. They do not configure the CPU backend's separate thread-affinity controls. Automatic @@ -174,7 +211,7 @@ order as the GPU indices. Each CPU list must use `numactl --physcpubind` syntax such as `0-3`, `0,2,4-7`, or `16-31,48-63`. ```bash -# Auto-detect NUMA nodes for visible GPUs +# Require binding, auto-detecting NUMA nodes for visible GPUs vllm serve meta-llama/Llama-3.1-8B-Instruct \ --tensor-parallel-size 4 \ --numa-bind @@ -195,10 +232,12 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \ Notes: -- CLI usage forces multiprocessing to use the `spawn` method automatically. If you enable NUMA binding through the Python API, also set `VLLM_WORKER_MULTIPROC_METHOD=spawn`. -- Automatic detection relies on NVML and NUMA support from the host. If it cannot determine the mapping reliably, pass `--numa-bind-nodes` explicitly. +- The CLI defaults to `spawn` when `VLLM_WORKER_MULTIPROC_METHOD` is unset. Automatic NUMA binding never changes the start method: if it is `fork`, default binding stays off. Python API usage normally defaults to `fork`; set `VLLM_WORKER_MULTIPROC_METHOD=spawn` to allow binding and follow Python's `if __name__ == "__main__":` requirements. +- Automatic memory preference permits allocations on other NUMA nodes when local memory is unavailable. Explicit `--numa-bind` still uses `--membind`, which restricts host allocations to the selected nodes; for host-memory-heavy setups (e.g. large CPU KV offload), avoid explicit strict binding if workers need memory from other nodes. The policy probes do not check allocation capacity or container memory limits. +- Automatic detection relies on NVML (CUDA) or amdsmi (ROCm) and NUMA support from the host. If it cannot determine the mapping reliably, default mode starts unbound; pass `--numa-bind-nodes` to provide a mapping explicitly. - Explicit `--numa-bind-nodes` and `--numa-bind-cpus` values must be valid `numactl` inputs. vLLM does a small amount of validation, but the effective binding semantics are still determined by `numactl`. - The current implementation binds GPU execution processes such as `EngineCore` and multiprocessing workers. It does not apply NUMA binding to frontend API server processes or the DP coordinator. +- Ray actors are not bound by the multiprocessing executable wrapper. Enabling automatic NUMA binding for Ray requires a separate actor-launch implementation. - In containerized environments, NUMA policy syscalls may require extra permissions, such as `--cap-add SYS_NICE` when running via `docker run`. ### CPU Backend Thread Affinity diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 938cf9b3999b..c183b32a8d32 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -804,6 +804,11 @@ def test_numa_bind_args(): assert engine_args.numa_bind_nodes == [0, 0, 1, 1] assert engine_args.numa_bind_cpus == ["0-3", "4-7", "8-11", "12-15"] + # Unset means auto; --no-numa-bind disables explicitly. + assert EngineArgs.from_cli_args(args=parser.parse_args([])).numa_bind is None + args = parser.parse_args(["--no-numa-bind"]) + assert EngineArgs.from_cli_args(args=args).numa_bind is False + def test_ir_op_priority(): from vllm.config.kernel import IrOpPriorityConfig, KernelConfig diff --git a/tests/utils_/test_numa_utils.py b/tests/utils_/test_numa_utils.py index 3c09722f82df..014e2f7da7dc 100644 --- a/tests/utils_/test_numa_utils.py +++ b/tests/utils_/test_numa_utils.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os +from copy import deepcopy from functools import partial from types import SimpleNamespace +from typing import Any import pytest @@ -39,20 +41,30 @@ def _no_pct_open(path, *args, **kwargs): def _make_config(**parallel_kwargs): - parallel_defaults = dict( + parallel_defaults: dict[str, Any] = dict( numa_bind=False, + _numa_bind_auto=False, numa_bind_nodes=None, numa_bind_cpus=None, + assigned_physical_gpu_ids=None, distributed_executor_backend="mp", data_parallel_backend="mp", nnodes_within_dp=1, data_parallel_rank_local=0, data_parallel_index=0, + data_parallel_size_local=1, pipeline_parallel_size=1, tensor_parallel_size=1, prefill_context_parallel_size=1, ) parallel_defaults.update(parallel_kwargs) + parallel_defaults.setdefault( + "local_world_size", + parallel_defaults["pipeline_parallel_size"] + * parallel_defaults["tensor_parallel_size"] + * parallel_defaults["prefill_context_parallel_size"] + // parallel_defaults["nnodes_within_dp"], + ) parallel_config = SimpleNamespace(**parallel_defaults) parallel_config.world_size = ( parallel_config.tensor_parallel_size @@ -214,12 +226,18 @@ def test_pct_binding_returns_none_when_node_cpulist_missing(monkeypatch): assert numa_utils._maybe_get_pct_cpu_binding([0]) is None -def test_get_numactl_args_uses_pct_when_user_did_not_specify_cpus(monkeypatch): +@pytest.mark.parametrize("automatic", [False, True]) +def test_get_numactl_args_uses_pct_when_user_did_not_specify_cpus( + monkeypatch, automatic +): _patch_pct_gates(monkeypatch, model_match=True, highest_perf=46) - vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[0, 1]) + vllm_config = _make_config( + numa_bind=True, _numa_bind_auto=automatic, numa_bind_nodes=[0, 1] + ) + memory_policy = "preferred" if automatic else "membind" assert ( numa_utils._get_numactl_worker_args(vllm_config.parallel_config, local_rank=1) - == "--physcpubind=0,1,16,17,64,65,80,81 --membind=1" + == f"--physcpubind=0,1,16,17,64,65,80,81 --{memory_policy}=1" ) @@ -249,7 +267,10 @@ def test_get_numactl_args_engine_core_baseline_spans_shard_numa_nodes(): ) -def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(monkeypatch): +@pytest.mark.parametrize("automatic", [False, True]) +def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes( + monkeypatch, automatic +): """PCT: EngineCore for a multi-NUMA shard binds to the union of priority cores across all shard nodes, so worker `--physcpubind` is always a subset of EngineCore's `cpus_allowed`.""" @@ -261,15 +282,17 @@ def test_get_numactl_args_engine_core_pct_spans_shard_numa_nodes(monkeypatch): ) vllm_config = _make_config( numa_bind=True, + _numa_bind_auto=automatic, numa_bind_nodes=[0, 0, 1, 1], tensor_parallel_size=4, ) + memory_policy = "preferred-many" if automatic else "membind" assert numa_utils._get_numactl_enginecore_args( vllm_config.parallel_config, local_rank=0 ) == ( "--physcpubind=" "0,1,16,17,64,65,80,81,128,129,144,145,192,193,208,209" - " --membind=0,1" + f" --{memory_policy}=0,1" ) @@ -543,6 +566,18 @@ def test_parallel_config_validates_numa_bind_nodes(): ParallelConfig(numa_bind_nodes=[0, -1]) +@pytest.mark.parametrize( + ("nodes", "cpus"), [([0, 1], None), (None, ["0-3"])], ids=["nodes", "cpus"] +) +def test_parallel_config_numa_lists_imply_numa_bind(nodes, cpus): + assert ParallelConfig().numa_bind is None + config = ParallelConfig(numa_bind_nodes=nodes, numa_bind_cpus=cpus) + assert config.numa_bind is True + assert config._numa_bind_auto is False + with pytest.raises(ValueError, match="cannot be combined with numa_bind=False"): + ParallelConfig(numa_bind=False, numa_bind_nodes=nodes, numa_bind_cpus=cpus) + + @pytest.mark.parametrize("cpuset", ["", "abc", "1-", "4-1", "1,,2", "1:2"]) def test_parallel_config_rejects_invalid_numa_bind_cpus(cpuset): with pytest.raises(ValueError, match="numa_bind_cpus"): @@ -564,7 +599,7 @@ def test_configure_subprocess_numa_fallback(monkeypatch): import multiprocessing monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/numactl") - monkeypatch.setattr(numa_utils.envs, "VLLM_WORKER_MULTIPROC_METHOD", "spawn") + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") node_config = _make_config(numa_bind=True, numa_bind_nodes=[0]) monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([])) @@ -593,3 +628,409 @@ def test_configure_subprocess_numa_fallback(monkeypatch): with numa_utils.configure_subprocess(node_config, local_rank=0): assert multiprocessing.spawn.get_executable() == before assert numa_utils._NUMACTL_ARGS_ENV not in os.environ + + +def _patch_auto_numa( + monkeypatch, start_method, detect, visible_gpus=1, node_cpus="0-7" +): + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/numactl") + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", start_method) + monkeypatch.setattr(numa_utils, "_is_auto_numa_available", lambda **kwargs: True) + monkeypatch.setattr(numa_utils, "get_auto_numa_nodes", detect) + monkeypatch.setattr( + "vllm.platforms.current_platform.device_count", lambda: visible_gpus + ) + real_read_text = numa_utils.Path.read_text + + def read_text(path, *args, **kwargs): + if str(path).startswith("/sys/devices/system/node/node"): + return node_cpus + return real_read_text(path, *args, **kwargs) + + monkeypatch.setattr(numa_utils.Path, "read_text", read_text) + + +@pytest.mark.parametrize("nodes", [[0], [0, 1]]) +def test_configure_subprocess_auto_binds_and_records_result(monkeypatch, nodes): + """Default mode resolves once in the launcher and stores the result, so + spawned children bind with it instead of re-checking eligibility.""" + _patch_auto_numa(monkeypatch, "spawn", lambda: nodes, visible_gpus=len(nodes)) + monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([])) + + vllm_config = _make_config(numa_bind=None, tensor_parallel_size=len(nodes)) + node_list = ",".join(map(str, nodes)) + preference = "preferred" if len(nodes) == 1 else "preferred-many" + with numa_utils.configure_subprocess( + vllm_config, local_rank=0, process_kind="EngineCore" + ): + assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == ( + f"--cpunodebind={node_list} --{preference}={node_list}" + ) + child_config = deepcopy(vllm_config) + assert vllm_config.parallel_config.numa_bind is True + assert child_config.parallel_config._numa_bind_auto is True + assert child_config.parallel_config.numa_bind_nodes == nodes + + monkeypatch.setattr(numa_utils, "get_auto_numa_nodes", _raise_probe_error) + monkeypatch.setattr(numa_utils, "_is_auto_numa_available", lambda **kwargs: False) + with numa_utils.configure_subprocess(child_config, local_rank=len(nodes) - 1): + assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == ( + f"--cpunodebind={nodes[-1]} --preferred={nodes[-1]}" + ) + + +def test_parallel_config_auto_policy_survives_copy_without_changing_hash(): + """Copying the launch decision must not change the DP configuration hash.""" + from vllm.config.utils import replace + + parallel_config = ParallelConfig() + original_hash = parallel_config.compute_hash() + parallel_config.numa_bind = True + parallel_config.numa_bind_nodes = [0] + parallel_config._numa_bind_auto = True + for copied in ( + replace(parallel_config), + deepcopy(parallel_config), + ): + assert copied._numa_bind_auto is True + assert copied.compute_hash() == original_hash + + +@pytest.mark.parametrize("dp_local_rank", [0, 1]) +def test_auto_numa_accepts_pcp_dp_topology(monkeypatch, dp_local_rank): + """Automatic eligibility must cover every PCP worker in both DP shards.""" + _patch_auto_numa(monkeypatch, "spawn", lambda: [0, 0, 1, 1], visible_gpus=4) + config = ParallelConfig( + distributed_executor_backend="mp", + prefill_context_parallel_size=2, + data_parallel_size=2, + data_parallel_size_local=2, + ) + assert numa_utils._can_auto_bind(config, local_rank=0, dp_local_rank=dp_local_rank) + assert config.numa_bind_nodes == [0, 0, 1, 1] + + +def test_auto_numa_redetects_when_visible_devices_change(monkeypatch): + """Independent launchers must not reuse another visible GPU set's mapping.""" + _patch_auto_numa( + monkeypatch, "spawn", numa_utils.get_auto_numa_nodes, visible_gpus=2 + ) + monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([])) + detections = [] + + def detect(): + devices = os.environ["CUDA_VISIBLE_DEVICES"] + detections.append(devices) + return {"0,1": [0, 0], "2,3": [1, 1]}[devices] + + monkeypatch.setattr( + "vllm.platforms.current_platform.get_all_device_numa_nodes", + detect, + raising=False, + ) + for devices, node in [("0,1", 0), ("2,3", 1)]: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", devices) + vllm_config = _make_config(numa_bind=None, tensor_parallel_size=2) + with numa_utils.configure_subprocess( + vllm_config, local_rank=0, process_kind="EngineCore" + ): + numactl_args = os.environ[numa_utils._NUMACTL_ARGS_ENV] + assert numactl_args == f"--cpunodebind={node} --preferred={node}" + assert vllm_config.parallel_config.numa_bind_nodes == [node, node] + assert detections == ["0,1", "2,3"] + + +@pytest.mark.parametrize("explicit", [False, True], ids=["automatic", "explicit"]) +@pytest.mark.parametrize( + "policy", + [0, 1, 2, 3, 4, 5, None], + ids=[ + "default", + "preferred", + "bind", + "interleave", + "local", + "preferred-many", + "unreadable", + ], +) +def test_configure_subprocess_respects_existing_mempolicy( + monkeypatch, policy, explicit +): + """Only explicit enablement may replace an existing non-default policy.""" + is_auto_numa_available = numa_utils._is_auto_numa_available + _patch_auto_numa(monkeypatch, "spawn", numa_utils.get_auto_numa_nodes) + monkeypatch.setattr(numa_utils, "_is_auto_numa_available", is_auto_numa_available) + monkeypatch.setattr("vllm.platforms.current_platform.is_cuda_alike", lambda: True) + monkeypatch.setattr(numa_utils.os.path, "isdir", lambda path: True) + monkeypatch.setattr(numa_utils.psutil, "cpu_count", lambda: 4) + monkeypatch.setattr( + numa_utils.psutil, + "Process", + lambda pid: SimpleNamespace(cpu_affinity=lambda: [0, 1, 2, 3]), + ) + monkeypatch.setattr( + "vllm.platforms.current_platform.get_all_device_numa_nodes", + lambda: [0], + raising=False, + ) + + def get_mempolicy(mode, *_args): + if policy is None: + return -1 + mode._obj.value = policy + return 0 + + monkeypatch.setattr( + numa_utils, + "get_libnuma", + lambda: SimpleNamespace(numa_available=lambda: 0, get_mempolicy=get_mempolicy), + ) + monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([])) + vllm_config = _make_config(numa_bind=True if explicit else None) + if explicit and policy is None: + with ( + pytest.raises(RuntimeError, match="could not detect"), + numa_utils.configure_subprocess(vllm_config, local_rank=0), + ): + pass + return + + expected_bind = policy is not None and (policy == 0 or explicit) + memory_arg = "--membind=0" if explicit else "--preferred=0" + with numa_utils.configure_subprocess(vllm_config, local_rank=0): + numactl_args = os.environ.get(numa_utils._NUMACTL_ARGS_ENV) + assert numactl_args == ( + f"--cpunodebind=0 {memory_arg}" if expected_bind else None + ) + assert vllm_config.parallel_config.numa_bind is expected_bind + + +def _raise_probe_error(): + raise RuntimeError("amdsmi init failed") + + +@pytest.mark.parametrize( + ("start_method", "detect", "assigned_gpus"), + [ + ("fork", lambda: [0], None), + ("spawn", lambda: None, None), + ("spawn", lambda: [], None), + ("spawn", _raise_probe_error, None), + ("spawn", lambda: [0, 1], [1]), + ], + ids=["fork", "no-topology", "no-visible-gpus", "probe-raises", "device-ids"], +) +def test_configure_subprocess_auto_falls_back_to_unbound( + monkeypatch, start_method, detect, assigned_gpus +): + """Default mode never fails startup or forces spawn; when binding cannot + engage (or would pick the wrong node) it starts unbound.""" + import multiprocessing + + _patch_auto_numa(monkeypatch, start_method, detect) + + vllm_config = _make_config(numa_bind=None, assigned_physical_gpu_ids=assigned_gpus) + before = multiprocessing.spawn.get_executable() + with numa_utils.configure_subprocess(vllm_config, local_rank=0): + assert multiprocessing.spawn.get_executable() == before + assert vllm_config.parallel_config.numa_bind is False + + +@pytest.mark.parametrize( + "parallel_kwargs", + [ + {"distributed_executor_backend": "ray"}, + {"data_parallel_backend": "ray"}, + {"distributed_executor_backend": "external_launcher"}, + ], + ids=["ray-executor", "ray-dp", "external-launcher"], +) +def test_auto_numa_skips_unsupported_launchers(monkeypatch, parallel_kwargs): + detections = [] + + def detect(): + detections.append(True) + return [0, 1] + + _patch_auto_numa(monkeypatch, "spawn", detect, visible_gpus=2) + vllm_config = _make_config(numa_bind=None, **parallel_kwargs) + with numa_utils.configure_subprocess(vllm_config, local_rank=0): + assert numa_utils._NUMACTL_ARGS_ENV not in os.environ + assert vllm_config.parallel_config.numa_bind is False + assert detections == [] + + +@pytest.mark.parametrize("nodes", [[0], [0, 1, 2], [0, -1], [0, None]]) +def test_auto_numa_rejects_incomplete_or_invalid_topology(monkeypatch, nodes): + _patch_auto_numa(monkeypatch, "spawn", lambda: nodes, visible_gpus=2) + vllm_config = _make_config(numa_bind=None, tensor_parallel_size=2) + with numa_utils.configure_subprocess(vllm_config, local_rank=0): + assert numa_utils._NUMACTL_ARGS_ENV not in os.environ + assert vllm_config.parallel_config.numa_bind is False + assert vllm_config.parallel_config.numa_bind_nodes is None + + +def test_auto_numa_checks_all_local_dp_workers_before_first_engine(monkeypatch): + _patch_auto_numa(monkeypatch, "spawn", lambda: [0, 1], visible_gpus=2) + vllm_config = _make_config( + numa_bind=None, tensor_parallel_size=2, data_parallel_size_local=2 + ) + with numa_utils.configure_subprocess( + vllm_config, local_rank=0, dp_local_rank=0, process_kind="EngineCore" + ): + assert numa_utils._NUMACTL_ARGS_ENV not in os.environ + assert vllm_config.parallel_config.numa_bind is False + + +def test_auto_numa_multi_node_topology_is_local(monkeypatch): + _patch_auto_numa(monkeypatch, "spawn", lambda: [2, 3], visible_gpus=2) + monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([])) + vllm_config = _make_config( + numa_bind=None, tensor_parallel_size=4, nnodes_within_dp=2 + ) + with numa_utils.configure_subprocess( + vllm_config, local_rank=0, process_kind="EngineCore" + ): + assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == ( + "--cpunodebind=2,3 --preferred-many=2,3" + ) + child_config = deepcopy(vllm_config) + with numa_utils.configure_subprocess(child_config, local_rank=1): + assert ( + os.environ[numa_utils._NUMACTL_ARGS_ENV] == "--cpunodebind=3 --preferred=3" + ) + + +@pytest.mark.parametrize("node_cpus", [None, "", "invalid"]) +def test_auto_numa_skips_nodes_without_usable_cpus(monkeypatch, node_cpus): + _patch_auto_numa(monkeypatch, "spawn", lambda: [0], node_cpus=node_cpus) + vllm_config = _make_config(numa_bind=None) + with numa_utils.configure_subprocess(vllm_config, local_rank=0): + assert numa_utils._NUMACTL_ARGS_ENV not in os.environ + assert vllm_config.parallel_config.numa_bind is False + + +@pytest.mark.parametrize( + "nodes, rejected_args, expected_args", + [ + ([0], "--preferred=", "--cpunodebind=0"), + ([0], "--cpunodebind=", None), + ([0, 1], "--preferred-many=", "--cpunodebind=0,1"), + ([0, 1], "--cpunodebind=", None), + ], +) +def test_auto_numa_records_effective_policy( + monkeypatch, caplog, nodes, rejected_args, expected_args +): + caplog.set_level("INFO", logger=numa_utils.logger.name) + _patch_auto_numa(monkeypatch, "spawn", lambda: nodes, visible_gpus=len(nodes)) + probes = [] + fake_run = _fake_numactl_run([rejected_args]) + + def run(cmd, *args, **kwargs): + probes.extend(cmd) + return fake_run(cmd, *args, **kwargs) + + monkeypatch.setattr(numa_utils.subprocess, "run", run) + vllm_config = _make_config(numa_bind=None, tensor_parallel_size=len(nodes)) + with numa_utils.configure_subprocess( + vllm_config, local_rank=0, process_kind="EngineCore" + ): + assert os.environ.get(numa_utils._NUMACTL_ARGS_ENV) == expected_args + assert vllm_config.parallel_config.numa_bind is (expected_args is not None) + launch_logs = [r.getMessage() for r in caplog.records if "Launching" in r.message] + if expected_args is None: + assert launch_logs == [] + assert vllm_config.parallel_config.numa_bind_nodes is None + assert vllm_config.parallel_config._numa_bind_auto is False + else: + assert len(launch_logs) == 1 + assert launch_logs[0].endswith(f"numactl {expected_args}") + assert not any(arg.startswith("--membind=") for arg in probes) + + +def test_auto_numa_does_not_swallow_launch_errors(monkeypatch): + import multiprocessing + + _patch_auto_numa(monkeypatch, "spawn", lambda: [0]) + monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([])) + monkeypatch.delenv(numa_utils._NUMACTL_ARGS_ENV, raising=False) + monkeypatch.delenv(numa_utils._NUMACTL_PYTHON_EXECUTABLE_ENV, raising=False) + original_executable = multiprocessing.spawn.get_executable() + with ( + pytest.raises(RuntimeError, match="worker startup failed"), + numa_utils.configure_subprocess(_make_config(numa_bind=None), local_rank=0), + ): + raise RuntimeError("worker startup failed") + assert multiprocessing.spawn.get_executable() == original_executable + assert numa_utils._NUMACTL_ARGS_ENV not in os.environ + assert numa_utils._NUMACTL_PYTHON_EXECUTABLE_ENV not in os.environ + + +@pytest.mark.parametrize("missing_numactl", [False, True]) +def test_auto_numa_skips_unavailable_environment(monkeypatch, missing_numactl): + _patch_auto_numa(monkeypatch, "spawn", lambda: [0]) + if missing_numactl: + monkeypatch.setattr("shutil.which", lambda name: None) + else: + monkeypatch.setattr( + numa_utils, "_is_auto_numa_available", lambda **kwargs: False + ) + vllm_config = _make_config(numa_bind=None) + with numa_utils.configure_subprocess(vllm_config, local_rank=0): + assert numa_utils._NUMACTL_ARGS_ENV not in os.environ + assert vllm_config.parallel_config.numa_bind is False + + +def test_explicit_numa_nodes_bypass_auto_eligibility(monkeypatch): + _patch_auto_numa(monkeypatch, "spawn", _raise_probe_error) + monkeypatch.setattr(numa_utils, "_is_auto_numa_available", lambda **kwargs: False) + monkeypatch.setattr(numa_utils.subprocess, "run", _fake_numactl_run([])) + vllm_config = _make_config(numa_bind=True, numa_bind_nodes=[7]) + with numa_utils.configure_subprocess(vllm_config, local_rank=0): + assert os.environ[numa_utils._NUMACTL_ARGS_ENV] == "--cpunodebind=7 --membind=7" + + +def test_explicit_numa_requires_detectable_topology(monkeypatch): + _patch_auto_numa(monkeypatch, "spawn", lambda: None) + with ( + pytest.raises(RuntimeError, match="could not detect"), + numa_utils.configure_subprocess(_make_config(numa_bind=True), local_rank=0), + ): + pass + + +@pytest.mark.parametrize("numa_bind", [False, True]) +def test_explicit_numa_handles_missing_numactl(monkeypatch, numa_bind): + monkeypatch.setattr("shutil.which", lambda name: None) + vllm_config = _make_config(numa_bind=numa_bind, numa_bind_nodes=[0]) + if numa_bind: + with ( + pytest.raises(RuntimeError, match="numactl is required"), + numa_utils.configure_subprocess(vllm_config, local_rank=0), + ): + pass + else: + with numa_utils.configure_subprocess(vllm_config, local_rank=0): + assert numa_utils._NUMACTL_ARGS_ENV not in os.environ + + +@pytest.mark.parametrize("affinity", [[0, 1], None, [3, 2, 1, 0]]) +def test_auto_numa_checks_cpu_affinity(monkeypatch, affinity): + monkeypatch.setattr("vllm.platforms.current_platform.is_cuda_alike", lambda: True) + monkeypatch.setattr(numa_utils.os.path, "isdir", lambda path: True) + monkeypatch.setattr(numa_utils.psutil, "cpu_count", lambda: 4) + process = SimpleNamespace() + if affinity is not None: + process.cpu_affinity = lambda: affinity + monkeypatch.setattr(numa_utils.psutil, "Process", lambda pid: process) + monkeypatch.setattr(numa_utils, "_get_mempolicy", lambda: 0) + monkeypatch.setattr( + "vllm.platforms.current_platform.get_all_device_numa_nodes", + lambda: [0, 1], + raising=False, + ) + assert numa_utils._is_auto_numa_available( + verbose=False, require_unbound_process=True + ) is (affinity == [3, 2, 1, 0]) diff --git a/tests/utils_/test_system_utils.py b/tests/utils_/test_system_utils.py index 592af714a5ba..25d4ef4b5d1d 100644 --- a/tests/utils_/test_system_utils.py +++ b/tests/utils_/test_system_utils.py @@ -11,6 +11,7 @@ from pathlib import Path import psutil +import pytest from vllm.utils import system_utils from vllm.utils.system_utils import ( @@ -32,11 +33,39 @@ def test_unique_filepath(): assert len(list(Path(temp_dir).glob("*.txt"))) == 10 -def test_numa_bind_forces_spawn(monkeypatch): - monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD", raising=False) - monkeypatch.setattr("sys.argv", ["vllm", "serve", "--numa-bind"]) +@pytest.mark.parametrize("start_method", [None, "fork"]) +@pytest.mark.parametrize( + "numa_args, force_spawn", + [ + (["--numa-bind"], True), + (["--numa-bind-nodes", "0"], True), + (["--numa-bind-nodes=0"], True), + (["--numa-bind-cpus", "0-3"], True), + (["--numa-bind-cpus=0-3"], True), + (["--numa_bind"], True), + (["--numa_bind_nodes", "0"], True), + (["--numa_bind_cpus=0-3"], True), + ([], False), + (["--no-numa-bind"], False), + ], +) +def test_numa_bind_cli_start_method(monkeypatch, start_method, numa_args, force_spawn): + if start_method is None: + monkeypatch.delenv("VLLM_WORKER_MULTIPROC_METHOD", raising=False) + else: + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", start_method) + monkeypatch.setattr(sys, "argv", ["vllm", "serve", *numa_args]) + for check in ( + "is_in_ray_actor", + "cuda_is_initialized", + "xpu_is_initialized", + "in_wsl", + ): + monkeypatch.setattr(system_utils, check, lambda: False) + _maybe_force_spawn() - assert os.environ["VLLM_WORKER_MULTIPROC_METHOD"] == "spawn" + expected_method = "spawn" if force_spawn else start_method + assert os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == expected_method def test_suppress_stdout_when_sys_stdout_has_no_fd(capfd): diff --git a/tests/v1/distributed/test_dense_dp_world_size.py b/tests/v1/distributed/test_dense_dp_world_size.py index 1fc37936463a..dd17b0234a23 100644 --- a/tests/v1/distributed/test_dense_dp_world_size.py +++ b/tests/v1/distributed/test_dense_dp_world_size.py @@ -1,17 +1,80 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import os +from typing import TYPE_CHECKING +import pytest import requests -from tests.utils import RemoteOpenAIServer +from tests.utils import VLLM_PATH, RemoteOpenAIServer from vllm.platforms import current_platform +if TYPE_CHECKING: + import torch + + from vllm.config import ParallelConfig + MODEL_NAME = "Qwen/Qwen3-0.6B" DP_SIZE = int(os.getenv("DP_SIZE", "2")) TP_SIZE = int(os.getenv("TP_SIZE", "1")) +class NumaBindingWorkerExtension: + parallel_config: "ParallelConfig" + local_rank: int + device: "torch.device" + + def check_numa_binding(self): + import subprocess + from pathlib import Path + + import psutil + import torch.distributed as dist + + from vllm.distributed import get_pcp_group, get_world_group + from vllm.utils.cpu_resource_utils import parse_id_list + + parallel = self.parallel_config + assert parallel.numa_bind and parallel._numa_bind_auto + assert parallel.numa_bind_nodes is not None + gpu_index = self.local_rank + node = current_platform.get_device_numa_node(self.device.index) + assert parallel.numa_bind_nodes[gpu_index] == node + node_cpus = parse_id_list( + Path(f"/sys/devices/system/node/node{node}/cpulist").read_text().strip() + ) + cpus = set(psutil.Process().cpu_affinity()) + assert cpus and cpus <= set(node_cpus) + result = subprocess.run( + ["numactl", "--show"], + check=True, + capture_output=True, + text=True, + timeout=10, + ) + policy = dict(line.split(":", 1) for line in result.stdout.splitlines()) + assert policy["policy"].strip() == "preferred" + assert policy["preferred node"].strip() == str(node) + parent = psutil.Process(os.getppid()) + parent_cpus = set(parent.cpu_affinity()) + assert cpus <= parent_cpus + state = { + "pid": os.getpid(), + "gpu_index": gpu_index, + "gpu_uuid": current_platform.get_device_uuid(self.device.index), + "dp_rank": parallel.data_parallel_rank, + "pcp_rank": get_pcp_group().rank_in_group, + "node": node, + "cpus": sorted(cpus), + "engine_core_pid": parent.pid, + "engine_core_cpus": sorted(parent_cpus), + } + world = get_world_group() + states = [None] * world.world_size + dist.all_gather_object(states, state, group=world.cpu_group) + return states + + def test_dense_dp_world_size(): server_args = [ "--dtype", @@ -50,3 +113,79 @@ def test_dense_dp_world_size(): ) response.raise_for_status() assert response.json() == {"world_size": TP_SIZE} + + +@pytest.mark.slow_test +def test_pcp_dp_with_auto_numa(): + """Automatic binding must cover both PCP workers in both local DP shards.""" + from shutil import which + + from vllm.utils import numa_utils + + if not current_platform.is_cuda() or current_platform.device_count() < 4: + pytest.skip("PCP+DP NUMA smoke test requires four NVIDIA GPUs") + if which("numactl") is None or not numa_utils._is_auto_numa_available( + verbose=False, require_unbound_process=True + ): + pytest.skip("Host cannot engage automatic NUMA binding") + + server_args = [ + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "8", + "--max-num-batched-tokens", + "2048", + "--kv-cache-memory-bytes", + str(256 << 20), + "--enforce-eager", + "--distributed-executor-backend", + "mp", + "--prefill-context-parallel-size", + "2", + "--data-parallel-size", + "2", + "--data-parallel-size-local", + "2", + "--worker-extension-cls", + "tests.v1.distributed.test_dense_dp_world_size.NumaBindingWorkerExtension", + ] + with RemoteOpenAIServer( + "deepseek-ai/DeepSeek-V2-Lite-Chat", + server_args, + env_dict={ + "VLLM_SERVER_DEV_MODE": "1", + "VLLM_WORKER_MULTIPROC_METHOD": "spawn", + "PYTHONPATH": os.pathsep.join( + filter(None, [str(VLLM_PATH), os.environ.get("PYTHONPATH")]) + ), + }, + max_wait_seconds=1200, + ) as server: + completion = server.get_client().completions.create( + model="deepseek-ai/DeepSeek-V2-Lite-Chat", + prompt="The first five positive integers are", + temperature=0, + max_tokens=16, + ) + assert completion.choices[0].text.strip() + assert completion.usage.completion_tokens > 0 + print(f"PCP+DP completion: {completion.choices[0].text!r}") + response = requests.post( + server.url_for("collective_rpc"), + json={"method": "check_numa_binding", "timeout": 30}, + timeout=60, + ) + response.raise_for_status() + states = response.json()["results"][0] + print(f"PCP+DP NUMA worker states: {states}") + assert sorted(s["gpu_index"] for s in states) == [0, 1, 2, 3] + assert len({s["gpu_uuid"] for s in states}) == 4 + assert sorted((s["dp_rank"], s["pcp_rank"]) for s in states) == [ + (0, 0), + (0, 1), + (1, 0), + (1, 1), + ] diff --git a/tests/v1/simple_kv_offload/test_integration.py b/tests/v1/simple_kv_offload/test_integration.py index c9119b7bff23..d0916a6e4406 100644 --- a/tests/v1/simple_kv_offload/test_integration.py +++ b/tests/v1/simple_kv_offload/test_integration.py @@ -10,8 +10,11 @@ from tests.models.utils import check_logprobs_close from vllm import LLM, SamplingParams, TokensPrompt, envs -from vllm.config import KVTransferConfig +from vllm.config import KVTransferConfig, ParallelConfig from vllm.platforms import current_platform +from vllm.utils import numa_utils +from vllm.v1.metrics.reader import Counter +from vllm.v1.simple_kv_offload.metrics import MetricName if not current_platform.is_cuda_alike(): pytest.skip("Requires CUDA or ROCm", allow_module_level=True) @@ -113,7 +116,7 @@ def reset_gpu_cache(): assert mixed_hit.num_cached_tokens == 16 -def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM: +def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int, **llm_kwargs) -> LLM: kv_transfer_config = KVTransferConfig( kv_connector="SimpleCPUOffloadConnector", kv_role="kv_both", @@ -122,12 +125,18 @@ def _make_llm(model: str, lazy: bool, cpu_bytes_to_use: int) -> LLM: "lazy_offload": lazy, }, ) + llm_kwargs = ( + dict( + kv_cache_memory_bytes=40 << 30, # 40 GiB + disable_hybrid_kv_cache_manager=False, + enable_prefix_caching=True, + ) + | llm_kwargs + ) return LLM( model=model, - kv_cache_memory_bytes=40 << 30, # 40 GiB - disable_hybrid_kv_cache_manager=False, - enable_prefix_caching=True, kv_transfer_config=kv_transfer_config, + **llm_kwargs, ) @@ -231,6 +240,90 @@ def _latency_test(llm: LLM, lazy: bool = False): ) +def _assert_worker_auto_numa(worker): + import subprocess + from pathlib import Path + + import psutil + + from vllm.utils.cpu_resource_utils import parse_id_list + + parallel_config = worker.parallel_config + assert parallel_config.numa_bind and parallel_config._numa_bind_auto + gpu_index = parallel_config.get_local_gpu_index(worker.local_rank) + node = parallel_config.numa_bind_nodes[gpu_index] + node_cpus = parse_id_list( + Path(f"/sys/devices/system/node/node{node}/cpulist").read_text().strip() + ) + cpus = set(psutil.Process().cpu_affinity()) + assert cpus and cpus <= set(node_cpus) + result = subprocess.run( + ["numactl", "--show"], check=True, capture_output=True, text=True, timeout=10 + ) + policy = dict(line.split(":", 1) for line in result.stdout.splitlines()) + assert policy["policy"].strip() == "preferred" + assert policy["preferred node"].strip() == str(node) + return {"node": node, "cpus": sorted(cpus), "policy": "preferred"} + + +@pytest.mark.slow_test +def test_simple_cpu_offload_with_auto_numa(monkeypatch, enable_pickle): + """Soft-bound NVIDIA workers must store and restore KV, not silently recompute.""" + if not current_platform.is_cuda(): + pytest.skip("Automatic NUMA KV-offload smoke test requires NVIDIA") + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") + preflight = ParallelConfig(distributed_executor_backend="mp") + preflight._numa_bind_auto = True + if not numa_utils._can_auto_bind(preflight, local_rank=0) or not ( + numa_utils._probe_numactl_args( + numa_utils._get_numactl_worker_args(preflight, local_rank=0) + ) + ): + pytest.skip("Host cannot engage automatic NUMA memory preference") + + model = "Qwen/Qwen3-0.6B" + llm = _make_llm( + model, + False, + 256 << 20, + kv_cache_memory_bytes=128 << 20, + max_model_len=1024, + distributed_executor_backend="mp", + enforce_eager=True, + disable_log_stats=False, + ) + + def loaded_blocks(llm: LLM): + return sum( + metric.value + for metric in llm.get_metrics() + if isinstance(metric, Counter) + and metric.name == MetricName.LOAD_BLOCKS.removesuffix("_total") + and metric.labels.get("model_name") == model + ) + + try: + states = llm.collective_rpc(_assert_worker_auto_numa, timeout=30) + print(f"NUMA worker state: {states}") + sampling_params = SamplingParams(max_tokens=8, temperature=0) + prompt = "hi " * 256 + "Let's count to ten. One, two, three, " + cold = llm.generate(prompt, sampling_params, use_tqdm=False)[0] + llm.generate(prompt, sampling_params, use_tqdm=False) + deadline = time.monotonic() + 30 + while not llm.reset_prefix_cache(reset_connector=False): + assert time.monotonic() < deadline, "CPU offload did not drain" + llm.generate("0", sampling_params, use_tqdm=False) + loaded_before = loaded_blocks(llm) + restored = llm.generate(prompt, sampling_params, use_tqdm=False)[0] + loaded = loaded_blocks(llm) - loaded_before + assert loaded > 0 + assert restored.outputs[0].token_ids == cold.outputs[0].token_ids + print(f"KV offload blocks restored={loaded}; tokens match") + finally: + del llm + + @pytest.mark.optional @pytest.mark.slow_test @pytest.mark.parametrize("model", SMALL_MODELS) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index f51a35c5c4a9..ca3daa163139 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -326,16 +326,32 @@ class is dynamically inherited by the worker class. This is used to inject nnodes: int = Field(default=1, ge=1) """num of nodes for multi-node distributed inference when distributed_executor_backend is mp.""" - numa_bind: bool = False + numa_bind: bool | None = None """Enable NUMA binding for GPU worker subprocesses. - By default, workers are pinned to their GPU's NUMA-local CPUs and - memory; on PCT-capable Xeons they also auto-bind to the SKU's - PCT priority cores. + Workers are pinned to their GPU's NUMA-local CPUs; on PCT-capable + Xeons they also auto-bind to the SKU's PCT priority cores. + + If None (default), vLLM binds only when the environment supports it + (multiple NUMA nodes, `numactl` installed, unrestricted CPU affinity, + a default NUMA memory policy, and the `spawn` start method with an + mp/uni executor and mp DP, including PCP). + Ray, external launchers and device-ID remapping are excluded + from automatic binding. Automatic mode prefers GPU-local memory but + permits allocations on other nodes. Unsupported environments start + unbound. With `True`, strict memory binding may override an existing + policy, and a missing `numactl` or undetectable GPU-to-NUMA topology is + a startup error; `False` disables binding. """ + _numa_bind_auto: bool = Field(default=False, repr=False) + """Internal launch decision: use soft memory preference in spawned children.""" numa_bind_nodes: list[int] | None = None """NUMA node to bind each GPU worker to. + Setting this list implies explicit `numa_bind=True` and requests strict + `--membind`, not automatic soft memory preference. It cannot be combined + with `numa_bind=False`. + Specify one NUMA node per visible GPU, for example `[0, 0, 1, 1]` for a 4-GPU system with GPUs 0-1 on NUMA node 0 and GPUs 2-3 on NUMA node 1. If unset and `numa_bind=True`, vLLM auto-detects the @@ -345,6 +361,10 @@ class is dynamically inherited by the worker class. This is used to inject numa_bind_cpus: list[str] | None = None """Optional CPU lists to bind each GPU worker to. + Setting this list implies explicit `numa_bind=True` and requests strict + `--membind`, not CPU-only binding or automatic soft memory preference. + It cannot be combined with `numa_bind=False`. + Specify one CPU list per visible GPU, for example `["0-3", "4-7", "8-11", "12-15"]`. When set, vLLM uses `numactl --physcpubind` instead of `--cpunodebind`. This is useful @@ -553,12 +573,12 @@ def _validate_parallel_config(self) -> Self: "data_parallel_external_lb can only be set when data_parallel_size > 1" ) - if not self.numa_bind and ( - self.numa_bind_nodes is not None or self.numa_bind_cpus is not None - ): - raise ValueError( - "numa_bind_nodes and numa_bind_cpus require numa_bind=True." - ) + if self.numa_bind_nodes is not None or self.numa_bind_cpus is not None: + if self.numa_bind is False: + raise ValueError( + "NUMA binding lists cannot be combined with numa_bind=False." + ) + self.numa_bind = True if self.enable_eplb: if not current_platform.is_cuda_alike() and not current_platform.is_xpu(): @@ -949,6 +969,7 @@ def compute_hash(self): # NUMA node in numa_bind_nodes (see vllm/utils/numa_utils.py # `_get_numa_node`), which would otherwise diverge the DP hash. "numa_bind", + "_numa_bind_auto", "numa_bind_nodes", "numa_bind_cpus", "assigned_physical_gpu_ids", diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 423544a2a135..316881e5d61b 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -541,7 +541,7 @@ class EngineArgs: cpu_distributed_timeout_seconds: int | None = ( ParallelConfig.cpu_distributed_timeout_seconds ) - numa_bind: bool = ParallelConfig.numa_bind + numa_bind: bool | None = ParallelConfig.numa_bind numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus device_ids: list[int | str] | None = None diff --git a/vllm/utils/numa_utils.py b/vllm/utils/numa_utils.py index 755d46e711fb..fefa30ea9ca7 100644 --- a/vllm/utils/numa_utils.py +++ b/vllm/utils/numa_utils.py @@ -26,6 +26,7 @@ logger = logging.getLogger(__name__) _NUMACTL_ARGS_ENV = "_VLLM_INTERNAL_NUMACTL_ARGS" _NUMACTL_PYTHON_EXECUTABLE_ENV = "_VLLM_INTERNAL_NUMACTL_PYTHON_EXECUTABLE" +_MPOL_DEFAULT = 0 @cache @@ -41,55 +42,84 @@ def get_libnuma(): return libnuma -def _can_set_mempolicy() -> bool: - """Check whether the current process can use NUMA memory policy syscalls.""" +def _get_mempolicy() -> int | None: + """Return the current memory-policy mode, or None if it cannot be read.""" try: libnuma = get_libnuma() if libnuma is None or libnuma.numa_available() < 0: - return False + return None mode = ctypes.c_int() ret = libnuma.get_mempolicy( ctypes.byref(mode), None, ctypes.c_ulong(0), None, ctypes.c_ulong(0) ) - return ret == 0 + return mode.value if ret == 0 else None except Exception: - return False + return None -def _is_auto_numa_available() -> bool: - """Check whether automatic GPU-to-NUMA detection should be attempted.""" +def _is_auto_numa_available( + verbose: bool = True, *, require_unbound_process: bool = False +) -> bool: + """Check whether automatic GPU-to-NUMA detection should be attempted. + + Args: + verbose: Log why detection is skipped at warning level; otherwise at + debug level, for the default mode where skipping is expected. + require_unbound_process: Require known unrestricted CPU affinity and + the default memory policy for default automatic binding. Explicit + topology discovery keeps its existing best-effort checks. + + """ from vllm.platforms import current_platform + log = logger.warning if verbose else logger.debug + if not current_platform.is_cuda_alike(): + log("Automatic NUMA binding requires a CUDA or ROCm platform.") return False if not os.path.isdir("/sys/devices/system/node/node1"): + log("Automatic NUMA binding requires multiple NUMA nodes.") return False try: process = psutil.Process(os.getpid()) cpu_affinity = process.cpu_affinity() cpu_count = psutil.cpu_count() - if cpu_count is not None and cpu_affinity != list(range(cpu_count)): - logger.warning( + if cpu_count is not None and set(cpu_affinity) != set(range(cpu_count)): + log( "CPU affinity is already constrained for this process. " "Skipping automatic NUMA binding; pass --numa-bind-nodes " "explicitly to override." ) return False + if cpu_count is None and require_unbound_process: + log("Could not determine CPU availability for automatic NUMA binding.") + return False except (AttributeError, NotImplementedError, psutil.Error): - pass + if require_unbound_process: + log("Could not check CPU affinity for automatic NUMA binding.") + return False - if not _can_set_mempolicy(): - logger.warning( - "User lacks permission to set NUMA memory policy. " + mempolicy = _get_mempolicy() + if mempolicy is None: + log( + "Could not read NUMA memory policy. " "Automatic NUMA detection may not work; if you are using Docker, " "try adding --cap-add SYS_NICE." ) return False + if require_unbound_process and mempolicy != _MPOL_DEFAULT: + log( + "NUMA memory policy is already set for this process (mode=%d). " + "Skipping automatic NUMA binding; pass --numa-bind explicitly " + "to override.", + mempolicy, + ) + return False if not hasattr(current_platform, "get_all_device_numa_nodes"): - logger.warning( + log( "Platform %s does not support automatic NUMA detection", type(current_platform).__name__, ) @@ -98,7 +128,6 @@ def _is_auto_numa_available() -> bool: return True -@cache def get_auto_numa_nodes() -> list[int] | None: """Auto-detect NUMA nodes for all visible GPUs.""" from vllm.platforms import current_platform @@ -112,6 +141,76 @@ def get_auto_numa_nodes() -> list[int] | None: return numa_nodes +def _can_auto_bind( + parallel_config, local_rank: int, dp_local_rank: int | None = None +) -> bool: + """Validate and store this launcher's automatic GPU-to-NUMA topology. + + Unlike ``--numa-bind``, this never fails startup or forces the ``spawn`` + start method: any unmet requirement means starting unbound. + """ + from shutil import which + + if envs.VLLM_WORKER_MULTIPROC_METHOD != "spawn": + logger.debug("Automatic NUMA binding requires the spawn start method.") + return False + if ( + parallel_config.distributed_executor_backend not in ("mp", "uni") + or parallel_config.data_parallel_backend != "mp" + ): + logger.debug("Automatic NUMA binding requires local mp/uni launchers.") + return False + # Topology is indexed in visible-device order, which --device-ids remaps. + if parallel_config.assigned_physical_gpu_ids is not None: + logger.debug("Automatic NUMA binding does not support --device-ids remapping.") + return False + if which("numactl") is None: + logger.debug("Automatic NUMA binding requires numactl on PATH.") + return False + try: + from vllm.platforms import current_platform + from vllm.utils.cpu_resource_utils import parse_id_list + + if not _is_auto_numa_available(verbose=False, require_unbound_process=True): + return False + numa_nodes = get_auto_numa_nodes() + if not numa_nodes: + logger.debug("No GPU-to-NUMA topology available for automatic binding.") + return False + if len(numa_nodes) != current_platform.device_count() or any( + not isinstance(node, int) or node < 0 for node in numa_nodes + ): + logger.debug("Incomplete or invalid GPU-to-NUMA topology: %s", numa_nodes) + return False + + last_local_rank = max(local_rank, parallel_config.local_world_size - 1) + last_dp_local_rank = (parallel_config.data_parallel_rank_local or 0) + max( + parallel_config.data_parallel_size_local - 1, 0 + ) + last_gpu_index = max( + parallel_config.get_local_gpu_index(last_local_rank, dp_local_rank), + parallel_config.get_local_gpu_index(last_local_rank, last_dp_local_rank), + ) + if last_gpu_index >= len(numa_nodes): + logger.debug( + "GPU-to-NUMA topology does not cover all local workers: " + "GPU index %d, mapping size %d.", + last_gpu_index, + len(numa_nodes), + ) + return False + for node in set(numa_nodes): + cpulist_path = Path(f"/sys/devices/system/node/node{node}/cpulist") + if not parse_id_list(cpulist_path.read_text().strip()): + logger.debug("NUMA node %d has no CPUs for automatic binding.", node) + return False + parallel_config.numa_bind_nodes = list(numa_nodes) + return True + except Exception: + logger.debug("Automatic GPU-to-NUMA topology validation failed", exc_info=True) + return False + + # PCT (Priority Core Turbo) auto-detection workaround for Granite Rapids # Xeon SKUs. # @@ -326,24 +425,13 @@ def _get_numactl_worker_args( gpu_index = parallel_config.get_local_gpu_index(local_rank, dp_local_rank) numa_node = _get_numa_node(parallel_config, gpu_index) cpu_binding = _get_cpu_binding(parallel_config, gpu_index, [numa_node]) + memory_policy = "preferred" if parallel_config._numa_bind_auto else "membind" + memory_arg = f"--{memory_policy}={numa_node}" if cpu_binding is not None: - logger.info( - "Binding worker subprocess (local_rank=%s, gpu_index=%s) to CPUs %s and NUMA node %s", # noqa: E501 - local_rank, - gpu_index, - cpu_binding, - numa_node, - ) - return f"--physcpubind={cpu_binding} --membind={numa_node}" + return f"--physcpubind={cpu_binding} {memory_arg}" - logger.info( - "Binding worker subprocess (local_rank=%s, gpu_index=%s) to NUMA node %s", - local_rank, - gpu_index, - numa_node, - ) - return f"--cpunodebind={numa_node} --membind={numa_node}" + return f"--cpunodebind={numa_node} {memory_arg}" def _get_enginecore_numa_nodes( @@ -387,7 +475,11 @@ def _get_numactl_enginecore_args( is also a safe superset by construction). """ shard_nodes = _get_enginecore_numa_nodes(parallel_config, dp_local_rank) - membind_arg = ",".join(str(n) for n in shard_nodes) + node_binding = ",".join(str(n) for n in shard_nodes) + memory_policy = "membind" + if parallel_config._numa_bind_auto: + memory_policy = "preferred" if len(shard_nodes) == 1 else "preferred-many" + memory_arg = f"--{memory_policy}={node_binding}" pct_cpus = ( None @@ -397,21 +489,9 @@ def _get_numactl_enginecore_args( if pct_cpus is not None: cpu_binding = ",".join(str(c) for c in pct_cpus) - logger.info( - "Binding EngineCore subprocess (local_rank=%s) to CPUs %s " - "and NUMA nodes %s", - local_rank, - cpu_binding, - membind_arg, - ) - return f"--physcpubind={cpu_binding} --membind={membind_arg}" + return f"--physcpubind={cpu_binding} {memory_arg}" - logger.info( - "Binding EngineCore subprocess (local_rank=%s) to NUMA nodes %s", - local_rank, - membind_arg, - ) - return f"--cpunodebind={membind_arg} --membind={membind_arg}" + return f"--cpunodebind={node_binding} {memory_arg}" def _log_numactl_show(label: str) -> bool: @@ -455,20 +535,22 @@ def _probe_numactl_args(numactl_args: str) -> bool: def _resolve_numactl_args(numactl_args: str) -> str: - """Drop ``--membind`` if the container rejects it, keeping CPU binding.""" + """Drop a rejected memory policy, keeping CPU binding when possible.""" cpu_only = " ".join( - t for t in numactl_args.split() if not t.startswith("--membind=") + t + for t in numactl_args.split() + if not t.startswith(("--membind=", "--preferred=", "--preferred-many=")) ) for candidate in (numactl_args, cpu_only, ""): if _probe_numactl_args(candidate): if candidate != numactl_args: logger.warning( - "numactl args %r rejected; falling back to %r. Add " - "--cap-add SYS_NICE for full NUMA binding.", + "numactl args %r rejected; falling back to %r.", numactl_args, candidate or "no binding", ) return candidate + logger.warning("numactl probes failed; starting without NUMA binding.") return "" @@ -481,6 +563,19 @@ def configure_subprocess( ): """Temporarily replace the multiprocessing executable with a numactl wrapper.""" parallel_config = vllm_config.parallel_config + if parallel_config.numa_bind is None: + # Resolve once in the launching process; spawned children inherit the + # result through the pickled config instead of re-checking (a bound + # EngineCore would fail the unconstrained-affinity check). + parallel_config.numa_bind = _can_auto_bind( + parallel_config, local_rank, dp_local_rank + ) + parallel_config._numa_bind_auto = parallel_config.numa_bind + if not parallel_config.numa_bind: + logger.info( + "Automatic NUMA binding unavailable on this host; starting " + "unbound. Enable debug logging for the skip reason." + ) if not parallel_config.numa_bind: yield return @@ -502,8 +597,20 @@ def configure_subprocess( numactl_args = _resolve_numactl_args(numactl_args) if not numactl_args: # No NUMA binding possible here; launch without the wrapper. + if parallel_config._numa_bind_auto: + parallel_config.numa_bind = False + parallel_config.numa_bind_nodes = None + parallel_config._numa_bind_auto = False yield return + if envs.VLLM_WORKER_MULTIPROC_METHOD == "spawn": + logger.info( + "Launching %s subprocess (local_rank=%s, dp_local_rank=%s) with numactl %s", + process_kind, + local_rank, + dp_local_rank, + numactl_args, + ) python_executable = os.fsdecode(multiprocessing.spawn.get_executable()) with ( _set_numa_wrapper_env(numactl_args, python_executable), diff --git a/vllm/utils/system_utils.py b/vllm/utils/system_utils.py index 5ae13232b069..cbea3c862d25 100644 --- a/vllm/utils/system_utils.py +++ b/vllm/utils/system_utils.py @@ -137,9 +137,12 @@ def _maybe_force_spawn(): os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address reasons.append("In a Ray actor and can only be spawned") - # Force spawn if NUMA binding is enabled via --numa-bind. - # NUMA binding uses executable hijacking which requires spawn - if "--numa-bind" in sys.argv: + # NUMA binding uses executable hijacking, which requires spawn. + if any( + arg.split("=", 1)[0].replace("_", "-") + in ("--numa-bind", "--numa-bind-nodes", "--numa-bind-cpus") + for arg in sys.argv + ): reasons.append("NUMA binding requires spawn method") if cuda_is_initialized(): From 0855defde8ea382781ab4211c4f67d17f6b96f50 Mon Sep 17 00:00:00 2001 From: Jie Hao Date: Tue, 27 Oct 2026 18:07:02 -0400 Subject: [PATCH 009/406] Revert "[NIXL] Bump NIXL version to 1.5.1" (#63817) (#64085) Signed-off-by: Jie Hao --- requirements/kv_connectors.txt | 2 +- requirements/kv_connectors_cu12.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index baddf1cd633c..aff67fce2ea1 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -2,7 +2,7 @@ lmcache >= 0.5.6 # CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime # images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded. cupy-cuda13x != 14.1.0 -nixl == 1.5.1 +nixl == 1.5.0 # CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for # the mooncake-transfer-engine-cuda13 variant of the same version. mooncake-transfer-engine >= 0.3.12 diff --git a/requirements/kv_connectors_cu12.txt b/requirements/kv_connectors_cu12.txt index 60bb194b32ec..4f5c26fc3bbc 100644 --- a/requirements/kv_connectors_cu12.txt +++ b/requirements/kv_connectors_cu12.txt @@ -3,5 +3,5 @@ lmcache == 0.5.6 cuda-python >= 12, < 13 cupy-cuda12x != 14.1.0 -nixl == 1.5.1 +nixl == 1.5.0 mooncake-transfer-engine >= 0.3.12 From c56706e88e0a699274e7ed807e072274048d1c33 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:21:39 -0400 Subject: [PATCH 010/406] [Perf] Use TMA for MOE up/down proj, 3.1~3.4% E2E Throughput Improvement for GLM5.3 (#63820) Signed-off-by: yewentao256 Co-authored-by: vllm-agent --- tests/kernels/moe/test_block_fp8.py | 55 +++++++++++++++++++ tests/kernels/moe/test_moe.py | 2 +- .../layers/fused_moe/fused_moe.py | 41 +++++++++++++- 3 files changed, 94 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index 77b638a76427..0a4fd90f093a 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.fused_moe import ( fused_experts, fused_topk, + override_config, ) from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -32,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, ) +from vllm.model_executor.layers.fused_moe.fused_moe import ( + try_get_optimal_moe_config, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) @@ -260,6 +264,57 @@ def test_w8a8_block_fp8_fused_moe( torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol) +@pytest.mark.parametrize("M", [24, 1024]) +@pytest.mark.parametrize("modular", [False, True]) +@torch.inference_mode() +def test_h200_block_fp8_up_and_down_tma_matches_pointer(M, modular, workspace_init): + if "H200" not in torch.cuda.get_device_name(): + pytest.skip("This block FP8 TMA test targets H200") + + E, N, K, topk = 256, 256, 128, 8 + torch.manual_seed(123) + a = torch.randn((M, K), dtype=torch.bfloat16) / 10 + score = torch.randn((M, E), dtype=torch.bfloat16) + w1, w2, quant_config = make_test_quant_config( + E, N, K, torch.bfloat16, torch.float8_e4m3fn, block_shape=[128, 128] + ) + topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False) + + config_args = (w1.shape, w2.shape, topk, "fp8_w8a8", M) + config = try_get_optimal_moe_config(*config_args, block_shape=[128, 128]) + assert w2.is_contiguous() + + if modular: + moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config) + + def run(): + return moe.apply( + a, + w1, + w2, + topk_weights, + topk_ids, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + expert_map=None, + global_num_experts=E, + ) + + else: + + def run(): + return fused_experts( + a, w1, w2, topk_weights, topk_ids, quant_config=quant_config + ) + + with set_current_vllm_config(vllm_config): + with override_config({**config, "USE_TMA": False}): + pointer_output = run() + tma_output = run() + + torch.testing.assert_close(tma_output, pointer_output, atol=0, rtol=0) + + @pytest.mark.parametrize("clamp_limit", [None, 10.0]) @torch.inference_mode() def test_w8a8_block_fp8_clamp_skips_fused_silu_quant( diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index d279c15c2c79..26e9f87f1014 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -75,7 +75,7 @@ def test_triton_moe_launcher_passes_scalar_scale_as_pointer(monkeypatch) -> None class FakeKernel: def __getitem__(self, grid): def launch(*args, **kwargs) -> None: - captured["a_scale"] = args[4] + captured["a_scale"] = args[5] return launch diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index e06288a3e92b..7e9f0e28f6ac 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -299,6 +299,7 @@ def fused_moe_kernel( # Pointers to matrices a_ptr, b_ptr, + b_desc, c_ptr, b_bias_ptr, a_scale_ptr, @@ -351,6 +352,7 @@ def fused_moe_kernel( SWAP_AB: tl.constexpr, # Tensor-descriptor path for the A gather and B load in the K-loop. USE_TD: tl.constexpr = False, + USE_TMA_B: tl.constexpr = False, ): """Implements the fused computation for a Mixture of Experts (MOE) using token and expert matrices. @@ -441,6 +443,7 @@ def fused_moe_kernel( offs_k = tl.arange(0, BLOCK_SIZE_K) # TD gather and the SWAP_AB accumulator layout are mutually exclusive. tl.static_assert(not (USE_TD and SWAP_AB)) + tl.static_assert(not (USE_TD and USE_TMA_B)) if USE_TD: # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx. m_td = num_valid_tokens // top_k @@ -523,16 +526,26 @@ def fused_moe_kernel( b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T elif SWAP_AB: a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :] - b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K a = tl.load(a_ptrs, mask=a_mask, other=0.0) - b = tl.load(b_ptrs, mask=b_mask, other=0.0) + if not USE_TMA_B: + b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K + b = tl.load(b_ptrs, mask=b_mask, other=0.0) else: a = tl.load( a_ptrs, mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), other=0.0, ) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + if not USE_TMA_B: + b = tl.load( + b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0 + ) + if USE_TMA_B: + b = b_desc.load( + [off_experts.to(tl.int32), pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K] + ).reshape((BLOCK_SIZE_N, BLOCK_SIZE_K)) + if not SWAP_AB: + b = b.T # We accumulate along the K dimension. if use_int8_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) @@ -842,6 +855,26 @@ def invoke_fused_moe_triton_kernel( BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K") if block_shape is not None: BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1])) + use_tma_b = ( + bool(config.pop("USE_TMA", True)) + and use_fp8_w8a8 + and block_shape is not None + and all(block_shape) + and current_platform.is_cuda() + and current_platform.is_device_capability(90) + and B.device.type == "cuda" + and B.is_contiguous() + and B.size(1) % config["BLOCK_SIZE_N"] == 0 + and B.size(2) % BLOCK_SIZE_K == 0 + ) + b_desc = None + if use_tma_b: + from triton.tools.tensor_descriptor import TensorDescriptor + + set_triton_allocator(B.device) + b_desc = TensorDescriptor.from_tensor( + B, [1, config["BLOCK_SIZE_N"], BLOCK_SIZE_K] + ) if use_td and A.size(1) % BLOCK_SIZE_K != 0: # TD gather/load feeding tl.dot with a non-block-aligned K # miscompiles (~74% of output elements wrong) on real HW; @@ -864,6 +897,7 @@ def invoke_fused_moe_triton_kernel( fused_moe_kernel[grid]( A, B, + b_desc, C, B_bias, A_scale, @@ -906,6 +940,7 @@ def invoke_fused_moe_triton_kernel( BLOCK_SIZE_K=BLOCK_SIZE_K, SWAP_AB=SWAP_AB, USE_TD=use_td, + USE_TMA_B=use_tma_b, **config, ) From e96b490b1204ff6f8ce856b500b368551fafd43b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20Bernl=C3=B6hr?= Date: Tue, 27 Oct 2026 23:35:07 +0100 Subject: [PATCH 011/406] [Bugfix][Kimi-K3] Reserve KDA prefill workspace before profiling (#63989) Signed-off-by: jbernloehr Co-authored-by: Codex Co-authored-by: Thien Tran Co-authored-by: vllm-agent --- vllm/models/kimi_k3/nvidia/kda.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index a921d1bd2879..d3329771181a 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -73,7 +73,10 @@ ) from vllm.v1.attention.backend import AttentionBackend from vllm.v1.kv_cache_interface import MambaSpec -from vllm.v1.worker.workspace import current_workspace_manager +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) logger = init_logger(__name__) @@ -857,6 +860,18 @@ def __init__( self.model_config.dtype, ) + # Profiling skips attention, so reserve its runtime layout before the + # KV cache is sized. Layers share the manager's grow-only allocation. + if is_workspace_manager_initialized(): + if self._flashkda_buffer_specs is not None: + current_workspace_manager().get_simultaneous( + *self._flashkda_buffer_specs + ) + elif self._flashinfer_kda_output_spec is not None: + current_workspace_manager().get_simultaneous( + self._flashinfer_kda_output_spec + ) + self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") decode_norm_weight = None if decode_conv1d_weight is not None: From 36f5ff3172eb75433e54e496d64c473e951f8ecc Mon Sep 17 00:00:00 2001 From: Moein Khazraee <33970824+mkhazraee@users.noreply.github.com> Date: Tue, 27 Oct 2026 16:41:00 -0700 Subject: [PATCH 012/406] [Perf][KV Offload] Reduce KVCR adapter overhead (#63889) Signed-off-by: Moein Khazraee Co-authored-by: vllm-agent --- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 50 ++++++++++++++----- vllm/v1/kv_offload/tiering/kvcr/manager.py | 32 ++++++------ 2 files changed, 52 insertions(+), 30 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index a6f67f80f3f4..beca574d8881 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -390,6 +390,7 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch, backpressure): assert list(blocks) == [key] (descriptor,) = blocks[key] assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2) + assert type(descriptor.element_index) is int assert list(tier.get_finished_jobs()) == [JobResult(7, True)] # Here we verify that request cleanup discards the request-scoped hint in @@ -479,13 +480,24 @@ def test_kvcr_tier_maps_query_status(monkeypatch, status, expected): 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.""" +@pytest.mark.parametrize( + ("job_keys", "chunk_ids", "success"), + [ + (None, [1, 3], True), + ((b"k2", b"k0"), [3, 1], True), + (None, [1], False), + ((b"k0", b"other"), [1, 3], False), + ], + ids=["original-keys", "reordered-keys", "short-chunks", "unexpected-key"], +) +def test_kvcr_tier_serves_primary_pin_request( + monkeypatch, job_keys, chunk_ids, success +): + """Validate pin jobs and hold accepted hits until KVCR releases the 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]: 3} lifecycle: list[str] = [] class Parent: @@ -499,8 +511,8 @@ def lookup(self, key, req_context): 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]), + keys=requested_keys if job_keys is None else job_keys, + chunk_ids=np.array(chunk_ids, dtype=np.int32), is_promotion=True, req_context=req_context, ) @@ -516,6 +528,11 @@ def on_request_finished(self, req_context): [(queued_request, result)] = bindings.poll_pin_results() assert queued_request == request + assert lifecycle == ["new", "finished"] + if not success: + assert result is None + assert list(tier.get_finished_jobs()) == [JobResult(11, False)] + return assert result is not None pin_handle, descriptors = result assert descriptors[keys[1]] is None @@ -523,7 +540,7 @@ def on_request_finished(self, req_context): assert block_descriptors is not None (descriptor,) = block_descriptors assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=3) - assert lifecycle == ["new", "finished"] + assert type(descriptor.element_index) is int polls = 0 @@ -626,6 +643,7 @@ def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): assert blocks[key] == [ MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2) ] + assert type(blocks[key][0].element_index) is int assert list(tier.get_finished_jobs()) == [JobResult(11, True)] assert kvcr.inventory_sink is not None @@ -655,7 +673,16 @@ def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch): ) -def test_kvcr_tier_accumulates_block_results(monkeypatch): +@pytest.mark.parametrize( + ("successes", "expected"), + [ + ((True, False), JobResult(13, False, {OffloadKey(b"k0")})), + ((False, True), JobResult(13, False, {OffloadKey(b"k1")})), + ((True, True), JobResult(13, True)), + ], + ids=["success-then-failure", "failure-then-success", "all-success"], +) +def test_kvcr_tier_accumulates_block_results(monkeypatch, successes, expected): """Wait for every block result and preserve partial success while draining.""" kvcr = RecordingKVCR() tier = _make_tier(monkeypatch, kvcr) @@ -672,13 +699,10 @@ def test_kvcr_tier_accumulates_block_results(monkeypatch): op_handle, _, _ = kvcr.deliver_calls[0] kvcr.completed.clear() - kvcr.complete(op_handle, {keys[0]: True}) + kvcr.complete(op_handle, {keys[0]: successes[0]}) assert list(tier.get_finished_jobs()) == [] - kvcr.complete(op_handle, {keys[1]: False}) + kvcr.complete(op_handle, {keys[1]: successes[1]}) tier.drain_jobs() - [result] = tier.get_finished_jobs() - assert result.job_id == 13 - assert not result.success - assert result.successful_keys == {keys[0]} + assert list(tier.get_finished_jobs()) == [expected] diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 1f3d9638b493..e24531120b5d 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -201,20 +201,17 @@ def _resolve_pin( 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): + if len(job_keys) != len(job.chunk_ids) or ( + job_keys is not hit_keys and set(job_keys) != set(hit_keys) + ): return None - descriptors: dict[BlockKey, list[MemoryRef] | 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) + descriptors: dict[BlockKey, list[MemoryRef] | None] = dict.fromkeys( + cast(tuple[BlockKey, ...], offload_keys) ) + make_descriptor = self._tier._make_descriptor + for key, chunk_id in zip(job_keys, job.chunk_ids.tolist()): + descriptors[BlockKey(key)] = [make_descriptor(chunk_id)] pin_handle = f"kvcr-job-{job.job_id}" self._pin_jobs[pin_handle] = job.job_id return pin_handle, descriptors @@ -476,6 +473,7 @@ def __init__( if local_mapping is not None: local_mapping.close() raise + self._nixl_agent_name = self._kvcr.config.nixl_agent_name self._local_dram_mmap = local_mapping self._finished_jobs: list[JobResult] = [] self._jobs_by_op: dict[OpHandle, _JobState] = {} @@ -493,9 +491,9 @@ def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: @override def submit_load(self, job_metadata: TransferJob) -> None: blocks = { - self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + self._key_adapter.encode(key): [self._make_descriptor(chunk_id)] for key, chunk_id in zip( - job_metadata.keys, job_metadata.chunk_ids, strict=True + job_metadata.keys, job_metadata.chunk_ids.tolist(), strict=True ) } if not blocks: @@ -517,9 +515,9 @@ def submit_load(self, job_metadata: TransferJob) -> None: @override def submit_store(self, job_metadata: TransferJob) -> None: blocks = { - self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + self._key_adapter.encode(key): [self._make_descriptor(chunk_id)] for key, chunk_id in zip( - job_metadata.keys, job_metadata.chunk_ids, strict=True + job_metadata.keys, job_metadata.chunk_ids.tolist(), strict=True ) } if not blocks: @@ -566,7 +564,7 @@ def _poll_finished_jobs(self) -> list[JobResult]: job, 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: + if successful_keys is not None and (remaining > 0 or not success): successful_keys.update( OffloadKey(bytes(key)) for key, entry in entries.items() @@ -706,6 +704,6 @@ def _record_inventory(self, event: InventoryEvent) -> None: def _make_descriptor(self, chunk_id: int) -> MemoryRef: return MemoryRef( - end_point_name=self._kvcr.config.nixl_agent_name, + end_point_name=self._nixl_agent_name, element_index=chunk_id, ) From 8697bbae4c7d4c15591bb85784835eeab2c8ecb1 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 27 Oct 2026 16:59:59 -0700 Subject: [PATCH 013/406] [Bugfix][Attention] Fix FlashInfer MLA DCP decode on ragged spec-decode batches (#59392) Signed-off-by: zjy0516 Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu Co-authored-by: Pavani Majety --- tests/v1/attention/test_flashinfer_mla_dcp.py | 38 +++++++++++-------- .../attention/backends/mla/flashinfer_mla.py | 26 +++++++++---- 2 files changed, 41 insertions(+), 23 deletions(-) diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py index a72dbc21d145..f14238fb38ec 100644 --- a/tests/v1/attention/test_flashinfer_mla_dcp.py +++ b/tests/v1/attention/test_flashinfer_mla_dcp.py @@ -87,7 +87,8 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch): @requires_flashinfer_mla @pytest.mark.parametrize("causal", [True, False], ids=["causal", "noncausal"]) -def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal): +@pytest.mark.parametrize("query_lens", [(3, 3), (3, 1)], ids=["uniform", "ragged"]) +def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal, query_lens): import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla impl = MagicMock() @@ -106,8 +107,11 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal): flashinfer_mla.FlashInferMLAImpl._flattened_decode_metadata, impl ) - num_reqs, query_len = 2, 3 - num_tokens = num_reqs * query_len + num_reqs = len(query_lens) + num_tokens = sum(query_lens) + query_len = max(query_lens) + split_rows = not causal or len(set(query_lens)) > 1 + lens = torch.tensor(query_lens) block_table = torch.tensor([[10], [20]], dtype=torch.int32) seq_lens = torch.tensor([4, 5], dtype=torch.int32) global_causal_seq_lens = torch.tensor([31, 35], dtype=torch.int32) @@ -120,15 +124,15 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal): attn_metadata.decode.seq_lens = seq_lens attn_metadata.decode.dcp_tot_seq_lens = global_causal_seq_lens attn_metadata.decode.max_query_len = query_len - attn_metadata.decode.query_start_loc = torch.arange( - 0, num_tokens + 1, query_len, dtype=torch.int32 + attn_metadata.decode.query_start_loc = torch.tensor( + [0, *lens.cumsum(0).tolist()], dtype=torch.int32 ) attn_metadata.decode.query_len = 0 query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16) kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16) - kernel_batch = num_reqs if causal else num_tokens - kernel_query_len = query_len if causal else 1 + kernel_batch = num_tokens if split_rows else num_reqs + kernel_query_len = 1 if split_rows else query_len kernel = MagicMock( return_value=( torch.ones( @@ -155,15 +159,17 @@ def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal): call = kernel.call_args.kwargs assert call["query"].shape == (kernel_batch, kernel_query_len, 24, 576) - expected_block_table = ( - block_table if causal else block_table.repeat_interleave(query_len, dim=0) - ) - expected_seq_lens = seq_lens if causal else seq_lens.repeat_interleave(query_len) - expected_global_seq_lens = ( - global_causal_seq_lens - if causal - else global_causal_seq_lens.repeat_interleave(query_len) - ) + expected_block_table = block_table + expected_seq_lens = seq_lens + expected_global_seq_lens = global_causal_seq_lens + if split_rows: + expected_block_table = block_table.repeat_interleave(lens, dim=0) + expected_seq_lens = seq_lens.repeat_interleave(lens) + expected_global_seq_lens = global_causal_seq_lens.repeat_interleave(lens) + if causal and split_rows: + # Each split row sees KV only up to its own position in its request. + rows_after = torch.cat([torch.arange(n - 1, -1, -1) for n in query_lens]) + expected_global_seq_lens = expected_global_seq_lens - rows_after.int() torch.testing.assert_close(call["block_tables"], expected_block_table) torch.testing.assert_close(call["seq_lens"], expected_seq_lens) assert call["backend"] == "cute-dsl" diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index 91492859d5d6..2179ab2a6ded 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -341,6 +341,8 @@ def forward_mqa( cum_seq_lens_q: torch.Tensor | None = None max_q_len: int | None = None row_req: torch.Tensor | None = None + # Per split row: query tokens of its request after it. + causal_rows_after: torch.Tensor | None = None if not attn_metadata.causal: # FlashInfer decode has no causal flag. Flatten each non-causal @@ -357,15 +359,23 @@ def forward_mqa( # do not return on the ragged path (flashinfer #3238). cum_seq_lens_q = attn_metadata.decode.query_start_loc max_q_len = attn_metadata.decode.max_query_len - # trtllm API requires extra dimension q_len_per_request for MTP - elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0: - logger.warning_once( - """FlashInferMLAImpl got a query of uneven length. - This usually indicates an issue in batch reordering - or incorrect setup in dummy_run.""" - ) + elif ( + attn_metadata.num_decode_tokens + != attn_metadata.num_decodes * attn_metadata.decode.max_query_len + ): + # Ragged causal DCP decode, e.g. requests with fewer drafts: split + # into single-token rows, each bounded at its own causal position. q = q.unsqueeze(1) + block_table, seq_lens, row_req = self._flattened_decode_metadata( + attn_metadata, q.shape[0] + ) + query_start_loc = attn_metadata.decode.query_start_loc + assert query_start_loc is not None + query_end = query_start_loc[1:] + rows = torch.arange(q.shape[0], device=q.device, dtype=query_end.dtype) + causal_rows_after = query_end[row_req] - 1 - rows else: + # trtllm API requires extra dimension q_len_per_request for MTP q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1]) if self.bmm1_scale is None: @@ -389,6 +399,8 @@ def forward_mqa( assert causal_seqlens_kv_global is not None if row_req is not None: causal_seqlens_kv_global = causal_seqlens_kv_global[row_req] + if causal_rows_after is not None: + causal_seqlens_kv_global = causal_seqlens_kv_global - causal_rows_after extra_kwargs.update( enable_dcp=True, cp_world=self.dcp_world_size, From 820be160ce0dc9098deaa2a4c6b5fc534c43881d Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:00:23 -0700 Subject: [PATCH 014/406] [Spec Decode] Disable batch-sharded draft logits for PCP (#64151) Signed-off-by: Giancarlo Delfin --- vllm/v1/worker/gpu/spec_decode/speculator.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 261c83a94727..517103949fee 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -277,8 +277,12 @@ def _init_draft_logits(self, target_model: nn.Module) -> None: rows of the requests it owns, at row req_state_idx // draft_logits_num_shards. """ + parallel_config = self.vllm_config.parallel_config if ( batch_sharded_sampling_enabled(self.vllm_config, target_model) + # PCPManager.sync_drafts indexes the draft logits by request state, + # so it must see every request's logits. + and parallel_config.prefill_context_parallel_size == 1 and self.draft_watermarker is None and self.supports_batch_sharded_draft_logits() ): From a91b939bfb3534d5f538ccaa331375c271000c69 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 27 Oct 2026 17:50:19 -0700 Subject: [PATCH 015/406] [Bugfix][Bench] Support Muse balanced routing (#63872) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- .buildkite/test_areas/quantization.yaml | 8 +- tests/benchmarks/test_forward_cli.py | 132 ++++++++++++++++++++++++ vllm/benchmarks/forward/routing.py | 4 + 3 files changed, 143 insertions(+), 1 deletion(-) diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 1d066c21f0a2..c9e54676d25d 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -90,6 +90,9 @@ steps: device: b200-k8s source_file_dependencies: - tests/quantization/test_blackwell_moe.py + - tests/benchmarks/test_forward_cli.py + - vllm/benchmarks/forward/ + - vllm/models/muse_spark/ - vllm/model_executor/models/deepseek_v2.py - vllm/model_executor/models/gpt_oss.py - vllm/model_executor/models/llama4.py @@ -99,7 +102,10 @@ steps: - vllm/model_executor/layers/quantization/mxfp4.py - vllm/v1/attention/backends/flashinfer.py commands: - - pytest -s -v tests/quantization/test_blackwell_moe.py + - >- + pytest -v -s + tests/quantization/test_blackwell_moe.py + tests/benchmarks/test_forward_cli.py::test_bench_forward_muse_balanced_routing mirror: amd: label: ":amd: (MI355) Quantized MoE" diff --git a/tests/benchmarks/test_forward_cli.py b/tests/benchmarks/test_forward_cli.py index 9b66086b76e0..32e1f427fc4a 100644 --- a/tests/benchmarks/test_forward_cli.py +++ b/tests/benchmarks/test_forward_cli.py @@ -16,6 +16,7 @@ from vllm.benchmarks.forward.kv_init import touched_kv_views, zero_kv_blocks from vllm.benchmarks.forward.main import run_case from vllm.benchmarks.forward.profile import summarize +from vllm.benchmarks.forward.routing import BalancedRouting from vllm.benchmarks.forward.shapes import jitter_kv_lens, parse_case from vllm.platforms import current_platform from vllm.v1.attention.backends.mla.cpu_mla import CPUMLABackend @@ -27,6 +28,137 @@ MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="Native Muse TRTLLM validation requires Blackwell.", +) +def test_bench_forward_muse_balanced_routing(tmp_path): + """Muse's native router supports finite, reproducible balanced graph replay.""" + config = dict( + architectures=["MuseSpark1xForCausalLM"], + model_type="muse_spark_text", + vocab_size=128, + hidden_size=512, + intermediate_size=512, + moe_hidden_size=512, + moe_intermediate_size=512, + num_hidden_layers=1, + num_attention_heads=8, + num_key_value_heads=8, + head_dim=64, + num_local_experts=8, + num_experts_per_tok=2, + num_shared_experts=0, + max_position_embeddings=64, + hidden_act="silu", + attention_bias=False, + layer_types=["full_attention"], + mlp_layer_types=["sparse"], + layer_rope_theta=[10000.0], + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + rms_norm_eps=1e-5, + post_norm_eps=1e-8, + residual_gate_temperature=0.3, + final_logit_softcapping=None, + output_multiplier=0.5, + routed_scaling_factor=1.0, + qk_scale_factor=1.0, + tie_word_embeddings=False, + torch_dtype="bfloat16", + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + entrypoint = ( + "from vllm.model_executor.models import ModelRegistry; " + "ModelRegistry.register_model('MuseSpark1xForCausalLM', " + "'vllm.models.muse_spark.nvidia.model:MuseSpark1xForCausalLM'); " + "from vllm.entrypoints.cli.main import main; main()" + ) + output = tmp_path / "results.json" + command = [ + sys.executable, + "-c", + entrypoint, + "bench", + "forward", + "--model", + str(tmp_path), + "--skip-tokenizer-init", + "--moe-backend", + "flashinfer_trtllm", + "--routing", + "balanced", + "--max-model-len", + "64", + "--max-num-seqs", + "4", + "--max-num-batched-tokens", + "64", + "--kv-cache-memory-bytes", + "16777216", + "--kv-len-jitter", + "0", + "--kernel-config.enable_jit_warmup", + "false", + "--kernel-config.enable_cutedsl_warmup", + "false", + "--disable-flashinfer-autotune", + "--case", + "prefill=32", + "--case", + "decode=2*1@32", + "--iters", + "2", + "--samples", + "2", + "--output-json", + str(output), + ] + result = subprocess.run(command, capture_output=True, text=True, timeout=600) + assert result.returncode == 0, result.stdout + result.stderr + rows = json.loads(output.read_text()) + assert len(rows) == 2 + for row in rows: + (rank,) = row["ranks"] + assert rank["finite"] and rank["graph_deterministic"] + assert rank["graph_vs_eager"]["relative_rms"] == 0 + + +def test_balanced_routing_restores_expert_bias(dist_init): + """Fixture routes stay balanced despite selection bias, then restore it.""" + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + from vllm.models.muse_spark.common.moe import MuseSpark1xRouter + + device = torch.device(current_platform.device_type) + layer = torch.nn.Module() + gate = GateLinear(4, 8, params_dtype=torch.float32).to(device) + layer.add_module("router", gate) + layer.top_k = 2 + bias = torch.tensor([100.0, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], device=device) + layer.register_buffer("expert_bias", bias.clone()) + selector = MuseSpark1xRouter( + top_k=2, + num_experts=8, + expert_bias=layer.expert_bias, + route_scale=1.0, + ) + hidden = torch.zeros(64, 4, device=device) + with torch.no_grad(): + gate.weight.zero_() + fixture = BalancedRouting(layer, 64, 0, 1, 4, 0, device) + try: + logits, _ = gate(hidden) + _, ids = selector.select_experts(hidden, logits) + fixture.check_applied() + counts = torch.bincount((ids // 2).flatten(), minlength=4) + assert counts.min() == counts.max() + finally: + fixture.remove() + torch.testing.assert_close(layer.expert_bias, bias, rtol=0, atol=0) + logits, _ = gate(hidden) + _, ids = selector.select_experts(hidden, logits) + assert set(ids.flatten().tolist()) == {0, 1} + + @pytest.mark.benchmark @pytest.mark.skipif( not current_platform.is_cuda(), reason="Forward graph timing requires CUDA." diff --git a/vllm/benchmarks/forward/routing.py b/vllm/benchmarks/forward/routing.py index 9d48a541e77d..ad9589835578 100644 --- a/vllm/benchmarks/forward/routing.py +++ b/vllm/benchmarks/forward/routing.py @@ -59,6 +59,8 @@ def _moe_layers(model: torch.nn.Module) -> list[tuple[torch.nn.Module, GateLinea layers: dict[GateLinear, torch.nn.Module] = {} for module in model.modules(): gate = getattr(module, "gate", None) + if not isinstance(gate, GateLinear): + gate = getattr(module, "router", None) if not isinstance(gate, GateLinear) or gate in layers: continue if _top_k(module) is not None: @@ -121,6 +123,8 @@ def __init__( bias = getattr(gate, "e_score_correction_bias", None) if bias is None: bias = getattr(module, "e_score_correction_bias", None) + if bias is None: + bias = getattr(module, "expert_bias", None) if isinstance(bias, torch.Tensor): self._biases.append((bias, bias.detach().clone())) bias.data.zero_() From 0106c8abce0f8fec21d2871e5b7404acd9c75065 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Tue, 27 Oct 2026 20:54:04 -0400 Subject: [PATCH 016/406] [CI] Shard FP8 MoE kernels and L4 Humming eval automatically (#61831) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: Thang Nguyen Co-authored-by: Codex Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/kernels.yaml | 2 ++ .buildkite/test_areas/lm_eval.yaml | 1 + 2 files changed, 3 insertions(+) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 6784ee1b28bd..f4fd3e7e1795 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -890,6 +890,7 @@ steps: - label: ":nvidia: (H100) [Nightly] DeepEP FP8 MoE Kernels" key: kernels-fp8-moe-test-2xh100 + automatic_shard: true timeout_in_minutes: 45 device: h100 num_devices: 2 @@ -900,6 +901,7 @@ steps: mirror: amd: label: ":amd: (MI355) DeepEP FP8 MoE Kernels" + automatic_shard: true dind: false device: mi355_2 timeout_in_minutes: 55 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 527203d0a818..377c8885d085 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -831,6 +831,7 @@ steps: - label: ":nvidia: (L4) [Nightly] LM Eval Humming TEMPORARY" key: lm-eval-humming-l4 + automatic_shard: true optional: true timeout_in_minutes: 90 device: l4 From 66f1d0fa169be86a531ef7fc488074673c242b8a Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 27 Oct 2026 20:57:29 -0400 Subject: [PATCH 017/406] [Bugfix][HiSparse] Read resident prefill rows by request state slot (#62739) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../test_indexer_deepseek_v4_slot_mapping.py | 1 + .../v1/attention/test_sparse_mla_backends.py | 89 ++++++++++++------- .../unit/test_hisparse_connector.py | 6 +- vllm/v1/attention/backends/mla/index_group.py | 6 +- .../v1/attention/backends/mla/sparse_utils.py | 44 +++++++++ vllm/v1/hisparse/binding.py | 4 +- vllm/v1/hisparse/runtime.py | 19 +--- 7 files changed, 116 insertions(+), 53 deletions(-) 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 f88c5bc0fb50..57164a4380ae 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -449,6 +449,7 @@ def test_compressed_slot_mapping_dcp_uses_rank_local_state_positions(): def test_index_conversion_warmup_uses_physical_block_stride(): config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=None), cache_config=SimpleNamespace(block_size=64), model_config=SimpleNamespace( max_model_len=1024, diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 78113b63a93a..9d33416984a5 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -4061,42 +4061,69 @@ def hook(rows, cache, indices, counts, *, is_decode, **kwargs): ) -@pytest.mark.parametrize("max_swap_rows", [1, 4], ids=["over_capacity", "mixed"]) -def test_hisparse_resident_prefill_uses_attention_block_stride( - monkeypatch, max_swap_rows +@pytest.mark.parametrize( + ("num_decode_tokens", "max_swap_rows"), + [(3, 1), (1, 4)], + ids=["over_capacity", "mixed"], +) +def test_hisparse_resident_prefill_reads_residency_by_state_row( + num_decode_tokens, max_swap_rows ): - """Resident batches with prefill rows read resident pages directly, whether - they exceed the swap-row capacity or fit it.""" - expected = torch.tensor([[19]], dtype=torch.int32) - cache_handle = _StubHiSparseCacheHandle( - all_context_pages_resident=True, - num_decode_tokens=1, - num_actual_tokens=2, - view=SimpleNamespace(block_size=64, attention_block_stride=832), - block_table=torch.tensor([[3]], dtype=torch.int32), - batch_block_table=lambda: torch.tensor([[3]], dtype=torch.int32), - runtime=SimpleNamespace(max_swap_rows=max_swap_rows), + """Resident prefill conversion reads each request's residency row through + its state slot, at the attention block stride. An unset slot reads slot 0.""" + topk = torch.zeros((3, 128), dtype=torch.int32, device=DEVICE_TYPE) + # [resident groups, state slots, pages]; the leader reads group 1. + residency = torch.tensor( + [[[7], [8]], [[3], [5]]], dtype=torch.int32, device=DEVICE_TYPE + ) + runtime = SimpleNamespace( + is_group_leader=True, + begin_forward=lambda: None, + eager_host_mirror=True, + max_swap_rows=max_swap_rows, + request_state_indices=torch.tensor( + [1, 0, -1], dtype=torch.int32, device=DEVICE_TYPE + ), ) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.caches = [cache_handle] - index_group.physical_topk_indices = torch.empty((2, 1), dtype=torch.int32) - convert_once = MagicMock(return_value=expected) - monkeypatch.setattr(index_group, "_convert_once", convert_once) - topk = torch.zeros((2, 1), dtype=torch.int32) - metadata = SimpleNamespace( - req_id_per_token=torch.zeros(2, dtype=torch.int32), + cache = HiSparseCacheHandle(runtime) + cache.view = SimpleNamespace(block_size=64, attention_block_stride=832) + cache.block_table = residency[1] + group = HiSparseMLAIndexGroup( + logical_topk_indices=topk, + physical_topk_indices=torch.empty_like(topk), + valid_topk_counts=torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE), + side_stream=torch.Stream(), + logical_topk_ready=torch.Event(), + physical_topk_ready=torch.Event(), + has_indexer=True, + caches=[cache], ) - - result = index_group.convert_logical_to_physical_topk( - 0, - topk, - metadata, - block_stride_rows=None, - return_valid_counts=False, + metadata = SimpleNamespace( + num_actual_tokens=3, + num_decode_tokens=num_decode_tokens, + max_query_len=1 if num_decode_tokens == 3 else 2, + num_reqs=3, + req_id_per_token=torch.tensor([0, 1, 2], dtype=torch.int32, device=DEVICE_TYPE), ) - assert result is expected - assert convert_once.call_args.kwargs["block_stride_rows"] == 832 + # Repeat after compilation so its host delay cannot hide a stream race, and + # change the residency so a stale table cannot pass. + for step in range(2): + residency.add_(step * 10) + expected = torch.tensor( + [[5 + step * 10], [3 + step * 10], [3 + step * 10]], device=DEVICE_TYPE + ) + expected *= 832 + torch.accelerator.synchronize() + group.prepare_for_batch(0, metadata) + group.set_logical_topk_ready(0) + # Compute-stream work between the indexer and the layer must not delay + # the inputs the side-stream converter reads. + torch.cuda._sleep(200_000_000) + physical = group.convert_logical_to_physical_topk( + 0, topk, metadata, block_stride_rows=None, return_valid_counts=False + ) + torch.testing.assert_close(physical, expected.int().expand_as(topk)) def test_hisparse_flashmla_reorders_full_speculative_window_as_decode(): diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index 5d6eaf7dd2f5..28a08d58c2f3 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -510,7 +510,7 @@ def test_residency_updates_persist_by_state_row(): An update for a lost page and an appended page must leave the other pages and the other request's row intact across steps that reorder the batch. """ - table = torch.zeros((4, 2, 4), dtype=torch.int32) + table = torch.zeros((2, 4, 4), dtype=torch.int32) update_hisparse_residency( table, { @@ -527,5 +527,5 @@ def test_residency_updates_persist_by_state_row(): torch.tensor([0, 2], dtype=torch.int32), ) - assert table[[0, 2], 0].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]] - assert table[[0, 2], 1].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]] + assert table[0, [0, 2]].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]] + assert table[1, [0, 2]].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]] diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 67069860e85c..c0f4d8894403 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -116,6 +116,7 @@ def _convert_once( *, block_stride_rows: int | None, return_valid_counts: bool, + idx_mapping: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: num_tokens = logical_topk_indices.shape[0] if num_tokens > self.physical_topk_indices.shape[0]: @@ -127,6 +128,7 @@ def _convert_once( BLOCK_STRIDE_ROWS=block_stride_rows, NUM_TOPK_TOKENS=logical_topk_indices.shape[1], return_valid_counts=return_valid_counts, + idx_mapping=idx_mapping, ) physical_topk_indices = self.physical_topk_indices[:num_tokens] @@ -150,6 +152,7 @@ def _convert_once( valid_counts_out=( valid_topk_counts if return_valid_counts else None ), + idx_mapping=idx_mapping, ) self.physical_topk_ready.record(self.side_stream) compute_stream.wait_event(self.physical_topk_ready) @@ -244,10 +247,11 @@ def convert_logical_to_physical_topk( layer_index, logical_topk_indices, req_id_per_token, - leader.batch_block_table(), + leader.block_table, leader.view.block_size, block_stride_rows=leader.view.attention_block_stride, return_valid_counts=return_valid_counts, + idx_mapping=leader.runtime.request_state_indices, ) source_block_table = cache.source_block_table assert source_block_table is not None diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 2447df5ceb71..9af7ad8d21b0 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -81,6 +81,7 @@ class CompileKey: block_stride_rows: int block_n: int has_prefill_workspace: bool + has_idx_mapping: bool count_valid: bool single_tile: bool compact_to_front: bool @@ -99,6 +100,7 @@ def sparse_mla_index_remap_kernel( token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] valid_count_ptr, # int32 [num_tokens] - output valid count per row + idx_mapping_ptr, # int32 [num_requests] - block table row per request prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr workspace_rank_stride, # rows per DCP rank of a gathered prefill workspace @@ -109,6 +111,7 @@ def sparse_mla_index_remap_kernel( BLOCK_N: tl.constexpr, # tile width along columns NUM_TOPK_TOKENS: tl.constexpr, HAS_PREFILL: tl.constexpr, + HAS_IDX_MAPPING: tl.constexpr, COUNT_VALID: tl.constexpr, # whether to count valid indices # BLOCK_N >= NUM_TOPK_TOKENS: one program owns the query, so its count # is an in-register reduction and needs no atomic. @@ -141,6 +144,8 @@ def sparse_mla_index_remap_kernel( # Load request id for this token (no mask: grid is exact) req = tl.load(req_id_ptr + token_id) + if HAS_IDX_MAPPING: + req = tl.maximum(tl.load(idx_mapping_ptr + req), 0) # Load token indices for this tile ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 @@ -239,6 +244,7 @@ def dispatch( # type: ignore[override] BLOCK_N: int, NUM_TOPK_TOKENS: int, HAS_PREFILL_WORKSPACE: bool, + HAS_IDX_MAPPING: bool, COUNT_VALID: bool, COMPACT_TO_FRONT: bool, DCP_SIZE: int, @@ -252,6 +258,7 @@ def dispatch( # type: ignore[override] block_stride_rows=BLOCK_STRIDE_ROWS, block_n=tiling[1], has_prefill_workspace=HAS_PREFILL_WORKSPACE, + has_idx_mapping=HAS_IDX_MAPPING, count_valid=COUNT_VALID, single_tile=tiling[0], compact_to_front=COMPACT_TO_FRONT, @@ -279,6 +286,20 @@ def get_warmup_keys( block_size * dcp_size, ) max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size) + hisparse_combos: list[dict[str, Any]] = [] + if vllm_config.attention_config.hisparse_config is not None: + hisparse_combos.extend( + dict( + HAS_PREFILL_WORKSPACE=False, + HAS_IDX_MAPPING=True, + COUNT_VALID=count_valid, + COMPACT_TO_FRONT=count_valid, + DCP_SIZE=1, + DCP_RANK=0, + DCP_INTERLEAVE=1, + ) + for count_valid in (False, True) + ) pcp_dcp_combos = [] if ( vllm_config.parallel_config.prefill_context_parallel_size > 1 @@ -288,6 +309,7 @@ def get_warmup_keys( pcp_dcp_combos.append( dict( HAS_PREFILL_WORKSPACE=True, + HAS_IDX_MAPPING=False, COUNT_VALID=True, COMPACT_TO_FRONT=True, DCP_SIZE=dcp_size, @@ -299,6 +321,7 @@ def get_warmup_keys( zip_inputs( dict( HAS_PREFILL_WORKSPACE=False, + HAS_IDX_MAPPING=False, COUNT_VALID=False, COMPACT_TO_FRONT=False, DCP_SIZE=1, @@ -307,6 +330,7 @@ def get_warmup_keys( ), dict( HAS_PREFILL_WORKSPACE=False, + HAS_IDX_MAPPING=False, COUNT_VALID=True, COMPACT_TO_FRONT=True, DCP_SIZE=1, @@ -315,6 +339,7 @@ def get_warmup_keys( ), dict( HAS_PREFILL_WORKSPACE=True, + HAS_IDX_MAPPING=False, COUNT_VALID=True, COMPACT_TO_FRONT=True, DCP_SIZE=1, @@ -323,6 +348,7 @@ def get_warmup_keys( ), dict( HAS_PREFILL_WORKSPACE=False, + HAS_IDX_MAPPING=False, COUNT_VALID=True, COMPACT_TO_FRONT=False, DCP_SIZE=dcp_size, @@ -331,6 +357,7 @@ def get_warmup_keys( ), dict( HAS_PREFILL_WORKSPACE=False, + HAS_IDX_MAPPING=False, COUNT_VALID=True, COMPACT_TO_FRONT=dcp_size > 1, DCP_SIZE=dcp_size, @@ -338,6 +365,7 @@ def get_warmup_keys( DCP_INTERLEAVE=dcp_interleave, ), *pcp_dcp_combos, + *hisparse_combos, ), BLOCK_SIZE=block_size, BLOCK_STRIDE_ROWS=block_stride_rows, @@ -373,6 +401,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: prefill_workspace_starts=( int32_ptr if compile_key.has_prefill_workspace else None ), + idx_mapping=(int32_ptr if compile_key.has_idx_mapping else None), max_num_blocks_per_req=1, workspace_rank_stride=1, BLOCK_SIZE=compile_key.block_size, @@ -380,6 +409,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: BLOCK_N=compile_key.block_n, NUM_TOPK_TOKENS=compile_key.num_topk_tokens, HAS_PREFILL_WORKSPACE=compile_key.has_prefill_workspace, + HAS_IDX_MAPPING=compile_key.has_idx_mapping, COUNT_VALID=compile_key.count_valid, COMPACT_TO_FRONT=compile_key.compact_to_front, DCP_SIZE=compile_key.dcp_size, @@ -397,6 +427,7 @@ def __call__( valid_counts: torch.Tensor | None, prefill_workspace_request_ids: torch.Tensor | None, prefill_workspace_starts: torch.Tensor | None, + idx_mapping: torch.Tensor | None, *, workspace_rank_stride: int, max_num_blocks_per_req: int, @@ -405,6 +436,7 @@ def __call__( BLOCK_N: int, NUM_TOPK_TOKENS: int, HAS_PREFILL_WORKSPACE: bool, + HAS_IDX_MAPPING: bool, COUNT_VALID: bool, COMPACT_TO_FRONT: bool, DCP_SIZE: int, @@ -419,6 +451,7 @@ def __call__( valid_count_ptr=valid_counts, prefill_request_id_ptr=prefill_workspace_request_ids, workspace_starts_ptr=prefill_workspace_starts, + idx_mapping_ptr=idx_mapping, BLOCK_N=block_n, HAS_PREFILL=HAS_PREFILL_WORKSPACE, SINGLE_TILE=single_tile, @@ -475,6 +508,7 @@ def triton_convert_req_index_to_global_index( return_valid_counts: bool = False, out: torch.Tensor | None = None, valid_counts_out: torch.Tensor | None = None, + idx_mapping: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """out[token_id, indice_id] = block_table[req_id[token_id], @@ -494,11 +528,17 @@ def triton_convert_req_index_to_global_index( prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace starts for each prefill request + When idx_mapping (int32 [num_requests]) is given, block_table is + indexed by idx_mapping[req_id[token_id]] instead of req_id[token_id]. + When return_valid_counts is True, also returns the count of valid (non -1) indices per row, computed during the same kernel pass (no extra overhead). """ assert req_id.dtype == torch.int32 assert block_table.dtype == torch.int32 + assert idx_mapping is None or ( + idx_mapping.dtype == torch.int32 and idx_mapping.is_contiguous() + ) assert token_indices.dtype == torch.int32 assert req_id.shape[0] == token_indices.shape[0], ( f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) " @@ -572,6 +612,7 @@ def triton_convert_req_index_to_global_index( valid_counts, prefill_workspace_request_ids, prefill_workspace_starts, + idx_mapping, # shapes / constexprs workspace_rank_stride=prefill_workspace_rank_stride or 0, max_num_blocks_per_req=max_num_blocks_per_req, @@ -582,6 +623,7 @@ def triton_convert_req_index_to_global_index( BLOCK_N=BLOCK_N, NUM_TOPK_TOKENS=NUM_TOPK_TOKENS, HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE, + HAS_IDX_MAPPING=idx_mapping is not None, COUNT_VALID=return_valid_counts, COMPACT_TO_FRONT=return_valid_counts, DCP_SIZE=dcp_size, @@ -680,6 +722,7 @@ def triton_filter_and_convert_dcp_index( # No prefill workspace on the DCP decode path. None, None, + None, workspace_rank_stride=0, max_num_blocks_per_req=max_num_blocks_per_req, BLOCK_SIZE=BLOCK_SIZE, @@ -689,6 +732,7 @@ def triton_filter_and_convert_dcp_index( BLOCK_N=BLOCK_N, NUM_TOPK_TOKENS=NUM_TOPK_TOKENS, HAS_PREFILL_WORKSPACE=False, + HAS_IDX_MAPPING=False, COUNT_VALID=count_valid, COMPACT_TO_FRONT=compact_valid_to_front, DCP_SIZE=dcp_size, diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 6d591cd66e40..9dd061ef97aa 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -211,7 +211,7 @@ def bind_hisparse_kv_caches( assert len(table_shapes) == 1 max_num_reqs, max_num_pages = table_shapes.pop() residency = torch.zeros( - (max_num_reqs, len(resident_groups), max_num_pages), + (len(resident_groups), max_num_reqs, max_num_pages), dtype=torch.int32, device=block_tables.input_block_tables[resident_groups[0][0]].device, ) @@ -229,7 +229,7 @@ def bind_hisparse_kv_caches( block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, - block_table=residency[:, resident_source_index], + block_table=residency[resident_source_index], slot_mapping=block_tables.slot_mappings[group_id], ) cache_handle.residency = residency diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index e9d0b09f37ed..50a18f84ead6 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -1145,13 +1145,13 @@ def update_hisparse_residency( ) -> None: """Apply scheduled requests' residency changes. - ``residency`` is ``[state rows, resident groups, pages]``: the GPU block of + ``residency`` is ``[resident groups, state rows, pages]``: the GPU block of each page, or block 0 when the page is read from the host. """ if not updates: return batch_rows = {request_id: row for row, request_id in enumerate(request_ids)} - num_groups = residency.shape[1] + num_groups = residency.shape[0] rows: list[int] = [] pages: list[int] = [] block_ids: list[list[int]] = [[] for _ in range(num_groups)] @@ -1172,7 +1172,7 @@ def update_hisparse_residency( host_values, device=request_state_indices.device, dtype=torch.int32 ) state_rows = request_state_indices[values[:, 0]] - residency[state_rows, :, values[:, 1]] = values[:, 2:] + residency[:, state_rows, values[:, 1]] = values[:, 2:].T class HiSparseCacheHandle: @@ -1197,7 +1197,6 @@ def __init__(self, runtime: HiSparseRuntime) -> None: # Speculator layers write their rows after the target forward. self.draft_layer = False self.index_group_caches: list[HiSparseCacheHandle] = [self] - self._batch_block_table: torch.Tensor | None = None def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: assert self.runtime.is_group_leader @@ -1206,7 +1205,6 @@ def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: def _prepare_for_batch(self, attn_metadata: Any | None) -> None: self.dummy_batch = attn_metadata is None - self._batch_block_table = None self.runtime.begin_forward() self.num_actual_tokens = ( attn_metadata.num_actual_tokens if attn_metadata is not None else 0 @@ -1222,17 +1220,6 @@ def _prepare_for_batch(self, attn_metadata: Any | None) -> None: not self.decode_batch or self.runtime.eager_host_mirror ) - def batch_block_table(self) -> torch.Tensor: - """Resident rows by this step's batch row, for the prefill paths.""" - if self._batch_block_table is None: - assert self.block_table is not None - indices = self.runtime.request_state_indices - assert indices is not None - self._batch_block_table = self.block_table.index_select( - 0, indices.clamp(min=0) - ) - return self._batch_block_table - def write_target( self, num_input_rows: int, num_slot_rows: int ) -> tuple[torch.Tensor, torch.Tensor, int]: From ebaa8124755b10b956dc0b69082d9a94c12285d6 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 27 Oct 2026 21:00:46 -0400 Subject: [PATCH 018/406] [Bugfix][HiSparse] Pad staged prefill KV to avoid a trtllm-gen out-of-bounds read (#61659) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../v1/attention/test_sparse_mla_backends.py | 27 +++++++++++++++++++ vllm/v1/hisparse/runtime.py | 13 ++++----- 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 9d33416984a5..1b053014a2e1 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -3659,6 +3659,33 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows(): torch.testing.assert_close(staged_flat[i], expected) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_hisparse_staged_prefill_cache_backs_min_attention_rows(): + """A staged cache of a few blocks still has 256 allocated rows behind it. + + trtllm-gen sparse MLA reads up to ~250 rows from the start of its KV cache + whatever its size, faulting when a smaller cache ends at an unmapped page. + """ + if not _has_hisparse_ops(): + pytest.skip("hisparse CUDA ops unavailable") + device = torch.device("cuda") + block_size, row_width = 64, 16 + block_table = torch.tensor([[3, 1]], dtype=torch.int32, device=device) + seq_lens = torch.tensor([100], dtype=torch.int32, device=device) + plan = build_hisparse_prefill_staging_plan( + block_table, seq_lens, block_size, staging_block_capacity=2 + ) + host_cache = torch.zeros(4, block_size, row_width).pin_memory() + + staged = HiSparseRuntime.gather_prefill_cache(None, host_cache, plan) + + assert staged.shape[0] * block_size < 256 + backing_rows = staged.untyped_storage().nbytes() // ( + row_width * staged.element_size() + ) + assert backing_rows >= 256 + + def test_hisparse_fp8_decode_resolves_rows_once_then_runs_batched_attention(): """Verification rows resolve in one swap_in; FP8 attention sees them batched.""" device = torch.device("cpu") diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 50a18f84ead6..01af2be565df 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -39,6 +39,9 @@ # fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE. FP8_DS_MLA_ROW_BYTES = 656 HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor] +# trtllm-gen sparse MLA reads up to ~250 rows from the start of its KV cache +# whatever the cache's size: https://github.com/flashinfer-ai/flashinfer/issues/6484 +_MIN_STAGED_PREFILL_KV_ROWS = 256 def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool: @@ -854,15 +857,13 @@ def gather_prefill_cache( ) row_width = kv_cache.shape[-1] + num_blocks = plan.row_ids.shape[1] // plan.block_size + min_blocks = cdiv(_MIN_STAGED_PREFILL_KV_ROWS, plan.block_size) staged = torch.empty( - ( - plan.row_ids.shape[1] // plan.block_size, - plan.block_size, - row_width, - ), + (max(num_blocks, min_blocks), plan.block_size, row_width), dtype=kv_cache.dtype, device=plan.block_table.device, - ) + )[:num_blocks] if plan.gpu_row_ids is not None and resident_cache is not None: self.backend.gather_resident_rows( resident_cache, staged.view(-1, row_width), plan.gpu_row_ids[0] From 458334f00f7cbff7c8a1feae449656aa7343c476 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Tue, 27 Oct 2026 21:15:00 -0400 Subject: [PATCH 019/406] [MRV2][Spec Decode] Close DFlash/DSpark and adaptive-verification JIT warmup gaps (#63819) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .../spec_decode/test_acceptance_estimator.py | 33 +++++++++++++++++++ .../spec_decode/test_dflash_prepare_inputs.py | 19 ++++++++++- tests/v1/worker/test_gpu_warmup_blocks.py | 1 + tests/v1/worker/test_mixed_warmup_gate.py | 2 ++ .../gpu/spec_decode/acceptance_estimator.py | 19 +++++++++-- .../gpu/spec_decode/dflash/speculator.py | 2 +- vllm/v1/worker/gpu/warmup.py | 5 +++ 7 files changed, 76 insertions(+), 5 deletions(-) diff --git a/tests/v1/spec_decode/test_acceptance_estimator.py b/tests/v1/spec_decode/test_acceptance_estimator.py index 022ed463879c..41e42a1b197f 100644 --- a/tests/v1/spec_decode/test_acceptance_estimator.py +++ b/tests/v1/spec_decode/test_acceptance_estimator.py @@ -15,6 +15,8 @@ from vllm.platforms import current_platform from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import ( OnlineAcceptanceEstimator, + _accumulate_kernel, + _refit_kernel, ) pytest.importorskip("triton") @@ -295,3 +297,34 @@ def test_feature_ignores_which_token_was_drawn(): atol=2e-3, rtol=2e-3, ) + + +def test_warmup_covers_every_batch_size_and_the_refit(): + """After warmup, no batch size and no refit may JIT-compile during serving.""" + device = current_platform.device_type + + def compiled_variants() -> set: + return { + (kernel.fn.__name__, key) + for kernel in (_accumulate_kernel, _refit_kernel) + for caches in kernel.device_caches.values() + for key in caches[0] + } + + estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device) + estimator.REFIT_INTERVAL = 2 + slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32) + num_accepted = torch.ones(REFIT_NUM_REQS, device=device) + features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) + _drive_round(estimator, features[:1], num_accepted[:1], slots[:1]) + estimator.warmup() + warmed = compiled_variants() + + for num_reqs in (8, 300): + _drive_round( + estimator, features[:num_reqs], num_accepted[:num_reqs], slots[:num_reqs] + ) + torch.accelerator.synchronize() + + assert estimator._refits == 1 + assert compiled_variants() == warmed diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index b40c5d2d906e..2c1ebc609044 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -12,6 +12,7 @@ from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( DFlashSpeculator, + _prepare_dflash_inputs_kernel, ) pytestmark = pytest.mark.skipif( @@ -28,6 +29,7 @@ def _run_prepare( cp_size: int = 1, cp_interleave: int = 1, draft_dcp_size: int | None = None, + num_scheduled_tokens: int = 4, ): device = torch.device(current_platform.device_type) max_num_reqs = 4 @@ -46,7 +48,7 @@ def _run_prepare( input_batch = SimpleNamespace( num_reqs=1, num_tokens=4, - num_scheduled_tokens=np.array([4], dtype=np.int32), + num_scheduled_tokens=np.array([num_scheduled_tokens], dtype=np.int32), seq_lens_cpu_upper_bound=torch.tensor([target_positions[-1] + 1]), positions=torch.tensor(target_positions, dtype=torch.int64, device=device), query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device), @@ -220,3 +222,18 @@ def test_prepare_dflash_inputs_never_writes_the_null_block(): PAD_SLOT_ID, PAD_SLOT_ID, ] + + +def test_prepare_dflash_inputs_compiles_once_for_all_query_lengths(): + def compiled_variants() -> set: + return { + key + for caches in _prepare_dflash_inputs_kernel.device_caches.values() + for key in caches[0] + } + + kwargs = dict(target_positions=[10, 11, 12, 13], block_table_values=[0, 7, 8]) + _run_prepare(**kwargs) + warmed = compiled_variants() + _run_prepare(**kwargs, num_scheduled_tokens=300) + assert compiled_variants() == warmed diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index ddfbdb7fac5c..cc30e3b40912 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -90,6 +90,7 @@ def _make_runner( num_speculative_steps=num_spec_steps, speculative_config=SimpleNamespace() if num_spec_steps else None, adaptive_verification=None, + speculator=None, rejection_sampler=None, decode_query_len=num_spec_steps + 1, is_pooling_model=False, diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index 01f303a27597..74547825d27d 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -76,6 +76,7 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_ rejection_sampler = SimpleNamespace(enable_adaptive_verification=True) runner = SimpleNamespace( adaptive_verification=manager, + speculator=None, rejection_sampler=rejection_sampler, speculative_config=SimpleNamespace(), ) @@ -116,6 +117,7 @@ def test_kernel_warmup_runs_same_rounds_on_every_pp_stage(monkeypatch, speculati rounds.clear() runner = SimpleNamespace( adaptive_verification=None, + speculator=None, speculative_config=SimpleNamespace() if speculative else None, rejection_sampler=( SimpleNamespace(enable_adaptive_verification=False) diff --git a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py index f7a4a2f2c1e3..c210e23a1787 100644 --- a/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py +++ b/vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py @@ -347,6 +347,7 @@ def __init__( self.device = device self._steps_since_refit = 0 self._refits = 0 + self._accumulate_block_r = max(256, triton.next_power_of_2(max_num_reqs)) # The learned coefficients used for predicting acceptance. self.slope = torch.full((1,), _INIT_SLOPE, dtype=torch.float32, device=device) @@ -422,8 +423,8 @@ def step( self.predictions.stride(0), self.counts, num_reqs, - BLOCK_R=max(256, triton.next_power_of_2(num_reqs)), - num_warps=1 if num_reqs <= 256 else 4, + BLOCK_R=self._accumulate_block_r, + num_warps=1 if self._accumulate_block_r <= 256 else 4, ) self._steps_since_refit += 1 @@ -433,6 +434,19 @@ def step( # Fit the coefficients to the accumulated statistics gathered over the # course of the last REFIT_INTERVAL steps. + self._refit() + self._refits += 1 + + def warmup(self) -> None: + """Drop warmup statistics and compile the refit kernel. A refit on + zeroed statistics leaves the coefficients unchanged.""" + self.info.zero_() + self.grad.zero_() + self.counts.zero_() + self._steps_since_refit = 0 + self._refit() + + def _refit(self) -> None: _refit_kernel[(1,)]( self.slope, self.intercepts, @@ -446,7 +460,6 @@ def step( DAMPING=self.DAMPING_OBSERVATIONS, BLOCK=triton.next_power_of_2(self.num_speculative_steps), ) - self._refits += 1 def predict( self, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index c694d99a4335..fa5b3365561e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -888,7 +888,7 @@ def prepare_dflash_inputs( # per-request query length, not the total token count across the batch. max_target_query_len = int(input_batch.num_scheduled_tokens[:num_reqs].max()) max_tokens_per_req = max_target_query_len + num_query_per_req - BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) + BLOCK_SIZE = 128 num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE) _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)]( input_buffers.input_ids, diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 13de688b5d5d..b1f5aed0b64b 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -252,6 +252,11 @@ def warmup_kernels( rejection_sampler.enable_adaptive_verification = True if adaptive_verification is not None: adaptive_verification.warmup_budget_kernel() + acceptance_estimator = getattr( + model_runner.speculator, "acceptance_estimator", None + ) + if acceptance_estimator is not None: + acceptance_estimator.warmup() def _warmup_kernels( From 6376fda7a95904abf8977011af526c693c4c3004 Mon Sep 17 00:00:00 2001 From: ZhenchengLin <130115215+ZhenchengLin@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:22:50 -0700 Subject: [PATCH 020/406] [Bugfix][Rust Frontend] Use server-reported prompt_tokens in vllm-bench (#59694) --- rust/src/bench/src/backends/mod.rs | 30 ++++++++ rust/src/bench/src/backends/openai_chat.rs | 56 +++++++++++++-- .../bench/src/backends/openai_completions.rs | 72 +++++++++++++++++-- .../bench/src/backends/openai_responses.rs | 26 ++++++- rust/src/bench/src/benchmark.rs | 4 +- rust/src/bench/src/metrics/calculator.rs | 47 ++++++------ rust/src/bench/src/metrics/steady_state.rs | 46 +++--------- 7 files changed, 203 insertions(+), 78 deletions(-) diff --git a/rust/src/bench/src/backends/mod.rs b/rust/src/bench/src/backends/mod.rs index b5400c5d0f49..da23eb46a0ab 100644 --- a/rust/src/bench/src/backends/mod.rs +++ b/rust/src/bench/src/backends/mod.rs @@ -52,6 +52,21 @@ pub struct ChatDelta { #[derive(Deserialize)] pub struct ChunkUsage { pub completion_tokens: Option, + pub prompt_tokens: Option, +} + +impl ChunkUsage { + /// Record the server-reported token counts on `output`. A zero + /// `prompt_tokens` keeps the client estimate: some OpenAI-compatible + /// servers send placeholder zeros in usage chunks. + pub fn apply(&self, output: &mut RequestFuncOutput) { + if let Some(ct) = self.completion_tokens { + output.output_tokens = ct as usize; + } + if let Some(pt) = self.prompt_tokens.filter(|&pt| pt > 0) { + output.prompt_len = pt as usize; + } + } } use crate::cli::BackendKind; @@ -246,6 +261,21 @@ pub fn build_headers( mod tests { use super::*; + /// A placeholder zero `prompt_tokens` must not replace the client estimate. + #[test] + fn chunk_usage_ignores_zero_prompt_tokens() { + let mut output = RequestFuncOutput { + prompt_len: 100, + ..Default::default() + }; + ChunkUsage { + completion_tokens: Some(3), + prompt_tokens: Some(0), + } + .apply(&mut output); + assert_eq!((output.prompt_len, output.output_tokens), (100, 3)); + } + #[test] fn test_update_server_metrics_converts_ms_to_seconds() { let chunk: CompletionChunk = serde_json::from_value(serde_json::json!({ diff --git a/rust/src/bench/src/backends/openai_chat.rs b/rust/src/bench/src/backends/openai_chat.rs index 30c6093296d0..36ce961abf81 100644 --- a/rust/src/bench/src/backends/openai_chat.rs +++ b/rust/src/bench/src/backends/openai_chat.rs @@ -130,10 +130,8 @@ impl OpenAIChatBackend { } // Separate `if` (not `else if`) — Dynamo may send // both choices and usage in the same chunk. - if let Some(ref usage) = data.usage - && let Some(ct) = usage.completion_tokens - { - output.output_tokens = ct as usize; + if let Some(usage) = &data.usage { + usage.apply(&mut output); } } } @@ -363,7 +361,11 @@ mod tests { /// Serve one streaming chat response over a raw socket: `num_tokens` /// content chunks back-to-back, then `usage_delay` of silence, then the /// choice-less usage chunk and `[DONE]`. Returns the endpoint URL. - fn spawn_sse_server(num_tokens: usize, usage_delay: std::time::Duration) -> String { + fn spawn_sse_server( + num_tokens: usize, + usage_delay: std::time::Duration, + prompt_tokens: Option, + ) -> String { use std::io::{BufRead, BufReader, Read, Write}; let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); @@ -393,8 +395,10 @@ mod tests { send("data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n"); } std::thread::sleep(usage_delay); + let prompt_tokens = + prompt_tokens.map_or(String::new(), |pt| format!(",\"prompt_tokens\":{pt}")); send(&format!( - "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}}}}}\n\n" + "data: {{\"choices\":[],\"usage\":{{\"completion_tokens\":{num_tokens}{prompt_tokens}}}}}\n\n" )); send("data: [DONE]\n\n"); }); @@ -411,7 +415,7 @@ mod tests { // real delay is needed here; there is no observable event to wait on. let usage_delay = std::time::Duration::from_millis(150); let input = RequestFuncInput { - api_url: spawn_sse_server(3, usage_delay), + api_url: spawn_sse_server(3, usage_delay, None), model: "test-model".to_string(), output_len: 3, ..Default::default() @@ -442,4 +446,42 @@ mod tests { // keys already set above must not be overridden by extra_body assert_eq!(v["stream"], true); } + + /// The server's `usage.prompt_tokens` replaces the client-side estimate, + /// so input-token totals count what the server actually processed + /// (chat template, multimodal tokens). Mirrors Python #38654. + #[tokio::test] + async fn test_usage_prompt_tokens_override_prompt_len() { + let input = RequestFuncInput { + api_url: spawn_sse_server(3, std::time::Duration::ZERO, Some(128)), + model: "test-model".to_string(), + prompt_len: 100, + output_len: 3, + ..Default::default() + }; + + let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap(); + + assert!(output.success, "{}", output.error); + assert_eq!(output.prompt_len, 128); + } + + /// Servers that omit `usage.prompt_tokens` keep the client-side estimate, + /// and `completion_tokens` is still recorded. + #[tokio::test] + async fn test_missing_prompt_tokens_keeps_client_prompt_len() { + let input = RequestFuncInput { + api_url: spawn_sse_server(3, std::time::Duration::ZERO, None), + model: "test-model".to_string(), + prompt_len: 100, + output_len: 3, + ..Default::default() + }; + + let output = OpenAIChatBackend.send_request(&input, &reqwest::Client::new()).await.unwrap(); + + assert!(output.success, "{}", output.error); + assert_eq!(output.prompt_len, 100); + assert_eq!(output.output_tokens, 3); + } } diff --git a/rust/src/bench/src/backends/openai_completions.rs b/rust/src/bench/src/backends/openai_completions.rs index 0b62d342bcde..049988191318 100644 --- a/rust/src/bench/src/backends/openai_completions.rs +++ b/rust/src/bench/src/backends/openai_completions.rs @@ -156,10 +156,8 @@ impl OpenAICompletionsBackend { } // Separate `if` (not `else if`) — Dynamo may send // both choices and usage in the same chunk. - if let Some(ref usage) = data.usage - && let Some(ct) = usage.completion_tokens - { - output.output_tokens = ct as usize; + if let Some(usage) = &data.usage { + usage.apply(&mut output); } } } @@ -201,3 +199,69 @@ impl OpenAICompletionsBackend { Ok(output) } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Serve `body` as one streaming completions response over a raw socket. + /// Returns the endpoint URL. + fn spawn_sse_server(body: &'static str) -> String { + use std::io::{BufRead, BufReader, Read, Write}; + + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut content_length = 0; + loop { + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + if let Some(v) = line.to_ascii_lowercase().strip_prefix("content-length:") { + content_length = v.trim().parse().unwrap(); + } + if line == "\r\n" { + break; + } + } + reader.read_exact(&mut vec![0; content_length]).unwrap(); + stream + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\n\ + connection: close\r\n\r\n{body}" + ) + .as_bytes(), + ) + .unwrap(); + }); + format!("http://{addr}/v1/completions") + } + + /// The server's `usage.prompt_tokens` replaces the client-side estimate, + /// as in the chat backend. + #[tokio::test] + async fn test_usage_prompt_tokens_override_prompt_len() { + let input = RequestFuncInput { + api_url: spawn_sse_server(concat!( + "data: {\"choices\":[{\"text\":\"a\"}]}\n\n", + "data: {\"choices\":[{\"text\":\"b\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"completion_tokens\":2,\"prompt_tokens\":128}}\n\n", + "data: [DONE]\n\n", + )), + model: "test-model".to_string(), + prompt_len: 100, + output_len: 2, + ..Default::default() + }; + + let output = OpenAICompletionsBackend + .send_request(&input, &reqwest::Client::new()) + .await + .unwrap(); + + assert!(output.success, "{}", output.error); + assert_eq!((output.prompt_len, output.output_tokens), (128, 2)); + } +} diff --git a/rust/src/bench/src/backends/openai_responses.rs b/rust/src/bench/src/backends/openai_responses.rs index 2a9982636137..3178b5f4b3bf 100644 --- a/rust/src/bench/src/backends/openai_responses.rs +++ b/rust/src/bench/src/backends/openai_responses.rs @@ -27,6 +27,7 @@ struct ResponsesEventResponse { #[derive(serde::Deserialize)] struct ResponsesUsage { + input_tokens: Option, output_tokens: Option, } @@ -180,8 +181,14 @@ impl ResponsesStream { // with `status: "incomplete"`; the spec also has `response.incomplete`. Some("response.completed" | "response.incomplete") => { self.saw_terminal_event = true; - if let Some(n) = event.response.and_then(|r| r.usage?.output_tokens) { - self.output.output_tokens = n as usize; + if let Some(usage) = event.response.and_then(|r| r.usage) { + if let Some(n) = usage.output_tokens { + self.output.output_tokens = n as usize; + } + // Same rule as `ChunkUsage::apply`: a zero keeps the client estimate. + if let Some(n) = usage.input_tokens.filter(|&n| n > 0) { + self.output.prompt_len = n as usize; + } } } kind if matches!(kind, Some("error" | "response.failed")) || event.error.is_some() => { @@ -253,6 +260,21 @@ mod tests { stream.finish() } + /// The terminal event's `usage.input_tokens` replaces the client estimate, + /// as in Python's `async_request_openai_responses`. + #[test] + fn test_terminal_usage_input_tokens_override_prompt_len() { + let output = run(&[ + text_delta("Hi"), + event( + "response.completed", + ",\"response\":{\"usage\":{\"input_tokens\":128,\"output_tokens\":1}}", + ), + ]); + assert!(output.success, "{}", output.error); + assert_eq!((output.prompt_len, output.output_tokens), (128, 1)); + } + #[test] fn test_reasoning_starts_ttft_and_latency_ends_at_last_token() { let output = run(&[ diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 5833a3bdcea5..117e9fe72d9c 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -645,7 +645,7 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { let total_output_tokens: usize = input_requests.iter().map(|r| r.expected_output_len).sum(); println!("Dry run stats:"); println!(" Total prompts: {}", input_requests.len()); - println!(" Total input tokens: {total_input_tokens}"); + println!(" Estimated total input tokens: {total_input_tokens}"); println!(" Total expected output tokens: {total_output_tokens}"); println!( " Avg input tokens: {:.1}", @@ -1118,7 +1118,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { (m, Vec::new()) } else { calculate_metrics( - &input_requests, &outputs, benchmark_duration, &config.selected_percentiles, @@ -1148,7 +1147,6 @@ pub async fn run_benchmark(config: &BenchConfig) -> Result { ) { let ss = steady_state::compute( &outputs, - &input_requests, &window, &config.selected_percentiles, config.backend.is_pooling(), diff --git a/rust/src/bench/src/metrics/calculator.rs b/rust/src/bench/src/metrics/calculator.rs index 7e3fc77565ce..b55ed1cbde6d 100644 --- a/rust/src/bench/src/metrics/calculator.rs +++ b/rust/src/bench/src/metrics/calculator.rs @@ -3,7 +3,6 @@ use crate::backends::RequestFuncOutput; use crate::config::GoodputConfig; -use crate::datasets::SampleRequest; use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics}; use crate::multi_turn::ConversationOutput; @@ -27,7 +26,6 @@ fn log_failed_requests(outputs: &[RequestFuncOutput]) { /// /// Mirrors Python's `calculate_metrics()` from serve.py:392-599. pub fn calculate_metrics( - input_requests: &[SampleRequest], outputs: &[RequestFuncOutput], dur_s: f64, selected_percentiles: &[f64], @@ -44,7 +42,7 @@ pub fn calculate_metrics( let mut ttfts: Vec = Vec::new(); let mut e2els: Vec = Vec::new(); - for (i, output) in outputs.iter().enumerate() { + for output in outputs.iter() { if output.success { let output_len = if output.output_tokens > 0 { output.output_tokens @@ -57,7 +55,7 @@ pub fn calculate_metrics( }; actual_output_lens.push(output_len); - total_input += input_requests[i].prompt_len; + total_input += output.prompt_len; if output_len > 1 { let latency_minus_ttft = output.latency - output.ttft; @@ -283,7 +281,6 @@ fn percentiles_from_sorted(sorted: &[f64], percentiles: &[f64]) -> Vec<(f64, f64 /// /// Mirrors Python's `calculate_metrics_for_embeddings()` from serve.py. /// Key differences from `calculate_metrics()`: -/// - Uses `output.prompt_len` (server-reported) instead of `input_requests[i].prompt_len` /// - Only computes E2EL (no TTFT/TPOT/ITL since pooling is non-streaming) /// - `total_output` is 0 (pooling produces no output tokens) /// - `total_token_throughput` = `total_input / dur_s` (input tokens only) @@ -380,7 +377,7 @@ pub fn calculate_embedding_metrics( /// Calculate multi-turn benchmark metrics with per-turn breakdown. /// -/// 1. Flatten all turns into SampleRequest/RequestFuncOutput pairs → overall metrics +/// 1. Flatten all turns' outputs → overall metrics /// 2. Group by turn_index → per-turn metrics /// 3. Conversation-level stats pub fn calculate_multi_turn_metrics( @@ -391,24 +388,15 @@ pub fn calculate_multi_turn_metrics( max_turn_count: usize, ) -> MultiTurnMetrics { // Flatten all turns for overall metrics - let mut all_requests: Vec = Vec::new(); let mut all_outputs: Vec = Vec::new(); for conv in conversation_outputs { for turn in &conv.turns { - all_requests.push(SampleRequest { - prompt: std::sync::Arc::from(""), - prompt_len: turn.request_output.prompt_len, - expected_output_len: turn.request_output.output_tokens, - request_id: None, - ..Default::default() - }); all_outputs.push(turn.request_output.clone()); } } let (overall, _) = calculate_metrics( - &all_requests, &all_outputs, dur_s, selected_percentiles, @@ -419,25 +407,16 @@ pub fn calculate_multi_turn_metrics( // Per-turn breakdown let mut per_turn: Vec = Vec::new(); for turn_idx in 0..max_turn_count { - let mut turn_requests: Vec = Vec::new(); let mut turn_outputs: Vec = Vec::new(); for conv in conversation_outputs { if let Some(turn) = conv.turns.get(turn_idx) { - turn_requests.push(SampleRequest { - prompt: std::sync::Arc::from(""), - prompt_len: turn.request_output.prompt_len, - expected_output_len: turn.request_output.output_tokens, - request_id: None, - ..Default::default() - }); turn_outputs.push(turn.request_output.clone()); } } if !turn_outputs.is_empty() { let (turn_metrics, _) = calculate_metrics( - &turn_requests, &turn_outputs, dur_s, selected_percentiles, @@ -509,6 +488,26 @@ mod tests { assert!((p99 - 98.01).abs() < 0.1); } + /// Input-token totals come from each output's `prompt_len` (server-reported + /// via `usage.prompt_tokens`), not the client-side request estimate. + /// Mirrors Python's `serve.py`. + #[test] + fn test_total_input_uses_output_prompt_len() { + let outputs = vec![ + RequestFuncOutput { + success: true, + prompt_len: 128, + output_tokens: 3, + ..Default::default() + }; + 2 + ]; + + let (metrics, _) = calculate_metrics(&outputs, 1.0, &[], false, &GoodputConfig::default()); + + assert_eq!(metrics.total_input, 256); + } + #[test] fn test_multi_turn_metrics_use_request_prompt_lengths() { let turn = |turn_index, prompt_len, output_tokens| crate::multi_turn::TurnOutput { diff --git a/rust/src/bench/src/metrics/steady_state.rs b/rust/src/bench/src/metrics/steady_state.rs index a403275b7ac6..f108bb7c8501 100644 --- a/rust/src/bench/src/metrics/steady_state.rs +++ b/rust/src/bench/src/metrics/steady_state.rs @@ -170,16 +170,12 @@ pub fn detect_window( }) } -use crate::datasets::SampleRequest; - -/// Compute steady-state metrics from per-request outputs, the detected window, -/// and the original request list (for input-token attribution). +/// Compute steady-state metrics from per-request outputs and the detected window. /// /// `is_pooling` gates off TTFT because pooling backends write `ttft = latency` /// as a placeholder; reporting TTFT over that would silently be E2EL. pub fn compute( outputs: &[RequestFuncOutput], - input_requests: &[SampleRequest], window: &SteadyStateWindow, percentiles: &[f64], is_pooling: bool, @@ -200,7 +196,7 @@ pub fn compute( let mut ttfts_in_window: Vec = Vec::new(); let mut tpots_in_window: Vec = Vec::new(); - for (idx, o) in outputs.iter().enumerate() { + for o in outputs.iter() { if !o.success { continue; } @@ -213,7 +209,7 @@ pub fn compute( } if started_in { - input_tokens_in_window += input_requests[idx].prompt_len; + input_tokens_in_window += o.prompt_len; if !is_pooling { ttfts_in_window.push(o.ttft); // Per-request TPOT: (latency - ttft) / (output_tokens - 1). @@ -417,10 +413,6 @@ mod tests { assert!((w.end_s - 10.0).abs() < 1e-9); } - use std::sync::Arc; - - use crate::datasets::SampleRequest; - fn mk_full( start: f64, latency: f64, @@ -440,16 +432,6 @@ mod tests { } } - fn mk_sample(prompt_len: usize) -> SampleRequest { - SampleRequest { - prompt: Arc::from(""), - prompt_len, - expected_output_len: 0, - request_id: None, - ..Default::default() - } - } - #[test] fn compute_constant_concurrency_matches_full_run() { // 3 identical requests, each 10s, each emitting 5 tokens uniformly via itl=2s. @@ -459,13 +441,10 @@ mod tests { mk_full(0.0, 10.0, 2.0, itl.clone(), 5), mk_full(0.0, 10.0, 2.0, itl.clone(), 5), ]; - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); let m = compute( &outputs, - &requests, &w, &[99.0], false, // is_pooling @@ -502,12 +481,10 @@ mod tests { let late_itl = vec![0.02; 999]; // 1000 tokens; emission starts after end_s let late = mk_full(9.9, 30.0, 0.05, late_itl, 1000); outputs.push(late); - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); // Target = 2 so threshold_abs = 2; window opens when both are in flight. let w = detect_window(&outputs, Some(2), 1.0, 0.1, 40.0).unwrap(); - let m = compute(&outputs, &requests, &w, &[99.0], false); + let m = compute(&outputs, &w, &[99.0], false); // Late request started in window but most of its tokens are emitted AFTER end_s. // Only a handful (≤3) should count. @@ -534,10 +511,8 @@ mod tests { mk_full(0.0, 5.0, 5.0, Vec::new(), 0), mk_full(0.0, 5.0, 5.0, Vec::new(), 0), ]; - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); let w = detect_window(&outputs, Some(2), 1.0, 0.1, 10.0).unwrap(); - let m = compute(&outputs, &requests, &w, &[99.0], /* is_pooling */ true); + let m = compute(&outputs, &w, &[99.0], /* is_pooling */ true); assert_eq!(m.mean_ttft_ms, 0.0); assert_eq!(m.median_ttft_ms, 0.0); assert!(m.percentiles_ttft_ms.is_empty()); @@ -560,10 +535,8 @@ mod tests { mk_full(0.0, 10.0, 2.0, itl.clone(), 5), mk_full(0.0, 10.0, 2.0, itl.clone(), 5), ]; - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); - let m = compute(&outputs, &requests, &w, &[99.0], false); + let m = compute(&outputs, &w, &[99.0], false); assert!((m.mean_tpot_ms - 2000.0).abs() < 1e-6); assert!((m.median_tpot_ms - 2000.0).abs() < 1e-6); @@ -574,10 +547,8 @@ mod tests { #[test] fn json_roundtrip_some() { let outputs = vec![mk_full(0.0, 10.0, 2.0, vec![2.0; 4], 5); 3]; - let requests: Vec = - outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); - let m = compute(&outputs, &requests, &w, &[99.0], false); + let m = compute(&outputs, &w, &[99.0], false); let s = serde_json::to_string(&m).unwrap(); let parsed: SteadyStateMetrics = serde_json::from_str(&s).unwrap(); assert!((parsed.output_throughput - m.output_throughput).abs() < 1e-12); @@ -605,14 +576,13 @@ mod tests { let itl = vec![1.0; 9]; let outputs: Vec = (0..5).map(|i| mk_full(i as f64, 20.0, 1.0, itl.clone(), 10)).collect(); - let requests: Vec = (0..5).map(|_| mk_sample(100)).collect(); let w = detect_window(&outputs, Some(5), 1.0, 1.0, 25.0).unwrap(); assert_eq!(w.threshold_abs, 5); assert!((w.start_s - 4.0).abs() < 1e-9); assert!((w.end_s - 20.0).abs() < 1e-9); - let m = compute(&outputs, &requests, &w, &[50.0, 99.0], false); + let m = compute(&outputs, &w, &[50.0, 99.0], false); // Round-trip through JSON. let s = serde_json::to_string(&m).unwrap(); From 06e172b5cf8845c80c40a6eaaecbcbde3e38e8df Mon Sep 17 00:00:00 2001 From: Shreyash Bhatkar Date: Tue, 27 Oct 2026 21:08:14 -0500 Subject: [PATCH 021/406] [Bugfix][Spec Decode] Align MRV2 lm_head LoRA with logits rows (#59341) Signed-off-by: Shreyash Bhatkar <60093808+Shr3yash@users.noreply.github.com> Co-authored-by: linitra24 --- tests/v1/worker/test_mrv2_lmhead_lora.py | 61 ++++++++++++++++++++++++ vllm/v1/worker/gpu/lora_utils.py | 5 +- vllm/v1/worker/gpu/model_runner.py | 1 + 3 files changed, 66 insertions(+), 1 deletion(-) create mode 100644 tests/v1/worker/test_mrv2_lmhead_lora.py diff --git a/tests/v1/worker/test_mrv2_lmhead_lora.py b/tests/v1/worker/test_mrv2_lmhead_lora.py new file mode 100644 index 000000000000..341854ad9875 --- /dev/null +++ b/tests/v1/worker/test_mrv2_lmhead_lora.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest + +from vllm.lora.request import LoRARequest +from vllm.v1.worker.gpu.lora_utils import LoraState + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + + +def test_prompt_lora_mapping_aligns_with_spec_decode_logits_rows(): + """lm_head LoRA ids follow logits rows, not requests. + + Under speculative decoding cu_num_logits is one bonus row plus each + draft row. A per-request mapping writes the adapter onto a neighbor. + """ + state = LoraState(max_num_reqs=4) + lora_b = LoRARequest("adapter-b", 5, "/tmp/adapter-b") + # State slots are not batch order: slot 1 is req-a, slot 0 is req-b. + state.add_request("req-a", 1, None) + state.add_request("req-b", 0, lora_b) + + idx_mapping = np.array([1, 0], dtype=np.int32) + # req-a: 1 bonus + 2 drafts. req-b: 1 bonus and no drafts. + cu_num_logits = np.array([0, 3, 4], dtype=np.int32) + num_logits = np.diff(cu_num_logits).astype(np.int32) + # Query length is independent: req-a is still prefilling. + num_scheduled = np.array([8, 1], dtype=np.int32) + + prompt, token, requests = state.make_lora_inputs( + ["req-a", "req-b"], + idx_mapping, + num_scheduled, + num_logits, + ) + + assert num_logits.tolist() == [3, 1] + assert prompt == (0, 0, 0, 5) + assert len(prompt) == int(cu_num_logits[-1]) + assert token == (0,) * 8 + (5,) + assert requests == {lora_b} + + +def test_prompt_lora_mapping_repeats_for_multiple_sampled_tokens(): + """Multiple sampled tokens per step still get one lm_head id per row.""" + state = LoraState(max_num_reqs=2) + lora = LoRARequest("adapter", 3, "/tmp/adapter") + state.add_request("req-a", 0, lora) + state.add_request("req-b", 1, None) + # Two logits rows per request, no drafts. + num_logits = np.array([2, 2], dtype=np.int32) + prompt, token, _requests = state.make_lora_inputs( + ["req-a", "req-b"], + np.array([0, 1], dtype=np.int32), + np.array([2, 2], dtype=np.int32), + num_logits, + ) + assert prompt == (3, 3, 0, 0) + assert token == (3, 3, 0, 0) diff --git a/vllm/v1/worker/gpu/lora_utils.py b/vllm/v1/worker/gpu/lora_utils.py index eb81a83662a1..738221317ef3 100644 --- a/vllm/v1/worker/gpu/lora_utils.py +++ b/vllm/v1/worker/gpu/lora_utils.py @@ -96,9 +96,12 @@ def make_lora_inputs( req_ids: list[str], idx_mapping: np.ndarray, num_scheduled_tokens: np.ndarray, + num_logits: np.ndarray, ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: lora_ids = self.lora_ids[idx_mapping] - prompt_lora_mapping = tuple(lora_ids) + # One lm_head id per logits row: 1 + num_drafts rows per request + # under spec decode. + prompt_lora_mapping = tuple(lora_ids.repeat(num_logits)) token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens)) active_lora_requests: set[LoRARequest] = self.get_activate_loras(req_ids) return prompt_lora_mapping, token_lora_mapping, active_lora_requests diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b8d34076b082..6509eaf52c20 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1904,6 +1904,7 @@ def execute_model( input_batch.req_ids, input_batch.idx_mapping_np, input_batch.num_scheduled_tokens, + np.diff(input_batch.cu_num_logits_np), ) self._set_active_loras(*lora_inputs) else: From 11a3634094a5b44fa01e66fe66e05650546f62d7 Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Tue, 27 Oct 2026 22:14:24 -0400 Subject: [PATCH 022/406] [CI] Shard Humming B200/H100 evals and two entrypoints AMD mirrors automatically (#62513) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/entrypoints.yaml | 2 ++ .buildkite/test_areas/lm_eval.yaml | 7 ++++--- .../gsm8k/configs/humming/config-h100-shard-0.txt | 6 ------ .../gsm8k/configs/humming/config-h100-shard-1.txt | 5 ----- .../gsm8k/configs/humming/config-h100-shard-2.txt | 4 ---- tests/evals/gsm8k/configs/humming/config-h100.txt | 15 +++++++++++++++ 6 files changed, 21 insertions(+), 18 deletions(-) delete mode 100644 tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt delete mode 100644 tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt delete mode 100644 tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt create mode 100644 tests/evals/gsm8k/configs/humming/config-h100.txt diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 8a568046eebc..3fda7836b81f 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -156,6 +156,7 @@ steps: - pytest -v -s entrypoints/openai/chat_completion mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Entrypoints Integration (OpenAI API chat_completion)" dind: false device: mi355_dpx @@ -185,6 +186,7 @@ steps: - pytest -v -s entrypoints/cohere mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Entrypoints Integration (API Server Generate)" dind: false device: mi355_dpx diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 377c8885d085..b55515d0beea 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -881,13 +881,13 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-act-int8-shard-$$BUILDKITE_PARALLEL_JOB.txt -- label: ":nvidia: (H100) [Nightly] LM Eval Humming Default TEMPORARY Shard %N" +- label: ":nvidia: (H100) [Nightly] LM Eval Humming Default TEMPORARY" key: lm-eval-humming-h100-default + automatic_shard: true optional: true timeout_in_minutes: 70 device: h100 num_devices: 1 - parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming/ @@ -895,7 +895,7 @@ steps: - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100.txt - label: ":nvidia: (H100) [Nightly] LM Eval Humming Activation FP8/INT8 TEMPORARY" key: lm-eval-humming-h100-act @@ -935,6 +935,7 @@ steps: - label: ":nvidia: (B200) [Nightly] LM Eval Humming Default TEMPORARY" key: lm-eval-humming-b200-default + automatic_shard: true optional: true timeout_in_minutes: 50 device: b200-k8s diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt deleted file mode 100644 index 737a92cf48ae..000000000000 --- a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt +++ /dev/null @@ -1,6 +0,0 @@ -Qwen3.5-35B-A3B-experts-int8-humming.yaml -Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml -Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml -Qwen3-30B-A3B-Fp8-v1-humming.yaml -Qwen3-30B-A3B-NVFP4-humming.yaml -Qwen2-1.5B-Instruct-FP8W8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt deleted file mode 100644 index 0676dc7623f5..000000000000 --- a/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt +++ /dev/null @@ -1,5 +0,0 @@ -Qwen3.6-35B-A3B-NVFP4-humming.yaml -Qwen3-30B-A3B-AWQ-humming.yaml -Qwen3-30B-A3B-GPTQ-Int4-humming.yaml -Qwen3-30B-A3B-FP8-block-humming.yaml -Qwen3-0.6B-MXFP8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt deleted file mode 100644 index 9d8fe128d0ff..000000000000 --- a/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt +++ /dev/null @@ -1,4 +0,0 @@ -NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-humming.yaml -Qwen3.5-35B-A3B-FP8-humming.yaml -Qwen3-30B-A3B-MXFP4A16-humming.yaml -gpt-oss-20b-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-h100.txt b/tests/evals/gsm8k/configs/humming/config-h100.txt new file mode 100644 index 000000000000..ba6a1e81da52 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-h100.txt @@ -0,0 +1,15 @@ +Qwen3.5-35B-A3B-experts-int8-humming.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml +Qwen3-30B-A3B-Fp8-v1-humming.yaml +Qwen3-30B-A3B-NVFP4-humming.yaml +Qwen2-1.5B-Instruct-FP8W8-humming.yaml +Qwen3.6-35B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming.yaml +Qwen3-30B-A3B-FP8-block-humming.yaml +Qwen3-0.6B-MXFP8-humming.yaml +NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-humming.yaml +Qwen3.5-35B-A3B-FP8-humming.yaml +Qwen3-30B-A3B-MXFP4A16-humming.yaml +gpt-oss-20b-humming.yaml From 1c787a25cc19bc6e6f7674d4f6f49873ca0ed755 Mon Sep 17 00:00:00 2001 From: Guancheng Fu <110874468+gc-fu@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:17:20 +0800 Subject: [PATCH 023/406] [Bugfix][LoRA] Reject non-linear expert placement with EP (#62797) Signed-off-by: Guancheng Fu Co-authored-by: OpenAI Codex --- tests/lora/test_moe_lora_ep_load.py | 57 +++++++++++++++++++++++++++++ vllm/lora/layers/fused_moe.py | 9 +++++ 2 files changed, 66 insertions(+) diff --git a/tests/lora/test_moe_lora_ep_load.py b/tests/lora/test_moe_lora_ep_load.py index 60d0c2c9ba7e..158c9375a208 100644 --- a/tests/lora/test_moe_lora_ep_load.py +++ b/tests/lora/test_moe_lora_ep_load.py @@ -4,11 +4,18 @@ load-time slicing """ +from contextlib import nullcontext +from dataclasses import replace +from types import SimpleNamespace + import pytest import torch +from vllm.lora.layers.fused_moe import FusedMoEWithLoRA from vllm.lora.lora_model import LoRAModel, MoEEPLoadSpec from vllm.lora.peft_helper import PEFTHelper +from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig +from vllm.model_executor.layers.fused_moe.expert_map_manager import ExpertMapManager NUM_LAYERS = 48 GLOBAL_NUM_EXPERTS = 128 @@ -17,6 +24,56 @@ NON_EXPERT_MODULES = ("q_proj", "k_proj", "v_proj", "o_proj", "gate") +@pytest.mark.parametrize( + "use_ep,requested_placement,num_expert_group,resolved_placement", + [ + (False, "linear", 2, "linear"), + (False, "round_robin", 2, "linear"), + (True, "linear", 2, "linear"), + (True, "round_robin", 1, "linear"), + (True, "round_robin", 2, "round_robin"), + ], +) +def test_moe_lora_ep_requires_contiguous_expert_placement( + use_ep, requested_placement, num_expert_group, resolved_placement +): + """Reject non-contiguous ownership, but allow requests resolved to linear.""" + parallel_config = replace( + FusedMoEParallelConfig.make_no_parallel(), + use_ep=use_ep, + ep_size=2 if use_ep else 1, + ) + expert_map = ExpertMapManager( + max_num_batched_tokens=16, + top_k=2, + global_num_experts=4, + num_redundant_experts=0, + num_expert_group=num_expert_group, + moe_parallel_config=parallel_config, + placement_strategy=requested_placement, + enable_eplb=False, + ) + assert expert_map.placement_strategy == resolved_placement + + layer = FusedMoEWithLoRA.__new__(FusedMoEWithLoRA) + torch.nn.Module.__init__(layer) + layer.moe_config = SimpleNamespace(moe_parallel_config=parallel_config) + layer.base_layer = SimpleNamespace( + expert_placement_strategy=expert_map.placement_strategy + ) + error = ( + pytest.raises( + ValueError, + match="Set expert_placement_strategy='linear' to use MoE LoRA " + "with expert parallelism", + ) + if use_ep and resolved_placement != "linear" + else nullcontext() + ) + with error: + layer._ep_check() + + def _expected_lora_modules() -> set[str]: """Replicate the set ``WorkerLoRAManager._load_adapter`` would build from this model's ``packed_modules_mapping``.""" diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 19a1128ce145..aa20290330c2 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -233,6 +233,15 @@ def _create_lora_b_weights(self, max_loras: int, lora_config: LoRAConfig): def _ep_check(self): if self.use_ep: + placement = self.base_layer.expert_placement_strategy + if placement != "linear": + raise ValueError( + "Fused MoE LoRA with EP requires linear expert placement: " + "adapter loading assumes a contiguous range of experts " + f"on each rank, got '{placement}'. " + "Set expert_placement_strategy='linear' to use MoE LoRA " + "with expert parallelism." + ) moe_config = self.moe_config all2all_backend = moe_config.moe_parallel_config.all2all_backend assert all2all_backend == "allgather_reducescatter", ( From e4c2afdc171f1ee1d501d35a4de7870838a905bc Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 28 Oct 2026 10:20:49 +0800 Subject: [PATCH 024/406] [Rust Frontend] Reject max_tokens above max_model_len like Python (#63964) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .../src/routes/openai/chat_completions.rs | 8 ++ .../server/src/routes/openai/completions.rs | 8 ++ .../src/server/src/routes/openai/utils/mod.rs | 21 +++++ rust/src/server/src/routes/render.rs | 11 +++ rust/src/server/src/routes/tests.rs | 81 +++++++++++++++++++ rust/src/text/src/lower.rs | 9 +++ 6 files changed, 138 insertions(+) diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index d243ffdfd1af..6ba84d867750 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -45,6 +45,7 @@ use crate::routes::openai::utils::types::{ ToolCallDelta, Usage, }; use crate::routes::openai::utils::usage::ContinuousUsage; +use crate::routes::openai::utils::validate_max_output_tokens; use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::routes::output_log::{LoggedOutput, chat_output_text, log_output}; use crate::state::AppState; @@ -74,6 +75,13 @@ pub async fn chat_completions( Ok(prepared) => prepared, Err(error) => return error.into_response(), }; + if let Err(error) = validate_max_output_tokens( + "max_completion_tokens", + prepared.chat_request.sampling_params.max_tokens, + state.chat.engine_core_client().max_model_len(), + ) { + return error.into_response(); + } let request_span = tracing::info_span!( "chat_completions", request_id = %prepared.request_id, diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 6400a9ced88a..228f39185136 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -44,6 +44,7 @@ use crate::routes::openai::completions::types::{ }; use crate::routes::openai::utils::types::LogProbs; use crate::routes::openai::utils::usage::ContinuousUsage; +use crate::routes::openai::utils::validate_max_output_tokens; use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::routes::output_log::{LoggedOutput, log_output}; use crate::state::AppState; @@ -81,6 +82,13 @@ pub async fn completions( Ok(prepared) => prepared, Err(error) => return error.into_response(), }; + if let Err(error) = validate_max_output_tokens( + "max_tokens", + prepared.text_request.sampling_params.max_tokens, + state.chat.engine_core_client().max_model_len(), + ) { + return error.into_response(); + } let request_span = tracing::info_span!( "completions", request_id = %prepared.request_id, diff --git a/rust/src/server/src/routes/openai/utils/mod.rs b/rust/src/server/src/routes/openai/utils/mod.rs index a5f00ae0eadf..f0159d6cddea 100644 --- a/rust/src/server/src/routes/openai/utils/mod.rs +++ b/rust/src/server/src/routes/openai/utils/mod.rs @@ -33,6 +33,27 @@ pub(crate) fn validate_generation_prompt_truncation( Ok(()) } +/// Reject an output budget larger than the whole context window, like +/// Python's `TokenizeParams` check for chat and completions. A budget that +/// fits the window is still clamped to the room the prompt leaves when the +/// request is lowered; `/inference/v1/generate` only clamps, as in Python. +pub(crate) fn validate_max_output_tokens( + param: &'static str, + max_tokens: Option, + max_model_len: u32, +) -> Result<(), ApiError> { + if let Some(max_tokens) = max_tokens + && max_tokens > max_model_len + { + bail_invalid_request!( + param = param, + "{param}={max_tokens} cannot be greater than max_model_len={max_model_len}. \ + Please request fewer output tokens." + ); + } + Ok(()) +} + pub(crate) fn resolve_generation_prompt_truncation( limit: Option, side: Option, diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs index a2de402dac2d..12acd90582d4 100644 --- a/rust/src/server/src/routes/render.rs +++ b/rust/src/server/src/routes/render.rs @@ -23,6 +23,7 @@ use crate::routes::inference::generate::{ GenerateRequest, GenerateSamplingParams, validate_request_compat as validate_generate_request, }; use crate::routes::openai::utils::types::{ListModelsResponse, ModelObject, StreamOptions}; +use crate::routes::openai::utils::validate_max_output_tokens; use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::routes::openai::{ ChatCompletionRequest, CompletionRequest, lower_chat_request, lower_completion_request, @@ -164,6 +165,11 @@ async fn render_chat( let return_token_ids = body.return_token_ids; let request_context = resolve_request_context(&headers, body.request_id.as_deref()); let chat_request = lower_chat_request(body, &model_resolution(&state), request_context)?; + validate_max_output_tokens( + "max_completion_tokens", + chat_request.sampling_params.max_tokens, + state.text.max_model_len(), + )?; let (text_request, _) = state .chat .prepare(chat_request, &state.text) @@ -196,6 +202,11 @@ async fn render_completion( request_context, tokenizer.as_ref(), )?; + validate_max_output_tokens( + "max_tokens", + text_request.sampling_params.max_tokens, + state.text.max_model_len(), + )?; Ok(Json(vec![lower_render_request( &state, text_request, diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 587810226290..3d015ea14b53 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -26,6 +26,7 @@ use vllm_chat::{ ChatRequestProcessor, ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, ParserSelection, }; +use vllm_engine_core_client::mock_engine::DEFAULT_MOCK_MAX_MODEL_LEN; use vllm_engine_core_client::mock_engine::default_ready_response; use vllm_engine_core_client::protocol::decode_value; use vllm_engine_core_client::protocol::dtype::TensorDtype; @@ -1721,6 +1722,44 @@ async fn render_list_models_reports_configured_max_model_len() { } } +/// Like Python's `TokenizeParams` check, a requested output budget above the +/// context window is a 400, while one that fits the window is accepted. +#[tokio::test] +async fn render_rejects_max_tokens_above_max_model_len() { + let mut app = test_render_app_with_max_model_len(Some(64)); + + let mut errors = Vec::new(); + for (uri, body) in [ + ( + "/v1/chat/completions/render", + json!({ + "model": "render-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 65, + }), + ), + ( + "/v1/completions/render", + json!({"model": "render-model", "prompt": "hello", "max_tokens": 65}), + ), + ] { + let (status, json) = post_json(&mut app, uri, body).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{uri}"); + errors.push(json.to_string()); + } + expect_test::expect![[r#" + {"error":{"message":"max_completion_tokens=65 cannot be greater than max_model_len=64. Please request fewer output tokens.","type":"BadRequestError","param":"max_completion_tokens","code":400}} + {"error":{"message":"max_tokens=65 cannot be greater than max_model_len=64. Please request fewer output tokens.","type":"BadRequestError","param":"max_tokens","code":400}}"#]].assert_eq(&errors.join("\n")); + + let (status, json) = post_json( + &mut app, + "/v1/completions/render", + json!({"model": "render-model", "prompt": "hello", "max_tokens": 64}), + ) + .await; + assert_eq!(status, StatusCode::OK, "{json}"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn render_chat_response_can_be_submitted_to_generate_unchanged() { @@ -8881,6 +8920,48 @@ async fn tokenize_allows_prompts_at_or_above_max_model_len() { assert_eq!(chat_json["max_model_len"], 4); } +/// Like Python's `TokenizeParams` check, chat and completions reject a +/// requested output budget above the context window before submitting. The +/// mock engine answers the first request, so a submitted one returns 200. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn chat_and_completions_reject_max_tokens_above_max_model_len() { + let (mut app, _engine_task) = test_app_with_engine_handle().await; + let max_tokens = DEFAULT_MOCK_MAX_MODEL_LEN + 1; + + let mut errors = Vec::new(); + for (uri, body) in [ + ( + "/v1/chat/completions", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": max_tokens, + }), + ), + ( + "/v1/chat/completions", + json!({ + "model": "Qwen/Qwen1.5-0.5B-Chat", + "messages": [{"role": "user", "content": "hello"}], + "max_completion_tokens": max_tokens, + }), + ), + ( + "/v1/completions", + json!({"model": "Qwen/Qwen1.5-0.5B-Chat", "prompt": "hello", "max_tokens": max_tokens}), + ), + ] { + let (status, json) = post_json(&mut app, uri, body).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{uri}: {json}"); + errors.push(json.to_string()); + } + expect_test::expect![[r#" + {"error":{"message":"max_completion_tokens=1048577 cannot be greater than max_model_len=1048576. Please request fewer output tokens.","type":"BadRequestError","param":"max_completion_tokens","code":400}} + {"error":{"message":"max_completion_tokens=1048577 cannot be greater than max_model_len=1048576. Please request fewer output tokens.","type":"BadRequestError","param":"max_completion_tokens","code":400}} + {"error":{"message":"max_tokens=1048577 cannot be greater than max_model_len=1048576. Please request fewer output tokens.","type":"BadRequestError","param":"max_tokens","code":400}}"#]].assert_eq(&errors.join("\n")); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn tokenize_chat_conflicting_generation_flags_returns_400() { diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 12b4cf254f39..5cbf58e7361b 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -1693,6 +1693,15 @@ mod tests { assert_eq!(result.unwrap(), 100); } + /// Shared lowering only clamps, even a budget above the whole context + /// window: `/inference/v1/generate` accepts it like Python, while chat and + /// completions reject it in their HTTP handlers. + #[test] + fn resolve_max_tokens_caps_budget_above_model_len() { + let result = resolve_max_tokens(Some(1000), None, 200, 100); + assert_eq!(result.unwrap(), 100); + } + #[test] fn lower_text_request_preserves_non_streaming_request_metadata() { let mut request = sample_request(); From 65de04f8b763d656c714c213dab54e3727b4bc4a Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 28 Oct 2026 10:21:38 +0800 Subject: [PATCH 025/406] [Rust Frontend] Accept list-form chat templates like Transformers (#63963) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- rust/src/chat/src/backend/hf.rs | 70 +++++++++++++++++++++- rust/src/chat/src/renderer/hf/mod.rs | 44 ++++++++++++-- rust/src/text/src/backend/hf/config.rs | 83 +++++++++++++++++++++++++- rust/src/text/src/backend/hf/mod.rs | 4 +- 4 files changed, 191 insertions(+), 10 deletions(-) diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index f0da863534c7..8b766e537c5e 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -237,8 +237,8 @@ mod tests { use super::HfChatBackend; use crate::backend::{ChatBackend, LoadModelBackendsOptions, NewChatOutputProcessorOptions}; - use crate::request::{ChatContent, ChatMessage, ChatRequest}; - use crate::{ParserSelection, RendererSelection}; + use crate::request::{ChatContent, ChatMessage, ChatRequest, ResolvedToolContext}; + use crate::{ChatTool, ParserSelection, RendererSelection}; fn request_with_user_text(text: &str) -> ChatRequest { ChatRequest { @@ -343,6 +343,72 @@ mod tests { assert_eq!(prompt, "hello"); } + #[test] + fn hf_renderer_selects_named_tokenizer_config_templates_like_transformers() { + let config = r#"{"model_type":"llama"}"#; + let named_templates = r#"{"chat_template":[ + {"name":"default","template":"default:{{ messages[0].content }}"}, + {"name":"tool_use","template":"tool_use:{{ tools[0].function.name }}"} + ]}"#; + let load = |files, chat_template: Option<&str>| { + HfChatBackend::from_resolved_model_files( + files, + "test-model".to_string(), + LoadModelBackendsOptions { + chat_template: chat_template.map(str::to_owned), + ..Default::default() + }, + test_tokenizer(), + ) + .unwrap() + }; + + let tokenizer_config_only = load(resolved_files(config, named_templates), None); + // A configured or dedicated template replaces every named template, + // `tool_use` included. + let configured = load( + resolved_files(config, named_templates), + Some("configured:{{ messages[0].content }}"), + ); + let mut files = resolved_files(config, named_templates); + let dedicated_path = + files.config_path.as_ref().unwrap().with_file_name("chat_template.jinja"); + std::fs::write(&dedicated_path, "dedicated:{{ messages[0].content }}").unwrap(); + files.chat_template_path = Some(dedicated_path); + let dedicated = load(files, None); + + let request = request_with_user_text("hello"); + let mut with_tools = request.clone(); + with_tools.tool_context = ResolvedToolContext::new( + &request.messages, + vec![ChatTool { + name: "lookup".to_string(), + description: None, + parameters: serde_json::json!({"type": "object"}), + strict: None, + defer_loading: None, + }], + None, + true, + ) + .unwrap(); + let rendered = [tokenizer_config_only, configured, dedicated].map(|backend| { + [&request, &with_tools].map(|request| { + backend.chat_renderer().render(request).unwrap().prompt.into_text().unwrap() + }) + }); + + assert_eq!( + rendered, + [ + ["default:hello", "tool_use:lookup"], + ["configured:hello", "configured:hello"], + ["dedicated:hello", "dedicated:hello"], + ] + .map(|row| row.map(str::to_owned)) + ); + } + #[test] fn native_renderer_inherits_deployment_reasoning_below_request_controls() { let backend = HfChatBackend::from_resolved_model_files( diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index dfa6a1d2165a..6979c1964494 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -9,7 +9,7 @@ use thiserror_ext::AsReport as _; use tracing::{info, trace, warn}; use vllm_text::Prompt; use vllm_text::backend::hf::{ - HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config, + HfChatTemplate, HfSpecialTokens, HfTokenizerConfig, ResolvedModelFiles, load_tokenizer_config, }; use self::format::{ @@ -73,6 +73,9 @@ pub struct MultimodalRenderInfo { /// state. pub struct HfChatRenderer { default_template: Option, + /// The model's template for requests with tools, from a list-form + /// `chat_template` entry named `tool_use`; falls back to the default. + tool_use_template: Option, default_template_kwargs: HashMap, content_format: ContentFormatOption, special_tokens: Option, @@ -97,6 +100,7 @@ impl HfChatRenderer { .map_err(|error| Error::ChatTemplate(error.to_report_string())) }) .transpose()?, + tool_use_template: None, default_template_kwargs, content_format, special_tokens: None, @@ -109,6 +113,18 @@ impl HfChatRenderer { }) } + /// Use a separate model template for requests with tools, as Transformers + /// selects a `tool_use` template when tools are passed. + pub fn with_tool_use_template(mut self, template: Option) -> Result { + self.tool_use_template = template + .map(|template| { + CompiledChatTemplate::new(template, self.content_format) + .map_err(|error| Error::ChatTemplate(error.to_report_string())) + }) + .transpose()?; + Ok(self) + } + pub fn with_special_tokens(mut self, special_tokens: Option) -> Self { self.special_tokens = special_tokens; self @@ -137,7 +153,16 @@ impl HfChatRenderer { chat_template, .. } = load_tokenizer_config(files.tokenizer_config_path.as_deref())?; - let mut template = chat_template; + let mut template = chat_template + .as_ref() + .and_then(HfChatTemplate::default_template) + .map(str::to_owned); + // A configured template or a dedicated template file replaces every + // tokenizer_config template, `tool_use` included. + let mut tool_use_template = chat_template + .as_ref() + .and_then(HfChatTemplate::tool_use_template) + .map(str::to_owned); let special_tokens = (!special_tokens.is_empty()).then_some(special_tokens); if let Some(configured_template) = options.chat_template.as_deref() { @@ -145,6 +170,7 @@ impl HfChatRenderer { resolve_chat_template(configured_template) .map_err(|error| Error::ChatTemplate(error.to_report_string()))?, ); + tool_use_template = None; info!("using configured chat template override"); } else if let Some(chat_template_path) = files.chat_template_path.as_deref() { // If independent chat template file(s) exist and contain non-empty content, @@ -158,6 +184,7 @@ impl HfChatRenderer { "loaded dedicated chat template file, overriding tokenizer_config chat_template" ); template = file_template; + tool_use_template = None; } else { warn!( path = %chat_template_path.display(), @@ -171,6 +198,7 @@ impl HfChatRenderer { options.default_chat_template_kwargs, options.chat_template_content_format, )? + .with_tool_use_template(tool_use_template)? .with_special_tokens(special_tokens) .with_multimodal(multimodal) .with_trust_request_chat_template(options.trust_request_chat_template)) @@ -183,7 +211,9 @@ impl HfChatRenderer { /// temporary template is compiled from that string and used instead of /// the model's default. Overrides are refused unless request templates /// are trusted, and render under a fuel budget; the model's default - /// template has none, as in the Python frontend. + /// template has none, as in the Python frontend. Without an override, a + /// request with tools uses the model's `tool_use` template if it has one, + /// as Transformers' `get_chat_template` does. fn apply_chat_template(&self, request: &ChatRequest) -> Result { if request.chat_options.chat_template.is_some() && !self.trust_request_chat_template { return Err(Error::UntrustedRequestChatTemplate); @@ -198,9 +228,15 @@ impl HfChatRenderer { .map_err(|error| Error::ChatTemplate(error.to_report_string())) }) .transpose()?; + let model_template = if request.tools().is_empty() { + None + } else { + self.tool_use_template.as_ref() + } + .or(self.default_template.as_ref()); let template = override_template .as_ref() - .or(self.default_template.as_ref()) + .or(model_template) .ok_or(Error::MissingChatTemplate)?; self.apply_chat_template_inner(template, request) diff --git a/rust/src/text/src/backend/hf/config.rs b/rust/src/text/src/backend/hf/config.rs index 9f22e7989329..d204757fa0ce 100644 --- a/rust/src/text/src/backend/hf/config.rs +++ b/rust/src/text/src/backend/hf/config.rs @@ -17,7 +17,7 @@ use crate::normalize_top_k; pub struct HfTokenizerConfig { #[serde(flatten)] pub special_tokens: HfSpecialTokens, - pub chat_template: Option, + pub chat_template: Option, /// The `tokenizer_class` field from HuggingFace tokenizer configs. Some /// tiktoken-based models (e.g. DeepSeek, Kimi K2) set this to a value /// containing "Tiktoken" which can be used as a hint for backend @@ -28,6 +28,53 @@ pub struct HfTokenizerConfig { pub response_template: Option, } +/// The `chat_template` field of `tokenizer_config.json`: one template, or the +/// named templates Transformers serializes as a list of `{name, template}` +/// (e.g. `default` and `tool_use`). +#[derive(Debug, Clone, Deserialize)] +#[serde(untagged)] +pub enum HfChatTemplate { + Single(String), + Named(Vec), +} + +/// One entry of a list-form `chat_template`. +#[derive(Debug, Clone, Deserialize)] +pub struct HfNamedChatTemplate { + pub name: String, + pub template: String, +} + +impl HfChatTemplate { + /// The template Transformers uses when no tools are passed: the single + /// template, or the one named `default`. + pub fn default_template(&self) -> Option<&str> { + match self { + Self::Single(template) => Some(template), + Self::Named(_) => self.named("default"), + } + } + + /// The template Transformers prefers when tools are passed: the one named + /// `tool_use`, if any. + pub fn tool_use_template(&self) -> Option<&str> { + self.named("tool_use") + } + + /// Look up a named template. Transformers collects the list into a dict, + /// so a later entry wins over an earlier one with the same name. + fn named(&self, name: &str) -> Option<&str> { + match self { + Self::Single(_) => None, + Self::Named(templates) => templates + .iter() + .rev() + .find(|template| template.name == name) + .map(|template| template.template.as_str()), + } + } +} + /// Hugging Face named special tokens may be serialized as a string or an /// object carrying the token content. #[derive(Debug, Clone, Deserialize)] @@ -321,7 +368,39 @@ where #[cfg(test)] mod tests { - use super::{GenerationConfig, ModelConfig}; + use super::{GenerationConfig, HfTokenizerConfig, ModelConfig}; + + #[test] + fn tokenizer_config_reads_single_and_named_chat_templates() { + let single: HfTokenizerConfig = + serde_json::from_str(r#"{"chat_template":"single"}"#).unwrap(); + let named: HfTokenizerConfig = serde_json::from_str( + r#"{"chat_template":[ + {"name":"default","template":"old default"}, + {"name":"tool_use","template":"tool use"}, + {"name":"default","template":"default"} + ]}"#, + ) + .unwrap(); + let without_default: HfTokenizerConfig = + serde_json::from_str(r#"{"chat_template":[{"name":"rag","template":"rag"}]}"#).unwrap(); + + let selected = [single, named, without_default].map(|config| { + let template = config.chat_template.unwrap(); + ( + template.default_template().map(str::to_owned), + template.tool_use_template().map(str::to_owned), + ) + }); + assert_eq!( + selected, + [ + (Some("single".to_string()), None), + (Some("default".to_string()), Some("tool use".to_string())), + (None, None), + ] + ); + } #[test] fn generation_config_normalizes_vllm_top_k_values() { diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index 77b47f62a31c..910fbf35ba1f 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -15,8 +15,8 @@ use vllm_tokenizer::{ use self::config::{GenerationConfig, load_generation_config}; pub use self::config::{ - HfSpecialTokens, HfTokenizerConfig, ModelConfig, NamedSpecialToken, load_model_config, - load_tokenizer_config, + HfChatTemplate, HfNamedChatTemplate, HfSpecialTokens, HfTokenizerConfig, ModelConfig, + NamedSpecialToken, load_model_config, load_tokenizer_config, }; pub use self::model_files::{ModelFileSources, ModelSource, ResolvedModelFiles, TokenizerSource}; pub use self::overrides::HfOverrides; From 00d4b782211b890f8ebb7a447bb51049cc771b8c Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 28 Oct 2026 10:23:27 +0800 Subject: [PATCH 026/406] [Rust Frontend] Count gpt-oss reasoning tokens like Python (#63960) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- rust/src/chat/src/output/harmony/mod.rs | 81 +++++++++++-- rust/src/chat/src/output/harmony/tests.rs | 132 +++++++++++++++++++++- 2 files changed, 196 insertions(+), 17 deletions(-) diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs index d4227041a7e3..4e38d8efb9c8 100644 --- a/rust/src/chat/src/output/harmony/mod.rs +++ b/rust/src/chat/src/output/harmony/mod.rs @@ -15,7 +15,7 @@ use vllm_text::output::{DecodedTextEvent, SampledDelta}; use crate::Result as ChatResult; use crate::error::{Error, Result}; -use crate::event::AssistantBlockKind; +use crate::event::{AssistantBlockKind, ChatTokenUsage}; use crate::output::{ AssistantEvent, ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream, generate_tool_call_id, @@ -47,6 +47,8 @@ struct HarmonyGroupKey { struct HarmonyGroup { key: HarmonyGroupKey, text: String, + /// Reasoning tokens attributed to this group's text. + reasoning_tokens: usize, } #[derive(Debug)] @@ -57,6 +59,8 @@ struct OpenHarmonyToolCall { struct HarmonyState { /// Incremental Harmony parser over assistant token IDs. parser: StreamableParser, + /// Harmony encoding, used to tell control tokens from content tokens. + encoding: &'static HarmonyEncoding, /// Whether tool-call content should surface as structured tool events. tool_calls_enabled: bool, /// Count of completed visible assistant messages for newline insertion. @@ -67,6 +71,13 @@ struct HarmonyState { current_text_group: Option, /// The currently open Harmony tool recipient, if any. open_tool_call: Option, + /// Total generated tokens counted as reasoning, reported in the final + /// usage. + reasoning_tokens: usize, + /// Reasoning tokens of the current message not yet attributed to an + /// emitted reasoning delta, such as a token whose content delta is still + /// empty. + unattributed_reasoning_tokens: usize, } impl HarmonyChatOutputProcessor { @@ -119,18 +130,48 @@ impl ChatOutputProcessor for HarmonyChatOutputProcessor { impl HarmonyState { /// Create one fresh Harmony streaming state for a new assistant response. - fn new(encoding: HarmonyEncoding, tool_calls_enabled: bool) -> Result { + fn new(encoding: &'static HarmonyEncoding, tool_calls_enabled: bool) -> Result { Ok(Self { - parser: StreamableParser::new(encoding, Some(Role::Assistant)) + parser: StreamableParser::new(encoding.clone(), Some(Role::Assistant)) .map_err(harmony_output_parsing_error)?, + encoding, tool_calls_enabled, completed_visible_messages: 0, completed_reasoning_messages: 0, current_text_group: None, open_tool_call: None, + reasoning_tokens: 0, + unattributed_reasoning_tokens: 0, }) } + /// Count the token the parser just consumed as reasoning if it is content + /// of an `analysis` message or of a `commentary` message with a recipient, + /// excluding Harmony control tokens, like Python's + /// `HarmonyParser._is_reasoning_token`. + fn count_reasoning_token(&mut self, token_id: u32) { + let is_reasoning_channel = match self.parser.current_channel().as_deref() { + Some("analysis") => true, + Some("commentary") => self.parser.current_recipient().is_some(), + _ => false, + }; + if is_reasoning_channel && !self.encoding.tokenizer().is_special_token(token_id) { + self.reasoning_tokens += 1; + self.unattributed_reasoning_tokens += 1; + } + } + + /// Take the reasoning tokens not yet attributed if `key` is a reasoning + /// text group. + fn take_reasoning_tokens_for(&mut self, key: &HarmonyGroupKey) -> usize { + match text_block_kind(key.channel.as_deref(), key.recipient.as_deref()) { + Some(AssistantBlockKind::Reasoning) => { + std::mem::take(&mut self.unattributed_reasoning_tokens) + } + _ => 0, + } + } + fn process_token_ids(&mut self, token_ids: &[u32]) -> Result> { let mut events = Vec::new(); let mut pending_group: Option = None; @@ -139,6 +180,7 @@ impl HarmonyState { let completed_before = self.parser.messages().len(); self.parser.process(token_id).map_err(harmony_output_parsing_error)?; let completed_after = self.parser.messages().len(); + self.count_reasoning_token(token_id); if let Some(delta) = self .parser @@ -151,14 +193,22 @@ impl HarmonyState { channel: self.parser.current_channel(), recipient: self.parser.current_recipient(), }; + let reasoning_tokens = self.take_reasoning_tokens_for(&key); match pending_group.as_mut() { - Some(group) if group.key == key => group.text.push_str(&delta), + Some(group) if group.key == key => { + group.text.push_str(&delta); + group.reasoning_tokens += reasoning_tokens; + } _ => { if let Some(group) = pending_group.take() { self.emit_group(group, &mut events); } - pending_group = Some(HarmonyGroup { key, text: delta }); + pending_group = Some(HarmonyGroup { + key, + text: delta, + reasoning_tokens, + }); } } } @@ -214,10 +264,12 @@ impl HarmonyState { let final_text = harmony_message_text(final_message); let tail = final_text.strip_prefix(&pending_content).unwrap_or(final_text).to_string(); if !tail.is_empty() { + let reasoning_tokens = self.take_reasoning_tokens_for(&pending_key); self.emit_group( HarmonyGroup { key: pending_key, text: tail, + reasoning_tokens, }, &mut events, ); @@ -258,8 +310,6 @@ impl HarmonyState { events.push(AssistantEvent::TextDelta { kind, delta: "\n".to_string(), - // TODO: measure reasoning tokens from native - // analysis-channel token counts. token_count: None, }); } @@ -267,12 +317,12 @@ impl HarmonyState { self.current_text_group = Some(group.key.clone()); } + let token_count = + (kind == AssistantBlockKind::Reasoning).then_some(group.reasoning_tokens); events.push(AssistantEvent::TextDelta { kind, delta: group.text, - // TODO: measure reasoning tokens from native - // analysis-channel token counts. - token_count: None, + token_count, }); return; } @@ -307,6 +357,10 @@ impl HarmonyState { /// Update newline and open-tool state after one Harmony message completes. fn handle_completed_message(&mut self, key: HarmonyGroupKey) { + // Reasoning tokens of a tool-call message are never attributed to a + // delta; they reach the client through the final usage only. + self.unattributed_reasoning_tokens = 0; + if self.current_text_group.as_ref() == Some(&key) { self.current_text_group = None; } @@ -334,7 +388,7 @@ async fn harmony_assistant_event_stream( tool_calls_enabled: bool, mut y: TryYielder, ) -> Result<()> { - let mut state = HarmonyState::new(encoding.clone(), tool_calls_enabled)?; + let mut state = HarmonyState::new(encoding, tool_calls_enabled)?; futures::pin_mut!(decoded); while let Some(event) = decoded.next().await.transpose()? { @@ -378,7 +432,10 @@ async fn harmony_assistant_event_stream( if let Some(finished) = finished { y.yield_ok(AssistantEvent::Done { - usage: finished.usage.into(), + usage: ChatTokenUsage { + engine: finished.usage, + reasoning_tokens: state.reasoning_tokens, + }, timestamps: finished.timestamps, spec_decode_metrics: finished.spec_decode_metrics, finish_reason: finished.finish_reason, diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index e05a15bfd53a..a985b2960765 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -11,6 +11,7 @@ use vllm_text::output::{ }; use super::*; +use crate::event::AssistantBlockKind; use crate::output::ChatOutputProcessor; use crate::request::{ChatRequest, ChatTool, ChatToolChoice, ResolvedToolContext}; use crate::{AssistantMessageExt, ChatEvent, ChatTokenUsage, FinishReason}; @@ -182,11 +183,14 @@ fn interrupted_analysis_message_is_preserved() { text: "think".to_string(), }], }, - usage: ChatTokenUsage::from(vllm_llm::TokenUsage { - prompt_token_count: 0, - output_token_count: 0, - prompt_cache_usage: Default::default(), - }), + usage: ChatTokenUsage { + engine: vllm_llm::TokenUsage { + prompt_token_count: 0, + output_token_count: 0, + prompt_cache_usage: Default::default(), + }, + reasoning_tokens: 1, + }, timestamps: Default::default(), spec_decode_metrics: None, finish_reason: FinishReason::stop_eos(), @@ -286,6 +290,124 @@ fn tool_calls_stream_arguments_and_finish_with_local_id_shape() { assert_eq!(end.arguments, r#"{"city":"Paris"}"#); } +fn content_token_count(text: &str) -> usize { + harmony_encoding().unwrap().tokenizer().encode_ordinary(text).len() +} + +fn reasoning_token_counts(events: &[ChatEvent]) -> (Vec, usize) { + let deltas = events + .iter() + .filter_map(|event| match event { + ChatEvent::BlockDelta { + kind: AssistantBlockKind::Reasoning, + token_count, + .. + // The newline separating two reasoning messages has no tokens. + } => Some(token_count.unwrap_or_default()), + _ => None, + }) + .collect(); + let ChatEvent::Done { usage, .. } = events.last().unwrap() else { + panic!("expected done"); + }; + (deltas, usage.reasoning_tokens) +} + +#[test] +fn reasoning_tokens_count_analysis_and_tool_call_content() { + let analysis = "Let me check the weather."; + let arguments = r#"{"city":"Paris"}"#; + let tokens = completion_tokens(&[ + text_message("analysis", analysis), + tool_message("get_weather", arguments, "commentary"), + ]); + let events = block_on(collect_events( + HarmonyChatOutputProcessor::new(&request_with_tools()).unwrap(), + vec![ + decoded_start(), + decoded_tokens(tokens, None, Some(finished())), + ], + )); + + // Like Python, content of `analysis` and of `commentary` with a recipient + // counts, Harmony control and header tokens do not. Tool-call tokens are + // only in the final usage. + let (deltas, total) = reasoning_token_counts(&events); + assert_eq!(deltas.iter().sum::(), content_token_count(analysis)); + assert_eq!( + total, + content_token_count(analysis) + content_token_count(arguments) + ); +} + +#[test] +fn reasoning_tokens_exclude_final_and_preamble_content() { + let tokens = completion_tokens(&[ + text_message("commentary", "Let me check."), + text_message("final", "Hello"), + ]); + let events = block_on(collect_events( + HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(), + vec![ + decoded_start(), + decoded_tokens(tokens, None, Some(finished())), + ], + )); + + assert_eq!(reasoning_token_counts(&events), (vec![], 0)); +} + +fn stream_one_token_per_update(messages: &[Message]) -> Vec { + let tokens = completion_tokens(messages); + let last = tokens.len() - 1; + let mut decoded = vec![decoded_start()]; + for (i, token) in tokens.into_iter().enumerate() { + decoded.push(decoded_tokens( + vec![token], + None, + (i == last).then(finished), + )); + } + block_on(collect_events( + HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap(), + decoded, + )) +} + +#[test] +fn reasoning_tokens_are_attributed_to_their_streamed_deltas() { + let first = "Let me think about this."; + let second = "Then answer."; + let events = stream_one_token_per_update(&[ + text_message("analysis", first), + text_message("analysis", second), + text_message("final", "Hi"), + ]); + + // One delta per content token, and the separating newline carries none. + let mut expected = vec![1; content_token_count(first)]; + expected.push(0); + expected.extend(vec![1; content_token_count(second)]); + let total = expected.iter().sum(); + assert_eq!(reasoning_token_counts(&events), (expected, total)); +} + +#[test] +fn reasoning_tokens_of_a_split_character_join_its_delta() { + // Each of these characters spans several byte-level tokens, so the first + // tokens of a character produce no text delta of their own. + let text = "\u{1D518}\u{1D52B}"; + let tokens = content_token_count(text); + assert!(tokens > text.chars().count()); + let events = + stream_one_token_per_update(&[text_message("analysis", text), text_message("final", "Hi")]); + + let (deltas, total) = reasoning_token_counts(&events); + assert_eq!(deltas.len(), text.chars().count()); + assert_eq!(deltas.iter().sum::(), tokens); + assert_eq!(total, tokens); +} + #[test] fn semantic_events_precede_same_update_logprobs() { let tokens = completion_tokens(&[text_message("final", "hello")]); From 4eb15aa369b194dc12034aff165b8f0d8d79e7ca Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 28 Oct 2026 10:30:16 +0800 Subject: [PATCH 027/406] [Rust Frontend] Return 404 for a missing LoRA adapter path like Python (#63961) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- rust/src/server/src/error.rs | 7 +++++- rust/src/server/src/grpc/control.rs | 17 +++++++++++++++ rust/src/server/src/lora.rs | 33 +++++++++++++++++++++++++---- rust/src/server/src/routes/lora.rs | 19 +++++++++++++++++ 4 files changed, 71 insertions(+), 5 deletions(-) diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index b53ca55a66e4..802741ed7af2 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -21,6 +21,8 @@ pub enum ApiError { }, /// The requested model name does not match the single configured model. ModelNotFound { model: String }, + /// A resource other than the model, such as a LoRA adapter, does not exist. + NotFound { message: String }, /// The request body could not be parsed as valid JSON. JsonParseError { message: String }, /// The requested operation conflicts with current server state. @@ -36,7 +38,7 @@ impl ApiError { pub fn status_code(&self) -> StatusCode { match self { Self::InvalidRequest { .. } => StatusCode::BAD_REQUEST, - Self::ModelNotFound { .. } => StatusCode::NOT_FOUND, + Self::ModelNotFound { .. } | Self::NotFound { .. } => StatusCode::NOT_FOUND, Self::Conflict { .. } => StatusCode::CONFLICT, Self::ServerError { .. } => StatusCode::INTERNAL_SERVER_ERROR, Self::JsonParseError { .. } => StatusCode::BAD_REQUEST, @@ -58,6 +60,7 @@ impl ApiError { format!("The model `{model}` does not exist."), Some("model".to_string()), ), + Self::NotFound { message } => ("NotFoundError", message.clone(), None), // Python reports a body that fails to parse into the request // model through FastAPI's validation handler, which names the // HTTP status instead of an error class. @@ -141,6 +144,7 @@ mod tests { let errors = [ ApiError::invalid_request("top_p must be in (0, 1].", Some("top_p")), ApiError::model_not_found("missing"), + ApiError::not_found("No adapter found for /not/an/adapter"), ApiError::json_parse_error("expected value at line 1 column 1"), ApiError::conflict("A profiling session is already active."), ApiError::server_error("Internal server error"), @@ -153,6 +157,7 @@ mod tests { expect![[r#" {"error":{"message":"top_p must be in (0, 1].","type":"BadRequestError","param":"top_p","code":400}} {"error":{"message":"The model `missing` does not exist.","type":"NotFoundError","param":"model","code":404}} + {"error":{"message":"No adapter found for /not/an/adapter","type":"NotFoundError","param":null,"code":404}} {"error":{"message":"expected value at line 1 column 1","type":"Bad Request","param":null,"code":400}} {"error":{"message":"A profiling session is already active.","type":"Conflict","param":null,"code":409}} {"error":{"message":"Internal server error","type":"InternalServerError","param":null,"code":500}} diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index 482a7d1d6301..85ef55779d72 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -159,6 +159,7 @@ fn load_lora_status(error: LoadLoraError) -> Status { let code = match &error { LoadLoraError::Disabled(_) => Code::FailedPrecondition, LoadLoraError::InvalidAdapter { .. } => Code::InvalidArgument, + LoadLoraError::PathAccess(LoraPathAccessError::NotFound { .. }) => Code::NotFound, LoadLoraError::PathAccess(LoraPathAccessError::InvalidPath { .. }) => Code::InvalidArgument, LoadLoraError::PathAccess(LoraPathAccessError::InvalidConfiguration { .. }) => { Code::Internal @@ -539,3 +540,19 @@ pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option LoraPathAccessError::NotFound { + lora_path: lora_path.to_string(), + }, + _ => invalid_path!("Local LoRA adapter path must exist and be accessible."), + })?; let canonical_prefixes = allowed_prefixes .iter() .map(|prefix| { @@ -347,7 +353,7 @@ mod tests { use std::path::PathBuf; use std::time::{SystemTime, UNIX_EPOCH}; - use super::validate_lora_path_access; + use super::{LoraPathAccessError, validate_lora_path_access}; fn temp_lora_dir(test_name: &str) -> PathBuf { let suffix = SystemTime::now() @@ -417,6 +423,25 @@ mod tests { fs::remove_dir_all(root).ok(); } + /// Python answers a missing adapter with 404 `NotFoundError`, so the + /// guard reports a missing path apart from an inaccessible one. + #[test] + fn lora_path_reports_missing_paths_as_not_found() { + let root = temp_lora_dir("missing-path"); + let missing = root.join("missing-adapter"); + + let prefixes = [root.clone()]; + let error = + validate_lora_path_access(missing.to_str().expect("utf-8 temp path"), Some(&prefixes)) + .expect_err("missing adapter path should be rejected"); + assert!( + matches!(&error, LoraPathAccessError::NotFound { lora_path } if *lora_path == missing.to_str().unwrap()), + "{error:?}" + ); + + fs::remove_dir_all(root).ok(); + } + #[test] fn lora_path_rejects_parent_escape_from_configured_prefixes() { let root = temp_lora_dir("parent-escape"); diff --git a/rust/src/server/src/routes/lora.rs b/rust/src/server/src/routes/lora.rs index ac0cb11e1a77..ec7b37a138b0 100644 --- a/rust/src/server/src/routes/lora.rs +++ b/rust/src/server/src/routes/lora.rs @@ -41,6 +41,9 @@ fn load_lora_api_error(error: LoadLoraError) -> ApiError { match error { LoadLoraError::Disabled(_) => ApiError::invalid_request(message, None), LoadLoraError::InvalidAdapter { .. } => ApiError::invalid_request(message, None), + LoadLoraError::PathAccess(LoraPathAccessError::NotFound { .. }) => { + ApiError::not_found(message) + } LoadLoraError::PathAccess(LoraPathAccessError::InvalidPath { .. }) => { ApiError::invalid_request(message, Some("lora_path")) } @@ -120,3 +123,19 @@ pub async fn unload_lora_adapter( lora_request.lora_name )) } + +#[cfg(test)] +mod tests { + use axum::http::StatusCode; + + use super::load_lora_api_error; + use crate::lora::{LoadLoraError, LoraPathAccessError}; + + #[test] + fn load_lora_reports_missing_adapter_path_as_not_found() { + let error = load_lora_api_error(LoadLoraError::PathAccess(LoraPathAccessError::NotFound { + lora_path: "/not/an/adapter".to_string(), + })); + assert_eq!(error.status_code(), StatusCode::NOT_FOUND); + } +} From 9df11cf6089c3ec2361b5220927770e5867afb0a Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 28 Oct 2026 10:31:05 +0800 Subject: [PATCH 028/406] [Rust Frontend] Export vllm:engine_sleep_state like Python (#63959) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 --- .../examples/external_engine_utility_call.rs | 5 +- rust/src/engine-core-client/src/client.rs | 16 +- rust/src/engine-core-client/src/client/imp.rs | 12 +- rust/src/engine-core-client/src/metrics.rs | 57 ++++++- .../src/protocol/utility.rs | 37 +++++ .../engine-core-client/src/tests/client.rs | 145 +++++++++++++++++- rust/src/metrics/src/scheduler.rs | 24 +++ rust/src/server/src/grpc/control.rs | 7 +- rust/src/server/src/grpc/tests.rs | 39 +++++ rust/src/server/src/routes/sleep.rs | 16 +- 10 files changed, 334 insertions(+), 24 deletions(-) diff --git a/rust/src/engine-core-client/examples/external_engine_utility_call.rs b/rust/src/engine-core-client/examples/external_engine_utility_call.rs index 8b73a12ce2f6..280f480de795 100644 --- a/rust/src/engine-core-client/examples/external_engine_utility_call.rs +++ b/rust/src/engine-core-client/examples/external_engine_utility_call.rs @@ -6,7 +6,7 @@ use std::time::Duration; use anyhow::{Context, Result, bail}; use clap::Parser; use tracing_subscriber::EnvFilter; -use vllm_engine_core_client::protocol::utility::PauseMode; +use vllm_engine_core_client::protocol::utility::{PauseMode, SleepLevel}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode}; #[derive(Debug, Parser)] @@ -111,7 +111,8 @@ async fn main() -> Result<()> { if args.skip_sleep_wake { println!("sleep_wake=skipped"); } else { - client.sleep(args.sleep_level, args.sleep_mode).await.with_context(|| { + let sleep_level = SleepLevel::try_from(args.sleep_level).map_err(anyhow::Error::msg)?; + client.sleep(sleep_level, args.sleep_mode).await.with_context(|| { format!( "failed to call sleep utility with level={} mode={}", args.sleep_level, args.sleep_mode diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 6eb9b9197260..fddf24306664 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -17,11 +17,12 @@ use tracing::{debug, info, trace}; use crate::client::imp::{ClientInner, run_abort_loop, run_output_dispatcher_loop}; use crate::coordinator::CoordinatorHandle; use crate::error::{Error, Result, bail_invalid_client_config}; +use crate::metrics::EngineSleepState; use crate::protocol::dtype::ModelDtype; use crate::protocol::handshake::EngineCoreReadyResponse; use crate::protocol::lora::LoraRequest; use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; -use crate::protocol::utility::{CheckpointState, EngineCoreUtilityRequest, PauseMode}; +use crate::protocol::utility::{CheckpointState, EngineCoreUtilityRequest, PauseMode, SleepLevel}; use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime}; use crate::transport::{self, ConnectedEngine}; @@ -953,25 +954,32 @@ impl EngineCoreClient { } /// Put the engine to sleep. - pub async fn sleep(&self, level: u32, mode: PauseMode) -> Result<()> { + pub async fn sleep(&self, level: SleepLevel, mode: PauseMode) -> Result<()> { self.call_utility::<(), _>("sleep", (level, mode)).await?; + self.inner.record_sleep_state(EngineSleepState::Asleep(level)); Ok(()) } /// Release KV cache memory while keeping model weights resident. pub async fn release_kv_cache_memory(&self) -> Result<()> { self.call_utility::<(), _>("release_kv_cache_memory", ()).await?; + self.inner.record_sleep_state(EngineSleepState::KvCacheReleased); Ok(()) } /// Wake the engine from sleep, optionally limiting the wake-up to specific /// tags, and return whether every engine is fully awake. pub async fn wake_up(&self, tags: Option>) -> Result { - Ok(self + let fully_awake = self .call_utility::("wake_up", (tags,)) .await? .into_iter() - .all(|fully_awake| fully_awake)) + .all(|fully_awake| fully_awake); + // A partial wake-up keeps the previous sleep state, as in Python. + if fully_awake { + self.inner.record_sleep_state(EngineSleepState::Awake); + } + Ok(fully_awake) } /// Pause generation (e.g. for weight updates). In `abort` and `wait` mode diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index fa82d9591aa8..f2b8e6f910a7 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -22,7 +22,9 @@ use crate::client::state::{ use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; -use crate::metrics::{IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder}; +use crate::metrics::{ + EngineSleepState, IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder, +}; use crate::protocol::encode_msgpack; use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs}; use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; @@ -103,6 +105,14 @@ impl ClientInner { &self.model_name } + /// Record a completed sleep or wake transition in + /// `vllm:engine_sleep_state`, exported only when engines record stats. + pub fn record_sleep_state(&self, state: EngineSleepState) { + if let Some(recorder) = &self.scheduler_stats_recorder { + recorder.record_sleep_state(state); + } + } + /// Register a newly added request. Return the selected engine id and the /// per-request output channel bound to its `request_id`. /// diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index dd2835972fdc..041e163c4d6a 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -6,8 +6,8 @@ use std::collections::BTreeSet; use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ - CacheHitSourceLabels, EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, - LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily, + CacheHitSourceLabels, EngineLabels, EnginePositionLabels, EngineSleepStateLabels, F64Gauge, + Family, HistogramMetric, LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, WaitingReasonLabels, }; @@ -16,11 +16,39 @@ use crate::protocol::stats::{ CachedTokensBySource, KvConnectorStats, MooncakeStats, MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats, }; +use crate::protocol::utility::SleepLevel; use crate::transport::ConnectedEngine; const WAITING_REASON_CAPACITY: &str = "capacity"; const WAITING_REASON_DEFERRED: &str = "deferred"; +/// `sleep_state` label values of `vllm:engine_sleep_state`, in the order of +/// [`EngineSleepState::gauge_values`]. +const SLEEP_STATES: [&str; 3] = ["awake", "weights_offloaded", "discard_all"]; + +/// Engine sleep state exported as `vllm:engine_sleep_state`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum EngineSleepState { + Awake, + Asleep(SleepLevel), + /// After `release_kv_cache_memory`, which Python records like a level-0 + /// sleep. + KvCacheReleased, +} + +impl EngineSleepState { + /// Gauge values for the `awake`, `weights_offloaded` and `discard_all` + /// series, as Python's `PrometheusStatLogger.record_sleep_state` sets them. + fn gauge_values(self) -> [u64; 3] { + match self { + Self::Awake => [1, 0, 0], + Self::Asleep(SleepLevel::PauseOnly) | Self::KvCacheReleased => [0, 0, 0], + Self::Asleep(SleepLevel::OffloadWeights) => [0, 1, 0], + Self::Asleep(SleepLevel::DiscardAll) => [0, 0, 1], + } + } +} + /// Cached output-batch metric handles for one model and engine index. pub(crate) struct IterationMetricHandles { pub iteration_tokens_total: HistogramMetric, @@ -55,6 +83,8 @@ struct SchedulerStatsHandles { scheduler_waiting_capacity: U64Gauge, scheduler_waiting_deferred: U64Gauge, kv_cache_usage: F64Gauge, + /// One handle per [`SLEEP_STATES`] entry. + engine_sleep_state: [U64Gauge; 3], // Prefix-cache counters, including the connector-backed external cache path. prefix_cache_queries: U64Counter, @@ -129,7 +159,21 @@ impl SchedulerStatsRecorder { }) .collect(); - Self { engines } + let recorder = Self { engines }; + // Engines start awake, as Python sets when it creates the stat logger. + recorder.record_sleep_state(EngineSleepState::Awake); + recorder + } + + /// Record the sleep state of every connected engine. Sleep and wake are + /// utility calls on all engines, so the state is not per engine. + pub(crate) fn record_sleep_state(&self, state: EngineSleepState) { + let values = state.gauge_values(); + for handles in self.engines.values() { + for (gauge, value) in handles.engine_sleep_state.iter().zip(values) { + gauge.set(value); + } + } } /// Record one scheduler-stats payload for the given engine index. @@ -171,6 +215,13 @@ fn resolve_scheduler_stats_handles( .scheduler_waiting_by_reason .get_or_create_owned(&deferred), kv_cache_usage: metrics.kv_cache_usage.get_or_create_owned(&labels), + engine_sleep_state: SLEEP_STATES.map(|sleep_state| { + metrics.engine_sleep_state.get_or_create_owned(&EngineSleepStateLabels { + model_name: model_name.to_string(), + engine, + sleep_state, + }) + }), prefix_cache_queries: metrics.prefix_cache_queries.get_or_create_owned(&labels), prefix_cache_hits: metrics.prefix_cache_hits.get_or_create_owned(&labels), external_prefix_cache_queries: metrics diff --git a/rust/src/engine-core-client/src/protocol/utility.rs b/rust/src/engine-core-client/src/protocol/utility.rs index 5379ff49e9e8..97a880c3acf2 100644 --- a/rust/src/engine-core-client/src/protocol/utility.rs +++ b/rust/src/engine-core-client/src/protocol/utility.rs @@ -118,6 +118,43 @@ impl FromStr for CheckpointState { } } +/// Level of the `sleep` utility call, sent as Python's integer level. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(into = "u32", try_from = "u32")] +pub enum SleepLevel { + /// Level 0: pause scheduling only, without changing GPU memory. + PauseOnly, + /// Level 1: offload model weights to CPU and discard the KV cache. + OffloadWeights, + /// Level 2: discard all GPU memory. + DiscardAll, +} + +impl From for u32 { + fn from(level: SleepLevel) -> Self { + match level { + SleepLevel::PauseOnly => 0, + SleepLevel::OffloadWeights => 1, + SleepLevel::DiscardAll => 2, + } + } +} + +impl TryFrom for SleepLevel { + type Error = String; + + fn try_from(level: u32) -> std::result::Result { + match level { + 0 => Ok(Self::PauseOnly), + 1 => Ok(Self::OffloadWeights), + 2 => Ok(Self::DiscardAll), + other => Err(format!( + "invalid sleep level `{other}`; expected one of: 0, 1, 2" + )), + } + } +} + /// Utility call id as carried on the engine-core MessagePack wire. /// /// Python emits utility ids as MessagePack integers, including values that may diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index f483892d4c94..b47dc08a20a0 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -38,7 +38,7 @@ use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; use crate::protocol::sampling::EngineCoreSamplingParams; use crate::protocol::stats::{KvConnectorStats, MooncakeOperation, SchedulerStats}; use crate::protocol::tensor::{WireArrayData, WireTensor}; -use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; +use crate::protocol::utility::{PauseMode, SleepLevel, UtilityOutput, UtilityResultEnvelope}; use crate::test_utils::{ IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets, setup_mock_engine_with_init, spawn_mock_engine_task, @@ -2416,6 +2416,149 @@ async fn wake_up_requires_every_engine_to_be_fully_awake() { client.shutdown().await.unwrap(); } +/// Render the `vllm:engine_sleep_state` series of one model as +/// `awake,weights_offloaded,discard_all` for engines 0 and 1. +fn engine_sleep_state_series(model_name: &str) -> [String; 2] { + let rendered = vllm_metrics::METRICS.render().unwrap(); + [0, 1].map(|engine| { + ["awake", "weights_offloaded", "discard_all"] + .map(|state| { + let prefix = format!( + "vllm:engine_sleep_state{{model_name=\"{model_name}\",engine=\"{engine}\",sleep_state=\"{state}\"}} " + ); + rendered + .lines() + .find_map(|line| line.strip_prefix(&prefix)) + .unwrap_or("missing") + .to_owned() + }) + .join(",") + }) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sleep_and_wake_calls_record_engine_sleep_state() { + init_tracing(); + // A model name of its own keeps other tests' clients off these series. + const MODEL: &str = "engine-sleep-state-model"; + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let mut engines = Vec::new(); + for engine_index in 0..2_u16 { + // Each engine answers the calls in order; `None` stands for a `()` + // result. On the partial wake-up only engine 1 reports fully awake. + let answers = [ + ("sleep", None), + ("wake_up", Some(engine_index == 1)), + ("wake_up", Some(true)), + ("release_kv_cache_memory", None), + ("wake_up", Some(true)), + ("sleep", None), + ]; + engines.push(spawn_mock_engine_task( + handshake_address.clone(), + EngineId::from_engine_index(engine_index).into_frame().to_vec(), + move |dealer, push| { + Box::pin(async move { + for (method, fully_awake) in answers { + let utility = recv_engine_message(dealer).await; + let payload = decode_value(&utility[1]); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call_id"); + assert_eq!(array[2], Value::from(method)); + send_outputs( + push, + UtilityCallOutput { + engine_index: u32::from(engine_index), + timestamp: 0.0, + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: fully_awake.map(utility_result_value), + }, + } + .into(), + ) + .await; + } + }) + }, + )); + } + let client = connect_client_with_ipc( + handshake_test_config(handshake_address, 2, MODEL, Duration::from_secs(2), 5, None), + &ipc, + ) + .await; + + let mut series = vec![("connected", engine_sleep_state_series(MODEL))]; + client.sleep(SleepLevel::OffloadWeights, PauseMode::default()).await.unwrap(); + series.push(("sleep level 1", engine_sleep_state_series(MODEL))); + assert!(!client.wake_up(Some(vec!["weights".to_owned()])).await.unwrap()); + series.push(("partial wake", engine_sleep_state_series(MODEL))); + assert!(client.wake_up(None).await.unwrap()); + series.push(("full wake", engine_sleep_state_series(MODEL))); + client.release_kv_cache_memory().await.unwrap(); + series.push(("release kv cache", engine_sleep_state_series(MODEL))); + assert!(client.wake_up(None).await.unwrap()); + client.sleep(SleepLevel::DiscardAll, PauseMode::default()).await.unwrap(); + series.push(("sleep level 2", engine_sleep_state_series(MODEL))); + + expect_test::expect![[r#" + [ + ( + "connected", + [ + "1,0,0", + "1,0,0", + ], + ), + ( + "sleep level 1", + [ + "0,1,0", + "0,1,0", + ], + ), + ( + "partial wake", + [ + "0,1,0", + "0,1,0", + ], + ), + ( + "full wake", + [ + "1,0,0", + "1,0,0", + ], + ), + ( + "release kv cache", + [ + "0,0,0", + "0,0,0", + ], + ), + ( + "sleep level 2", + [ + "0,0,1", + "0,0,1", + ], + ), + ] + "#]] + .assert_debug_eq(&series); + + for (shutdown_tx, engine_task) in engines { + let _ = shutdown_tx.send(()); + engine_task.await.unwrap(); + } + client.shutdown().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn collective_rpc_flattens_results_from_all_engines() { init_tracing(); diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs index 05392a2cf56e..d5327a8cea0b 100644 --- a/rust/src/metrics/src/scheduler.rs +++ b/rust/src/metrics/src/scheduler.rs @@ -91,6 +91,15 @@ pub struct WaitingReasonLabels { pub reason: &'static str, } +/// Labels for `vllm:engine_sleep_state`: one series per engine and state, +/// `awake`, `weights_offloaded` (sleep level 1) or `discard_all` (level 2). +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] +pub struct EngineSleepStateLabels { + pub model_name: String, + pub engine: u32, + pub sleep_state: &'static str, +} + /// Labels for per-operation Mooncake store connector telemetry. `operation` /// and `status` are dynamic (not known ahead of time), unlike the other /// fixed-label families in this module. @@ -218,6 +227,10 @@ pub struct SchedulerMetrics { pub scheduler_waiting_by_reason: Family, pub kv_cache_usage: Family, + /// `vllm:engine_sleep_state`, set by the engine client's sleep and wake + /// calls rather than from `SchedulerStats`. + pub engine_sleep_state: Family, + /// `vllm:lora_requests_info`. Value is the emit-time unix timestamp in /// seconds. pub lora_info: Family, @@ -310,6 +323,16 @@ impl SchedulerMetrics { kv_cache_usage.clone(), ); + let engine_sleep_state = Family::default(); + registry.register( + "vllm:engine_sleep_state", + "Engine sleep state; awake = 0 means engine is sleeping; \ + awake = 1 means engine is awake; \ + weights_offloaded = 1 means sleep level 1; \ + discard_all = 1 means sleep level 2.", + engine_sleep_state.clone(), + ); + let lora_info = Family::default(); registry.register( "vllm:lora_requests_info", @@ -596,6 +619,7 @@ impl SchedulerMetrics { scheduler_waiting, scheduler_waiting_by_reason, kv_cache_usage, + engine_sleep_state, lora_info, prefix_cache_queries, prefix_cache_hits, diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index 85ef55779d72..d522c7cc8e35 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -10,7 +10,7 @@ use tokio_util::sync::CancellationToken; use tonic::{Code, Request, Response, Status}; use vllm_engine_core_client::EngineCoreClient; use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; -use vllm_engine_core_client::protocol::utility::PauseMode as EnginePauseMode; +use vllm_engine_core_client::protocol::utility::{PauseMode as EnginePauseMode, SleepLevel}; use super::{ControlServer, pb}; use crate::config::LoraModulePath; @@ -378,7 +378,10 @@ impl pb::control_server::Control for ControlServiceImpl { self.require_sleep_mode()?; let request = request.into_inner(); let mode = pause_mode(request.mode)?; - let level = request.level.unwrap_or(1); + let level = match request.level { + Some(level) => SleepLevel::try_from(level).map_err(Status::invalid_argument)?, + None => SleepLevel::OffloadWeights, + }; let _guard = self.rl_lock.lock().await; self.client() .sleep(level, mode) diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index d1590c2129e6..b4c0cfebca92 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -2392,6 +2392,45 @@ async fn control_list_loras_requires_lora_enabled_engine() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn control_sleep_rejects_unknown_level() { + let mut ready = default_ready_response(); + ready.enable_sleep_mode = true; + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service_with_engine_script( + b"engine-grpc-sleep-level".to_vec(), + ready, + Arc::new(FakeTextBackend::default()), + None, + |_, _| boxed_test_future(async move {}), + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut control_client = ControlClient::new(channel); + let status = control_client + .sleep(pb::SleepRequest { + level: Some(3), + mode: pb::PauseMode::Unspecified.into(), + }) + .await + .expect_err("level 3 should be rejected"); + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert_eq!( + status.message(), + "invalid sleep level `3`; expected one of: 0, 1, 2" + ); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_forwards_weight_update_without_pause_guard() { diff --git a/rust/src/server/src/routes/sleep.rs b/rust/src/server/src/routes/sleep.rs index b6239dc2d4b5..33b78169f55b 100644 --- a/rust/src/server/src/routes/sleep.rs +++ b/rust/src/server/src/routes/sleep.rs @@ -7,7 +7,7 @@ use axum::Json; use axum::extract::rejection::QueryRejection; use axum::extract::{Query, RawQuery, State}; use serde::{Deserialize, Serialize}; -use vllm_engine_core_client::protocol::utility::PauseMode; +use vllm_engine_core_client::protocol::utility::{PauseMode, SleepLevel}; use vllm_metrics::METRICS; use crate::error::ApiError; @@ -22,7 +22,7 @@ pub(crate) struct IsSleepingResponse { #[derive(Serialize)] pub struct SleepResponse { status: &'static str, - level: u32, + level: SleepLevel, } #[derive(Serialize)] @@ -40,13 +40,13 @@ pub struct WakeUpResponse { #[derive(Debug, Deserialize)] pub(crate) struct SleepParams { #[serde(default = "default_sleep_level")] - level: u32, + level: SleepLevel, #[serde(default)] mode: PauseMode, } -const fn default_sleep_level() -> u32 { - 1 +const fn default_sleep_level() -> SleepLevel { + SleepLevel::OffloadWeights } fn invalid_query(error: QueryRejection) -> ApiError { @@ -59,12 +59,6 @@ pub async fn sleep( params: Result, QueryRejection>, ) -> Result, ApiError> { let Query(SleepParams { level, mode }) = params.map_err(invalid_query)?; - if level > 2 { - return Err(ApiError::invalid_request( - "level must be between 0 and 2", - Some("level"), - )); - } let _recorder = METRICS.api_server.record_sleep_mode_operation("sleep"); state From 6282c4c5cf7c58b1b0dfeefb0ff78bc32e6d8a9e Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 27 Oct 2026 19:33:38 -0700 Subject: [PATCH 029/406] [Bench] Default to Zipf routing for forward benchmarks with dummy weights (#64167) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- .agents/skills/model-microbenchmark/SKILL.md | 5 +- docs/benchmarking/cli.md | 23 ++++++ docs/cli/README.md | 6 +- tests/benchmarks/test_forward_cli.py | 38 ++++++++-- vllm/benchmarks/forward/main.py | 17 +++-- vllm/benchmarks/forward/routing.py | 77 +++++++++++++++----- 6 files changed, 130 insertions(+), 36 deletions(-) diff --git a/.agents/skills/model-microbenchmark/SKILL.md b/.agents/skills/model-microbenchmark/SKILL.md index 6da792aba4de..b1bfabeb578f 100644 --- a/.agents/skills/model-microbenchmark/SKILL.md +++ b/.agents/skills/model-microbenchmark/SKILL.md @@ -45,7 +45,8 @@ Key flags (plus all `EngineArgs`): - `--kv-len-jitter 0.5` (default): prefixes drawn within ±50%, same on every DP rank, capped at `max_model_len`. Use `0` for exact lengths. -- `--routing auto|real|balanced`: `auto` = balanced with dummy weights. +- `--routing auto|real|balanced|zipf`: `auto` = Zipf with dummy weights; + `zipf` uses fixed alpha=1.25 skew (target only, no PCP or grouped top-k). - `--iters 20 --samples 5`, `--output-json out.json` (per-rank latency stats, `samples_us`, graph-vs-eager check). - Defaults: `--load-format dummy`, prefix caching off, no torch.compile. @@ -95,7 +96,7 @@ that the JSON summary excludes. ## Caveats - **Dummy weights lie for data-dependent kernels.** Dummy gates route every - token to the same experts (hence balanced routing). Cross-check suspicious + token to the same experts (hence synthetic routing). Cross-check suspicious kernels with real weights. - **torch.compile models (e.g. Qwen3/3.5) look slow**: no compile here. Use them for correctness only; trust perf for models under `vllm/models/`. diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 1ea961fefa4f..533b89816971 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -14,6 +14,29 @@ It's a living document, updated as new features and datasets become available. For exact-batch model timing, including steady-state pipeline-parallel forwards, see [vllm bench forward](../cli/README.md#forward). +For skewed synthetic MoE expert loads, use `--routing zipf`: + +```bash +vllm bench forward --model Qwen/Qwen3-30B-A3B \ + --routing zipf --case decode=6*1@8192 --case prefill=2048 +``` + +Each token independently selects distinct experts without replacement, weighted +by popularity rank to the power of `-1.25`. Expert popularity is randomly +permuted per layer to avoid systematically assigning hot experts to the lowest +EP ranks. +The fixture is seeded by `--seed`, fixed across timing replays, and generated +outside the timed forward. Sequence-parallel TP ranks use their corresponding +token rows. Zipf allows EP load imbalance; `--routing balanced` gives equal +load per EP rank. Neither synthetic mode supports grouped top-k routing; +Zipf also does not support prefill context parallelism. + +The default `--routing auto` selects Zipf with dummy weights and real routing +otherwise. When the model or parallel layout does not support the synthetic +fixture, `auto` warns and falls back to real routing. Use `--routing balanced` +to explicitly request equal EP load. Zipf is a synthetic workload, not a fit +to a particular model's routing or speculative-token correlations. + ## Attention Kernel Benchmark Run the attention kernel smoke benchmark from the repository root using the diff --git a/docs/cli/README.md b/docs/cli/README.md index a8880f493d3c..116cd12f9c58 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -183,8 +183,10 @@ Use `--exclude-drafter` to measure only the target forward. Verification assumes all draft tokens are accepted and uses a fixed bonus token; target sampling and rejection sampling are excluded. Both paths use CUDA graphs and restore the batch's state before each replay. -Balanced MoE routing fixtures apply to the target model; the drafter uses its -model routing, so use real weights when measuring a MoE drafter. +Synthetic MoE routing fixtures (`--routing balanced` or `--routing zipf`) apply +to the target model; the drafter uses its model routing, so use real weights +when measuring a MoE drafter. For the Zipf distribution and its limitations, +see [Model Forward Benchmark](../benchmarking/cli.md#model-forward-benchmark). ```bash vllm bench forward --model Qwen/Qwen3-1.7B \ diff --git a/tests/benchmarks/test_forward_cli.py b/tests/benchmarks/test_forward_cli.py index 32e1f427fc4a..9f1012414870 100644 --- a/tests/benchmarks/test_forward_cli.py +++ b/tests/benchmarks/test_forward_cli.py @@ -16,7 +16,7 @@ from vllm.benchmarks.forward.kv_init import touched_kv_views, zero_kv_blocks from vllm.benchmarks.forward.main import run_case from vllm.benchmarks.forward.profile import summarize -from vllm.benchmarks.forward.routing import BalancedRouting +from vllm.benchmarks.forward.routing import RoutingFixture, zipf_topk_ids from vllm.benchmarks.forward.shapes import jitter_kv_lens, parse_case from vllm.platforms import current_platform from vllm.v1.attention.backends.mla.cpu_mla import CPUMLABackend @@ -125,6 +125,8 @@ def test_bench_forward_muse_balanced_routing(tmp_path): def test_balanced_routing_restores_expert_bias(dist_init): """Fixture routes stay balanced despite selection bias, then restore it.""" + from vllm.config import get_current_vllm_config + from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.models.muse_spark.common.moe import MuseSpark1xRouter @@ -144,13 +146,14 @@ def test_balanced_routing_restores_expert_bias(dist_init): hidden = torch.zeros(64, 4, device=device) with torch.no_grad(): gate.weight.zero_() - fixture = BalancedRouting(layer, 64, 0, 1, 4, 0, device) + fixture = RoutingFixture(layer, 64, 0, 1, 4, 0, device, mode="balanced") try: - logits, _ = gate(hidden) - _, ids = selector.select_experts(hidden, logits) - fixture.check_applied() - counts = torch.bincount((ids // 2).flatten(), minlength=4) - assert counts.min() == counts.max() + with set_forward_context(None, get_current_vllm_config(), num_tokens=64): + logits, _ = gate(hidden) + _, ids = selector.select_experts(hidden, logits) + fixture.check_applied() + counts = torch.bincount((ids // 2).flatten(), minlength=4) + assert counts.min() == counts.max() finally: fixture.remove() torch.testing.assert_close(layer.expert_bias, bias, rtol=0, atol=0) @@ -163,7 +166,8 @@ def test_balanced_routing_restores_expert_bias(dist_init): @pytest.mark.skipif( not current_platform.is_cuda(), reason="Forward graph timing requires CUDA." ) -def test_bench_forward_hybrid_verify(tmp_path): +@pytest.mark.parametrize("routing", ["auto", "balanced"]) +def test_bench_forward_hybrid_verify(tmp_path, routing): """Repeated target verification starts from the same recurrent state.""" config = dict( architectures=["Qwen3NextForCausalLM"], @@ -204,6 +208,8 @@ def test_bench_forward_hybrid_verify(tmp_path): "--model", str(tmp_path), "--skip-tokenizer-init", + "--routing", + routing, "--exclude-drafter", "--speculative-config", '{"method":"ngram","num_speculative_tokens":5}', @@ -239,11 +245,27 @@ def test_bench_forward_hybrid_verify(tmp_path): assert [row["case"] for row in rows] == ["verify", "mixed"] for row in rows: (rank,) = row["ranks"] + assert rank["routing"] == ("zipf" if routing == "auto" else routing) assert rank["finite"] and rank["graph_deterministic"] assert rank["timing_scope"] == "target_only" assert row["median_us"] > 0 +def test_zipf_expert_selection(): + """Popularity follows Zipf, with distinct experts and reproducible rows.""" + rng = np.random.default_rng(0) + ids = zipf_topk_ids(50000, 32, 1, rng) + frequencies = np.sort(np.bincount(ids.ravel(), minlength=32))[::-1] / len(ids) + expected = np.arange(1, 33, dtype=float) ** -1.25 + expected /= expected.sum() + np.testing.assert_allclose(frequencies, expected, atol=0.005, rtol=0) + + ids = zipf_topk_ids(2048, 384, 6, np.random.default_rng(0)) + assert ((ids >= 0) & (ids < 384)).all() + assert (np.diff(np.sort(ids, axis=1), axis=1) > 0).all() + assert np.array_equal(ids[:6], zipf_topk_ids(6, 384, 6, np.random.default_rng(0))) + + @pytest.mark.benchmark @pytest.mark.skipif( not current_platform.is_cuda(), reason="Forward graph timing requires CUDA." diff --git a/vllm/benchmarks/forward/main.py b/vllm/benchmarks/forward/main.py index 57d37a7b2c68..22622e6697bf 100644 --- a/vllm/benchmarks/forward/main.py +++ b/vllm/benchmarks/forward/main.py @@ -66,10 +66,13 @@ def add_cli_args(parser: FlexibleArgumentParser) -> None: ) parser.add_argument( "--routing", - choices=["auto", "real", "balanced"], + choices=["auto", "real", "balanced", "zipf"], default="auto", help="MoE routing. balanced: fixed routing with equal load per EP rank. " - "auto: balanced with dummy weights, else real.", + "zipf: fixed weighted sampling without replacement (alpha=1.25), " + "with expert popularity shuffled per layer. " + "auto: zipf with dummy weights, else real; falls back to real when " + "the synthetic fixture is unsupported.", ) parser.add_argument( "--exclude-drafter", @@ -209,17 +212,17 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any profile_prefix, include_drafter=include_drafter, ) - routing, balanced = install_routing( + routing, fixture = install_routing( runner, args.routing, args.load_format == "dummy", batch.num_tokens, args.seed ) try: result = bench.run() - if balanced is not None: - balanced.check_applied() + if fixture is not None: + fixture.check_applied() finally: runner.execute_model(batch.finish()) - if balanced is not None: - balanced.remove() + if fixture is not None: + fixture.remove() return { "num_tokens": batch.num_tokens, "kv_blocks": batch.num_blocks, diff --git a/vllm/benchmarks/forward/routing.py b/vllm/benchmarks/forward/routing.py index ad9589835578..e68c1344a1de 100644 --- a/vllm/benchmarks/forward/routing.py +++ b/vllm/benchmarks/forward/routing.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Balanced MoE routing for benchmarks with dummy weights. +"""Synthetic MoE routing for forward benchmarks. Dummy gate weights give degenerate routing, which skews MoE and all-to-all cost. The routing still runs; its result is replaced by the fixture at @@ -18,10 +18,13 @@ """ import inspect +import zlib import numpy as np import torch +from vllm.distributed.parallel_state import get_tp_group +from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear @@ -54,6 +57,34 @@ def balanced_topk_ids( return ranks * local + (offset + repeat) % local +def zipf_topk_ids( + num_tokens: int, num_experts: int, top_k: int, rng: np.random.Generator +) -> np.ndarray: + """Sample distinct experts per token with Zipf(alpha=1.25) popularity.""" + expert_ids = rng.permutation(num_experts) + inverse_weights = np.arange(1, num_experts + 1, dtype=np.float64) ** 1.25 + ids = np.empty((num_tokens, top_k), dtype=np.int64) + for start in range(0, num_tokens, 1024): + count = min(1024, num_tokens - start) + # Exponential races sample weighted top-k without replacement. + scores = rng.exponential(size=(count, num_experts)) * inverse_weights + selected = np.argpartition(scores, top_k - 1, axis=1)[:, :top_k] + ids[start : start + count] = expert_ids[selected] + return ids + + +def _fixture_rows(fixture: torch.Tensor, num_tokens: int) -> torch.Tensor: + """Select this rank's rows when the model shards tokens across TP ranks.""" + descriptor = get_forward_context().batch_descriptor + offset = 0 + if descriptor is not None: + tp = get_tp_group() + full = descriptor.num_tokens + if num_tokens < full and num_tokens == -(-full // tp.world_size): + offset = tp.rank_in_group * num_tokens + return fixture[offset : offset + num_tokens] + + def _moe_layers(model: torch.nn.Module) -> list[tuple[torch.nn.Module, GateLinear]]: """(MoE module, gate) per layer, in model order.""" layers: dict[GateLinear, torch.nn.Module] = {} @@ -78,8 +109,8 @@ def _top_k(module: torch.nn.Module) -> int | None: return None -class BalancedRouting: - """Installs balanced routing for one case; `remove()` restores the model.""" +class RoutingFixture: + """Installs synthetic routing for one case; `remove()` restores the model.""" def __init__( self, @@ -90,6 +121,7 @@ def __init__( ep_size: int, seed: int, device: torch.device, + mode: str, ): layers = _moe_layers(model) if not layers: @@ -103,14 +135,22 @@ def __init__( rows = -(-num_tokens // 256) * 256 + 256 for layer, (module, gate) in enumerate(layers): if (getattr(module, "num_expert_group", None) or 1) > 1: - raise NotImplementedError("Balanced routing with grouped top-k.") + raise NotImplementedError("Synthetic routing with grouped top-k.") num_experts = gate.weight.shape[0] top_k = _top_k(module) assert top_k is not None - rng = np.random.default_rng([seed, layer]) - ids = balanced_topk_ids(rows * dp_size, num_experts, top_k, ep_size, rng)[ - dp_rank * rows : (dp_rank + 1) * rows - ] + layer_seed = layer + if mode == "zipf" and gate.prefix: + # Preserve layer popularity when pipeline stages change. + layer_seed = zlib.crc32(gate.prefix.encode()) + rng = np.random.default_rng([seed, layer_seed]) + if mode == "zipf": + ids = zipf_topk_ids(rows * dp_size, num_experts, top_k, rng) + else: + ids = balanced_topk_ids( + rows * dp_size, num_experts, top_k, ep_size, rng + ) + ids = ids[dp_rank * rows : (dp_rank + 1) * rows] logits = torch.full((rows, num_experts), _OTHER_LOGIT, dtype=torch.float32) logits.scatter_(1, torch.from_numpy(ids), _CHOSEN_LOGIT) # Distinct values: no ties for any top-k implementation. @@ -153,7 +193,7 @@ def hook(module, inputs, output): logits = output[0] if isinstance(output, tuple) else output if cast[0].dtype != logits.dtype: cast[0] = fixture.to(logits.dtype) # Once, at warmup. - replaced = cast[0][: logits.shape[0]] + replaced = _fixture_rows(cast[0], logits.shape[0]) return (replaced, *output[1:]) if isinstance(output, tuple) else replaced return hook @@ -167,7 +207,7 @@ def hook(module, args, kwargs): ids = kwargs["topk_ids"] if "topk_ids" in kwargs else args[index] if cast[0].dtype != ids.dtype: cast[0] = fixture.to(ids.dtype) # Once, at warmup. - replaced = cast[0][: ids.shape[0]] + replaced = _fixture_rows(cast[0], ids.shape[0]) if "topk_ids" in kwargs: return args, {**kwargs, "topk_ids": replaced} return (*args[:index], replaced, *args[index + 1 :]), kwargs @@ -183,7 +223,7 @@ def check_applied(self) -> None: ] if missed: raise RuntimeError( - f"Balanced routing didn't reach {len(missed)} MoE layers " + f"Synthetic routing didn't reach {len(missed)} MoE layers " f"({missed[0]}): their gate is fused away. Use --routing real." ) @@ -196,15 +236,15 @@ def remove(self) -> None: def install_routing( runner, requested: str, dummy_weights: bool, num_tokens: int, seed: int -) -> tuple[str, BalancedRouting | None]: - """Resolve `--routing` and install balanced routing if selected. +) -> tuple[str, RoutingFixture | None]: + """Resolve `--routing` and install a synthetic fixture if selected. - `auto` is balanced with dummy weights and real otherwise; it falls back to + `auto` is Zipf with dummy weights and real otherwise; it falls back to real routing when the model doesn't support the fixture. """ mode = requested if mode == "auto": - mode = "balanced" if dummy_weights else "real" + mode = "zipf" if dummy_weights else "real" if mode == "real": return mode, None parallel = runner.parallel_config @@ -219,7 +259,9 @@ def install_routing( if not parallel.enable_expert_parallel: num_tokens *= parallel.prefill_context_parallel_size try: - routing = BalancedRouting( + if mode == "zipf" and parallel.prefill_context_parallel_size > 1: + raise NotImplementedError("Zipf routing with prefill context parallelism.") + routing = RoutingFixture( runner.model, num_tokens, runner.dp_rank, @@ -227,9 +269,10 @@ def install_routing( ep_size, seed, runner.device, + mode, ) except (ValueError, NotImplementedError) as e: - if requested == "balanced": + if requested != "auto": raise logger.warning("Using real routing: %s", e) return "real", None From 2b678ac739166bd0356883048673bbd1c34a3b0e Mon Sep 17 00:00:00 2001 From: Giannis <145396613+giannis2two@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:37:09 +0200 Subject: [PATCH 030/406] [RL] Support per-update source ranks for NCCL weight transfer (#62557) Signed-off-by: Dharak Kharod Co-authored-by: Dharak Rajeshbhai Kharod Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test_areas/distributed.yaml | 2 + docs/training/weight_transfer/README.md | 9 +- docs/training/weight_transfer/nccl.md | 61 +++- examples/rl/rlhf_nccl_multi_source.py | 207 ++++++++++++ tests/distributed/test_weight_transfer.py | 318 ++++++++++++++++-- .../weight_transfer/nccl_common.py | 5 +- .../weight_transfer/nccl_engine.py | 212 ++++++++---- 7 files changed, 694 insertions(+), 120 deletions(-) create mode 100644 examples/rl/rlhf_nccl_multi_source.py diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 883aa03aba05..777f812ece5f 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -475,6 +475,7 @@ steps: # see https://github.com/vllm-project/vllm/pull/5689 for details - pytest -v -s distributed/test_custom_all_reduce.py - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s distributed/test_weight_transfer.py -k multiple_sources - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py mirror: @@ -496,6 +497,7 @@ steps: commands: - pytest -v -s distributed/test_custom_all_reduce.py - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s distributed/test_weight_transfer.py -k multiple_sources - TARGET_TEST_SUITE=MI355 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)' diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index 9c946eba0f13..73b56c6cc2fc 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -131,10 +131,11 @@ parameter is usually itself a collective among its owners (an FSDP `full_tensor()` all-gather, a Megatron export) that would deadlock if some of them skipped it. Which ranks own which parameter is declared in [`ParamMeta.owners`](base.md#parammeta-ownership-and-layout). -Only [sharded RDT](sharded_rdt.md) accepts ranks that own different parameters, -such as pipeline stages. Pipeline-parallel support for NCCL and IPC is in -progress; until then they, like NCCL M2N, need every parameter owned by rank 0 -or by every rank (see each backend's Limitations). +Only [sharded RDT](sharded_rdt.md) and [NCCL](nccl.md#broadcasting-ranks) +accept ranks that own different parameters, such as pipeline stages. +Pipeline-parallel support for IPC is in progress; until then it, like NCCL M2N, +needs every parameter owned by rank 0 or by every rank (see each backend's +Limitations). Pass each process's own `rank` on the init info. It is explicit rather than read from a global process group, which is ambiguous once several groups (FSDP / TP / diff --git a/docs/training/weight_transfer/nccl.md b/docs/training/weight_transfer/nccl.md index 8ada8ed870f9..279ba3f21e93 100644 --- a/docs/training/weight_transfer/nccl.md +++ b/docs/training/weight_transfer/nccl.md @@ -10,7 +10,7 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b ## How It Works -1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction). The trainer is rank 0; workers start at `rank_offset` (1). +1. The trainer's broadcasting ranks and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction). Trainer rank 0 is group rank 0; with [several broadcasters](#broadcasting-ranks) they take `[0, rank_offset)` and workers follow them (`rank_offset` is 1 otherwise). 2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights. 3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py). @@ -52,7 +52,7 @@ engine = WeightTransferTrainerFactory.trainer_init( init_info=NCCLTrainerInitInfo( master_address=master_address, master_port=master_port, - world_size=world_size, # trainer + all inference workers + world_size=world_size, # broadcasting trainer ranks + all inference workers rank=0, # this trainer rank; rank 0 is the sender packed=True, ), @@ -66,7 +66,8 @@ engine.send_weights() # once per sync `trainer_init` drives the whole handshake: it kicks off the inference side's `init_weight_transfer_engine` (through the client) on a side thread while opening the trainer's own rank-0 endpoint, since both ends must rendezvous together. It -builds the worker's init info itself, with `rank_offset=1` and the same packed +builds the worker's init info itself, with `rank_offset` set to the number of +broadcasting ranks (1 unless the source is partitioned) and the same packed params, so the two sides cannot disagree. `send_weights()` then drives `start_weight_update`, `update_weights` — run @@ -79,7 +80,7 @@ once every transfer has drained. | ----- | ------- | ----------- | | `master_address` | — | Rendezvous host | | `master_port` | — | Rendezvous port | -| `world_size` | — | Full trainer + worker NCCL group size | +| `world_size` | — | Transfer NCCL group size: the [broadcasting trainer ranks](#broadcasting-ranks) plus the inference workers (`1 + num_workers` for a single sender). Other trainer ranks are not counted | | `rank` | — | Keyword-only. This trainer process's rank; **0 is the sender** | | `packed` | `True` | Use packed broadcasting | | `packed_buffer_size_bytes` | 1 GiB | Packed buffer size | @@ -124,13 +125,37 @@ rather than letting it reach the wire. `ModuleSource` satisfies this by construction. If you write a custom source — a Megatron export, an MoE re-fusing pass — this is the invariant to test first. -### One sender - -Each parameter is sent by its first owner, and this backend has a single -sender, trainer rank 0. Every parameter's `owners` must therefore start at rank -0 or be unset; a source whose pipeline stages own different parameters is -rejected at `trainer_init`. Other owners (FSDP peers of rank 0) only -drive `iter_full()` to take part in the gathers. +### Broadcasting ranks + +Each parameter is broadcast by its first owner (rank 0 when `owners` is unset). +When that is rank 0 for every parameter, rank 0 is the only broadcaster. Other +owners (FSDP peers of rank 0) only drive `iter_full()` to take part in the +gathers. + +A partitioned source, such as pipeline stages declaring `owners=(stage +ranks...)`, has several broadcasters. They join the transfer group at group +ranks `[0, rank_offset)` in trainer-rank order, rank 0 always included. Every +rank derives this plan from `metadata`, which is identical on every rank, so +the broadcasters need no extra coordination. In `send_weights()`, rank 0 sends +one `update_weights` per broadcaster, in order, each tagged with that +broadcaster's group rank as `src`. Each broadcaster sends its own parameters +during its turn and receives and discards the others' to stay in the +collective. Initialize and call `send_weights()` on every trainer rank +concurrently. + +Every parameter a rank owns must share its first owner, since each broadcaster +gathers and sends in its own turn. So with several broadcasters every parameter +must declare `owners`. Rank 0 is always in the transfer group and only gathers +in its own turn, so if it owns a parameter it must be that parameter's first +owner. `trainer_init` raises on every rank otherwise. + +A failure on any broadcaster during an update is fatal for the transfer. That +rank raises immediately, but the other trainer ranks and the inference workers +can stay blocked in that update's broadcasts, so tear down the trainer job and +restart or re-initialize the inference engine before transferring again. + +[`rlhf_nccl_multi_source.py`](../../../examples/rl/rlhf_nccl_multi_source.py) +shows two pipeline stages updating one tensor-parallel inference engine. ### Metadata is sent once, then reused @@ -142,7 +167,7 @@ round. So the worker remembers the list: | `update_info` you send | What the worker loads | | ---------------------- | --------------------- | | `names`, `dtype_names`, `shapes` filled in | Exactly these parameters. The worker also saves this list. | -| All three lists empty (or `{}`) | The list saved by the previous `update_weights` call on the same model. | +| All three lists empty (or `{}`) | The list saved by the previous `update_weights` call on the same model from the same `src`. | | All three lists empty, nothing saved yet | Nothing. | Rules to remember: @@ -195,6 +220,9 @@ while rank 0 broadcasts in `names` order, then `POST /finish_weight_update`. checkpoint shapes. From the second sync on, send `{"update_info": {}}` if the parameters did not change; see [Metadata is sent once, then reused](#metadata-is-sent-once-then-reused). +An optional `src` (default `0`, below `rank_offset`) selects the trainer rank +that broadcasts that update, so one start/finish lifecycle can carry updates +from several trainer ranks. - Unpacked: one broadcast per tensor. - Packed: concatenate `t.contiguous().view(-1).view(torch.uint8)` in `names` @@ -275,11 +303,9 @@ Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still ## Limitations -- **Pipeline-parallel trainers are not supported yet (work in progress).** Rank 0 - is the only sender, so it must own every parameter (see - [One sender](#one-sender)). A source whose ranks own different parameters is - rejected at `trainer_init`. -- **Full tensors on the wire.** Rank 0 materializes each parameter in full and +- **A broadcaster failure is fatal.** See + [Broadcasting ranks](#broadcasting-ranks). +- **Full tensors on the wire.** Each broadcaster materializes its parameters in full and every worker receives all of it, whatever its tensor-parallel shard. For large or expert-parallel models, see [sharded RDT](sharded_rdt.md) or [NCCL M2N](m2n.md). @@ -288,6 +314,7 @@ Sparse NCCL sends only `O(nnz)` indices and values. Checkpoint application still ## Examples - [RLHF with NCCL weight syncing (`vllm serve`, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - **Start here.** Trainer on one GPU, 2x tensor-parallel fp8 server on two others; HTTP control plane, NCCL data plane. Launches and tears down its own server +- [Multi-source NCCL weight syncing](../../../examples/rl/rlhf_nccl_multi_source.py) - Two pipeline stages each broadcast their own parameters to one TP2 inference engine within a single lifecycle - [NCCL weight syncing from a trainer that does not import vLLM](../../../examples/rl/rlhf_http_nccl_no_vllm.py) - `torch` + `requests` only, `nccl_unique_id_b64` rendezvous; see [Trainer Without vLLM over HTTP](#trainer-without-vllm-over-http) - [RLHF with NCCL + FSDP2 and expert parallelism](../../../examples/rl/rlhf_nccl_fsdp_ep.py) - Multi-rank trainer: every FSDP rank builds an engine and joins the `full_tensor()` gather while only rank 0 touches the wire - [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Qwen3 MoE per-expert checkpoint updates with one trainer GPU and two TP2/EP2 inference GPUs diff --git a/examples/rl/rlhf_nccl_multi_source.py b/examples/rl/rlhf_nccl_multi_source.py new file mode 100644 index 000000000000..730d4443528c --- /dev/null +++ b/examples/rl/rlhf_nccl_multi_source.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dense NCCL weight transfer from two disjoint trainer sources. + +The example uses four GPUs: + +* GPUs 0-1 run two trainer actors. Each actor exposes the parameters owned by + one logical pipeline stage. +* GPUs 2-3 run one tensor-parallel vLLM engine. + +Each stage declares itself the owner of its parameters in `ParamMeta.owners`, +so both trainer ranks join the weight-transfer communicator and each +broadcasts its own parameters. Rank 0 drives the control plane: one +`update_weights` per stage, within one start/finish lifecycle. + +For simplicity, both trainer actors load the small model before filtering their +parameter streams. A real pipeline-parallel trainer would already hold only its +stage-local parameters; the transfer behavior is the same. + +This example was written for four GPUs with at least 8 GiB each. +""" + +import os +import time +from functools import cached_property + +import ray +import torch +from huggingface_hub import snapshot_download +from transformers import AutoModelForCausalLM + +from vllm import LLM, SamplingParams +from vllm.config import WeightTransferConfig +from vllm.distributed.weight_transfer import ( + ParamMeta, + RayVLLMWeightSyncClient, + WeightSource, + WeightTransferTrainerFactory, +) +from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo +from vllm.utils.network_utils import get_ip, get_open_port + +MODEL_NAME = "Qwen/Qwen3-0.6B" +NUM_SOURCES = 2 +INFERENCE_TP_SIZE = 2 + + +class StageSource(WeightSource): + """Expose the checkpoint parameters assigned to one logical stage.""" + + def __init__( + self, + model: torch.nn.Module, + stage_rank: int, + num_stages: int, + num_layers: int, + ) -> None: + self.model = model + self.stage_rank = stage_rank + self.num_stages = num_stages + self.num_layers = num_layers + + def _owner(self, name: str) -> int: + parts = name.split(".") + for index, part in enumerate(parts[:-1]): + if part == "layers" and parts[index + 1].isdigit(): + layer = int(parts[index + 1]) + return min( + layer * self.num_stages // self.num_layers, self.num_stages - 1 + ) + return 0 if "embed_tokens" in name else self.num_stages - 1 + + @cached_property + def metadata(self) -> list[ParamMeta]: + return [ + ParamMeta( + name, param.dtype, tuple(param.shape), owners=(self._owner(name),) + ) + for name, param in self.model.named_parameters() + ] + + def iter_full(self): + for name, param in self.model.named_parameters(): + if self._owner(name) == self.stage_rank: + yield name, param + + +class MyLLM(LLM): + """Allow the vLLM Ray executor to allocate its own worker GPUs.""" + + def __init__(self, *args, **kwargs): + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + super().__init__(*args, **kwargs) + + def ready(self): + return True + + +@ray.remote(num_cpus=1, num_gpus=1) +class TrainerSource: + def __init__(self, model_path: str, source_rank: int): + self.source_rank = source_rank + torch.accelerator.set_device_index(0) + self.model = AutoModelForCausalLM.from_pretrained( + model_path, torch_dtype=torch.bfloat16 + ).cuda() + self.engine = None + + def ready(self): + return self.source_rank + + def setup_engine( + self, + llm_handle, + master_address: str, + master_port: int, + world_size: int, + ): + source = StageSource( + self.model, + self.source_rank, + NUM_SOURCES, + self.model.config.num_hidden_layers, + ) + self.engine = WeightTransferTrainerFactory.trainer_init( + init_info=NCCLTrainerInitInfo( + master_address=master_address, + master_port=master_port, + world_size=world_size, + rank=self.source_rank, + packed=True, + ), + client=RayVLLMWeightSyncClient(llm_handle), + source=source, + ) + + def send_weights(self): + self.engine.send_weights() + + +def generate(llm, prompts: list[str]) -> list[list[int]]: + outputs = ray.get( + llm.generate.remote( + prompts, + SamplingParams(temperature=0, max_tokens=16), + ) + ) + return [output.outputs[0].token_ids for output in outputs] + + +def main(): + ray.init() + model_path = snapshot_download(MODEL_NAME) + + sources = [ + TrainerSource.remote(model_path, source_rank) + for source_rank in range(NUM_SOURCES) + ] + ray.get([source.ready.remote() for source in sources]) + + llm = ray.remote(num_cpus=0, num_gpus=0)(MyLLM).remote( + model=model_path, + enforce_eager=True, + tensor_parallel_size=INFERENCE_TP_SIZE, + distributed_executor_backend="ray", + weight_transfer_config=WeightTransferConfig(backend="nccl"), + load_format="dummy", + ) + ray.get(llm.ready.remote()) + num_workers = ray.get(llm.get_world_size.remote()) + + prompts = ["The capital of France is", "The future of AI is"] + before = generate(llm, prompts) + + master_address = get_ip() + master_port = get_open_port() + world_size = NUM_SOURCES + num_workers + t0 = time.perf_counter() + ray.get( + [ + source.setup_engine.remote( + llm, + master_address, + master_port, + world_size, + ) + for source in sources + ] + ) + print(f"[sync] engine setup took {time.perf_counter() - t0:.2f}s") + + t0 = time.perf_counter() + ray.get([source.send_weights.remote() for source in sources]) + print(f"[sync] iter 0 took {time.perf_counter() - t0:.3f}s") + after = generate(llm, prompts) + t0 = time.perf_counter() + ray.get([source.send_weights.remote() for source in sources]) + print(f"[sync] iter 1 took {time.perf_counter() - t0:.3f}s") + repeated = generate(llm, prompts) + + assert after != before, "generation did not change after weight transfer" + assert repeated == after, "repeating the same transfer changed generation" + print("Multi-source NCCL weight transfer completed successfully.") + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index fefd9059ee01..eba3ccc20eed 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -8,6 +8,7 @@ import builtins import importlib.util +import math import pickle import runpy import threading @@ -281,12 +282,14 @@ def test_parse_update_info_valid(self): "names": ["w1", "w2"], "dtype_names": ["float32", "bfloat16"], "shapes": [[100, 100], [50]], + "src": 1, } ) assert isinstance(update_info, NCCLWeightTransferUpdateInfo) assert update_info.names == ["w1", "w2"] assert update_info.dtype_names == ["float32", "bfloat16"] assert update_info.shapes == [[100, 100], [50]] + assert update_info.src == 1 # --- Unit Tests: Engine Registry --- @@ -627,6 +630,197 @@ def test_nccl_weight_transfer_between_processes(batch_invariant): ) +@ray.remote(num_gpus=1) +class MultiSourceInferenceWorker: + def __init__(self, worker_rank: int, num_workers: int): + import contextlib + + device = _set_ray_assigned_device() + + from vllm.config.weight_transfer import WeightTransferConfig + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, + ) + + class Recorder(torch.nn.Module): + def __init__(self): + super().__init__() + self.received: dict[str, float] = {} + self.load_counts: dict[str, int] = {} + + def load_weights(self, weights): + for name, tensor in weights: + self.received[name] = tensor.sum().item() + self.load_counts[name] = self.load_counts.get(name, 0) + 1 + + self.recorder = Recorder() + self.engine = NCCLWeightTransferEngine( + WeightTransferConfig(backend="nccl"), + create_mock_vllm_config(rank=worker_rank, world_size=num_workers), + device, + self.recorder, + ) + import vllm.config as vllm_config_mod + + vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext() + + def ready(self): + return True + + def initialize(self, init_info): + self.engine.init_transfer_engine(self.engine.parse_init_info(init_info)) + + def update(self, update_info): + self.engine.receive_weights(self.engine.parse_update_info(update_info)) + + def results(self): + torch.accelerator.synchronize() + results = { + "received": self.recorder.received, + "load_counts": self.recorder.load_counts, + } + self.engine.shutdown() + return results + + +@ray.remote(num_gpus=1) +def trainer_send_from_source( + receivers, + master_address: str, + master_port: int, + world_size: int, + num_sources: int, + rank: int, + tensor_shape: list[int], + num_tensors_per_source: int, + packed: bool, +) -> bool: + from functools import cached_property + + device = _set_ray_assigned_device() + + from vllm.distributed.weight_transfer import ( + ParamMeta, + WeightSource, + WeightTransferTrainerFactory, + ) + from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo + + class Source(WeightSource): + """Rank r owns (and broadcasts) `source_r.*`, like a pipeline stage.""" + + @cached_property + def metadata(self): + return [ + ParamMeta( + f"source_{r}.weight_{i}", + torch.float32, + tuple(tensor_shape), + owners=(r,), + ) + for r in range(num_sources) + for i in range(num_tensors_per_source) + ] + + def iter_full(self): + for i in range(num_tensors_per_source): + value = rank * num_tensors_per_source + i + 1 + yield ( + f"source_{rank}.weight_{i}", + torch.full(tensor_shape, value, dtype=torch.float32, device=device), + ) + + class Client: + def call(self, method, *args): + return ray.get( + [getattr(receiver, method).remote(*args) for receiver in receivers] + ) + + def init_weight_transfer_engine(self, init_info): + self.call("initialize", init_info) + + def start_weight_update(self): + pass + + def update_weights(self, update_info): + self.call("update", update_info) + + def finish_weight_update(self, weight_version=None): + pass + + engine = WeightTransferTrainerFactory.trainer_init( + NCCLTrainerInitInfo( + master_address=master_address, + master_port=master_port, + world_size=world_size, + rank=rank, + packed=packed, + packed_buffer_size_bytes=1024, + ), + client=Client(), + source=Source(), + ) + for _ in range(2): + engine.send_weights() + engine.shutdown() + return True + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 4, + reason="Need 4 GPUs for two NCCL sources and two inference workers.", +) +@pytest.mark.parametrize("packed", [False, True], ids=["unpacked", "packed"]) +def test_nccl_weight_transfer_from_multiple_sources(packed): + """Two trainer ranks each own and broadcast half the parameters (like + pipeline stages); both workers receive all of them, twice (the second + round reusing metadata).""" + _init_ray_for_weight_transfer() + + master_address = "127.0.0.1" + master_port = get_open_port() + num_sources = 2 + num_workers = 2 + world_size = num_sources + num_workers + tensor_shape = [8, 8] + num_tensors_per_source = 10 + + receivers = [ + MultiSourceInferenceWorker.remote( # type: ignore[attr-defined] + worker_rank, num_workers + ) + for worker_rank in range(num_workers) + ] + ray.get([receiver.ready.remote() for receiver in receivers]) + futures = [ + trainer_send_from_source.remote( + receivers, + master_address, + master_port, + world_size, + num_sources, + rank, + tensor_shape, + num_tensors_per_source, + packed, + ) + for rank in range(num_sources) + ] + assert ray.get(futures, timeout=300) == [True] * num_sources + expected_numel = math.prod(tensor_shape) + expected = { + f"source_{r}.weight_{i}": (r * num_tensors_per_source + i + 1) * expected_numel + for r in range(num_sources) + for i in range(num_tensors_per_source) + } + results = ray.get([receiver.results.remote() for receiver in receivers]) + for result in results: + assert result["received"] == expected + assert result["load_counts"] == dict.fromkeys(expected, 2) + for receiver in receivers: + ray.kill(receiver) + + def test_sparse_nccl_checkpoint_chunks_to_ep_local_experts_cpu(monkeypatch): """Replay global expert patches through two EP-local loaders on CPU.""" @@ -1702,7 +1896,9 @@ def test_nccl_trainer_init_ships_worker_init_info(monkeypatch, packed): # Bypass the real NCCL rendezvous. group = object() - monkeypatch.setattr(nccl_engine_mod, "open_trainer_endpoint", lambda info: group) + monkeypatch.setattr( + nccl_engine_mod, "open_trainer_endpoint", lambda info, rank: group + ) client = RecordingClient() engine = WeightTransferTrainerFactory.trainer_init( @@ -1849,7 +2045,7 @@ def test_nccl_trainer_send_weights_drives_client_in_order(): packed=False, ) # Bypass the real NCCL rendezvous; broadcast is a no-op. - engine.model_update_group = MagicMock() + engine.model_update_group = MagicMock(rank=0) engine.send_weights() @@ -1873,7 +2069,7 @@ def test_nccl_trainer_sends_metadata_only_when_changed(): source=ModuleSource(_module_with(("w", torch.zeros(4, device="cuda")))), packed=False, ) - engine.model_update_group = MagicMock() + engine.model_update_group = MagicMock(rank=0) engine.send_weights() assert client.last_update_info is not None @@ -1886,8 +2082,9 @@ def test_nccl_trainer_sends_metadata_only_when_changed(): def test_nccl_worker_reuses_metadata_per_model(monkeypatch): - """Empty update info reuses the previous call's metadata for the same model; - target and draft keep separate copies, and re-init clears them.""" + """Empty update info reuses the previous call's metadata for the same + (model, src) and receives from that src; re-init clears every entry, and + src must be a trainer rank.""" import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod monkeypatch.setattr( @@ -1896,7 +2093,9 @@ def test_nccl_worker_reuses_metadata_per_model(monkeypatch): target, draft = MagicMock(spec=torch.nn.Module), MagicMock(spec=torch.nn.Module) received: list = [] for model in (target, draft): - model.load_weights = lambda w, m=model: received.append((m, w[0][0])) + model.load_weights = lambda w, m=model: received.append( + (m, w[0][0], tuple(w[0][1].shape)) + ) engine = NCCLWeightTransferEngine( WeightTransferConfig(backend="nccl"), create_mock_vllm_config(), @@ -1904,29 +2103,53 @@ def test_nccl_worker_reuses_metadata_per_model(monkeypatch): target, ) init_info = NCCLWeightTransferInitInfo( - master_address="127.0.0.1", master_port=29500, rank_offset=1, world_size=2 + master_address="127.0.0.1", master_port=29500, rank_offset=2, world_size=3 ) engine.init_transfer_engine(init_info) monkeypatch.setattr(torch.cuda, "current_stream", MagicMock()) - def update(model, *names): + def update(model, src, name=None, shape=None): engine.set_weight_update_target(model, None) engine.receive_weights( NCCLWeightTransferUpdateInfo( - names=list(names), - dtype_names=["float32"] * len(names), - shapes=[[2]] * len(names), + names=[name] if name else [], + dtype_names=["float32"] if name else [], + shapes=[shape] if name else [], + src=src, ) ) - update(target, "t") - update(draft, "d") - update(target) - update(draft) + seeded = { + (target, 0): ("t0", [2]), + (target, 1): ("t1", [3]), + (draft, 0): ("d0", [4]), + (draft, 1): ("d1", [5]), + } + for (model, src), (name, shape) in seeded.items(): + update(model, src, name, shape) + reused = [(draft, 1), (target, 0), (draft, 0), (target, 1)] + for model, src in reused: + update(model, src) + + order = [*seeded, *reused] + assert received == [ + (model, seeded[model, src][0], tuple(seeded[model, src][1])) + for model, src in order + ] + group = engine.model_update_group + assert [c.kwargs["src"] for c in group.broadcast.call_args_list] == [ + src for _, src in order + ] + engine.init_transfer_engine(init_info) - update(target) + for model, src in seeded: + update(model, src) + assert len(received) == len(order) + engine.model_update_group.broadcast.assert_not_called() - assert received == [(target, "t"), (draft, "d"), (target, "t"), (draft, "d")] + for bad in (-1, 2): + with pytest.raises(ValueError, match=r"sender rank in \[0, 2\)"): + update(target, bad, "c", [2]) class _ScriptedSource(WeightSource): @@ -1950,7 +2173,7 @@ def _mock_group_engine(source, monkeypatch, **kwargs): engine = NCCLTrainerWeightTransferEngine( client=RecordingClient(), source=source, packed=False, **kwargs ) - engine.model_update_group = MagicMock() + engine.model_update_group = MagicMock(rank=0) monkeypatch.setattr(torch.cuda, "current_stream", MagicMock()) return engine @@ -2007,29 +2230,62 @@ def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch): engine.send_weights() -def test_nccl_trainer_rejects_a_second_sender(): - """Each parameter is sent by its first owner and broadcast has one sender, - so a pipeline stage that does not start at rank 0 cannot be sent; every - rank sees the same metadata and raises alike.""" - meta = [ - ParamMeta("w", torch.float32, (4,), owners=(0, 1)), - ParamMeta("b", torch.float32, (2,), owners=(2, 3)), - ] - source = _ScriptedSource(meta, []) +_TWO_STAGES = [ + ParamMeta("w", torch.float32, (4,), owners=(0, 1)), + ParamMeta("b", torch.float32, (2,), owners=(2, 3)), +] + + +@pytest.mark.parametrize( + ("meta", "match"), + [ + ( + [ParamMeta("w", torch.float32, (4,)), _TWO_STAGES[1]], + "'w' has no owners", + ), + ( + [*_TWO_STAGES, ParamMeta("x", torch.float32, (2,), owners=(1, 2))], + "rank 1 owns parameters broadcast by ranks 0 and 1", + ), + ( + [ParamMeta("w", torch.float32, (4,), owners=(1, 0))], + r"'w' has owners \(1, 0\)", + ), + ( + [ParamMeta("w", torch.float32, (4,), owners=(1, 0)), _TWO_STAGES[1]], + r"'w' has owners \(1, 0\)", + ), + ], +) +def test_nccl_trainer_rejects_ranks_spanning_broadcasters(meta, match): + """Each broadcaster sends in its own turn, so a rank owning parameters of + two broadcasters would gather out of turn; every rank sees the same + metadata and raises alike.""" for rank in (0, 2): - with pytest.raises(ValueError, match="'b' is sent by its first owner, rank 2"): + with pytest.raises(ValueError, match=match): NCCLTrainerWeightTransferEngine.trainer_init( NCCLTrainerInitInfo( master_address="127.0.0.1", master_port=29500, - world_size=2, + world_size=3, rank=rank, ), client=RecordingClient(), - source=source, + source=_ScriptedSource(meta, []), ) +def test_nccl_trainer_rejects_world_size_without_workers(): + with pytest.raises(ValueError, match="leaves no inference workers"): + NCCLTrainerWeightTransferEngine.trainer_init( + NCCLTrainerInitInfo( + master_address="127.0.0.1", master_port=29500, world_size=2, rank=0 + ), + client=RecordingClient(), + source=_ScriptedSource(_TWO_STAGES, []), + ) + + def test_nccl_trainer_non_sender_owner_drives_its_gathers(monkeypatch): """A non-sender owner (an FSDP peer) still iterates what it owns, since the sender's gathers need it; it never touches the client or the group.""" @@ -2098,7 +2354,7 @@ def iter_full(self): engine = NCCLTrainerWeightTransferEngine( client=_BlockingClient(), source=_FailingSource(), packed=False ) - engine.model_update_group = MagicMock() + engine.model_update_group = MagicMock(rank=0) monkeypatch.setattr(torch.cuda, "current_stream", MagicMock()) started = time.perf_counter() diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py index b7de9f8becfc..70cc77fcb944 100644 --- a/vllm/distributed/weight_transfer/nccl_common.py +++ b/vllm/distributed/weight_transfer/nccl_common.py @@ -305,9 +305,8 @@ def trainer_init( - master_address: str - master_port: int - world_size: int - rank: This trainer process's rank in the group. The broadcast backends - have a single trainer at rank 0; multi-rank trainers (m2n) occupy - ranks `[0, num_trainer_ranks)` and the workers start after them. + rank: This trainer process's rank in the group. Trainer ranks occupy + `[0, num_trainer_ranks)` and the workers start after them. Returns: PyNcclCommunicator for weight transfer. diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 7ec98ef9969c..9b1cf998c429 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/vllm/distributed/weight_transfer/nccl_engine.py @@ -55,8 +55,7 @@ class NCCLTrainerInitInfo(TrainerInitInfo): """Trainer-side init info for the dense NCCL weight transfer backend. The sender opens its endpoint as NCCL rank 0, so it needs no `rank_offset`. - `world_size` is the full trainer+worker NCCL group size. `rank` (from - `TrainerInitInfo`) identifies this trainer process; rank 0 is the sender. + `rank` (from `TrainerInitInfo`) identifies this trainer process. The trainer joins over a TCPStore rendezvous (`master_address` + `master_port`). Torch-free trainers (e.g. JAX) that cannot join a TCPStore @@ -76,6 +75,8 @@ class NCCLTrainerInitInfo(TrainerInitInfo): master_address: str master_port: int world_size: int + """Transfer group size: rank 0 and each distinct `ParamMeta.owners[0]`, plus + the inference workers (`1 + num_workers` for a single sender).""" packed: bool = True packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS @@ -85,19 +86,16 @@ class NCCLTrainerInitInfo(TrainerInitInfo): class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo): """Per-round update info for the dense NCCL weight transfer backend. - Whether the transfer is packed (and the buffer geometry) is a must-agree - wire param carried on the init info (`NCCLTrainerInitInfo` / - `NCCLWeightTransferInitInfo`), not here; this carries only the per-round - parameter metadata. + Carries the parameter metadata and `src`, the broadcasting group rank; the + packed wire params live on the init info. - Leave all three lists empty to reuse the metadata of the previous - `update_weights` call on the same model, so a trainer that sends the same - parameters every round ships the names, dtypes and shapes only once. + Empty lists reuse the previous update's metadata for the same model and `src`. """ names: list[str] = field(default_factory=list) dtype_names: list[str] = field(default_factory=list) shapes: list[list[int]] = field(default_factory=list) + src: int = 0 def __post_init__(self): """Validate that all lists have the same length.""" @@ -119,10 +117,8 @@ class NCCLWeightTransferEngine( ): """Weight transfer engine using NCCL for communication between trainer and workers. - This implementation uses NCCL broadcast operations to transfer dense - checkpoint-format weights from the trainer (rank 0) to all inference workers - in a process group. Received weights are loaded via the model's - `load_weights` using the layerwise reload lifecycle. + Receives dense checkpoint-format weights broadcast from the trainer ranks + below `rank_offset` and loads them with the layerwise reload lifecycle. """ # Define backend-specific dataclass types @@ -140,11 +136,14 @@ def __init__( self.model_update_group: PyNcclCommunicator | None = None # Set from the trainer-supplied init info at the handshake; defaults are # only for the (unreachable) receive-before-init case. + self.rank_offset = 1 self.packed = False self.packed_buffer_size_bytes = DEFAULT_PACKED_BUFFER_SIZE_BYTES self.packed_num_buffers = DEFAULT_PACKED_NUM_BUFFERS - # Metadata of the last update per model, reused by empty update info. - self._last_update_info: dict[torch.nn.Module, NCCLWeightTransferUpdateInfo] = {} + # Metadata of the last update per (model, src), reused by empty update info. + self._last_update_info: dict[ + tuple[torch.nn.Module, int], NCCLWeightTransferUpdateInfo + ] = {} def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None: """Initialize NCCL process group with the trainer and record the @@ -156,6 +155,7 @@ def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None: rank offset, world size, and the packed wire params """ + self.rank_offset = init_info.rank_offset self.packed = init_info.packed self.packed_buffer_size_bytes = init_info.packed_buffer_size_bytes self.packed_num_buffers = init_info.packed_num_buffers @@ -189,8 +189,8 @@ def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None: encoded. Args: - update_info: NCCL update info containing parameter names, dtypes, - and shapes; empty to reuse the previous call's + update_info: Parameter names, dtypes, shapes and `src`; empty lists + reuse the previous call's from the same `src` """ if self.model_update_group is None: @@ -198,10 +198,17 @@ def receive_weights(self, update_info: NCCLWeightTransferUpdateInfo) -> None: "NCCL weight transfer not initialized. " "Call init_transfer_engine() first." ) + src = update_info.src + if not 0 <= src < self.rank_offset: + raise ValueError( + f"NCCL update src must be a sender rank in [0, {self.rank_offset}), " + f"got {src!r}." + ) + key = (self.model, src) if update_info.names: - self._last_update_info[self.model] = update_info + self._last_update_info[key] = update_info else: - update_info = self._last_update_info.get(self.model, update_info) + update_info = self._last_update_info.get(key, update_info) from vllm.model_executor.model_loader.mtp_validation import ( disable_mtp_completeness_check, @@ -220,7 +227,7 @@ def state_dict_info_iterator(): packed_nccl_broadcast_consumer( iterator=state_dict_info_iterator(), group=self.model_update_group, - src=0, + src=src, post_unpack_func=self.model.load_weights, buffer_size_bytes=self.packed_buffer_size_bytes, num_buffers=self.packed_num_buffers, @@ -234,7 +241,7 @@ def state_dict_info_iterator(): dtype = getattr(torch, dtype_name) weight = torch.empty(shape, dtype=dtype, device=self.device) self.model_update_group.broadcast( - weight, src=0, stream=torch.cuda.current_stream() + weight, src=src, stream=torch.cuda.current_stream() ) self.model.load_weights([(name, weight)]) del weight @@ -248,17 +255,11 @@ def shutdown(self) -> None: class NCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[NCCLTrainerInitInfo]): """Trainer-side NCCL weight transfer engine. - On the sender (rank 0) holds the NCCL communicator and drives the full - update round trip: it runs the inference-side `update_weights` concurrently - with the trainer-side broadcast (both rendezvous inside the same NCCL - calls), then finishes the update. Non-sender trainer ranks hold no - communicator; they only drive `iter_full()` over the names they own, to stay - in the trainer-side gathers (e.g. FSDP `full_tensor()`), and skip the client - RPCs and the broadcast (all guarded on `is_sender`). + Each parameter is broadcast by its first owner (rank 0 when unset); rank 0 + also drives the client RPCs. See "Broadcasting ranks" in `nccl.md`. - Each parameter is sent by its first owner, so every parameter's owners must - start at rank 0 (or be unset): a partitioned source with several senders - (pipeline stages) is rejected. + Other trainer ranks hold no communicator and only drive `iter_full()` to + stay in the trainer-side gathers (e.g. FSDP `full_tensor()`). `packed` / buffer sizes come from `NCCLTrainerInitInfo`; the sender propagates them to the worker at `trainer_init` (on the worker-side init @@ -283,6 +284,8 @@ def __init__( self.packed_buffer_size_bytes = packed_buffer_size_bytes self.packed_num_buffers = packed_num_buffers self.model_update_group: PyNcclCommunicator | None = None + # This rank's rank in the transfer group, set by trainer_init. + self._group_rank = 0 # Metadata the worker already holds; an unchanged round sends none. self._sent_meta: list[ParamMeta] | None = None @@ -296,7 +299,19 @@ def _build_worker_init_payload(cls, init_info: NCCLTrainerInitInfo) -> dict: ) def _open_trainer_endpoint(self, init_info: NCCLTrainerInitInfo) -> None: - self.model_update_group = open_trainer_endpoint(init_info) + self.model_update_group = open_trainer_endpoint(init_info, self._group_rank) + + @staticmethod + def _broadcast_plan( + meta: list[ParamMeta], + ) -> tuple[list[int], list[tuple[int, list[ParamMeta]]]]: + """The transfer group's trainer ranks, in group-rank order, and each + broadcaster's parameters (in metadata order), in broadcaster order.""" + slices: dict[int, list[ParamMeta]] = {} + for m in meta: + slices.setdefault(0 if m.owners is None else m.owners[0], []).append(m) + group_ranks = sorted({0, *slices}) + return group_ranks, sorted(slices.items()) @classmethod def trainer_init( @@ -314,13 +329,39 @@ def trainer_init( "does not implement iter_full()." ) # Identical on every rank, so every rank raises alike. + group_ranks, slices = cls._broadcast_plan(source.metadata) for m in source.metadata: - if m.owners is not None and m.owners[0] != 0: + if m.owners is not None and 0 in m.owners[1:]: raise ValueError( - f"parameter {m.name!r} is sent by its first owner, rank " - f"{m.owners[0]}, but NCCL broadcast has a single sender, " - "rank 0. Every parameter's owners must start at rank 0." + f"parameter {m.name!r} has owners {m.owners}: rank 0 only " + "gathers in its own turn, so it must be the first owner of " + "every parameter it owns." ) + if len(slices) > 1: + broadcaster_of: dict[int, int] = {} + for broadcaster, metas in slices: + for m in metas: + if m.owners is None: + raise ValueError( + f"parameter {m.name!r} has no owners, so every rank " + "owns it, but the source has several broadcasting " + f"ranks {[r for r, _ in slices]}. Every parameter " + "must declare its owners." + ) + for r in m.owners: + prev = broadcaster_of.setdefault(r, broadcaster) + if prev != broadcaster: + raise ValueError( + f"rank {r} owns parameters broadcast by ranks " + f"{prev} and {broadcaster}, but each broadcaster " + "sends in its own turn, so every parameter a " + "rank owns must share its first owner." + ) + if init_info.world_size <= len(group_ranks): + raise ValueError( + f"world_size {init_info.world_size} leaves no inference workers " + f"after the {len(group_ranks)} broadcasting trainer ranks." + ) engine = cls( client=client, source=source, @@ -330,13 +371,17 @@ def trainer_init( packed_buffer_size_bytes=init_info.packed_buffer_size_bytes, packed_num_buffers=init_info.packed_num_buffers, ) + if engine.rank not in group_ranks: + # Ranks that broadcast nothing stay out of the transfer group; they + # only join the trainer-side gathers in send_weights. + return engine + engine._group_rank = group_ranks.index(engine.rank) if not engine.is_sender: - # Non-sender trainer ranks aren't part of the transfer NCCL group and - # don't drive the inference side; they only participate in the - # trainer-side gather during send_weights. + engine._open_trainer_endpoint(init_info) return engine payload = cls._build_worker_init_payload(init_info) + payload["rank_offset"] = len(group_ranks) run_trainer_rendezvous( init_worker=lambda: engine.client.init_weight_transfer_engine(payload), init_trainer=lambda: engine._open_trainer_endpoint(init_info), @@ -351,34 +396,57 @@ def send_weights(self) -> None: # Metadata is declared without gathering. For Megatron it is itself a # collective, so every rank runs it; only the sender ships it. meta = source.metadata + group_ranks, slices = self._broadcast_plan(meta) - if not self.is_sender: - # Non-sender ranks only join the trainer-side gather collective. - self._broadcast(source, meta) + if self.rank not in group_ranks: + # Ranks outside the transfer group only join the trainer-side + # gather collective. + for _ in source.iter_full(): + pass self._post_send_sync() return - if meta == self._sent_meta: - update_info = NCCLWeightTransferUpdateInfo() + reuse = self.is_sender and meta == self._sent_meta + if self.is_sender: + self.client.start_weight_update() + for broadcaster, metas in slices: + src = group_ranks.index(broadcaster) + if self.is_sender: + self._send_update(metas, src, reuse) + else: + self._broadcast(metas, src) + if self.is_sender: + self.client.finish_weight_update() + self._sent_meta = meta + self._post_send_sync() + + def _send_update(self, meta: list[ParamMeta], src: int, reuse: bool) -> None: + """Rank 0's part of one broadcaster's turn: run its `update_weights` + concurrently with `_broadcast`.""" + if reuse: + update_info = NCCLWeightTransferUpdateInfo(src=src) else: update_info = NCCLWeightTransferUpdateInfo( names=[m.name for m in meta], dtype_names=[str(m.dtype).split(".")[-1] for m in meta], shapes=[list(m.shape) for m in meta], + src=src, ) - - self.client.start_weight_update() + payload = asdict(update_info) + if not src: + # Keep single-broadcaster payloads readable by older servers. + del payload["src"] # update_weights (workers receive) must run concurrently with the # trainer-side broadcast — both rendezvous inside the same NCCL calls. exe = ThreadPoolExecutor(max_workers=1) try: - future = exe.submit(self.client.update_weights, asdict(update_info)) + future = exe.submit(self.client.update_weights, payload) # Cheap best-effort: if update_weights already failed (e.g. a bad # request rejected before any NCCL call), surface it now instead of # hanging in broadcast waiting for a peer that will never arrive. if future.done(): future.result() - self._broadcast(source, meta) + self._broadcast(meta, src) future.result() # surface inference-side errors finally: # Never wait for the RPC thread here. If the broadcast raised, the @@ -389,28 +457,23 @@ def send_weights(self) -> None: # thread is still joined at interpreter exit, so a caller that wants # to exit cleanly after such a failure must not wait on it.) exe.shutdown(wait=False) - self.client.finish_weight_update() - self._sent_meta = meta - self._post_send_sync() - - def _broadcast(self, source: WeightSource, meta: list[ParamMeta]) -> None: - """Drive `iter_full()` (gathering among each parameter's owners) and, on - the sender, broadcast from rank 0, packed or one-by-one. Other owners only - join the gathers.""" - if not self.is_sender: - for _ in source.iter_full(): - pass - return + def _broadcast(self, meta: list[ParamMeta], src: int) -> None: + """Broadcast `meta` from group rank `src`: that rank sends from + `iter_full()`, every other group member receives and discards.""" assert self.model_update_group is not None, ( "trainer_init() must be called before _broadcast()." ) - pairs = self._checked_iter(source, meta) + if src != self.model_update_group.rank: + self._discard(meta, src) + return + assert self.source is not None + pairs = self._checked_iter(self.source, meta) if self.packed: packed_nccl_broadcast_producer( iterator=pairs, group=self.model_update_group, - src=0, + src=src, post_iter_func=lambda item: item[1], buffer_size_bytes=self.packed_buffer_size_bytes, num_buffers=self.packed_num_buffers, @@ -423,7 +486,27 @@ def _broadcast(self, source: WeightSource, meta: list[ParamMeta]) -> None: # pointer. Keep the copy referenced until the broadcast is # enqueued. (The packed path linearizes in `pack_tensors`.) send = tensor if tensor.is_contiguous() else tensor.contiguous() - self.model_update_group.broadcast(send, src=0, stream=stream) + self.model_update_group.broadcast(send, src=src, stream=stream) + + def _discard(self, meta: list[ParamMeta], src: int) -> None: + """Receive another broadcaster's parameters and drop them.""" + assert self.model_update_group is not None + device = self.model_update_group.device + if self.packed: + packed_nccl_broadcast_consumer( + iterator=((m.name, (list(m.shape), m.dtype)) for m in meta), + group=self.model_update_group, + src=src, + post_unpack_func=lambda weights: None, + buffer_size_bytes=self.packed_buffer_size_bytes, + num_buffers=self.packed_num_buffers, + device=device, + ) + else: + stream = torch.cuda.current_stream() + for m in meta: + weight = torch.empty(m.shape, dtype=m.dtype, device=device) + self.model_update_group.broadcast(weight, src=src, stream=stream) @staticmethod def _checked_iter( @@ -437,8 +520,7 @@ def _checked_iter( iteration disagrees with it, the two sides split the stream differently and the transfer hangs in NCCL or loads garbage. Checking here costs one comparison per parameter and turns that into an error naming the first - divergent parameter. The sender owns every parameter, so it yields all - of `meta`. + divergent parameter. """ sent = 0 for name, tensor in source.iter_full(): From b71c1d5f14e593b6c5af9a689c6f6afead33f7e1 Mon Sep 17 00:00:00 2001 From: seungrokj <144636725+seungrokj@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:46:23 -0700 Subject: [PATCH 031/406] [ROCm][Bugfix][Kimi-K3] Build the folded q_b_proj / W_UK weight on the root model (#64158) Signed-off-by: seungrokj Co-authored-by: Claude Opus 5.5 --- tests/models/kimi_k3/test_amd_qb_uk_fold.py | 13 +++++++++++++ vllm/models/kimi_k3/amd/model.py | 5 +++++ 2 files changed, 18 insertions(+) diff --git a/tests/models/kimi_k3/test_amd_qb_uk_fold.py b/tests/models/kimi_k3/test_amd_qb_uk_fold.py index 8a2ab1f360bb..abacd4a8993e 100644 --- a/tests/models/kimi_k3/test_amd_qb_uk_fold.py +++ b/tests/models/kimi_k3/test_amd_qb_uk_fold.py @@ -93,3 +93,16 @@ def test_fold_through_aiter_gemm(num_tokens: int) -> None: ql_nope_ref, q_pe_ref = _split_fold_output(x.float() @ w_fold.float().T, num_heads) torch.testing.assert_close(ql_nope.float(), ql_nope_ref, rtol=2e-2, atol=2e-2) torch.testing.assert_close(q_pe.float(), q_pe_ref, rtol=2e-2, atol=2e-2) + + +def test_root_model_builds_the_fold() -> None: + """The loader calls only the root model's post-load hook, so + KimiK3ForConditionalGeneration must forward it to the text model, whose hook + builds each MLA wrapper's W_fold (else _fused_decode asserts at DCP1).""" + from unittest.mock import MagicMock + + from vllm.models.kimi_k3.amd.model import KimiK3ForConditionalGeneration + + root = MagicMock() + KimiK3ForConditionalGeneration.process_weights_after_loading(root) + root.language_model.process_weights_after_loading.assert_called_once_with() diff --git a/vllm/models/kimi_k3/amd/model.py b/vllm/models/kimi_k3/amd/model.py index aacbd086b511..927214852901 100644 --- a/vllm/models/kimi_k3/amd/model.py +++ b/vllm/models/kimi_k3/amd/model.py @@ -252,3 +252,8 @@ def get_mamba_state_copy_func(cls): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def process_weights_after_loading(self) -> None: + # The loader calls only the root model's hook; forward it so the text + # model's MLA layers build their folded q_b_proj / W_UK weight. + self.language_model.process_weights_after_loading() From 84e455367336b62547d1f1273e9b50bc7f8edb47 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Wed, 28 Oct 2026 11:13:43 +0800 Subject: [PATCH 032/406] [CI][XPU] Give test_mrv2_weight_offloading more GPU memory headroom (#63895) Signed-off-by: Chaojun Zhang Co-authored-by: vllm-agent --- tests/basic_correctness/cpu_offload/test_cpu_offload.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/basic_correctness/cpu_offload/test_cpu_offload.py b/tests/basic_correctness/cpu_offload/test_cpu_offload.py index 46bdb7218bd2..39a162b1ee6f 100644 --- a/tests/basic_correctness/cpu_offload/test_cpu_offload.py +++ b/tests/basic_correctness/cpu_offload/test_cpu_offload.py @@ -66,7 +66,7 @@ def test_mrv2_weight_offloading( with vllm_runner( "hmellor/tiny-random-LlamaForCausalLM", enforce_eager=True, - gpu_memory_utilization=0.02, + gpu_memory_utilization=0.03, max_model_len=128, max_num_seqs=1, **offload_kwargs, From 68eca9af01b8b323667cb4b87148547e8b98eef6 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 28 Oct 2026 11:20:51 +0800 Subject: [PATCH 033/406] [Rust Frontend] Export vllm:num_requests_kv_fetch_by_stage like Python (#63962) Signed-off-by: Bugen Zhao Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- rust/src/engine-core-client/src/metrics.rs | 72 +++++++++++++++++-- .../engine-core-client/src/protocol/stats.rs | 61 +++++++++++++++- rust/src/llm/src/log_stats.rs | 43 ++++++++++- rust/src/metrics/src/scheduler.rs | 19 +++++ 4 files changed, 185 insertions(+), 10 deletions(-) diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index 041e163c4d6a..22680427970c 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -7,14 +7,15 @@ use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ CacheHitSourceLabels, EngineLabels, EnginePositionLabels, EngineSleepStateLabels, F64Gauge, - Family, HistogramMetric, LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily, - MooncakeOperationHistogramFamily, MooncakeOperationLabels, RequestMetrics, - SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, WaitingReasonLabels, + Family, HistogramMetric, KvFetchStageLabels, LoraAdapterNames, LoraInfoLabels, + MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, MooncakeOperationLabels, + RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, + WaitingReasonLabels, }; use crate::protocol::stats::{ - CachedTokensBySource, KvConnectorStats, MooncakeStats, MooncakeTransferStats, - MultiConnectorStats, NixlStats, SchedulerStats, + CachedTokensBySource, KvConnectorStats, KvFetchReqsByStage, MooncakeStats, + MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats, }; use crate::protocol::utility::SleepLevel; use crate::transport::ConnectedEngine; @@ -82,6 +83,8 @@ struct SchedulerStatsHandles { scheduler_waiting: U64Gauge, scheduler_waiting_capacity: U64Gauge, scheduler_waiting_deferred: U64Gauge, + /// One handle per `KvFetchReqsByStage::STAGES` entry. + scheduler_kv_fetch_by_stage: [U64Gauge; 3], kv_cache_usage: F64Gauge, /// One handle per [`SLEEP_STATES`] entry. engine_sleep_state: [U64Gauge; 3], @@ -214,6 +217,13 @@ fn resolve_scheduler_stats_handles( scheduler_waiting_deferred: metrics .scheduler_waiting_by_reason .get_or_create_owned(&deferred), + scheduler_kv_fetch_by_stage: KvFetchReqsByStage::STAGES.map(|stage| { + metrics.scheduler_kv_fetch_by_stage.get_or_create_owned(&KvFetchStageLabels { + model_name: model_name.to_string(), + engine, + stage, + }) + }), kv_cache_usage: metrics.kv_cache_usage.get_or_create_owned(&labels), engine_sleep_state: SLEEP_STATES.map(|sleep_state| { metrics.engine_sleep_state.get_or_create_owned(&EngineSleepStateLabels { @@ -300,6 +310,13 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: & .set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs); handles.scheduler_waiting_capacity.set(stats.num_waiting_reqs); handles.scheduler_waiting_deferred.set(stats.num_skipped_waiting_reqs); + for (handle, count) in handles + .scheduler_kv_fetch_by_stage + .iter() + .zip(stats.num_kv_fetch_reqs_by_stage.counts()) + { + handle.set(count); + } handles.kv_cache_usage.set(stats.kv_cache_usage); // Prefix-cache counters, including the connector-backed external cache path. @@ -538,12 +555,13 @@ mod tests { use std::collections::{BTreeMap, BTreeSet}; use expect_test::expect; + use itertools::Itertools as _; use vllm_metrics::Metrics; use crate::metrics::LoraInfoExporter; use crate::protocol::stats::{ - KvConnectorStats, MooncakeOperation, MooncakeRecord, MooncakeStats, MooncakeStatus, - MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats, + KvConnectorStats, KvFetchReqsByStage, MooncakeOperation, MooncakeRecord, MooncakeStats, + MooncakeStatus, MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats, }; fn names(values: &[&str]) -> BTreeSet { @@ -673,6 +691,46 @@ mod tests { .assert_eq(&mooncake_transfer_series); } + /// Exports every engine's KV fetch stage gauges from startup and sets them + /// from each `SchedulerStats` update. + #[test] + fn kv_fetch_reqs_by_stage_are_recorded_into_stage_gauges() { + let metrics = Metrics::new(); + let handles = super::resolve_scheduler_stats_handles(&metrics.scheduler, "model", 0); + // Series of one family render in hash order, so sort them. + let kv_fetch_series = || { + metrics + .render() + .unwrap() + .lines() + .filter(|line| line.starts_with("vllm:num_requests_kv_fetch_by_stage{")) + .sorted() + .join("\n") + }; + + expect![[r#" + vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="completed_waiting"} 0 + vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="in_progress"} 0 + vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="waiting_to_start"} 0"#]] + .assert_eq(&kv_fetch_series()); + + let stats = SchedulerStats { + num_kv_fetch_reqs_by_stage: KvFetchReqsByStage { + waiting_to_start: 3, + in_progress: 1, + completed_waiting: 2, + }, + ..Default::default() + }; + super::record_scheduler_stats_with_handles(&handles, &stats); + + expect![[r#" + vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="completed_waiting"} 2 + vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="in_progress"} 1 + vllm:num_requests_kv_fetch_by_stage{model_name="model",engine="0",stage="waiting_to_start"} 3"#]] + .assert_eq(&kv_fetch_series()); + } + /// Records one MultiConnector payload into Mooncake and NIXL metrics. #[test] fn kv_connector_stats_are_recorded_into_mooncake_and_nixl_metrics() { diff --git a/rust/src/engine-core-client/src/protocol/stats.rs b/rust/src/engine-core-client/src/protocol/stats.rs index c56f32eb244e..b0491e6a021c 100644 --- a/rust/src/engine-core-client/src/protocol/stats.rs +++ b/rust/src/engine-core-client/src/protocol/stats.rs @@ -49,6 +49,32 @@ impl CachedTokensBySource { } } +/// Waiting requests with an async KV load, per load stage. Python sends a +/// `dict[str, int]` keyed by stage, so absent stages default to zero. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default)] +pub struct KvFetchReqsByStage { + /// Needs an async KV load that has not started. + pub waiting_to_start: u64, + /// Load started and not yet reported finished. + pub in_progress: u64, + /// Load finished, request not running yet. + pub completed_waiting: u64, +} + +impl KvFetchReqsByStage { + /// Label values, in the same order as [`Self::counts`]. + pub const STAGES: [&'static str; 3] = ["waiting_to_start", "in_progress", "completed_waiting"]; + + pub fn counts(&self) -> [u64; 3] { + [ + self.waiting_to_start, + self.in_progress, + self.completed_waiting, + ] + } +} + /// Stores prefix cache hit statistics. /// - `reset`: Whether `reset_prefix_cache` was invoked. /// - `queries`: Refers to the number of tokens that were queried. @@ -349,6 +375,9 @@ pub struct SchedulerStats { /// Number of waiting requests deferred by transient constraints or a blocked status. #[serde(default)] pub num_skipped_waiting_reqs: u64, + /// Waiting requests with an async KV load, by load stage. + #[serde(default)] + pub num_kv_fetch_reqs_by_stage: KvFetchReqsByStage, /// Internal DP load-balancing step counter. pub step_counter: u64, /// Internal DP load-balancing wave number. @@ -373,7 +402,10 @@ pub struct SchedulerStats { #[cfg(test)] mod tests { - use super::{CachedTokensBySource, KvConnectorStats, PrefixCacheStats}; + use super::{ + CachedTokensBySource, KvConnectorStats, KvFetchReqsByStage, PrefixCacheStats, + SchedulerStats, + }; #[test] fn prefix_cache_stats_hits_by_source_defaults_and_decodes() { @@ -432,4 +464,31 @@ mod tests { if stats.nixl.is_some() && stats.mooncake_transfer.is_some() )); } + + #[test] + fn scheduler_stats_kv_fetch_reqs_by_stage_defaults_and_decodes() { + let mut payload = serde_json::json!({ + "num_running_reqs": 1, "num_waiting_reqs": 2, + "step_counter": 0, "current_wave": 0, "kv_cache_usage": 0.5, + "prefix_cache_stats": { + "reset": false, "requests": 0, "queries": 0, "hits": 0, + "preempted_requests": 0, "preempted_queries": 0, "preempted_hits": 0 + }, + "kv_cache_eviction_events": [] + }); + + // Engines before vllm-project/vllm#58874 omit the field entirely. + let wire = rmp_serde::to_vec_named(&payload).unwrap(); + let stats: SchedulerStats = rmp_serde::from_slice(&wire).unwrap(); + assert_eq!( + stats.num_kv_fetch_reqs_by_stage, + KvFetchReqsByStage::default() + ); + + payload["num_kv_fetch_reqs_by_stage"] = + serde_json::json!({"waiting_to_start": 3, "in_progress": 1}); + let wire = rmp_serde::to_vec_named(&payload).unwrap(); + let stats: SchedulerStats = rmp_serde::from_slice(&wire).unwrap(); + assert_eq!(stats.num_kv_fetch_reqs_by_stage.counts(), [3, 1, 0]); + } } diff --git a/rust/src/llm/src/log_stats.rs b/rust/src/llm/src/log_stats.rs index d16c8bcbe3a3..70e575706479 100644 --- a/rust/src/llm/src/log_stats.rs +++ b/rust/src/llm/src/log_stats.rs @@ -6,9 +6,11 @@ use std::time::{Duration, Instant}; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, info}; +use vllm_engine_core_client::protocol::stats::KvFetchReqsByStage; use vllm_metrics::{ - EngineLabels, F64Gauge, METRICS, PromptTokenSourceLabels, SchedulerLogStatsAccumulator, - SchedulerLogStatsInterval, U64Counter, U64Gauge, WaitingReasonLabels, + EngineLabels, F64Gauge, KvFetchStageLabels, METRICS, PromptTokenSourceLabels, + SchedulerLogStatsAccumulator, SchedulerLogStatsInterval, U64Counter, U64Gauge, + WaitingReasonLabels, }; const WAITING_REASON_DEFERRED: &str = "deferred"; @@ -37,6 +39,8 @@ struct EngineMetrics { scheduler_running: U64Gauge, scheduler_waiting: U64Gauge, scheduler_deferred: U64Gauge, + /// One handle per `KvFetchReqsByStage::STAGES` entry. + scheduler_kv_fetch_by_stage: [U64Gauge; 3], kv_cache_usage: F64Gauge, } @@ -162,6 +166,15 @@ fn resolve_engine_metrics(model_name: &str, engine_indices: &[u32]) -> Vec, interval // Read scheduler gauges (aggregate across engines). let (num_running, num_waiting, kv_cache_usage) = read_scheduler_gauges(&engines); let num_deferred = read_deferred_waiting(&engines); + let [ + kv_fetch_waiting_to_start, + kv_fetch_in_progress, + kv_fetch_completed_waiting, + ] = read_kv_fetch_by_stage(&engines); let delta_preemptions = curr.num_preemptions.wrapping_sub(prev.num_preemptions); // Compute prefix cache hit rate over this interval. @@ -252,6 +270,15 @@ async fn run_stats_logger(model_name: String, engine_indices: Vec, interval if num_deferred > 0 { write!(msg, ", Deferred: {num_deferred} reqs").unwrap(); } + if kv_fetch_waiting_to_start + kv_fetch_in_progress + kv_fetch_completed_waiting > 0 { + write!( + msg, + ", KV fetch: {kv_fetch_waiting_to_start} waiting to start, \ + {kv_fetch_in_progress} in progress, \ + {kv_fetch_completed_waiting} completed waiting" + ) + .unwrap(); + } if delta_preemptions > 0 { write!(msg, ", Preemptions: {delta_preemptions}").unwrap(); } @@ -366,6 +393,18 @@ fn read_deferred_waiting(engines: &[EngineMetrics]) -> u64 { engines.iter().map(|e| e.scheduler_deferred.get()).sum() } +/// Read waiting requests per async KV load stage across all engines, in +/// `KvFetchReqsByStage::STAGES` order. +fn read_kv_fetch_by_stage(engines: &[EngineMetrics]) -> [u64; 3] { + let mut counts = [0u64; 3]; + for e in engines { + for (count, gauge) in counts.iter_mut().zip(&e.scheduler_kv_fetch_by_stage) { + *count += gauge.get(); + } + } + counts +} + /// Return the cache hit rate as a percentage for a counter delta. fn cache_hit_rate(hits: u64, queries: u64) -> f64 { if queries > 0 { diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs index d5327a8cea0b..f8b00310058d 100644 --- a/rust/src/metrics/src/scheduler.rs +++ b/rust/src/metrics/src/scheduler.rs @@ -100,6 +100,13 @@ pub struct EngineSleepStateLabels { pub sleep_state: &'static str, } +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] +pub struct KvFetchStageLabels { + pub model_name: String, + pub engine: u32, + pub stage: &'static str, +} + /// Labels for per-operation Mooncake store connector telemetry. `operation` /// and `status` are dynamic (not known ahead of time), unlike the other /// fixed-label families in this module. @@ -225,6 +232,7 @@ pub struct SchedulerMetrics { pub scheduler_running: Family, pub scheduler_waiting: Family, pub scheduler_waiting_by_reason: Family, + pub scheduler_kv_fetch_by_stage: Family, pub kv_cache_usage: Family, /// `vllm:engine_sleep_state`, set by the engine client's sleep and wake @@ -316,6 +324,16 @@ impl SchedulerMetrics { scheduler_waiting_by_reason.clone(), ); + let scheduler_kv_fetch_by_stage = Family::default(); + registry.register( + "vllm:num_requests_kv_fetch_by_stage", + "Number of waiting requests by async KV load stage. \ + Stage labels: 'waiting_to_start' = needs an async KV load that has not started; \ + 'in_progress' = load started and not yet reported finished; \ + 'completed_waiting' = load finished, request not running yet.", + scheduler_kv_fetch_by_stage.clone(), + ); + let kv_cache_usage = Family::default(); registry.register( "vllm:kv_cache_usage_perc", @@ -618,6 +636,7 @@ impl SchedulerMetrics { scheduler_running, scheduler_waiting, scheduler_waiting_by_reason, + scheduler_kv_fetch_by_stage, kv_cache_usage, engine_sleep_state, lora_info, From b766f15983bd7b926c93015dd8551a6101a2ca20 Mon Sep 17 00:00:00 2001 From: Hritik Raj Date: Wed, 28 Oct 2026 09:12:11 +0530 Subject: [PATCH 034/406] [Bugfix][Frontend] fix Anthropic inline-system probe crashing startup (#64020) Signed-off-by: Hritik Raj Co-authored-by: vllm-agent --- .../test_anthropic_messages_conversion.py | 21 +++++++++++++++++++ vllm/entrypoints/anthropic/serving.py | 4 +++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index aa50a16b2365..0525272576dc 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -1454,6 +1454,27 @@ def test_no_template_defaults_merge(self): """No chat_template → conservative default: merge.""" assert AnthropicServingMessages._detect_merge_inline_system(None) is True + def test_tools_none_check_template_merges(self): + """OLMo 3-style ``tools | tojson`` on undefined tools falls back to merge.""" + template = ( + "{% if tools is none %}{% else %}{{ tools | tojson }}{% endif %}" + "{% for m in messages %}{{ m.content }}{% endfor %}" + ) + assert AnthropicServingMessages._detect_merge_inline_system(template) is True + + def test_tools_loop_template_no_merge(self): + """Hermes-style ``for tool in tools`` must not raise in the probe.""" + template = ( + "{% for tool in tools %}{{ tool | tojson }}{% endfor %}" + "{% for m in messages %}{{ m.content }}{% endfor %}" + ) + assert AnthropicServingMessages._detect_merge_inline_system(template) is False + + def test_non_template_error_defaults_merge(self): + """A non-TemplateError render failure must not abort startup.""" + template = "{{ undefined_var | tojson }}" + assert AnthropicServingMessages._detect_merge_inline_system(template) is True + @pytest.mark.parametrize( ("resolved", "expected"), [(PERMISSIVE_TEMPLATE, False), (Q35_TEMPLATE, True)], diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index e1b362ab4136..6d653616b148 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -324,7 +324,9 @@ def _detect_merge_inline_system(chat_template: str | None) -> bool: add_generation_prompt=False, ) return False - except jinja2.TemplateError: + except Exception: + # Non-TemplateError failures (e.g. TypeError from filters on + # undefined variables) must not abort server startup. return True @staticmethod From 9bc1c8b4fede8d009b308fecf74e560a1b4f7eca Mon Sep 17 00:00:00 2001 From: yiminghub2024 <482890@qq.com> Date: Wed, 28 Oct 2026 11:54:50 +0800 Subject: [PATCH 035/406] [Doc] Add autotrust/JEV-27B-VL multimodal recipe (vLLM >= 0.32.0) (#64180) Signed-off-by: Cursor Agent Co-authored-by: Cursor Agent Co-authored-by: yiminghub2024 Co-authored-by: vllm-agent --- recipes/models/autotrust/JEV-27B-VL.yaml | 200 +++++++++++++++++++++++ 1 file changed, 200 insertions(+) create mode 100644 recipes/models/autotrust/JEV-27B-VL.yaml diff --git a/recipes/models/autotrust/JEV-27B-VL.yaml b/recipes/models/autotrust/JEV-27B-VL.yaml new file mode 100644 index 000000000000..61a0ae1327ac --- /dev/null +++ b/recipes/models/autotrust/JEV-27B-VL.yaml @@ -0,0 +1,200 @@ +# yaml-language-server: $schema=https://recipes.vllm.ai/recipe.schema.json + +meta: + title: "JEV-27B-VL" + slug: "jev-27b-vl" + provider: "AutoTrust" + description: "Dense 27B vision-language decision model on a Qwen3.8-27B backbone, with a rank-32 LoRA (jev-decision) for calibrated choices over text and images." + date_added: 2026-10-28 + date_updated: 2026-10-28 + tasks: + - multimodal + - text + performance_headline: "Qwen3.8-27B VL backbone; System 1 decisions through the jev-decision LoRA" + related_recipes: + - "Qwen/Qwen3.8-27B" + +model: + model_id: "autotrust/JEV-27B-VL" + # Model author: this checkpoint only works on vLLM >= 0.32.0. + # Qwen3_5ForConditionalGeneration is registered in that release. + min_vllm_version: "0.32.0" + architecture: dense + # safetensors parameter total is 27,781,427,952 BF16 (vision tower included). + parameter_count: "27B" + active_parameters: "27B" + # text_config.max_position_embeddings. The validated serve line caps the + # window at 32768 via --max-model-len; 256K is the backbone native length. + context_length: 262144 + base_args: + - "--served-model-name" + - "autotrust/JEV-27B-VL" + - "--enable-lora" + - "--max-lora-rank" + - "32" + - "--lora-modules" + - "jev-decision=JEV-27B-VL/adapter_vllm" + - "--logprobs-mode" + - "processed_logprobs" + - "--max-model-len" + - "32768" + - "--enable-prefix-caching" + - "--mamba-cache-mode" + - "align" + - "--limit-mm-per-prompt" + - '{"image": 8}' + - "--max-num-seqs" + - "8" + - "--trust-request-chat-template" + base_env: {} + +# Host checkout for the LoRA. vLLM serve still pulls weights by model_id; +# --lora-modules needs adapter_vllm on disk. Same subdirectory on the +# ModelScope checkout. install_modes includes docker because the image +# cannot fetch this relative path for you. +dependencies: + - note: "Local layout for LoRA module jev-decision. After this download the adapter is JEV-27B-VL/adapter_vllm; a ModelScope checkout uses the same adapter_vllm directory under whatever folder you chose." + command: "hf download autotrust/JEV-27B-VL --local-dir JEV-27B-VL" + install_modes: [pip, docker] + +features: + text_only: + description: "Skip the vision encoder for text-only workloads and free its memory for KV cache." + args: + - "--language-model-only" + encoder_parallel: + description: "Run the vision encoder data-parallel across the tensor-parallel ranks instead of sharding it." + args: + - "--mm-encoder-tp-mode" + - "data" + +opt_in_features: + - text_only + - encoder_parallel + +variants: + default: + precision: bf16 + # 55,562,855,904 bytes on disk; ceil(55.563 GB × 1.2) = 67. + # Uniform BF16, so params × 2 × 1.2 lands on the same budget. + vram_minimum_gb: 67 + description: "BF16 weights (about 52–56 GB). The model card's quick start targets one GPU with 80 GB or more; this recipe does not mark any GPU verified." + +compatible_strategies: + - single_node_tp + - multi_node_tp + - multi_node_tp_pp + +# serve.sh does not set tensor-parallel-size: one GPU. 67 GB fits a single +# high-memory GPU, so the single-node default stays TP 1 instead of the +# whole node. +strategy_overrides: + single_node_tp: + tp: 1 + +guide: | + ## Overview + + [JEV-27B-VL](https://huggingface.co/autotrust/JEV-27B-VL) is AutoTrust's + vision-language decision model. The checkpoint is a dense BF16 + `Qwen3_5ForConditionalGeneration` model (about 27.8B parameters, vision + tower included) continued from + [Qwen/Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B). System 2 is + the served checkpoint, named `autotrust/JEV-27B-VL`. System 1 is the + rank-32 LoRA `jev-decision` (`adapter_vllm/`, `r` 32, `lora_alpha` 64), + used for calibrated choice probabilities over text or images. + + This recipe has no `verified` hardware entry. The model card's quick start + asks for one GPU with 80 GB or more, and the card's published latency + figures are from a B200. Those runs are the author's, not a badge on this + page. + + ## Prerequisites + + - **vLLM:** >= 0.32.0. Earlier releases are not supported for this + checkpoint. + - **Weights:** `hf download autotrust/JEV-27B-VL --local-dir JEV-27B-VL` + so the LoRA path below exists. The same `adapter_vllm` directory is in + the [ModelScope](https://modelscope.cn/models/autotrust/JEV-27B-VL) + checkout; if that folder is not named `JEV-27B-VL`, point the module at + `/adapter_vllm` (the HF snapshot path ends in + `snapshots//adapter_vllm`). + - **Docker:** the image only bind-mounts the Hugging Face cache. Mount the + checkout into the container and keep `jev-decision` pointed at that + `adapter_vllm` directory. + - **Context:** `text_config.max_position_embeddings` is 262144. The + validated command uses `--max-model-len 32768` (the `serve.sh` default). + The card says a full 256K prompt needs about 17 GB of KV cache on top of + the weights, and suggests `--max-model-len 131072` on an 80 GB GPU. + + ## Decision LoRA and serve flags + + These flags are the model card's validated `vllm serve` line. They are + base arguments, not optional toggles. + + - **`jev-decision`** — `--enable-lora`, `--max-lora-rank 32`, and + `--lora-modules jev-decision=JEV-27B-VL/adapter_vllm`. The module name + stays `jev-decision`. The path is the local checkout + `JEV-27B-VL/adapter_vllm`, or the `adapter_vllm` directory inside an + HF or ModelScope checkout. + - **`--mamba-cache-mode align`** — the Qwen3.8 hybrid backbone mixes + full attention with Mamba-style layers; the card serves it in `align` + mode. + - **`--limit-mm-per-prompt '{"image": 8}'`** — at most 8 images per + request. + - **`--max-num-seqs 8`** — required. The card says that with more than 8 + sequences in one batch, vLLM's LoRA path for this multimodal class + returns wrong System 1 probabilities. Extra requests queue. + - **`--logprobs-mode processed_logprobs`** — System 1 reads post-processor + logprobs. On `/v1/chat/completions` or `/v1/completions` for + `jev-decision`, the card also says to send `top_k: 0` and `top_p: 1.0`: + `generation_config.json` sets `top_k` 20 and `top_p` 0.95, and those + defaults would truncate the returned distribution. + - **`--trust-request-chat-template`** — image decisions send a raw + decision chat template on the request. System 2 keeps the checkpoint's + own template. + - **`--served-model-name autotrust/JEV-27B-VL`** — public name for + System 2. `--enable-prefix-caching` matches the card's serve line. + + `POST /v1/decide` is not part of stock `vllm serve`. The checkpoint ships + `serve_decide.py`, which starts the OpenAI server with these same flags + and adds that route. Stock `vllm serve` still exposes System 2 as + `autotrust/JEV-27B-VL` and System 1 as the LoRA model `jev-decision`. + + ## Client usage + + System 2 (the base checkpoint), with an image: + + ```python + client.chat.completions.create( + model="autotrust/JEV-27B-VL", + messages=[{ + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": image_url}}, + {"type": "text", "text": "What does this image show?"}, + ], + }], + ) + ``` + + System 1 through the checkpoint's `/v1/decide` route (`serve_decide.py`): + + ```bash + curl localhost:8000/v1/decide -H 'Content-Type: application/json' -d '{ + "kind": "choice", + "state": "Customer: my card was charged twice for one coffee.", + "question": "Which team should handle this?", + "options": ["billing", "shipping", "tech support"]}' + ``` + + Text-only System 1 calls the LoRA directly (`model` `jev-decision`). The + card says those text decisions match + [autotrust/JEV-27B](https://huggingface.co/autotrust/JEV-27B). + + ## References + + - [Hugging Face model card](https://huggingface.co/autotrust/JEV-27B-VL) + - [ModelScope (autotrust/JEV-27B-VL)](https://modelscope.cn/models/autotrust/JEV-27B-VL) + - [ModelScope mirror listed on the card (AutoTrust-Labs/JEV-27B-VL)](https://www.modelscope.ai/models/AutoTrust-Labs/JEV-27B-VL) + - [Base model: Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B) From 7e3e3f6ab652550fdefecd8e9823394fd8e51950 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 27 Oct 2026 23:11:19 -0500 Subject: [PATCH 036/406] [CI][Bugfix] Pin the ColUltraVBERT test checkpoint to the revision vLLM supports (#64189) Signed-off-by: Andreas Karatzas Co-authored-by: Andreas Karatzas Co-authored-by: Claude --- tests/models/multimodal/pooling/test_colultravbert.py | 3 +++ tests/models/registry.py | 6 +++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/models/multimodal/pooling/test_colultravbert.py b/tests/models/multimodal/pooling/test_colultravbert.py index 81deedcb0334..d26d5011c715 100644 --- a/tests/models/multimodal/pooling/test_colultravbert.py +++ b/tests/models/multimodal/pooling/test_colultravbert.py @@ -9,6 +9,8 @@ from vllm.model_executor.models.colultravbert import get_capped_longest_edge MODEL_NAME = "athrael-soju/ColUltraVBERT" +# Later revisions declare the UltraVBertModel architecture instead. +REVISION = "509adbab61b71da0484b20f3baf718e15219d761" DTYPE = "bfloat16" PAGE_PROMPT = "User:Describe the image.\nAssistant:" IMAGE_SEQ_LEN = 64 @@ -53,6 +55,7 @@ def test_colultravbert_tile_cap(width, height, grid): def colultravbert_model(vllm_runner): with vllm_runner( MODEL_NAME, + revision=REVISION, runner="pooling", dtype=DTYPE, max_model_len=2048, diff --git a/tests/models/registry.py b/tests/models/registry.py index 7b8223ccad97..8ff1654ed0c9 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -701,7 +701,11 @@ def check_available_online( "ColModernVBertForRetrieval": _HfExamplesInfo( "ModernVBERT/colmodernvbert-merged", ), - "ColUltraVBertForRetrieval": _HfExamplesInfo("athrael-soju/ColUltraVBERT"), + "ColUltraVBertForRetrieval": _HfExamplesInfo( + "athrael-soju/ColUltraVBERT", + # Later revisions declare the UltraVBertModel architecture instead. + revision="509adbab61b71da0484b20f3baf718e15219d761", + ), "ColPaliForRetrieval": _HfExamplesInfo("vidore/colpali-v1.3-hf"), "ModernVBertModel": _HfExamplesInfo("ModernVBERT/colmodernvbert-merged"), "ColQwen3": _HfExamplesInfo( From 88d3995ea18de47ba5c306789aa615474ad74592 Mon Sep 17 00:00:00 2001 From: AlanFokCo <49893220+AlanFokCo@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:18:16 +0800 Subject: [PATCH 037/406] [Feature][Core] Add opt-in multi-worker CUDA checkpoint sleep (#64186) Signed-off-by: AlanFokCo Co-authored-by: aoshen02 --- docs/features/sleep_mode.md | 131 ++++++ .../sleep_mode/test_process_checkpoint.py | 133 ++++++ .../executor/test_cuda_checkpoint_executor.py | 417 ++++++++++++++++++ .../executor/test_cuda_process_checkpoint.py | 231 ++++++++++ vllm/utils/cuda_process_checkpoint.py | 248 +++++++++++ vllm/v1/executor/cuda_checkpoint_executor.py | 297 +++++++++++++ 6 files changed, 1457 insertions(+) create mode 100644 tests/basic_correctness/memory/sleep_mode/test_process_checkpoint.py create mode 100644 tests/v1/executor/test_cuda_checkpoint_executor.py create mode 100644 tests/v1/executor/test_cuda_process_checkpoint.py create mode 100644 vllm/utils/cuda_process_checkpoint.py create mode 100644 vllm/v1/executor/cuda_checkpoint_executor.py diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index bb6c6559228c..8c73dd16a624 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -242,3 +242,134 @@ curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache' ## Limitation On ROCm, the virtual memory allocation on ROCm is done through chunked memory allocation. You can control the chunk size through `VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE` (in MB). The default value is set at 256MB. The larger the chunk size the faster the performance. However, setting it too large will cause OOM. So if you encounter OOM when using sleep mode. Try reducing the chunk size. It is recommended to define the chunk size as a power of 2. + +## Experimental CUDA checkpoint executor + +This capability is experimental and disabled by default. Selecting the executor +below checkpoints each CUDA worker after ordinary level-1 sleep, +then restores all workers before waking allocator-managed memory. This preserves +the remaining CUDA state, including captured graph executable objects, without +recapturing graphs. Ordinary level-1 semantics still apply: KV cache contents are +discarded. This is not a durable application checkpoint or model switching API. + +Set `VLLM_USE_V2_MODEL_RUNNER=1` before starting the engine: + +```python +from vllm import LLM + +llm = LLM( + "Qwen/Qwen3-0.6B", + enable_sleep_mode=True, + distributed_executor_backend=( + "vllm.v1.executor.cuda_checkpoint_executor.CUDACheckpointExecutor" + ), +) +llm.sleep(level=1) +llm.wake_up() +``` + +This experimental implementation requires Linux, NVIDIA driver checkpoint APIs, +and dedicated worker processes. It uses an external helper to call +`cuCheckpointProcessLock`, `cuCheckpointProcessCheckpoint`, +`cuCheckpointProcessRestore`, and `cuCheckpointProcessUnlock`. CPU threads remain +alive, but worker CUDA calls cannot execute while checkpointed. The helper uses +the same process namespace and user as the executor. Driver permission and +resource restrictions still apply; see the +[CUDA checkpoint API documentation](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-driver-api/group__CUDA__CHECKPOINT.html). + +The initial validation used NVIDIA H20 with driver 580.126.09, NCCL 2.29.7, +and FlashInfer 0.6.18.post1. See [PR #58183](https://github.com/vllm-project/vllm/pull/58183) +for versioned test results and communication performance comparisons. Support +for other GPU, driver, and communication-library combinations is not established. + +Only single-node tensor parallelism is accepted. Tensor parallel sizes 1, 2, +and 4 have been GPU-tested; other sizes remain unvalidated. Use +`VLLM_USE_V2_MODEL_RUNNER=1`, the standard GPU worker and the `cumem` sleep +backend. Pipeline, data, context and expert parallelism, external cache or +weight transfer, custom worker classes, and sleep level 2 are rejected before +worker startup (sleep level 2 is rejected when requested). Worker extensions +that allocate additional CUDA IPC resources are outside the validated scope. +Level 0 remains +a scheduler pause and does not perform a CUDA checkpoint. + +For tensor parallel sizes greater than 1, use NCCL 2.29.7 or later with memory +suspension support, set `enable_nccl_comm_suspend=True`, and pass +`disable_custom_all_reduce=True`. The following explicit configuration retains +NCCL GPU peer-to-peer or legacy shared-memory communication, selected by NCCL. +`NCCL_P2P_DISABLE` and `NCCL_SHM_DISABLE` need not be set when communicator +suspension is enabled. NCCL CUDA host-memory allocation remains disabled; it +is not equivalent to legacy shared-memory transport. Before driver checkpoint, +ordinary worker +sleep suspends NCCL mappings; worker wake restores them at their reserved +addresses. The executor checks each rank's actual suspension state. + +```bash +export NCCL_IB_DISABLE=1 +export NCCL_CUMEM_ENABLE=1 NCCL_CUMEM_HOST_ENABLE=0 NCCL_NVLS_ENABLE=0 +export VLLM_ALLREDUCE_USE_SYMM_MEM=0 VLLM_USE_NCCL_SYMM_MEM=0 +export VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC=0 +export VLLM_ALLREDUCE_USE_FLASHINFER=1 +export VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm +``` + +FlashInfer must provide stable-address workspace `checkpoint_prepare` and +`checkpoint_restore` methods. The executor calls existing worker hooks around +the driver checkpoint. Disabling standalone FlashInfer all-reduce with +`VLLM_ALLREDUCE_USE_FLASHINFER=0` does not disable the compiler +`fuse_allreduce_rms` pass, which can allocate FlashInfer workspaces independently. +When that pass is enabled, the same backend restriction and workspace hooks +apply. A configuration using only NCCL must also disable that compiler pass. +The FlashInfer `mnnvl` backend is not accepted, and the NCCL suspend configuration +requires NVLS to be disabled. Live multicast +allocations can block the combined sleep/checkpoint flow even on newer +user-mode drivers. Keep these restrictions when upgrading the driver. The +helper must use the same user-mode driver as the workers; mixed versions are +unsupported. Performance costs of the restrictions must be measured against +the workload's default communication +configuration; retaining GPU peer-to-peer communication alone does not establish +unchanged inference performance. + +A more restrictive Socket-only configuration is also accepted without +`enable_nccl_comm_suspend`: set `NCCL_P2P_DISABLE=1`, `NCCL_SHM_DISABLE=1`, +`NCCL_CUMEM_ENABLE=0`, and +`VLLM_ALLREDUCE_USE_FLASHINFER=0`, retaining the other communication restrictions +above. This configuration is primarily useful for compatibility comparisons. +Do not create additional CUDA IPC resources or fall back to an independently +created PyTorch NCCL process group while using the checkpoint executor. + +Provide host memory for both the ordinary weight backup and the driver's +remaining CUDA state backup. Restore needs sufficient free GPU memory; release +any temporary GPU workload before waking. Checkpoint and restore add latency to +ordinary sleep/wake. The helper itself can retain a small driver allocation, so +worker offload does not imply a zero device-wide memory reading. + +All ranks are locked before any is checkpointed. All ranks are restored before +any is unlocked. Operations within each stage run concurrently across workers. +The whole stage has a 30-second response deadline (each lock also has a +10-second driver timeout). Worker sleep/wake and FlashInfer workspace +hooks have a 120-second response deadline. Failure terminates the workers rather than +serving with a partially restored rank set; restart the engine after failure. +Shutdown terminates the dedicated worker and helper processes without restoring +their CUDA state or requiring free GPU memory for the model. + +Repeated level-1 sleep is idempotent. A first partial wake, such as +`wake_up(tags=["weights"])`, restores the **whole CUDA process** before restoring +the selected allocator tags. It cannot selectively restore graph or driver +state. Ordinary worker RPCs are rejected while checkpointed. A checkpointed +worker cannot execute another model until restored. + +Serialize sleep, wake, pause, and resume calls across all clients, including +while level-1 sleep is draining requests. Wait for each call to complete before +issuing the next. A partial wake does not automatically resume scheduling; wake +all allocator tags before explicitly resuming or submitting new inference. +Do not use an explicit resume call to bypass sleeping memory. + +If EngineCore exits during a management call, the client may remain waiting even +after the workers terminate. Use an external timeout and restart the failed +engine; a timeout does not establish that the operation completed. Do not retry +control operations on an engine whose checkpoint outcome is unknown. + +As with the ordinary executor, `sleep()` is a no-op while weights remain asleep. +For example, after `wake_up(tags=["kv_cache"])`, another `sleep()` does not +checkpoint the process again. Wake the remaining weights before starting a new +full sleep/checkpoint cycle. diff --git a/tests/basic_correctness/memory/sleep_mode/test_process_checkpoint.py b/tests/basic_correctness/memory/sleep_mode/test_process_checkpoint.py new file mode 100644 index 000000000000..b60394353bf2 --- /dev/null +++ b/tests/basic_correctness/memory/sleep_mode/test_process_checkpoint.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in Linux/NVIDIA driver integration test, with exclusive GPU access. + +Run with VLLM_TEST_CUDA_PROCESS_CHECKPOINT=1. The driver must expose CUDA +process checkpoint APIs. Override VLLM_TEST_MODEL and VLLM_TEST_SLEEP_CYCLES +for larger-model validation; this test requires enough RAM for weight backups. +The selected model must support batch invariance. Hybrid GDN models require +separate paired accuracy evaluation because that mode is unsupported for GDN. +""" + +import os + +import pytest + +from vllm import LLM, SamplingParams + + +def _graph_identity(worker): + import torch + + if not hasattr(worker, "_checkpoint_test_replays"): + worker._checkpoint_test_replays = 0 + worker._checkpoint_test_captures = 0 + replay = torch.cuda.CUDAGraph.replay + capture = torch.cuda.CUDAGraph.capture_begin + + def counted_replay(graph): + worker._checkpoint_test_replays += 1 + return replay(graph) + + def counted_capture(graph, *args, **kwargs): + worker._checkpoint_test_captures += 1 + return capture(graph, *args, **kwargs) + + torch.cuda.CUDAGraph.replay = counted_replay + torch.cuda.CUDAGraph.capture_begin = counted_capture + torch.accelerator.synchronize() + graphs = worker.model_runner.cudagraph_manager.graphs + return { + "pid": os.getpid(), + "graphs": sorted(g.raw_cuda_graph_exec() for g in graphs.values()), + "replays": worker._checkpoint_test_replays, + "captures": worker._checkpoint_test_captures, + } + + +@pytest.mark.skipif( + os.environ.get("VLLM_TEST_CUDA_PROCESS_CHECKPOINT") != "1", + reason="requires exclusive GPUs and a checkpoint-capable NVIDIA driver", +) +@pytest.mark.parametrize( + "tp,transport", [(1, "p2p"), (2, "p2p"), (4, "p2p"), (2, "shm"), (4, "shm")] +) +def test_process_checkpoint_preserves_outputs_and_graphs(tp, transport, monkeypatch): + import torch + + if torch.accelerator.device_count() < tp: + pytest.skip(f"requires {tp} GPUs") + config = { + "VLLM_USE_V2_MODEL_RUNNER": "1", + "VLLM_WORKER_MULTIPROC_METHOD": "spawn", + "VLLM_ALLOW_INSECURE_SERIALIZATION": "1", # Test observation RPC only. + "VLLM_BATCH_INVARIANT": "1", + "NCCL_P2P_DISABLE": "1" if transport == "shm" else "0", + "NCCL_SHM_DISABLE": "0" if transport == "shm" else "1", + "NCCL_IB_DISABLE": "1", + "NCCL_CUMEM_ENABLE": "1", + "NCCL_CUMEM_HOST_ENABLE": "0", + "NCCL_NVLS_ENABLE": "0", + "VLLM_ALLREDUCE_USE_FLASHINFER": "0", + "VLLM_ALLREDUCE_USE_SYMM_MEM": "0", + "VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC": "0", + "VLLM_USE_NCCL_SYMM_MEM": "0", + } + for key, value in config.items(): + monkeypatch.setenv(key, value) + llm = LLM( + os.environ.get("VLLM_TEST_MODEL", "Qwen/Qwen3-0.6B"), + dtype="bfloat16", + tensor_parallel_size=tp, + distributed_executor_backend=( + "vllm.v1.executor.cuda_checkpoint_executor.CUDACheckpointExecutor" + ), + disable_custom_all_reduce=tp > 1, + enable_nccl_comm_suspend=tp > 1, + enable_sleep_mode=True, + enable_prefix_caching=False, + kv_cache_memory_bytes=1024**3, + max_model_len=256, + max_num_seqs=4, + compilation_config={ + "cudagraph_capture_sizes": [1, 2, 4], + "pass_config": {"fuse_allreduce_rms": False}, + }, + ) + try: + params = SamplingParams(temperature=0, max_tokens=32, ignore_eos=True) + + def generate(): + outputs = llm.generate( + ["The capital of France is", "Compute 7 times 8:"], + params, + use_tqdm=False, + ) + return [output.outputs[0].token_ids for output in outputs] + + reference = generate() + graphs = llm.collective_rpc(_graph_identity) + assert all(rank["graphs"] for rank in graphs) + for cycle in range(int(os.environ.get("VLLM_TEST_SLEEP_CYCLES", "3"))): + assert generate() == reference + llm.sleep(level=1) + llm.sleep(level=1) + order = (("weights", "kv_cache"), ("kv_cache", "weights"), (None,))[ + cycle % 3 + ] + for tag in order: + llm.wake_up(tags=[tag] if tag else None) + awake = llm.collective_rpc(_graph_identity) + assert generate() == reference + after = llm.collective_rpc(_graph_identity) + for original, before_generate, current in zip( + graphs, awake, after, strict=True + ): + assert current["pid"] == original["pid"] + assert current["graphs"] == original["graphs"] + assert current["replays"] > before_generate["replays"] + assert current["captures"] == 0 + graphs = after + llm.sleep(level=1) + finally: + llm.llm_engine.engine_core.shutdown() diff --git a/tests/v1/executor/test_cuda_checkpoint_executor.py b/tests/v1/executor/test_cuda_checkpoint_executor.py new file mode 100644 index 000000000000..b5ad8d88948d --- /dev/null +++ b/tests/v1/executor/test_cuda_checkpoint_executor.py @@ -0,0 +1,417 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU lifecycle tests for the opt-in process checkpoint executor.""" + +import multiprocessing +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from typing import Any +from unittest.mock import Mock, patch + +import pytest + +from vllm.v1.executor.cuda_checkpoint_executor import CUDACheckpointExecutor +from vllm.v1.executor.multiproc_executor import MultiprocExecutor + + +@pytest.fixture +def executor(): + calls: list[tuple[Any, ...]] = [] + obj = CUDACheckpointExecutor.__new__(CUDACheckpointExecutor) + obj._checkpoint_lock = threading.RLock() + obj._checkpoint_active = False + obj._use_nccl_suspend = False + obj._checkpoint_flashinfer = False + obj.is_failed = False + obj.sleeping_tags = set() + obj.failure_callback = None + obj.workers = [SimpleNamespace(proc=SimpleNamespace(pid=p)) for p in (11, 22)] + obj._checkpoint = SimpleNamespace( + suspend=lambda pids: calls.append(("checkpoint", pids)), + resume=lambda: calls.append(("restore",)), + close=lambda: calls.append(("close",)), + ) + + def sleep(self, level=1): + if "weights" in self.sleeping_tags: + return + calls.append(("sleep", level)) + self.sleeping_tags = {"weights", "kv_cache"} + + def wake(self, tags=None): + calls.append(("wake", tags)) + self.sleeping_tags -= set(tags) if tags else self.sleeping_tags.copy() + + with ( + patch.object(MultiprocExecutor, "sleep", sleep), + patch.object(MultiprocExecutor, "wake_up", wake), + patch.object( + MultiprocExecutor, "shutdown", lambda _: calls.append(("shutdown",)) + ), + ): + yield obj, calls + obj._checkpoint = None + + +def test_sleep_then_restore_precedes_worker_wake(executor): + obj, calls = executor + obj.sleep() + obj.wake_up() + assert calls == [ + ("sleep", 1), + ("checkpoint", [11, 22]), + ("restore",), + ("wake", None), + ] + + +def test_repeated_sleep_does_not_checkpoint_twice(executor): + obj, calls = executor + obj.sleep() + obj.sleep() + assert calls == [("sleep", 1), ("checkpoint", [11, 22])] + + +def test_partial_wake_restores_process_only_once(executor): + obj, calls = executor + obj.sleep() + obj.wake_up(["weights"]) + obj.wake_up(["kv_cache"]) + assert calls.count(("restore",)) == 1 + assert not obj.is_sleeping + + +def test_invalid_tag_does_not_restore(executor): + obj, calls = executor + obj.sleep() + with pytest.raises(ValueError): + obj.wake_up(["invalid"]) + assert ("restore",) not in calls + + +def test_failed_checkpoint_stops_workers_and_blocks_wake(executor): + obj, calls = executor + + def fail(_): + raise RuntimeError("partial checkpoint") + + obj._checkpoint.suspend = fail + with pytest.raises(RuntimeError, match="partial"): + obj.sleep() + assert obj.is_failed and ("shutdown",) in calls and ("close",) in calls + with pytest.raises(RuntimeError, match="failed"): + obj.wake_up() + assert not any(c[0] == "wake" for c in calls) + + +def test_shutdown_does_not_restore_checkpointed_memory(executor): + obj, calls = executor + obj.sleep() + + def fail(): + raise RuntimeError("restore out of memory") + + obj._checkpoint.resume = fail + obj.shutdown() + assert not obj.is_failed + assert ("shutdown",) in calls and ("close",) in calls + assert ("wake", None) not in calls + + +def test_level_two_rejected_before_sleep(executor): + obj, calls = executor + with pytest.raises(ValueError): + obj.sleep(2) + assert not calls + + +def test_checkpointed_worker_rejects_cuda_rpc_without_blocking_wake(executor): + obj, _ = executor + obj.sleep() + with pytest.raises(RuntimeError, match="checkpointed"): + obj.collective_rpc("profile") + obj.wake_up() + + +def test_monitor_shutdown_waits_for_inflight_checkpoint(executor): + obj, calls = executor + entered, release, closed = threading.Event(), threading.Event(), threading.Event() + + def suspend(pids): + entered.set() + assert release.wait(10) + + lock, contended = obj._checkpoint_lock, threading.Event() + + class ObservedLock: + def __enter__(self): + if not lock.acquire(blocking=False): + contended.set() + assert lock.acquire(timeout=10) + + def __exit__(self, *exc): + lock.release() + + obj._checkpoint_lock = ObservedLock() + obj._checkpoint.suspend = suspend + obj._checkpoint.close = lambda: closed.set() + with ThreadPoolExecutor(max_workers=2) as pool: + sleeper = pool.submit(obj.sleep) + try: + assert entered.wait(10) + closer = pool.submit(obj.shutdown) + assert contended.wait(5) + assert not closed.is_set() + finally: + release.set() + sleeper.result(timeout=10) + closer.result(timeout=10) + assert closed.is_set() + + +def test_repeated_pause_acknowledges_already_completed_device_work(executor): + obj, _ = executor + obj.sleep() + for method in ("synchronize_device", "reset_mm_cache", "reset_encoder_cache"): + assert obj.collective_rpc(method) == [None, None] + obj.wake_up() + + +def test_nonblocking_pause_returns_completed_future(executor): + obj, _ = executor + obj.sleep() + result = obj.collective_rpc("synchronize_device", non_block=True) + assert result.done() and result.result() == [None, None] + + +def test_failed_executor_never_acknowledges_pause(executor): + obj, _ = executor + obj.sleep() + obj.is_failed = True + with pytest.raises(RuntimeError, match="failed"): + obj.collective_rpc("synchronize_device") + + +@pytest.mark.parametrize( + "p2p,shm,tp,nvls,suspend,fusion,backend,accepted", + [ + pytest.param("0", "1", 2, "0", True, False, None, True, id="p2p-only"), + pytest.param("0", "0", 3, "0", True, False, None, True, id="p2p-and-shm"), + pytest.param("1", "0", 8, "0", True, False, None, True, id="shm-only"), + pytest.param(None, None, 2, "0", True, False, None, True, id="defaults"), + pytest.param("0", "0", 2, "0", True, True, "trtllm", True, id="trtllm"), + pytest.param("0", "0", 2, "1", True, False, None, False, id="nvls"), + pytest.param("0", "0", 2, "0", False, False, None, False, id="no-suspend"), + pytest.param("0", "0", 2, "0", True, True, None, False, id="fusion-default"), + pytest.param("0", "0", 2, "0", True, True, "mnnvl", False, id="mnnvl"), + ], +) +def test_transport_requires_suspension_and_no_nvls( + monkeypatch, nvls, suspend, accepted, fusion, backend, p2p, shm, tp +): + # Transport options are independent of the rejection rules; avoid their + # Cartesian product while retaining each option and validation branch. + monkeypatch.setattr("vllm.platforms.current_platform.is_cuda", lambda: True) + env = { + "NCCL_P2P_DISABLE": p2p, + "NCCL_SHM_DISABLE": shm, + "NCCL_IB_DISABLE": "1", + "NCCL_CUMEM_ENABLE": "1", + "NCCL_CUMEM_HOST_ENABLE": "0", + "NCCL_NVLS_ENABLE": nvls, + "VLLM_ALLREDUCE_USE_FLASHINFER": "0", + "VLLM_ALLREDUCE_USE_SYMM_MEM": "0", + "VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC": "0", + "VLLM_USE_NCCL_SYMM_MEM": "0", + } + for key, value in env.items(): + if value is None: + monkeypatch.delenv(key, raising=False) + else: + monkeypatch.setenv(key, value) + monkeypatch.delenv("VLLM_FLASHINFER_ALLREDUCE_BACKEND", raising=False) + if backend is not None: + monkeypatch.setenv("VLLM_FLASHINFER_ALLREDUCE_BACKEND", backend) + config = SimpleNamespace( + compilation_config=SimpleNamespace( + mode=3 if fusion else 0, + pass_config=SimpleNamespace(fuse_allreduce_rms=fusion), + ), + parallel_config=SimpleNamespace( + tensor_parallel_size=tp, + nnodes=1, + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + world_size=tp, + disable_custom_all_reduce=True, + enable_expert_parallel=False, + worker_cls="vllm.v1.worker.gpu_worker.Worker", + ), + model_config=SimpleNamespace( + enable_sleep_mode=True, + enable_nccl_comm_suspend=suspend, + sleep_mode_backend="cumem", + ), + use_v2_model_runner=True, + weight_transfer_config=None, + kv_transfer_config=None, + ec_transfer_config=None, + ) + with ( + patch.object(MultiprocExecutor, "__init__", return_value=None), + patch.object(MultiprocExecutor, "shutdown"), + ): + if accepted: + obj = CUDACheckpointExecutor(config) + assert obj._checkpoint_flashinfer == fusion + obj.shutdown() + else: + with pytest.raises(ValueError, match="NCCL checkpoint configuration"): + CUDACheckpointExecutor(config) + + +@pytest.mark.parametrize( + "has_symbols,actual,expected", + [(False, True, True), (True, False, True), (True, True, False)], +) +def test_validate_rejects_missing_or_inconsistent_nccl( + monkeypatch, has_symbols, actual, expected +): + from vllm.v1.executor.cuda_checkpoint_executor import _validate_nccl_checkpoint + + comm = SimpleNamespace( + available=True, + disabled=False, + _suspended=actual, + nccl=SimpleNamespace(has_symbol=lambda _: has_symbols), + ) + group = SimpleNamespace(device_communicator=SimpleNamespace(pynccl_comm=comm)) + monkeypatch.setattr("vllm.distributed.get_tp_group", lambda: group) + with pytest.raises(RuntimeError, match="active suspend/resume"): + _validate_nccl_checkpoint(None, expected) + + +def test_flashinfer_hooks_wrap_driver_checkpoint_once(executor): + obj, calls = executor + obj._checkpoint_flashinfer = True + with patch.object( + obj, "collective_rpc", side_effect=lambda method, **kw: calls.append((method,)) + ): + obj.sleep() + obj.wake_up(["weights"]) + obj.wake_up(["kv_cache"]) + assert calls == [ + ("sleep", 1), + ("checkpoint_prepare",), + ("checkpoint", [11, 22]), + ("restore",), + ("checkpoint_restore",), + ("wake", ["weights"]), + ("wake", ["kv_cache"]), + ] + + +def test_sleep_worker_rpc_has_deadline(executor): + obj, _ = executor + with patch.object(MultiprocExecutor, "collective_rpc", return_value=[]) as rpc: + obj.collective_rpc("sleep") + assert rpc.call_args.kwargs["timeout"] == 120 + obj.collective_rpc("wake_up", timeout=7) + assert rpc.call_args.kwargs["timeout"] == 7 + + +def test_worker_termination_reaps_real_processes(): + ctx = multiprocessing.get_context("spawn") + processes = [ctx.Process(target=time.sleep, args=(60,)) for _ in range(2)] + try: + for proc in processes: + proc.start() + CUDACheckpointExecutor._ensure_worker_termination(processes) + assert all(proc.exitcode is not None for proc in processes) + assert not any(proc.is_alive() for proc in processes) + finally: + for proc in processes: + if proc.is_alive(): + proc.kill() + proc.join(timeout=5) + proc.close() + + +def test_kv_first_partial_wake_then_sleep_preserves_parent_noop(executor): + obj, calls = executor + obj._use_nccl_suspend = True + obj.collective_rpc = Mock() + obj.sleep() + obj.wake_up(["kv_cache"]) + assert obj.sleeping_tags == {"weights"} + before = list(calls) + rpc_count = obj.collective_rpc.call_count + obj.sleep() + assert calls == before + assert obj.collective_rpc.call_count == rpc_count + assert not obj.is_failed + obj.wake_up(["weights"]) + assert not obj.is_sleeping + + +@pytest.mark.parametrize( + "component,field,value", + [ + ("parallel_config", "tensor_parallel_size", 0), + ("parallel_config", "tensor_parallel_size", -1), + ("parallel_config", "nnodes", 2), + ("parallel_config", "enable_expert_parallel", True), + ("parallel_config", "worker_cls", "custom.Worker"), + ("model_config", "sleep_mode_backend", "custom"), + (None, "use_v2_model_runner", False), + (None, "weight_transfer_config", object()), + ("platform", "is_cuda", False), + ("platform", "linux", False), + ], +) +def test_unsupported_checkpoint_configuration_fails_before_workers( + component, field, value, monkeypatch +): + config = SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=1, + world_size=1, + nnodes=1, + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=False, + worker_cls="vllm.v1.worker.gpu_worker.Worker", + ), + model_config=SimpleNamespace( + enable_sleep_mode=True, + enable_nccl_comm_suspend=False, + sleep_mode_backend="cumem", + ), + use_v2_model_runner=True, + weight_transfer_config=None, + kv_transfer_config=None, + ec_transfer_config=None, + ) + monkeypatch.setattr("vllm.platforms.current_platform.is_cuda", lambda: True) + if component == "platform": + if field == "is_cuda": + monkeypatch.setattr( + "vllm.platforms.current_platform.is_cuda", lambda: value + ) + else: + monkeypatch.setattr("sys.platform", "darwin") + else: + setattr(getattr(config, component) if component else config, field, value) + with ( + patch.object(MultiprocExecutor, "__init__") as initialize, + patch.object(MultiprocExecutor, "shutdown"), + pytest.raises(ValueError, match="checkpoint"), + ): + CUDACheckpointExecutor(config) + initialize.assert_not_called() diff --git a/tests/v1/executor/test_cuda_process_checkpoint.py b/tests/v1/executor/test_cuda_process_checkpoint.py new file mode 100644 index 000000000000..1e7aafb90ce1 --- /dev/null +++ b/tests/v1/executor/test_cuda_process_checkpoint.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Controller ordering and helper cleanup without CUDA initialization.""" + +import io +import json +import subprocess +import sys +from unittest.mock import Mock, patch + +import pytest + +from vllm.utils.cuda_process_checkpoint import CudaProcessCheckpoint, Driver + + +@pytest.mark.parametrize("rank_count", [1, 2]) +def test_all_ranks_restore_before_any_unlock(rank_count): + events: list[tuple[int, str]] = [] + with patch("vllm.utils.cuda_process_checkpoint.Driver") as driver: + driver.return_value.run_stage.side_effect = lambda pids, op: events.extend( + (pid, op) for pid in pids + ) + checkpoint = CudaProcessCheckpoint(lambda **event: None) + checkpoint.suspend(list(range(1, rank_count + 1))) + checkpoint.resume() + assert events == [ + (pid, op) + for op in ("Lock", "Checkpoint", "Restore", "Unlock") + for pid in range(1, rank_count + 1) + ] + assert checkpoint.pids == [] + + +def test_partial_failure_prevents_further_rank_operations(): + events = [] + + def invoke(pid, op): + events.append((pid, op)) + if pid == 22 and op == "Restore": + raise RuntimeError("injected restore failure") + + with patch("vllm.utils.cuda_process_checkpoint.Driver") as driver: + driver.return_value.run_stage.side_effect = lambda pids, op: [ + invoke(pid, op) for pid in pids + ] + checkpoint = CudaProcessCheckpoint(lambda **event: None) + checkpoint.suspend([11, 22]) + with pytest.raises(RuntimeError, match="injected"): + checkpoint.resume() + count = len(events) + with pytest.raises(RuntimeError, match="restart required"): + checkpoint.resume() + assert len(events) == count + assert not any(op == "Unlock" for _, op in events) + assert checkpoint.failed + + +def test_helper_initialization_failure_reaps_process(): + proc = Mock() + proc.poll.return_value = None + with ( + patch("subprocess.Popen", return_value=proc), + patch.object(Driver, "_read", side_effect=RuntimeError("init")), + pytest.raises(RuntimeError, match="init"), + ): + Driver(lambda **event: None) + proc.kill.assert_called_once() + proc.communicate.assert_called_once_with(timeout=5) + + +def test_helper_read_timeout_kills_process(): + driver = Driver.__new__(Driver) + driver.proc = Mock() + selector = Mock() + selector.select.return_value = [] + with ( + patch("selectors.DefaultSelector", return_value=selector), + pytest.raises(TimeoutError, match="deadline"), + ): + driver._read(0) + driver.proc.kill.assert_called_once() + selector.close.assert_called_once() + + +def test_helper_close_reaps_real_process(): + driver = Driver.__new__(Driver) + driver.proc = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(60)"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + text=True, + ) + try: + driver.close() + assert driver.proc.returncode is not None + assert driver.proc.stdin is not None and driver.proc.stdout is not None + assert driver.proc.stdin.closed and driver.proc.stdout.closed + finally: + if driver.proc.poll() is None: + driver.proc.kill() + driver.proc.wait(timeout=5) + + +def test_driver_stage_runs_ranks_concurrently(): + import threading + + from vllm.utils.cuda_process_checkpoint import _run_parallel_stage + + pids = [11, 22, 33, 44] + barrier = threading.Barrier(len(pids), timeout=2) + + def invoke(pid, op): + barrier.wait() + return {"pid": pid, "op": op, "rc": 0} + + rows = _run_parallel_stage(pids, "Restore", invoke) + assert rows == [{"pid": pid, "op": "Restore", "rc": 0} for pid in pids] + + +def test_driver_stage_collects_rank_failure(): + from vllm.utils.cuda_process_checkpoint import _run_parallel_stage + + def invoke(pid, op): + if pid == 22: + raise RuntimeError("injected") + return {"pid": pid, "op": op, "rc": 0} + + rows = _run_parallel_stage([11, 22], "Restore", invoke) + assert rows[0]["rc"] == 0 + assert rows[1]["pid"] == 22 and rows[1]["op"] == "Restore" + assert "injected" in rows[1]["error"] + + +@pytest.fixture +def stage_driver(monkeypatch): + driver = Driver.__new__(Driver) + driver.emit = Mock() + driver.proc = Mock(stdin=io.StringIO()) + response = Mock() + monkeypatch.setattr(driver, "_read", response) + return driver, response + + +@pytest.mark.parametrize( + "fault", ["missing", "duplicate", "wrong_op", "bad_state", "failed"] +) +def test_driver_rejects_invalid_stage_response(fault, stage_driver): + driver, response = stage_driver + rows = [ + {"pid": pid, "op": "Restore", "rc": 0, "after": {"rc": 0, "value": 1}} + for pid in [11, 22] + ] + if fault == "missing": + rows.pop() + elif fault == "duplicate": + rows[1]["pid"] = 11 + elif fault == "wrong_op": + rows[1]["op"] = "Unlock" + elif fault == "bad_state": + rows[1]["after"] = {"rc": 0, "value": 2} + else: + rows[1]["rc"] = 2 + response.return_value = {"kind": "stage", "op": "Restore", "results": rows} + with pytest.raises(RuntimeError): + driver.run_stage([11, 22], "Restore") + + +@pytest.mark.parametrize("response", [None, [], "bad", 0]) +def test_driver_rejects_non_object_response(response, stage_driver): + driver, read = stage_driver + read.return_value = response + with pytest.raises(RuntimeError, match="Invalid CUDA checkpoint stage response"): + driver.run_stage([11], "Restore") + + +def test_driver_accepts_complete_stage_response(stage_driver): + driver, response = stage_driver + rows = [ + {"pid": pid, "op": "Restore", "rc": 0, "after": {"rc": 0, "value": 1}} + for pid in [11, 22] + ] + response.return_value = {"kind": "stage", "op": "Restore", "results": rows} + driver.run_stage([11, 22], "Restore") + assert json.loads(driver.proc.stdin.getvalue()) == [[11, 22], "Restore"] + for row in rows: + driver.emit.assert_any_call(**row) + + +def test_rank_failure_with_blocked_peer_times_out_and_reaps_helper(): + import textwrap + + driver = Driver.__new__(Driver) + driver.emit = Mock() + driver.proc = subprocess.Popen( + [ + sys.executable, + "-u", + "-c", + textwrap.dedent(""" + import json + import sys + import time + from vllm.utils.cuda_process_checkpoint import _run_parallel_stage + + def invoke(pid, op): + if pid == 11: + raise RuntimeError("injected rank failure") + time.sleep(60) + + print(json.dumps({"kind": "ready"}), flush=True) + pids, op = json.loads(sys.stdin.readline()) + rows = _run_parallel_stage(pids, op, invoke) + print(json.dumps({"kind": "stage", "op": op, "results": rows})) + """), + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + text=True, + ) + try: + assert driver._read(10) == {"kind": "ready"} + read = driver._read + with ( + patch.object(driver, "_read", side_effect=lambda timeout: read(0.2)), + pytest.raises(TimeoutError, match="deadline"), + ): + driver.run_stage([11, 22], "Restore") + driver.proc.wait(timeout=5) + assert driver.proc.returncode != 0 + finally: + driver.close() diff --git a/vllm/utils/cuda_process_checkpoint.py b/vllm/utils/cuda_process_checkpoint.py new file mode 100644 index 000000000000..57ad0559d591 --- /dev/null +++ b/vllm/utils/cuda_process_checkpoint.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Linux CUDA process checkpoint calls run in a separate interpreter.""" + +import ctypes as C +import time +from collections.abc import Sequence +from typing import Literal, TextIO, cast + +CheckpointOperation = Literal["Lock", "Checkpoint", "Restore", "Unlock"] +_OPERATIONS: tuple[CheckpointOperation, ...] = ( + "Lock", + "Checkpoint", + "Restore", + "Unlock", +) + + +class LockArgs(C.Structure): + _fields_ = [ + ("timeoutMs", C.c_uint), + ("reserved0", C.c_uint), + ("reserved1", C.c_uint64 * 7), + ] + + +class DirectDriver: + def __init__(self, emit): + self.emit = emit + self.lib = C.CDLL("libcuda.so.1") + self.lib.cuInit.argtypes = [C.c_uint] + rc = self.lib.cuInit(0) + if rc: + raise RuntimeError(f"cuInit: {rc}") + self.lib.cuCheckpointProcessGetState.argtypes = [C.c_int, C.POINTER(C.c_int)] + self.lib.cuGetErrorName.argtypes = [C.c_int, C.POINTER(C.c_char_p)] + for op in _OPERATIONS: + getattr(self.lib, "cuCheckpointProcess" + op).argtypes = [ + C.c_int, + C.c_void_p, + ] + + def state(self, pid): + state = C.c_int(-1) + rc = self.lib.cuCheckpointProcessGetState(pid, C.byref(state)) + return {"rc": rc, "value": state.value} + + def invoke(self, pid, op: CheckpointOperation): + """Return a result without writing to the helper protocol stream.""" + before = self.state(pid) + args = LockArgs(10000) if op == "Lock" else None + start = time.perf_counter() + rc = getattr(self.lib, "cuCheckpointProcess" + op)( + pid, C.byref(args) if args is not None else None + ) + elapsed = time.perf_counter() - start + name = C.c_char_p() + self.lib.cuGetErrorName(rc, C.byref(name)) + after = self.state(pid) + return dict( + kind="control", + pid=pid, + op=op, + rc=rc, + error=name.value.decode() if name.value else None, + before=before, + after=after, + time_s=elapsed, + ) + + def __call__(self, pid, op: CheckpointOperation): + row = self.invoke(pid, op) + self.emit(**row) + if row["rc"] or row["after"] != { + "rc": 0, + "value": {"Lock": 1, "Checkpoint": 2, "Restore": 1, "Unlock": 0}[op], + }: + raise RuntimeError(str(row)) + + +def _run_parallel_stage(pids, op: CheckpointOperation, invoke): + """Run one stage in the helper; never emit partial protocol responses.""" + from concurrent.futures import ThreadPoolExecutor + + def run(pid): + try: + return invoke(pid, op) + except Exception as exc: + return {"pid": pid, "op": op, "error": repr(exc)} + + # Waiting threads live only in the helper. The parent's stage deadline + # terminates this process if any driver call blocks. + with ThreadPoolExecutor(max_workers=len(pids)) as pool: + return list(pool.map(run, pids)) + + +class Driver: + """Persistent helper isolates potentially blocking driver calls.""" + + def __init__(self, emit): + import subprocess + import sys + from pathlib import Path + + self.emit = emit + self.proc = subprocess.Popen( + [sys.executable, str(Path(__file__)), "--service"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + text=True, + ) + try: + self._read(30) + except Exception: + if self.proc.poll() is None: + self.proc.kill() + self.proc.communicate(timeout=5) + raise + + def _read(self, timeout): + import json + import selectors + + sel = selectors.DefaultSelector() + try: + sel.register(cast(TextIO, self.proc.stdout), selectors.EVENT_READ) + if not sel.select(timeout): + self.proc.kill() + raise TimeoutError("CUDA controller deadline exceeded") + line = cast(TextIO, self.proc.stdout).readline() + if not line: + raise RuntimeError("CUDA controller exited unexpectedly") + return json.loads(line) + finally: + sel.close() + + def __call__(self, pid, op: CheckpointOperation): + self.run_stage([pid], op) + + def run_stage(self, pids, op: CheckpointOperation): + import json + + for pid in pids: + self.emit(kind="control_start", pid=pid, op=op) + cast(TextIO, self.proc.stdin).write(json.dumps([pids, op]) + "\n") + cast(TextIO, self.proc.stdin).flush() + try: + response = self._read(30) + except Exception as exc: + self.emit(kind="control_failed", pids=pids, op=op, error=repr(exc)) + raise + if not isinstance(response, dict): + raise RuntimeError(f"Invalid CUDA checkpoint stage response: {response}") + rows = response.get("results", []) + if ( + response.get("kind") != "stage" + or response.get("op") != op + or not isinstance(rows, list) + or len(rows) != len(pids) + or any( + not isinstance(row, dict) + or row.get("pid") != pid + or row.get("op") != op + for pid, row in zip(pids, rows) + ) + ): + raise RuntimeError(f"Invalid CUDA checkpoint stage response: {response}") + for row in rows: + self.emit(**row) + failed = [ + row + for row in rows + if row.get("rc") != 0 + or row.get("after") + != { + "rc": 0, + "value": {"Lock": 1, "Checkpoint": 2, "Restore": 1, "Unlock": 0}[op], + } + ] + if failed: + raise RuntimeError(str(failed)) + + def close(self): + if self.proc.poll() is None: + self.proc.kill() + self.proc.communicate(timeout=1) + + +class CudaProcessCheckpoint: + """Suspend all local worker processes, or resume all before returning.""" + + def __init__(self, emit): + self.emit = emit + self.driver = Driver(emit) + self.pids = [] + self.failed = False + + def _run(self, operations: Sequence[CheckpointOperation]): + if self.failed: + raise RuntimeError("CUDA process checkpoint failed; restart required") + try: + for op in operations: + self.driver.run_stage(self.pids, op) + self.emit(kind="boundary", stage=op, pids=self.pids) + except Exception: + self.failed = True + raise + + def suspend(self, pids): + if self.pids: + raise RuntimeError("CUDA processes are already checkpointed") + if not pids or len(set(pids)) != len(pids): + raise ValueError("Expected distinct worker PIDs") + self.pids = list(pids) + self._run(("Lock", "Checkpoint")) + + def resume(self): + if not self.pids: + return + self._run(("Restore", "Unlock")) + self.pids = [] + + def close(self): + self.driver.close() + + +if __name__ == "__main__": + import json + import sys + + def output(**kw): + print(json.dumps(kw), flush=True) + + driver = DirectDriver(output) + output(kind="controller_ready") + for line in sys.stdin: + try: + pids, op = json.loads(line) + if isinstance(pids, int): + driver(pids, op) + else: + output( + kind="stage", + op=op, + results=_run_parallel_stage(pids, op, driver.invoke), + ) + except Exception: + break diff --git a/vllm/v1/executor/cuda_checkpoint_executor.py b/vllm/v1/executor/cuda_checkpoint_executor.py new file mode 100644 index 000000000000..9b743e49fd29 --- /dev/null +++ b/vllm/v1/executor/cuda_checkpoint_executor.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in, single-node sleep executor preserving complete CUDA worker state.""" + +import json +import os +import sys +import threading +import time +from concurrent.futures import Future +from functools import wraps +from multiprocessing.process import BaseProcess +from typing import Any, cast + +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.cuda_process_checkpoint import CudaProcessCheckpoint +from vllm.v1.executor.multiproc_executor import MultiprocExecutor + +logger = init_logger(__name__) + + +def _serialized(method): + @wraps(method) + def call(self, *args, **kwargs): + with self._checkpoint_lock: + return method(self, *args, **kwargs) + + return call + + +def _validate_nccl_checkpoint(worker, suspended): + from vllm.distributed import get_tp_group + + comm = cast(Any, get_tp_group().device_communicator).pynccl_comm + if ( + comm is None + or not comm.available + or comm.disabled + or not comm.nccl.has_symbol("ncclCommSuspend") + or not comm.nccl.has_symbol("ncclCommResume") + or comm._suspended != suspended + ): + raise RuntimeError( + "NCCL checkpoint requires an active suspend/resume implementation" + ) + + +class CUDACheckpointExecutor(MultiprocExecutor): + """Use dedicated workers even at TP=1 so failed restoration is terminable.""" + + def __init__(self, vllm_config, monitor_workers=True): + self._checkpoint_lock = threading.RLock() + self._checkpoint = None + self._checkpoint_active = False + parallel = vllm_config.parallel_config + if sys.platform != "linux" or not current_platform.is_cuda(): + raise ValueError("CUDA checkpoint requires Linux and NVIDIA CUDA") + if ( + not vllm_config.use_v2_model_runner + or vllm_config.model_config.sleep_mode_backend != "cumem" + or parallel.enable_expert_parallel + or parallel.worker_cls not in ("auto", "vllm.v1.worker.gpu_worker.Worker") + ): + raise ValueError( + "CUDA checkpoint requires the V2 GPU model runner, cumem sleep " + "backend and standard GPU worker, without expert parallelism" + ) + compilation = getattr(vllm_config, "compilation_config", None) + passes = getattr(compilation, "pass_config", None) + fused_allreduce = bool(getattr(compilation, "mode", 0)) and bool( + getattr(passes, "fuse_allreduce_rms", False) + ) + standalone_flashinfer = os.environ.get("VLLM_ALLREDUCE_USE_FLASHINFER") == "1" + # The compiler fusion allocates FlashInfer workspaces independently + # of the standalone all-reduce environment switch. + self._checkpoint_flashinfer = parallel.tensor_parallel_size > 1 and ( + standalone_flashinfer or fused_allreduce + ) + self._use_nccl_suspend = ( + parallel.tensor_parallel_size > 1 + and vllm_config.model_config.enable_nccl_comm_suspend + ) + if ( + parallel.tensor_parallel_size < 1 + or parallel.nnodes != 1 + or parallel.decode_context_parallel_size != 1 + or parallel.prefill_context_parallel_size != 1 + or parallel.pipeline_parallel_size != 1 + or parallel.data_parallel_size != 1 + or parallel.world_size != parallel.tensor_parallel_size + ): + raise ValueError("CUDA checkpoint requires single-node tensor parallelism") + if ( + vllm_config.kv_transfer_config + or vllm_config.ec_transfer_config + or vllm_config.weight_transfer_config + ): + raise ValueError( + "CUDA checkpoint does not support external cache or weight transfer" + ) + if not vllm_config.model_config.enable_sleep_mode: + raise ValueError("CUDA checkpoint requires enable_sleep_mode=True") + if parallel.tensor_parallel_size > 1: + required = { + "NCCL_P2P_DISABLE": "1", + "NCCL_SHM_DISABLE": "1", + "NCCL_IB_DISABLE": "1", + "NCCL_CUMEM_ENABLE": "0", + "NCCL_CUMEM_HOST_ENABLE": "0", + "VLLM_ALLREDUCE_USE_FLASHINFER": "0", + "VLLM_ALLREDUCE_USE_SYMM_MEM": "0", + "VLLM_ALLREDUCE_USE_FLASHINFER_PCIE_IPC": "0", + "VLLM_USE_NCCL_SYMM_MEM": "0", + } + if vllm_config.model_config.enable_nccl_comm_suspend: + # Suspension releases P2P mappings while retaining addresses. + # Legacy SHM also survives CUDA checkpoint; let NCCL select + # P2P or SHM rather than requiring one transport topology. + required.pop("NCCL_P2P_DISABLE") + required.pop("NCCL_SHM_DISABLE") + # Live multicast allocations can block checkpoint even on + # newer user-mode drivers. Keep NVLS disabled until the sleep + # path can release those allocations before checkpoint. + if os.environ.get("NCCL_DISABLE_MEM_MANAGER", "0") != "0": + raise ValueError("NCCL checkpoint requires the memory manager") + required.update( + { + "NCCL_CUMEM_ENABLE": "1", + "NCCL_NVLS_ENABLE": "0", + "VLLM_ALLREDUCE_USE_FLASHINFER": "1" + if standalone_flashinfer + else "0", + } + ) + if self._checkpoint_flashinfer: + required["VLLM_FLASHINFER_ALLREDUCE_BACKEND"] = "trtllm" + if not parallel.disable_custom_all_reduce or any( + os.environ.get(k) != v for k, v in required.items() + ): + raise ValueError( + "CUDA checkpoint currently requires disabled custom " + f"all-reduce and NCCL checkpoint configuration: {required}" + ) + super().__init__(vllm_config, monitor_workers) + + def _record(self, **event): + event["monotonic_s"] = time.monotonic() + logger.info("CUDA_PROCESS_CHECKPOINT %s", json.dumps(event)) + + def collective_rpc( + self, + method, + timeout=None, + args=(), + kwargs=None, + non_block=False, + unique_reply_rank=None, + kv_output_aggregator=None, + ec_output_aggregator=None, + ): + if self.is_failed: + raise RuntimeError("CUDA checkpoint executor failed") + if self._checkpoint_active: + # EngineCore's first pause synchronized and cleared these caches. + # No new worker work can run while checkpointed, so a repeated + # pause can acknowledge that state without calling locked CUDA APIs. + if ( + method + in ("synchronize_device", "reset_mm_cache", "reset_encoder_cache") + and not args + and not kwargs + and unique_reply_rank is None + and kv_output_aggregator is None + and ec_output_aggregator is None + ): + result = [None] * len(self.workers) + if non_block: + future: Future[Any] = Future() + future.set_result(result) + return future + return result + raise RuntimeError( + "Workers are checkpointed; wake them before issuing RPCs" + ) + if method in ("sleep", "wake_up") and timeout is None: + timeout = 120 + return super().collective_rpc( + method, + timeout=timeout, + args=args, + kwargs=kwargs, + non_block=non_block, + unique_reply_rank=unique_reply_rank, + kv_output_aggregator=kv_output_aggregator, + ec_output_aggregator=ec_output_aggregator, + ) + + def check_health(self): + if self.is_failed: + raise RuntimeError("CUDA checkpoint executor failed") + if not self._checkpoint_active: + super().check_health() + + @staticmethod + def _ensure_worker_termination(worker_procs: list[BaseProcess]): + # Dedicated workers can be inside locked CUDA calls. Destroy their + # processes without restoring GPU memory or waiting for CUDA cleanup. + # Reap within the outer process manager's shutdown window. + for proc in worker_procs: + if proc.is_alive(): + proc.kill() + deadline = time.monotonic() + 3 + for proc in worker_procs: + proc.join(timeout=max(0, deadline - time.monotonic())) + if proc.exitcode is None: + logger.error( + "Worker pid=%s did not exit within cleanup deadline", proc.pid + ) + + def _close_checkpoint(self): + checkpoint, self._checkpoint = self._checkpoint, None + if checkpoint is not None: + try: + checkpoint.close() + except Exception: + logger.exception("Failed to close CUDA checkpoint helper") + + def _stop_failed(self): + self.is_failed = True + # Terminate worker processes even if a CUDA API is blocked. Never try + # to continue serving a partially checkpointed rank set. + super().shutdown() + self._close_checkpoint() + callback = self.failure_callback + if callback is not None: + self.failure_callback = None + callback() + + @_serialized + def sleep(self, level=1): + if self.is_failed: + raise RuntimeError("CUDA checkpoint executor failed") + if level != 1: + raise ValueError("CUDA checkpoint currently supports sleep level 1 only") + if self._checkpoint_active: + return + # Match the parent's no-op after KV-first partial wake: weights are + # still asleep, but CUDA/NCCL have resumed and must not be assumed + # suspended. Fully wake all tags before starting another checkpoint. + if "weights" in self.sleeping_tags: + return + if self._checkpoint is None: + self._checkpoint = CudaProcessCheckpoint(self._record) + try: + if self._use_nccl_suspend: + self.collective_rpc( + _validate_nccl_checkpoint, args=(False,), timeout=30 + ) + super().sleep(level) + if self._use_nccl_suspend: + self.collective_rpc(_validate_nccl_checkpoint, args=(True,), timeout=30) + if self._checkpoint_flashinfer: + self.collective_rpc("checkpoint_prepare", timeout=120) + self._checkpoint.suspend([w.proc.pid for w in self.workers]) + self._checkpoint_active = True + except Exception: + self._stop_failed() + raise + + @_serialized + def wake_up(self, tags=None): + if self.is_failed: + raise RuntimeError("CUDA checkpoint executor failed") + if tags and any(tag not in self.sleeping_tags for tag in tags): + raise ValueError("Invalid or already awake sleep tag") + try: + if self._checkpoint_active: + cast(CudaProcessCheckpoint, self._checkpoint).resume() + self._checkpoint_active = False + if self._checkpoint_flashinfer: + self.collective_rpc("checkpoint_restore", timeout=120) + super().wake_up(tags) + if self._use_nccl_suspend: + self.collective_rpc( + _validate_nccl_checkpoint, args=(False,), timeout=30 + ) + except Exception: + self._stop_failed() + raise + + @_serialized + def shutdown(self): + try: + super().shutdown() + finally: + self._close_checkpoint() From dbcf568821827d2a32223b46d4bc95862dc6c185 Mon Sep 17 00:00:00 2001 From: Gabriel Wu <13583761+lucifer1004@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:32:12 +0800 Subject: [PATCH 038/406] [Perf][SM120][MLA] Use compact FP8 KV rows for NoPE models (#58305) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Kimi Code Co-authored-by: Claude Opus 5.5 --- csrc/cache_kernels.cu | 28 +++--- tests/kernels/attention/test_cache.py | 45 +++++++-- tests/v1/attention/test_mla_backends.py | 35 ++++++- .../v1/attention/test_sparse_mla_backends.py | 92 ++++++++++++++----- .../test_sparse_mla_kv_cache_layout.py | 1 + .../layers/attention/mla_attention.py | 7 +- .../layers/attention/sparse_mla_attention.py | 6 +- vllm/v1/attention/backend.py | 3 + .../mla/flashinfer_mla_sparse_sm120.py | 8 +- vllm/v1/attention/backends/mla/index_group.py | 13 +-- vllm/v1/hisparse/runtime.py | 2 - 11 files changed, 176 insertions(+), 64 deletions(-) diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index 9567d87f443a..3541d1050128 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -540,6 +540,7 @@ __global__ void concat_and_cache_ds_mla_kernel( const int kv_lora_rank, // const int pe_dim, // const int block_size, // + const bool has_rope_tail, // const float* scale // ) { const int64_t token_idx = blockIdx.x; @@ -556,15 +557,15 @@ __global__ void concat_and_cache_ds_mla_kernel( // For the NoPE part, each tile of 128 elements is handled by half of one warp // (16 threads). There are 4 total tiles, so 2 warps (64 threads). // Lanes 0 and 16 of each warp write the scale values for that warp's tiles. - // The RoPE part (last 64 elements) is handled by another 1 warp (32 threads). + // The optional RoPE tail is handled by another 1 warp (32 threads). // So in total, we use 3 warps (96 threads) per block. - // Cast kv_cache to 16_bit for RoPE values - scalar_t* kv_cache_16bit = - reinterpret_cast(&kv_cache[dst_idx_start]); - - // Zero the reserved RoPE tail for NoPE rows. if (threadIdx.x >= 64) { + if (!has_rope_tail) { + return; + } + scalar_t* kv_cache_16bit = + reinterpret_cast(&kv_cache[dst_idx_start]); // Each thread handles two elements of RoPE const int8_t pe_idx_start = (threadIdx.x - 64) * 2; // RoPE values start after the packed 8-bit NoPE values and the @@ -941,7 +942,7 @@ void reshape_and_cache_flash( reinterpret_cast(kv_cache.data_ptr()), \ slot_mapping.const_data_ptr(), block_stride, entry_stride, \ kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ - reinterpret_cast(scale.data_ptr())); + has_rope_tail, reinterpret_cast(scale.data_ptr())); void concat_and_cache_mla( torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] @@ -965,6 +966,8 @@ void concat_and_cache_mla( int kv_lora_rank = kv_c.size(1); int pe_dim = k_pe.size(1); int block_size = kv_cache.size(1); + const int64_t row_bytes = kv_cache.size(2) * kv_cache.element_size(); + const bool has_rope_tail = row_bytes == 656; const bool is_nvfp4_ds_mla = kv_cache_dtype == "nvfp4_ds_mla"; if (kv_cache_dtype == "fp8_ds_mla") { @@ -972,8 +975,9 @@ void concat_and_cache_mla( "kv_lora_rank must be 512 for fp8_ds_mla"); STD_TORCH_CHECK(pe_dim == 64 || pe_dim == 0, "pe_dim must be 64 or 0 for fp8_ds_mla"); - STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(), - "kv_cache.size(2) must be 656 bytes for fp8_ds_mla"); + STD_TORCH_CHECK(has_rope_tail || (pe_dim == 0 && row_bytes == 528), + "fp8_ds_mla requires 656-byte rows, or 528-byte rows with " + "pe_dim == 0"); STD_TORCH_CHECK(kv_c.element_size() == 2, "kv_c.element_size() must be 2 for fp8_ds_mla"); STD_TORCH_CHECK(k_pe.element_size() == 2, @@ -1011,8 +1015,8 @@ void concat_and_cache_mla( // For the NoPE part, each tile of 128 elements is handled by half of one // warp (16 threads). There are 4 total tiles, so 2 warps (64 threads). // Lanes 0 and 16 of each warp write the scale values for that warp's tiles. - // The RoPE part (last 64 elements) is handled by another 1 warp (32 - // threads). So in total, we use 3 warps (96 threads) per block. + // The optional RoPE tail is handled by another 1 warp (32 threads). + // So in total, we use 3 warps (96 threads) per block. dim3 block(96); DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, CALL_CONCAT_AND_CACHE_DS_MLA); @@ -1837,6 +1841,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16, "dst must be bfloat16"); STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA"); + STD_TORCH_CHECK(src_cache.size(2) * src_cache.element_size() >= 656, + "src_cache rows must contain at least 656 bytes for MLA"); int64_t block_table_stride = block_table.stride(0); int64_t cache_block_stride = src_cache.stride(0); diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 149aedd30fc1..116835e93223 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -1014,9 +1014,15 @@ def test_concat_and_cache_ds_mla( @pytest.mark.parametrize("device", CUDA_DEVICES) @pytest.mark.parametrize("block_size", [64, 256]) +@pytest.mark.parametrize( + "entry_size,row_padding,page_padding", + [(528, 0, 0), (656, 0, 0), (528, 128, 1), (656, 16, 1)], +) @torch.inference_mode() -def test_concat_and_cache_ds_mla_nope(device: str, block_size: int) -> None: - """NoPE matches zero RoPE, clears valid tails, and preserves unused slots.""" +def test_concat_and_cache_ds_mla_nope( + device: str, block_size: int, entry_size: int, row_padding: int, page_padding: int +) -> None: + """NoPE matches zero RoPE and preserves unused slots and stride padding.""" dtype = torch.bfloat16 if current_platform.is_rocm(): pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm") @@ -1032,28 +1038,51 @@ def test_concat_and_cache_ds_mla_nope(device: str, block_size: int) -> None: k_pe = torch.empty(num_tokens, 0, dtype=dtype, device=device) zero_k_pe = torch.zeros(num_tokens, 64, dtype=dtype, device=device) scale = torch.tensor(1.0, dtype=torch.float32, device=device) - kv_cache = torch.full( + storage = torch.full( + (num_blocks, block_size + page_padding, entry_size + row_padding), + 0xFF, + dtype=torch.uint8, + device=device, + ) + kv_cache = storage[:, :block_size, :entry_size] + ref_cache = torch.full( (num_blocks, block_size, 656), 0xFF, dtype=torch.uint8, device=device ) - ref_cache = kv_cache.clone() opcheck( torch.ops._C_cache_ops.concat_and_cache_mla, (kv_c, k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale), test_utils=DEFAULT_OPCHECK_TEST_UTILS, ) - kv_cache.fill_(0xFF) + storage.fill_(0xFF) ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale) ops.concat_and_cache_mla( kv_c, zero_k_pe, ref_cache, slot_mapping, "fp8_ds_mla", scale ) - torch.testing.assert_close(kv_cache, ref_cache, atol=0, rtol=0) - rows = kv_cache.view(num_blocks * block_size, 656) - assert (rows[slot_mapping[:2], 528:] == 0).all() + torch.testing.assert_close(kv_cache, ref_cache[..., :entry_size], atol=0, rtol=0) + rows = kv_cache.reshape(num_blocks * block_size, entry_size) + if entry_size == 656: + assert (rows[slot_mapping[:2], 528:] == 0).all() untouched = torch.ones(num_blocks * block_size, dtype=torch.bool, device=device) untouched[slot_mapping[:2]] = False assert (rows[untouched] == 0xFF).all() + assert (storage[..., entry_size:] == 0xFF).all() + assert (storage[:, block_size:] == 0xFF).all() + + if entry_size == 528: + with pytest.raises(RuntimeError, match="528-byte rows with pe_dim == 0"): + ops.concat_and_cache_mla( + kv_c, zero_k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale + ) + with pytest.raises(RuntimeError, match="at least 656 bytes"): + ops.cp_gather_and_upconvert_fp8_kv_cache( + kv_cache, + torch.empty(1, 576, dtype=dtype, device=device), + torch.zeros(1, 1, dtype=torch.int32, device=device), + torch.zeros(1, dtype=torch.int32, device=device), + 1, + ) # Bytes per token for the nvfp4_ds_mla cache layout (see flashmla_sparse.py). diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 6883dbc4b363..ddafc37ba647 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -10,7 +10,7 @@ import sys from types import MethodType, SimpleNamespace -from typing import Literal +from typing import Any, Literal import pytest import torch @@ -219,13 +219,39 @@ def test_masked_mha_routing_is_dimension_specific(): ("fp8_ds_mla", KVQuantMode.FP8_PER_TENSOR), ], ) +@pytest.mark.parametrize( + ("sm120", "kv_lora_rank", "rope_dim", "row_bytes"), + [ + (False, 512, 64, 656), + (False, 512, 0, 656), + (True, 512, 64, 656), + (True, 512, 0, 528), + (True, 256, 0, 656), + ], +) def test_mla_kv_cache_spec_uses_layer_cache_dtype( - cache_dtype: str, expected_quant_mode: KVQuantMode + cache_dtype: str, + expected_quant_mode: KVQuantMode, + sm120: bool, + kv_lora_rank: int, + rope_dim: int, + row_bytes: int, ): + from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import ( + FlashInferMLASparseSM120Impl, + ) + from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseImpl + + impl: Any = object.__new__( + FlashInferMLASparseSM120Impl if sm120 else FlashMLASparseImpl + ) + impl.kv_lora_rank = kv_lora_rank + impl.qk_rope_head_dim = rope_dim layer = SimpleNamespace( attn_backend=flashmla_module.FlashMLABackend, + impl=impl, kv_cache_dtype=cache_dtype, - head_size=576, + head_size=kv_lora_rank + rope_dim, topk_buffer_width=None, indexer=None, non_causal_multi_token_decode=False, @@ -242,7 +268,8 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype( assert spec.cache_dtype_str == cache_dtype assert spec.kv_quant_mode == expected_quant_mode if cache_dtype == "fp8_ds_mla": - assert spec.page_size_bytes == 64 * 656 + assert spec.state_content_bytes == row_bytes + assert spec.page_size_bytes == 64 * row_bytes @pytest.mark.cpu_test diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 1b053014a2e1..e15dc7ab5402 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -2157,6 +2157,38 @@ def test_sparse_mla_index_groups_own_distinct_physical_buffers(): ) +@pytest.mark.parametrize( + ("kv_cache_dtype", "row_width", "kv_dtype"), + [ + ("fp8_ds_mla", 528, torch.uint8), + ("fp8_ds_mla", 656, torch.uint8), + ("bfloat16", 512, torch.bfloat16), + ], +) +def test_hisparse_index_group_preserves_row_width( + monkeypatch, kv_cache_dtype, row_width, kv_dtype +): + monkeypatch.setattr(index_group_module, "_create_side_stream", MagicMock()) + monkeypatch.setattr(index_group_module, "_create_event", MagicMock()) + create_handle = MagicMock(side_effect=lambda *args, **kwargs: SimpleNamespace()) + monkeypatch.setattr( + index_group_module, "create_hisparse_cache_handle", create_handle + ) + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=object()), model_config=None + ) + builder = SparseMLAIndexGroupBuilder(torch.empty((2, 4), dtype=torch.int32)) + for leader in (True, False): + group, layer_index = builder.register_layer( + leader, config, row_width=row_width, kv_cache_dtype=kv_cache_dtype + ) + assert isinstance(group, HiSparseMLAIndexGroup) + assert layer_index == (0 if leader else 1) + assert create_handle.call_args.kwargs["row_width"] == row_width + assert create_handle.call_args.kwargs["kv_dtype"] == kv_dtype + assert create_handle.call_args.kwargs["is_index_group_leader"] == leader + + # HiSparse is host-resident-only and kernel-only: runtime construction # raises without the compiled CUDA ops. requires_hisparse_ops = pytest.mark.skipif( @@ -2175,6 +2207,7 @@ def _make_hisparse_runtime( max_swap_rows: int | None = None, index_group: hisparse_runtime.HiSparseIndexGroup | None = None, max_union_rows: int | None = None, + kv_dtype: torch.dtype = torch.float32, ) -> HiSparseRuntime: runtime = HiSparseRuntime( config=ResolvedHiSparseConfig( @@ -2184,7 +2217,7 @@ def _make_hisparse_runtime( ), max_num_reqs=max_num_reqs, row_width=row_width, - kv_dtype=torch.float32, + kv_dtype=kv_dtype, device=DEVICE_TYPE, max_swap_rows=max_swap_rows, index_group=index_group, @@ -2195,7 +2228,7 @@ def _make_hisparse_runtime( num_blocks = max_num_reqs * blocks_per_request + 1 raw = torch.zeros( num_blocks * block_size * row_width, - dtype=torch.float32, + dtype=kv_dtype, device=DEVICE_TYPE, ).view(torch.int8) block_table = torch.arange( @@ -2206,7 +2239,7 @@ def _make_hisparse_runtime( runtime.bind_hot_cache( raw, byte_offset=0, - block_stride=block_size * row_width * torch.float32.itemsize, + block_stride=block_size * row_width * kv_dtype.itemsize, num_blocks=num_blocks, block_size=block_size, block_table=block_table, @@ -2668,17 +2701,20 @@ def test_hisparse_resident_rows_bypass_hot_lru(): @requires_hisparse_ops -def test_hisparse_swap_in_preserves_rows_across_eviction(): +@pytest.mark.parametrize( + ("row_width", "kv_dtype"), + [(64, torch.float32), (528, torch.uint8), (656, torch.uint8)], +) +def test_hisparse_swap_in_preserves_rows_across_eviction(row_width, kv_dtype): device = torch.device(DEVICE_TYPE) block_size = 64 - row_width = 64 num_blocks = 64 top_k = 128 buf = 256 num_reqs = 4 - kv_pool = torch.randn( - (num_blocks + 1, block_size, row_width), dtype=torch.float32 + kv_pool = torch.randint( + 0, 256, (num_blocks + 1, block_size, row_width), dtype=kv_dtype ).pin_memory() flat_pool = kv_pool.reshape(-1, row_width) @@ -2687,6 +2723,7 @@ def test_hisparse_swap_in_preserves_rows_across_eviction(): device_buffer_size=buf, max_num_reqs=num_reqs, row_width=row_width, + kv_dtype=kv_dtype, ) runtime.bind_source_cache(kv_pool) cache = HiSparseCacheHandle(runtime) @@ -2700,7 +2737,7 @@ def test_hisparse_swap_in_preserves_rows_across_eviction(): req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device) seq_len = blocks_per_req * block_size base = torch.arange(top_k, dtype=torch.int32, device=device) - for step in range(3): + for step in (0, 0, 1, 2, 0): cache.runtime.begin_forward() topk = torch.stack( [(base + step * top_k + row * 17) % seq_len for row in range(num_reqs)] @@ -2731,7 +2768,13 @@ def test_hisparse_swap_in_preserves_rows_across_eviction(): hot_indices[valid].to(torch.long) ].cpu() expected = flat_pool[global_ref[valid].cpu().to(torch.long)] - torch.testing.assert_close(gathered, expected) + torch.testing.assert_close(gathered, expected, rtol=0, atol=0) + + hits, misses = runtime.index_group.swap_stats.tolist() + valid_rows_per_step = num_reqs * (top_k - 1) + assert hits >= valid_rows_per_step + assert misses > 3 * valid_rows_per_step + assert hits + misses == 5 * valid_rows_per_step @requires_hisparse_ops @@ -3596,12 +3639,16 @@ def test_hisparse_prefill_resident_sources_follow_each_resident_group(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_hisparse_gather_prefill_cache_prefers_resident_rows(): +@pytest.mark.parametrize( + ("row_width", "kv_dtype"), + [(16, torch.float32), (528, torch.uint8), (656, torch.uint8)], +) +def test_hisparse_gather_prefill_cache_prefers_resident_rows(row_width, kv_dtype): """Staged rows come from the resident cache when a shadow page exists.""" if not _has_hisparse_ops(): pytest.skip("hisparse CUDA ops unavailable") device = torch.device("cuda") - block_size, resident_block_size, row_width = 4, 2, 16 + block_size, resident_block_size = 4, 2 block_table = torch.tensor([[5, 2, 0], [9, 3, 0]], dtype=torch.int32, device=device) seq_lens = torch.tensor([5, 8], dtype=torch.int32, device=device) plan = build_hisparse_prefill_staging_plan( @@ -3619,20 +3666,15 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows(): ) num_host_blocks, num_res_blocks = 10, 24 - host_cache = ( - torch.arange(num_host_blocks * block_size * row_width, dtype=torch.float32) - .view(num_host_blocks, block_size, row_width) - .pin_memory() - ) - resident_cache = ( - ( - -torch.arange( - num_res_blocks * resident_block_size * row_width, dtype=torch.float32 - ) - - 1.0 - ) - .view(num_res_blocks, resident_block_size, row_width) - .to(device) + host_cache = torch.randint( + 0, 128, (num_host_blocks, block_size, row_width), dtype=kv_dtype + ).pin_memory() + resident_cache = torch.randint( + 128, + 256, + (num_res_blocks, resident_block_size, row_width), + dtype=kv_dtype, + device=device, ) runtime = _make_hisparse_runtime(row_width=row_width, block_size=block_size) diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index b3dd64800c24..7f61e235d567 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -80,6 +80,7 @@ def test_allocation_and_warmup_follow_addressing_mode( layer = SimpleNamespace( _vllm_config=config, attn_backend=backend, + impl=object.__new__(backend.get_impl_cls()), kv_cache_dtype=cache_dtype, head_size=576, topk_buffer_width=None, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 173c40dc89de..5eef73e49c4a 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1382,9 +1382,10 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), # ds_mla layouts pack NoPE + RoPE + scales into one opaque per-token # blob, so the size is not derivable from head_size. - # See flashmla_sparse.py. - state_content_bytes={"fp8_ds_mla": 656, "nvfp4_ds_mla": 352}.get( - self.kv_cache_dtype + state_content_bytes=( + cast(MLAAttentionImpl, self.impl).get_fp8_ds_mla_row_bytes() + if self.kv_cache_dtype == "fp8_ds_mla" + else {"nvfp4_ds_mla": 352}.get(self.kv_cache_dtype) ), ) if self.sliding_window is not None: diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 2f6583f1fcac..cfc3009f52e9 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -689,7 +689,11 @@ def __init__( index_group_builder.register_layer( indexer is not None, vllm_config, - head_size=head_size, + row_width=( + self.get_fp8_ds_mla_row_bytes() + if kv_cache_dtype == "fp8_ds_mla" + else head_size + ), kv_cache_dtype=kv_cache_dtype, ) ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 5c6ce28a956d..761a582b20ca 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -1045,6 +1045,9 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): hisparse_cache: "HiSparseCacheHandle | None" = None supports_pcp: bool = True + def get_fp8_ds_mla_row_bytes(self) -> int: + return 656 + @abstractmethod def __init__( self, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 04b6edf1256b..7fc81ded2299 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -6,6 +6,7 @@ import torch +from vllm.config import get_current_vllm_config from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, ) @@ -36,6 +37,11 @@ class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetada is_sparse = True supports_dense_mha_prefill = False + def get_fp8_ds_mla_row_bytes(self) -> int: + if self.kv_lora_rank == 512 and self.qk_rope_head_dim == 0: + return 528 + return super().get_fp8_ds_mla_row_bytes() + def __init__( self, num_heads: int, @@ -83,8 +89,6 @@ def __init__( topk_indices_buffer=topk_indices_buffer, **mla_args, ) - from vllm.config import get_current_vllm_config - vllm_config = get_current_vllm_config() model_type = None if vllm_config.model_config is not None: diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index c0f4d8894403..60a966e8ebf8 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -17,7 +17,6 @@ triton_convert_req_index_to_global_index, ) from vllm.v1.hisparse.runtime import ( - FP8_DS_MLA_ROW_BYTES, HiSparseCacheHandle, HiSparsePrefillStagingPlan, create_hisparse_cache_handle, @@ -44,7 +43,7 @@ def register_layer( self, vllm_config: VllmConfig | None = None, *, - head_size: int | None = None, + row_width: int | None = None, kv_cache_dtype: str | None = None, ) -> int: layer_index = self.num_layers @@ -178,22 +177,20 @@ def register_layer( self, vllm_config: VllmConfig | None = None, *, - head_size: int | None = None, + row_width: int | None = None, kv_cache_dtype: str | None = None, ) -> int: assert vllm_config is not None - assert head_size is not None + assert row_width is not None assert kv_cache_dtype is not None layer_index = super().register_layer() if self.prefill_stream is None: self.prefill_stream = torch.Stream(device=self.logical_topk_indices.device) if kv_cache_dtype == "fp8_ds_mla": - row_width = FP8_DS_MLA_ROW_BYTES kv_dtype = torch.uint8 else: from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype - row_width = head_size kv_dtype = kv_cache_dtype_str_to_dtype( kv_cache_dtype, vllm_config.model_config ) @@ -363,7 +360,7 @@ def register_layer( is_index_producing_layer: bool, vllm_config: VllmConfig | None = None, *, - head_size: int | None = None, + row_width: int | None = None, kv_cache_dtype: str | None = None, ) -> tuple[SparseMLAIndexGroup, int]: if is_index_producing_layer or self.current_group is None: @@ -395,7 +392,7 @@ def register_layer( group = self.current_group group_index = group.register_layer( vllm_config, - head_size=head_size, + row_width=row_width, kv_cache_dtype=kv_cache_dtype, ) return group, group_index diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 01af2be565df..966631def6cc 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -36,8 +36,6 @@ from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.kv_cache_interface import KVCacheConfig -# fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE. -FP8_DS_MLA_ROW_BYTES = 656 HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor] # trtllm-gen sparse MLA reads up to ~250 rows from the start of its KV cache # whatever the cache's size: https://github.com/flashinfer-ai/flashinfer/issues/6484 From edad8b03c5285a1c12e7e4f83c1af709b319a01b Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:01:19 +0400 Subject: [PATCH 039/406] [Bugfix] Merge supported KV cache layouts across pipeline stages instead of asserting they match (#61294) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude Co-authored-by: Thien Tran --- tests/v1/core/test_kv_cache_utils.py | 25 +++++++++++++++++++ vllm/v1/attention/backends/utils.py | 36 ++++++++++++++++++++-------- 2 files changed, 51 insertions(+), 10 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 5d0d9cd5a900..baba01b043eb 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1375,6 +1375,31 @@ def test_dcp_replicated_kpool_tail_keeps_block_interior_layout(monkeypatch): assert layout == KVCacheLayout.LBHNC +def test_pipeline_stages_resolve_a_layout_every_stage_supports(monkeypatch): + """PP stages hold different layers, so their workers can report different + supported layouts (e.g. only one stage has FlashInfer layers on SM100).""" + from vllm.v1.attention.backends.utils import ( + get_supported_kv_cache_layouts, + resolve_kv_cache_layout, + ) + + monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False) + config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) + config.cache_config.kv_cache_layout = None + any_layout = [layout.name for layout in get_supported_kv_cache_layouts([])] + layout = resolve_kv_cache_layout(config, [["LBHNC", "BLHNC"], any_layout]) + assert layout == KVCacheLayout.LBHNC + + # A stage that declares nothing must not outvote a declared preference. + config.cache_config.kv_cache_layout = None + layout = resolve_kv_cache_layout(config, [["BLHNC", "LBNHC"], any_layout]) + assert layout == KVCacheLayout.BLHNC + + config.cache_config.kv_cache_layout = None + with pytest.raises(ValueError, match="No KV cache layout satisfies"): + resolve_kv_cache_layout(config, [["LBHNC"], ["LBNHC"]]) + + @pytest.mark.parametrize("use_mla", [False, True]) def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla): config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 40b283a3f1c5..47a25ab240fc 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -251,12 +251,21 @@ def get_supported_kv_cache_layouts( their order; otherwise the layout the most backends put first wins, ties keeping the enum order. An empty intersection is a hard error. """ - supported_layouts_lists: list[Sequence[KVCacheLayout]] = [ - layouts - for backend in backends - if (layouts := backend.supported_kv_cache_layouts()) is not None - ] or [_DEFAULT_LAYOUT_PREFERENCE] + return _merge_supported_layouts( + [ + layouts + for backend in backends + if (layouts := backend.supported_kv_cache_layouts()) is not None + ] + or [_DEFAULT_LAYOUT_PREFERENCE] + ) + +def _merge_supported_layouts( + supported_layouts_lists: Sequence[Sequence[KVCacheLayout]], +) -> list[KVCacheLayout]: + """Layouts in every supported set, merged as in + ``get_supported_kv_cache_layouts``.""" first = supported_layouts_lists[0] if all(layouts == first for layouts in supported_layouts_lists[1:]): return list(first) @@ -298,8 +307,10 @@ def resolve_kv_cache_layout( """Resolve one KV cache layout for the whole model. Runs once in the engine core. Every worker reports the layouts its backends - support, most preferred first (``get_supported_kv_cache_layouts``); all - ranks run the same backends, so their lists must agree. Specs mixing HNC + support, most preferred first (``get_supported_kv_cache_layouts``). Pipeline + stages hold different layers, so their lists may differ; the candidates are + the layouts every worker supports, and the layout the most workers put first + wins. Workers whose backends declare nothing do not vote. Specs mixing HNC shapes narrow the candidates to block-compact layouts. An explicit ``VLLM_KV_CACHE_LAYOUT`` must be one of the candidates or resolution fails, with the legacy ``NHD``/``HND`` names as aliases for ``LBNHC``/``LBHNC``; the @@ -316,10 +327,15 @@ def resolve_kv_cache_layout( assert supported_layouts and all(supported_layouts), ( "No worker reported supported KV cache layouts." ) - assert all(names == supported_layouts[0] for names in supported_layouts[1:]), ( - f"Workers disagree on supported KV cache layouts: {supported_layouts}." + reported = [ + tuple(_layout_from_name(name) for name in names) for names in supported_layouts + ] + # Workers where no backend declares a set report the default preference; + # like undeclared backends of one worker, they don't vote. + candidates = _merge_supported_layouts( + [layouts for layouts in reported if layouts != _DEFAULT_LAYOUT_PREFERENCE] + or reported ) - candidates = [_layout_from_name(name) for name in supported_layouts[0]] # A block-compact layout means the block is densely packed in memory, so any mix of # specs can re-interpret HNC with different sizes as long as the total number of From f60edccd0079428503f9e9f66e3b651aa4da73db Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:44:17 +0400 Subject: [PATCH 040/406] [Bugfix] Omit null `annotations` and `usage.prompt_tokens_details` in chat completions (#64177) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- .../completion/test_completion_text_offset.py | 16 +++++++++++ .../openai/test_tool_choice_content_none.py | 27 ++++++++++++++++++- .../openai/chat_completion/protocol.py | 4 ++- vllm/entrypoints/serve/engine/protocol.py | 4 ++- 4 files changed, 48 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai/completion/test_completion_text_offset.py b/tests/entrypoints/openai/completion/test_completion_text_offset.py index 5c66b81ff7fc..a1205b442fd1 100644 --- a/tests/entrypoints/openai/completion/test_completion_text_offset.py +++ b/tests/entrypoints/openai/completion/test_completion_text_offset.py @@ -97,3 +97,19 @@ async def results(): assert response.choices[0].logprobs is not None assert response.choices[0].logprobs.text_offset == [0, 3, 6] assert stream_offsets == [0, 3, 6] + + +def test_non_stream_response_omits_null_prompt_tokens_details(): + """The OpenAI spec types ``usage.prompt_tokens_details`` as an object, not + nullable, so it is omitted when ``--enable-prompt-tokens-details`` is off.""" + response = _serving().request_output_to_completion_response( + [_output("abc", [0], finished=True)], + _request(stream=False), + "r", + 0, + "m", + None, + RequestResponseMetadata(request_id="r"), + ) + + assert "prompt_tokens_details" not in response.model_dump()["usage"] diff --git a/tests/entrypoints/openai/test_tool_choice_content_none.py b/tests/entrypoints/openai/test_tool_choice_content_none.py index f08915a93dca..3b37321d8a6c 100644 --- a/tests/entrypoints/openai/test_tool_choice_content_none.py +++ b/tests/entrypoints/openai/test_tool_choice_content_none.py @@ -23,7 +23,7 @@ ChatMessage, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.entrypoints.serve.engine.protocol import UsageInfo +from vllm.entrypoints.serve.engine.protocol import PromptTokenUsageInfo, UsageInfo from vllm.parser.abstract_parser import DelegatingParser pytestmark = pytest.mark.skip_global_cleanup @@ -152,6 +152,31 @@ def test_chat_completion_response_keeps_non_empty_tool_calls_payload(): assert message["tool_calls"][0]["function"]["name"] == "get_weather" +def test_chat_completion_response_omits_null_annotations_and_prompt_tokens_details(): + """The OpenAI spec types both as array/object, not nullable.""" + payload = _chat_response(ChatMessage(role="assistant", content="done")).model_dump() + + assert "annotations" not in payload["choices"][0]["message"] + assert "prompt_tokens_details" not in payload["usage"] + + +def test_chat_completion_response_keeps_annotations_and_prompt_tokens_details(): + citation = {"start_index": 0, "end_index": 4, "title": "t", "url": "https://x"} + response = _chat_response( + ChatMessage( + role="assistant", + content="done", + annotations=[{"type": "url_citation", "url_citation": citation}], + ) + ) + response.usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=1) + + parsed = OpenAIChatCompletion.model_validate(response.model_dump()) + + assert parsed.choices[0].message.annotations[0].url_citation.url == "https://x" + assert parsed.usage.prompt_tokens_details.cached_tokens == 1 + + def _stream_response(delta: DeltaMessage) -> ChatCompletionStreamResponse: return ChatCompletionStreamResponse( id="chatcmpl-test", diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 1341ef2b86fc..ca9e8a01a07b 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -77,7 +77,9 @@ class ChatMessage(OpenAIBaseModel): role: str content: str | None = None refusal: str | None = None - annotations: OpenAIAnnotation | None = None + annotations: list[OpenAIAnnotation] | None = Field( + default=None, exclude_if=lambda value: value is None + ) audio: OpenAIChatCompletionAudio | None = None function_call: FunctionCall | None = None tool_calls: list[ToolCall] = Field(default_factory=list) diff --git a/vllm/entrypoints/serve/engine/protocol.py b/vllm/entrypoints/serve/engine/protocol.py index 704d6b2455a8..ef1592199625 100644 --- a/vllm/entrypoints/serve/engine/protocol.py +++ b/vllm/entrypoints/serve/engine/protocol.py @@ -136,5 +136,7 @@ class UsageInfo(OpenAIBaseModel): prompt_tokens: int = 0 total_tokens: int = 0 completion_tokens: int | None = 0 - prompt_tokens_details: PromptTokenUsageInfo | None = None + prompt_tokens_details: PromptTokenUsageInfo | None = Field( + default=None, exclude_if=lambda value: value is None + ) completion_tokens_details: CompletionTokenUsageInfo | None = None From 665bfc8aaeb7eca008ca4102cbba68efdc4efabb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Wed, 28 Oct 2026 06:46:42 +0100 Subject: [PATCH 041/406] [Security] Fix ReDoS in validation error path sanitizer (#52120) Signed-off-by: jperezde --- .../test_error_sanitization.py | 54 +++++++++++++++++++ .../serve/exception_handling/utils.py | 3 +- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py index c783fb7db67c..0f06c3cf9bed 100644 --- a/tests/entrypoints/serve/exception_handling/test_error_sanitization.py +++ b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py @@ -153,3 +153,57 @@ def test_engine_dead_error_response_hides_log_oriented_message(): assert "stack trace" not in err.error.message assert "log" not in err.error.message assert "stack trace (above)" in str(exc) + + +class TestSanitizeMessageReDoSResistance: + """Regression tests for GHSA-f2g9-pmwr-xwc7. + + Verifies that sanitize_message performs bounded near-linear work + on adversarial slash-delimited input with no filename extension. + """ + + def test_slash_segments_no_extension_completes_quickly(self): + """50k slash-segments without dot-extension must complete in < 0.1s.""" + import time + + message = "/a" * 50000 + t0 = time.perf_counter() + result = sanitize_message(message) + elapsed = time.perf_counter() - t0 + assert elapsed < 0.1, f"Took {elapsed:.2f}s (expected < 0.1s)" + assert result == message + + def test_alphanumeric_control_completes(self): + """100k alphanumeric chars must complete near-instantly.""" + import time + + message = "a" * 100000 + t0 = time.perf_counter() + result = sanitize_message(message) + elapsed = time.perf_counter() - t0 + assert elapsed < 0.1, f"Took {elapsed:.2f}s" + assert result == message + + def test_linear_scaling(self): + """Verify near-linear scaling across input sizes.""" + import time + + times = [] + for n in [1000, 10000, 50000]: + message = "/a" * n + t0 = time.perf_counter() + sanitize_message(message) + times.append(time.perf_counter() - t0) + + ratio = times[-1] / max(times[0], 1e-9) + assert ratio < 100, ( + f"50k/1k time ratio is {ratio:.0f}x (expected < 100x for linear)" + ) + + def test_path_redaction_still_works(self): + """Legitimate paths are still redacted after the fix.""" + assert "" in sanitize_message("Error in /app/server.py") + assert "" in sanitize_message("at /workspace/vllm/engine.py") + assert "" in sanitize_message( + "error at /usr/lib/python3.12/dist-packages/vllm/core.py" + ) diff --git a/vllm/entrypoints/serve/exception_handling/utils.py b/vllm/entrypoints/serve/exception_handling/utils.py index 6c1dd205167c..1c26c90cd5c2 100644 --- a/vllm/entrypoints/serve/exception_handling/utils.py +++ b/vllm/entrypoints/serve/exception_handling/utils.py @@ -11,5 +11,6 @@ def sanitize_message(message: str) -> str: message = re.sub( r"/(?:home|usr|opt|var|tmp|root|lib|mnt|srv)(?:/[\w.\-]+)+", "", message ) - message = re.sub(r"(?:/[\w\-]+)+/[\w\-]+\.\w+", "", message) + if "." in message: + message = re.sub(r"(?>/[\w\-]+)+/[\w\-]+\.\w+", "", message) return message.strip() From 26302eb938966fbe2d17c2ad9cbbd75bbdc7bf9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20P=C3=A9rez=20de=20Algaba?= <124347725+jperezdealgaba@users.noreply.github.com> Date: Wed, 28 Oct 2026 07:10:09 +0100 Subject: [PATCH 042/406] [Security] Bound FunAudioChat runtime audio to the profiled duration (#57411) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Juan Pérez de Algaba --- .../processing/test_funaudiochat.py | 129 ++++++++++++++++++ vllm/model_executor/models/funaudiochat.py | 44 ++++-- 2 files changed, 162 insertions(+), 11 deletions(-) create mode 100644 tests/models/multimodal/processing/test_funaudiochat.py diff --git a/tests/models/multimodal/processing/test_funaudiochat.py b/tests/models/multimodal/processing/test_funaudiochat.py new file mode 100644 index 000000000000..b090c6f38bdb --- /dev/null +++ b/tests/models/multimodal/processing/test_funaudiochat.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch + +from vllm.model_executor.models.funaudiochat import ( + FunAudioChatDummyInputsBuilder, + FunAudioChatMultiModalProcessor, + FunAudioChatProcessingInfo, +) +from vllm.transformers_utils.configs.funaudiochat import FunAudioChatConfig + + +class RecordingFeatureExtractor: + sampling_rate = 16000 + n_fft = 400 + hop_length = 160 + + def __init__(self) -> None: + self.last_kwargs: dict[str, object] | None = None + + def __call__(self, wavs, **kwargs): + self.last_kwargs = kwargs + n_frames = [max(1, int(w.shape[0]) // self.hop_length) for w in wavs] + max_frames = max(n_frames) + features = torch.zeros((len(wavs), 128, max_frames)) + mask = torch.zeros((len(wavs), max_frames), dtype=torch.long) + for i, frames in enumerate(n_frames): + mask[i, :frames] = 1 + return {"input_features": features, "attention_mask": mask} + + +class MockSpeechTokenizer: + pad_token = "<|audio_pad|>" + + def __call__(self, speech_strs, **kwargs): + token_len = max(1, len(self.pad_token)) + lengths = [len(s) // token_len for s in speech_strs] + max_len = max(lengths) + pad_to = kwargs.get("pad_to_multiple_of") + if pad_to: + max_len = ((max_len + pad_to - 1) // pad_to) * pad_to + ids = torch.zeros((len(speech_strs), max_len), dtype=torch.long) + mask = torch.zeros((len(speech_strs), max_len), dtype=torch.long) + for i, n in enumerate(lengths): + mask[i, :n] = 1 + return {"input_ids": ids, "attention_mask": mask} + + +class MockTokenizer: + def encode(self, text, **kwargs): + return [1, 2, 3] + + +@pytest.fixture +def funaudiochat_processor(monkeypatch): + feature_extractor = RecordingFeatureExtractor() + speech_tokenizer = MockSpeechTokenizer() + tokenizer = MockTokenizer() + + monkeypatch.setattr( + "vllm.model_executor.models.funaudiochat.WhisperFeatureExtractor.from_pretrained", + lambda *args, **kwargs: feature_extractor, + ) + monkeypatch.setattr( + "vllm.model_executor.models.funaudiochat.TokenizersBackend.from_pretrained", + lambda *args, **kwargs: speech_tokenizer, + ) + + ctx = MagicMock() + ctx.get_hf_config.return_value = FunAudioChatConfig() + ctx.model_config.hf_config = FunAudioChatConfig() + ctx.model_config.model = "funaudiochat" + ctx.model_config.revision = None + ctx.model_config.tokenizer_revision = None + mm_config = MagicMock() + mm_config.enable_mm_embeds = False + mm_config.allow_missing_mm_embeddings = False + ctx.model_config.multimodal_config = mm_config + ctx.model_config.get_multimodal_config.return_value = mm_config + ctx.tokenizer = tokenizer + ctx.get_tokenizer.return_value = tokenizer + + info = FunAudioChatProcessingInfo(ctx) + processor = FunAudioChatMultiModalProcessor( + info, FunAudioChatDummyInputsBuilder(info) + ) + return processor, feature_extractor + + +def test_max_audio_samples_match_profiled_300s_budget(funaudiochat_processor): + processor, _ = funaudiochat_processor + sampling_rate = 16000 + + assert processor.info.get_max_speech_frames() == 7500 + assert processor.info.get_max_audio_samples(sampling_rate) == 4_800_000 + assert processor.info.get_max_audio_samples(sampling_rate) / sampling_rate == 300.0 + + +def test_over_budget_audio_is_rejected(funaudiochat_processor): + processor, feature_extractor = funaudiochat_processor + sampling_rate = 16000 + over_budget = np.zeros( + processor.info.get_max_audio_samples(sampling_rate) + 1, dtype=np.float32 + ) + mm_items = processor.info.parse_mm_data({"audio": [over_budget]}, validate=False) + + with pytest.raises(ValueError, match="too long for FunAudioChat"): + processor._apply_hf_processor_main(mm_items, {}) + + assert feature_extractor.last_kwargs is None + + +def test_profiled_max_audio_is_accepted(funaudiochat_processor): + processor, feature_extractor = funaudiochat_processor + sampling_rate = 16000 + max_samples = processor.info.get_max_audio_samples(sampling_rate) + audio = np.zeros(max_samples, dtype=np.float32) + mm_items = processor.info.parse_mm_data({"audio": [audio]}, validate=False) + + processed = processor._apply_hf_processor_main(mm_items, {}) + + assert processed["speech_attention_mask"].sum().item() == 7500 + assert feature_extractor.last_kwargs is not None + assert feature_extractor.last_kwargs.get("padding") == "max_length" diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 8e64fe03e22d..ac31c20d6b4c 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -532,7 +532,7 @@ def speech_tokenizer(self) -> TokenizersBackend: revision=self.ctx.model_config.revision, ) - def get_feature_extractor(self) -> WhisperFeatureExtractor: + def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor: return self.feature_extractor def get_speech_tokenizer(self) -> TokenizersBackend: @@ -568,6 +568,26 @@ def get_audio_group_size(self) -> int: audio_cfg = getattr(cfg, "audio_config", None) return int(getattr(audio_cfg, "group_size", 5)) + def get_max_speech_frames(self) -> int: + """Maximum 25 Hz speech frames admitted for one audio item. + + The discrete encoder emits ``max_source_positions`` replacement tokens + after grouping frames by ``group_size`` (1500 * 5 = 7500 frames, 300s). + """ + cfg = self.get_hf_config() + audio_cfg = getattr(cfg, "audio_config", None) + max_audio_tokens = int(getattr(audio_cfg, "max_source_positions", 1500)) + return max(1, max_audio_tokens) * max(1, self.get_audio_group_size()) + + def get_max_audio_samples(self, sampling_rate: int) -> int: + """Waveform samples that fill the profiled per-item speech-frame budget.""" + token_fps = int(self.token_fps) + target_num_frames = self.get_max_speech_frames() + return max( + 1, + (target_num_frames * sampling_rate + token_fps - 1) // token_fps, + ) + class FunAudioChatDummyInputsBuilder( BaseDummyInputsBuilder[FunAudioChatProcessingInfo] @@ -587,16 +607,7 @@ def get_dummy_mm_data( # Dummy inputs are used for profiling; construct the worst-case audio # length that maximizes the number of encoder tokens. - cfg = self.info.get_hf_config() - audio_cfg = getattr(cfg, "audio_config", None) - max_audio_tokens = int(getattr(audio_cfg, "max_source_positions", 1500)) - group_size = self.info.get_audio_group_size() - token_fps = int(getattr(self.info, "token_fps", 25)) - target_num_frames = max(1, max_audio_tokens) * max(1, group_size) - audio_len = max( - 1, - (target_num_frames * sampling_rate + token_fps - 1) // token_fps, - ) + audio_len = self.info.get_max_audio_samples(sampling_rate) return { "audio": self._get_dummy_audios( @@ -638,6 +649,8 @@ def _apply_hf_processor_main( feature_extractor = self.info.get_feature_extractor(**hf_kwargs) sr = int(feature_extractor.sampling_rate) min_samples = int(getattr(feature_extractor, "n_fft", 400) or 400) + max_samples = self.info.get_max_audio_samples(sr) + max_speech_frames = self.info.get_max_speech_frames() wavs: list[np.ndarray] = [] speech_strs: list[str] = [] @@ -654,6 +667,15 @@ def _apply_hf_processor_main( audio_np, (0, min_samples - audio_np.shape[0]), mode="constant" ) + if audio_np.shape[0] > max_samples: + duration_s = float(audio_np.shape[0]) / float(sr) + max_duration_s = float(max_samples) / float(sr) + raise ValueError( + "Audio is too long for FunAudioChat: " + f"{duration_s:.1f}s exceeds the profiled maximum of " + f"{max_duration_s:.1f}s ({max_speech_frames} speech frames)." + ) + wavs.append(audio_np) num_frames = int( (float(audio_np.shape[0]) / float(sr)) * float(self.info.token_fps) From e82e6d422fb065aa09c2e1b8efb210d2629b13f3 Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Wed, 28 Oct 2026 02:17:17 -0400 Subject: [PATCH 043/406] [Bugfix][Pooling] Reject prompt-wide truncation for Jina ranking (#55200) Signed-off-by: Bruno Volpato Co-authored-by: OpenAI Codex --- docs/models/pooling_models/token_embed.md | 9 + .../test_jina_ranking_io_processor_unit.py | 54 ------ .../language/pooling/test_jina_reranker_v3.py | 166 ++++++++++++++++++ .../pooling/scoring/io_processor.py | 27 ++- 4 files changed, 193 insertions(+), 63 deletions(-) diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index b0a09c0f887d..3823d5560527 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -82,6 +82,15 @@ If your model is not in the above list, we will try to automatically convert the jina-reranker-v3 is a listwise document reranker model with a novel `last but not late interaction` architecture. More information can be found at: [examples/pooling/token_embed/jina_reranker_v3_offline.py](../../../examples/pooling/token_embed/jina_reranker_v3_offline.py) +For `jinaai/jina-reranker-v3`, `truncate_prompt_tokens` is unsupported. This applies to +`LLM.encode(..., pooling_task="token_embed")`, `LLM.score`, `/pooling`, `/score`, and +`/rerank`, even when the limit would leave the prompt unchanged. Requests that set +it fail before inference; HTTP APIs return status 400 with +`param="truncate_prompt_tokens"`. For scoring, use `max_tokens_per_query` and +`max_tokens_per_doc` to truncate the query and documents separately. With `LLM.score`, +pass these limits through `PoolingParams(extra_kwargs={...})`. For `/pooling` and +`LLM.encode`, truncate each input text on the client before submitting the request. + --8<-- [end:supported-token-embed-models] ## Offline Inference diff --git a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py index e69dd6d45a5b..7fc2ab226c2b 100644 --- a/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_jina_ranking_io_processor_unit.py @@ -3,16 +3,13 @@ """Unit tests for JinaRankingIOProcessor request building.""" from types import SimpleNamespace -from unittest.mock import MagicMock import pytest from tokenizers import Tokenizer, models, pre_tokenizers from transformers import TokenizersBackend from vllm import PoolingParams -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessor -from vllm.entrypoints.pooling.scoring.protocol import RerankRequest from vllm.entrypoints.pooling.scoring.typing import ScoringData from vllm.entrypoints.pooling.typing import OfflineScoringInputsContext from vllm.renderers import TokenizeParams @@ -20,57 +17,6 @@ pytestmark = pytest.mark.skip_global_cleanup -def test_online_forwards_truncate_prompt_tokens_to_proxy(monkeypatch): - """The proxy request handed to the base factory must carry - truncate_prompt_tokens/truncation_side from the real request. - - JinaRankingIOProcessor swaps ctx.request for a proxy - PoolingCompletionRequest before delegating to the base factory, which - reads truncation off ctx.request. Dropping the fields on the proxy - silently disables truncate_prompt_tokens for Jina rerank/score. - """ - proc = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor) - monkeypatch.setattr( - proc, - "valid_inputs_online", - MagicMock(return_value=ScoringData(data_1=["query"], data_2=["doc"])), - ) - monkeypatch.setattr(proc, "_get_token_limits", MagicMock(return_value=(0, 0))) - monkeypatch.setattr( - proc, "ensure_str", MagicMock(side_effect=lambda data: list(data)) - ) - monkeypatch.setattr( - proc, "format_docs_prompts_func", MagicMock(return_value="formatted prompt") - ) - - captured: dict[str, object] = {} - - def _spy_base(self, ctx): - captured["truncate_prompt_tokens"] = ctx.request.truncate_prompt_tokens - captured["truncation_side"] = ctx.request.truncation_side - return [] - - monkeypatch.setattr(PoolingIOProcessor, "get_request_factory_online", _spy_base) - - request = RerankRequest( - model="m", - query="query", - documents=["doc"], - truncate_prompt_tokens=512, - truncation_side="left", - ) - ctx = MagicMock() - ctx.request = request - ctx.prompt_extras = None - - proc.get_request_factory_online(ctx) - - assert captured["truncate_prompt_tokens"] == 512 - assert captured["truncation_side"] == "left" - # The real request is restored after delegating. - assert ctx.request is request - - @pytest.fixture def offline_processor_and_context(): backend = Tokenizer(models.WordLevel({"[UNK]": 0}, unk_token="[UNK]")) diff --git a/tests/models/language/pooling/test_jina_reranker_v3.py b/tests/models/language/pooling/test_jina_reranker_v3.py index 3fc5c48c5a59..a53689cfbd31 100644 --- a/tests/models/language/pooling/test_jina_reranker_v3.py +++ b/tests/models/language/pooling/test_jina_reranker_v3.py @@ -7,8 +7,10 @@ import torch.nn.functional as F from tests.utils import RemoteOpenAIServer +from vllm import PoolingParams from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse from vllm.entrypoints.pooling.scoring.protocol import RerankResponse, ScoreResponse +from vllm.exceptions import VLLMValidationError model_name = "jinaai/jina-reranker-v3" query = "What are the health benefits of green tea?" @@ -53,6 +55,8 @@ def test_offline(vllm_runner, monkeypatch): _test_offline_1_v_1(llm) _test_offline_1_v_n(llm) _test_offline_n_v_n(llm) + _test_offline_truncation(llm) + _test_offline_nested_markers(llm) _test_offline_token_embed_illegal_inputs(llm) assert llm.model_config.embedding_size == EMBEDDING_SIZE @@ -65,6 +69,8 @@ def test_online(): _test_online_1_v_n(server) _test_online_n_v_n(server) _test_online_instruction(server) + _test_online_truncation(server) + _test_online_nested_markers(server) _test_online_token_embed_illegal_inputs(server) @@ -145,6 +151,81 @@ def _test_offline_token_embed_illegal_inputs(llm): llm.encode([1, 2, 3], pooling_task="token_embed") +def _test_offline_truncation(llm): + prompts = documents[:2] + [query] + (output,) = llm.encode(prompts, pooling_task="token_embed") + for side in ("left", "right"): + for limit in ( + -1, + 1, + len(output.prompt_token_ids) - 1, + len(output.prompt_token_ids) + 1, + ): + kwargs = {"truncate_prompt_tokens": limit, "truncation_side": side} + for call, args, extra in ( + (llm.encode, (prompts,), {"pooling_task": "token_embed"}), + (llm.score, (query, documents[:2]), {}), + ): + with pytest.raises( + VLLMValidationError, + match="JinaForRanking does not support truncate_prompt_tokens", + ) as exc_info: + call(*args, **extra, tokenization_kwargs=kwargs) + assert exc_info.value.parameter == "truncate_prompt_tokens" + + limited = llm.score( + "one two three four five", + ["six seven eight nine ten", "red green blue yellow white"], + pooling_params=PoolingParams( + extra_kwargs={"max_tokens_per_query": 3, "max_tokens_per_doc": 4} + ), + ) + expected = llm.score( + "one two three", ["six seven eight nine", "red green blue yellow"] + ) + assert len(limited) == len(expected) == 2 + for actual, reference in zip(limited, expected): + assert actual.prompt_token_ids == reference.prompt_token_ids + assert actual.outputs.score == pytest.approx(reference.outputs.score, abs=TOL) + + +def _test_offline_nested_markers(llm): + prompts = documents[:2] + [query] + (expected,) = llm.encode(prompts, pooling_task="token_embed") + instruction = "Rank relevant passages." + scores = llm.score( + query, + documents[:2], + pooling_params=PoolingParams( + extra_kwargs={"chat_template_kwargs": {"instruction": instruction}} + ), + ) + for token in ("embed", "rerank"): + nested = f"<|{token}<|{token}_token|>_token|>" + (actual,) = llm.encode( + [text + nested for text in prompts], pooling_task="token_embed" + ) + assert actual.outputs.data.shape == (len(prompts), EMBEDDING_SIZE), ( + actual.outputs.data.shape + ) + assert actual.prompt_token_ids == expected.prompt_token_ids + actual_scores = llm.score( + query + nested, + [doc + nested for doc in documents[:2]], + pooling_params=PoolingParams( + extra_kwargs={ + "chat_template_kwargs": {"instruction": instruction + nested} + } + ), + ) + assert len(actual_scores) == len(scores) + for actual_score, reference in zip(actual_scores, scores): + assert actual_score.prompt_token_ids == reference.prompt_token_ids + assert actual_score.outputs.score == pytest.approx( + reference.outputs.score, abs=TOL + ) + + def _get_score_response(server, query, document, **extra_body): payload = { "model": model_name, @@ -306,6 +387,91 @@ def _test_online_instruction(server): ) +def _test_online_truncation(server): + docs = documents[:2] + for endpoint, inputs in ( + ("pooling", {"input": docs + [query], "task": "token_embed"}), + ("score", {"queries": query, "documents": docs}), + ("rerank", {"query": query, "documents": docs}), + ): + payload = {"model": model_name, **inputs} + response = requests.post(server.url_for(endpoint), json=payload) + response.raise_for_status() + num_tokens = response.json()["usage"]["prompt_tokens"] + for side in ("left", "right"): + for limit in (-1, 1, num_tokens - 1, num_tokens + 1): + payload.update(truncate_prompt_tokens=limit, truncation_side=side) + response = requests.post(server.url_for(endpoint), json=payload) + assert response.status_code == 400, response.status_code + error = response.json()["error"] + assert ( + "JinaForRanking does not support truncate_prompt_tokens" + in error["message"] + ) + assert error["param"] == "truncate_prompt_tokens" + + for get_response in (_get_score_response, _get_rerank_response): + limited = get_response( + server, + "one two three four five", + ["six seven eight nine ten", "red green blue yellow white"], + max_tokens_per_query=3, + max_tokens_per_doc=4, + ) + expected = get_response( + server, "one two three", ["six seven eight nine", "red green blue yellow"] + ) + assert limited.usage.prompt_tokens == expected.usage.prompt_tokens + if isinstance(limited, ScoreResponse): + assert len(limited.data) == len(expected.data) == 2 + assert [d.score for d in limited.data] == pytest.approx( + [d.score for d in expected.data], abs=TOL + ) + else: + assert [d.index for d in limited.results] == [ + d.index for d in expected.results + ] + assert [d.relevance_score for d in limited.results] == pytest.approx( + [d.relevance_score for d in expected.results], abs=TOL + ) + + +def _test_online_nested_markers(server): + docs = documents[:2] + expected = _get_embeds(server, docs + [query]) + instruction = "Rank relevant passages." + for token in ("embed", "rerank"): + nested = f"<|{token}<|{token}_token|>_token|>" + actual = _get_embeds(server, [text + nested for text in docs + [query]]) + assert actual.shape == expected.shape == (len(docs) + 1, EMBEDDING_SIZE), ( + actual.shape + ) + torch.testing.assert_close(actual, expected) + for get_response in (_get_score_response, _get_rerank_response): + reference = get_response(server, query, docs, instruction=instruction) + actual_scores = get_response( + server, + query + nested, + [doc + nested for doc in docs], + instruction=instruction + nested, + ) + assert actual_scores.usage.prompt_tokens == reference.usage.prompt_tokens + if isinstance(actual_scores, ScoreResponse): + assert len(actual_scores.data) == len(reference.data) == len(docs) + assert [d.score for d in actual_scores.data] == pytest.approx( + [d.score for d in reference.data], abs=TOL + ) + else: + assert [d.index for d in actual_scores.results] == [ + d.index for d in reference.results + ] + assert [ + d.relevance_score for d in actual_scores.results + ] == pytest.approx( + [d.relevance_score for d in reference.results], abs=TOL + ) + + def _test_online_token_embed_illegal_inputs(server): response = requests.post( server.url_for("pooling"), diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index cd1b3df61b6a..483f0737aefb 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt -from vllm.exceptions import ChatTemplateResolutionError +from vllm.exceptions import ChatTemplateResolutionError, VLLMValidationError from vllm.logger import init_logger from vllm.renderers import TokenizeParams, merge_kwargs from vllm.renderers.hf import resolve_chat_template, safe_apply_chat_template @@ -833,7 +833,17 @@ def default_tokenizer_encode(): return full_prompt, engine_prompt -class JinaRankingIOProcessorMixin: +class JinaRankingIOProcessorMixin(PoolingIOProcessor): + def render(self, render_params: AnyRenderParam) -> PoolingEngineInput: + if render_params["tok_params"].truncate_prompt_tokens is not None: + raise VLLMValidationError( + "JinaForRanking does not support truncate_prompt_tokens. Use " + "max_tokens_per_query and max_tokens_per_doc for scoring, or " + "truncate each input text on the client.", + parameter="truncate_prompt_tokens", + ) + return super().render(render_params) + @staticmethod def format_docs_prompts_func( query: str, @@ -852,9 +862,12 @@ def format_docs_prompts_func( special_tokens = default_special_tokens def sanitize_input(text: str) -> str: - for token in special_tokens.values(): - text = text.replace(token, "") - return text + while True: + previous = text + for token in special_tokens.values(): + text = text.replace(token, "") + if text == previous: + return text query = sanitize_input(query) docs = [sanitize_input(doc) for doc in docs] @@ -950,10 +963,6 @@ def get_request_factory_online( for q, d in zip(queries, docs) ] - # Forward truncation from the real request: the base factory reads - # these off ctx.request, so omitting them here silently drops - # truncate_prompt_tokens for Jina rerank/score (unlike the embed and - # bi/cross-encoder paths, which read them from the real request). ctx.request = PoolingCompletionRequest( task="token_embed", input=prompts, From 08a236fa43296db31b702c71608833e9fd7fbc41 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Wed, 28 Oct 2026 14:21:27 +0800 Subject: [PATCH 044/406] [CI][XPU] Deselect Mamba deep-sleep case in Intel CI (#64057) Signed-off-by: Chaojun Zhang Co-authored-by: Kunshang Ji --- .buildkite/intel_jobs/basic_correctness_intel.yaml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.buildkite/intel_jobs/basic_correctness_intel.yaml b/.buildkite/intel_jobs/basic_correctness_intel.yaml index 69059228179b..da6d2ec0f41b 100644 --- a/.buildkite/intel_jobs/basic_correctness_intel.yaml +++ b/.buildkite/intel_jobs/basic_correctness_intel.yaml @@ -20,6 +20,8 @@ steps: - vllm/ - tests/basic_correctness/memory/ - tests/basic_correctness/cpu_offload/ + # MRV2 model state remains in the weights pool on XPU (vllm-project/vllm#63548). + # Re-enable Mamba deep sleep once XPU nested memory pools are supported. commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh @@ -28,4 +30,5 @@ steps: pytest -v -s basic_correctness/cpu_offload && pytest -v -s basic_correctness/memory - --deselect="tests/basic_correctness/memory/cumem/test_cumem.py::test_basic_cumem"' + --deselect="tests/basic_correctness/memory/cumem/test_cumem.py::test_basic_cumem" + --deselect="tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py::test_deep_sleep[yujiepan/mamba2-codestral-v0.1-tiny-random]"' From 41dd82e6a7392ead33b1e8a7d175442511e488ca Mon Sep 17 00:00:00 2001 From: Yejing Lai Date: Wed, 28 Oct 2026 14:31:01 +0800 Subject: [PATCH 045/406] [XPU][CI] Fix test_tp1_fp8_fusions rms_quant_fusion on XPU (#63975) Signed-off-by: Lai, Yejing Co-authored-by: vllm-agent --- tests/compile/fusions_e2e/test_tp1_quant.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/compile/fusions_e2e/test_tp1_quant.py b/tests/compile/fusions_e2e/test_tp1_quant.py index 6e675b3b660e..12cd841a09fe 100644 --- a/tests/compile/fusions_e2e/test_tp1_quant.py +++ b/tests/compile/fusions_e2e/test_tp1_quant.py @@ -133,6 +133,13 @@ def test_tp1_fp8_fusions( # enabling AITER. matches_check.remove("norm_rope_fusion") + if current_platform.is_xpu(): + # XPU defaults to the weight-only W8A16 FP8 path, which leaves no + # activation quant op for rms+quant fusion; opt into W8A8 instead. + model_kwargs["linear_backend"] = "xpu" + # Attention + quant fusion is not supported on XPU. + matches = matches._replace(attn_quant_fusion=0) + run_e2e_fusion_test( model_name, matches, From fd3469f4533d0517b0c8a6b72459d0fe7fd58d2f Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Wed, 28 Oct 2026 14:53:32 +0800 Subject: [PATCH 046/406] [XPU] bump xpu-kernels to v0.1.16.2 (#61956) Signed-off-by: Yan Ma Co-authored-by: vllm-agent --- .buildkite/intel_jobs/lora_intel.yaml | 2 +- requirements/xpu.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml index 366aaf7109aa..3ff650c99752 100644 --- a/.buildkite/intel_jobs/lora_intel.yaml +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -57,7 +57,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && - pytest -v -s lora/test_fused_moe_lora_kernel.py && + pytest -v -s lora/test_fused_moe_lora_kernel.py --deselect="lora/test_fused_moe_lora_kernel.py::test_fused_moe_lora_kernel_large_rank_fallback[xpu:0-dtype0-1-True-256]" && pytest -v -s lora/test_moe_lora_align_sum.py --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"' diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 14372071f0b5..11fd573a81ba 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -20,4 +20,4 @@ torchvision torchcodec >= 0.14 # Required for the torchcodec video decoding backend auto_round_lib==0.16.0 -vllm_xpu_kernels==0.1.15.4 +vllm_xpu_kernels==0.1.16.2 From 6abdff59301f96f6c8733c38691bd36bd2dd0066 Mon Sep 17 00:00:00 2001 From: nbommi19 Date: Wed, 28 Oct 2026 00:22:14 -0700 Subject: [PATCH 047/406] [Perf][DSv4.1] Tile the NVFP4 compressed-KV dequant-gather (#63018) Signed-off-by: Nitin Sai Bommi Co-authored-by: Claude Opus 5.5 --- tests/kernels/test_compressor_kv_cache.py | 152 ++++++++++++++++++ .../deepseek_v41/common/ops/cache_utils.py | 80 +++++---- 2 files changed, 205 insertions(+), 27 deletions(-) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 3cd998a60ec4..9d4756424a13 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -31,6 +31,7 @@ ) from vllm.models.deepseek_v4.compressor import _get_c128_boundary from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.mla.compressor_utils import ( get_dspark_swa_index_width, ) @@ -799,6 +800,157 @@ def check(shorter_by): check(1) +@triton.jit +def _per_token_nvfp4_gather_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, + cache_block_size: tl.constexpr, + block_stride: tl.constexpr, +): + """Reference for the tiled NVFP4 gather: one token per iteration, with + the e2m1 grid decoded arithmetically.""" + batch_idx = tl.program_id(0) + 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, 512) + for i in range(tl.program_id(1), gather_len, tl.num_programs(1)): + pos = start_pos + i + physical_block_idx = tl.load( + block_table_ptr + batch_idx * max_blocks_per_seq + pos // cache_block_size + ) + page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + token_ptr = page + (pos % cache_block_size) * 288 + packed = tl.load(token_ptr + tl.arange(0, 256)) + codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32)) + mag_code = codes & 7 + e = (mag_code >> 1).to(tl.float32) + m = (mag_code & 1).to(tl.float32) + mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0)) + vals = tl.where(codes >= 8, -mag, mag) + sf = tl.load(token_ptr + 256 + tl.arange(0, 32)) + scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) + dequant = tl.reshape(tl.reshape(vals, (32, 16)) * scale[:, None], (512,)) + out_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + tl.store(out_row_ptr + d, dequant.to(tl.bfloat16)) + + +def _per_token_nvfp4_gather( + out, k_cache, seq_lens, gather_lens, block_table, block_size, offset +): + _per_token_nvfp4_gather_kernel[(seq_lens.shape[0], 128)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + cache_block_size=block_size, + block_stride=k_cache.stride(0), + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize( + "block_size,num_reqs", [(64, 1), (128, 3), (64, 17), (128, 64)] +) +@pytest.mark.parametrize("partial", [False, True]) +def test_v41_nvfp4_gather_ragged_batch(block_size, num_reqs, partial): + """The NVFP4 gather writes the same bits as the per-token reference on + arbitrary cache bytes (NaN, subnormal and negative scales included) for + ragged batches: empty requests, partial last blocks and a long context, a + shuffled block table holding -1 in every slot the gather must not read, + and an offset into a wider strided workspace.""" + from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache + + torch.manual_seed(19) + device = "cuda" + if num_reqs == 1: + seq_lens_list = [65_541] + else: + edges = [0, 1, block_size - 1, block_size, block_size + 1, 3 * block_size + 5] + seq_lens_list = (edges + torch.randint(0, 5000, (num_reqs,)).tolist())[ + :num_reqs + ] + gather_lens_list = ( + [int(torch.randint(0, n + 1, ())) for n in seq_lens_list] if partial else None + ) + offset = 7 if partial else 0 + + blocks_per_req = [math.ceil(n / block_size) for n in seq_lens_list] + num_blocks = sum(blocks_per_req) + 1 + # Padded pages, so the block stride is not block_size * 288. + page_bytes = block_size * 288 + 512 + backing = torch.randint( + 0, 256, (num_blocks, page_bytes), dtype=torch.uint8, device=device + ) + cache = backing.as_strided((num_blocks, block_size, 288), (page_bytes, 288, 1)) + pages = iter(torch.randperm(num_blocks).tolist()) + block_table = torch.full((num_reqs, max(blocks_per_req) + 2), -1, dtype=torch.int32) + for req, seq_len in enumerate(seq_lens_list): + start = 0 if gather_lens_list is None else seq_len - gather_lens_list[req] + for block in range(start // block_size, blocks_per_req[req]): + block_table[req, block] = next(pages) + block_table = block_table.to(device) + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=device) + gather_lens = ( + None + if gather_lens_list is None + else torch.tensor(gather_lens_list, dtype=torch.int32, device=device) + ) + + width = max(seq_lens_list) + offset + storage = torch.full( + (num_reqs, width + 9, 512), -123.0, dtype=torch.bfloat16, device=device + ) + args = (cache, seq_lens, gather_lens, block_table) + dequantize_and_gather_k_cache( + storage[:, :width], *args, block_size=block_size, offset=offset + ) + expected = torch.full_like(storage, -123.0) + _per_token_nvfp4_gather(expected[:, :width], *args, block_size, offset) + assert torch.equal(storage.view(torch.int16), expected.view(torch.int16)) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_v41_nvfp4_gather_decodes_every_code_and_scale(): + """Every (packed byte, e4m3 scale byte) pair: token t holds the bytes + 0..255, and scale byte t for all of its 32 groups.""" + from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache + + device, block_size = "cuda", 64 + codes = torch.arange(256, device=device).to(torch.uint8) + records = torch.empty(256, 288, dtype=torch.uint8, device=device) + records[:, :256] = codes + records[:, 256:] = codes[:, None] + args = ( + records.view(256 // block_size, block_size, 288), + torch.tensor([256], dtype=torch.int32, device=device), + None, + torch.arange(4, dtype=torch.int32, device=device)[None], + block_size, + 0, + ) + out = torch.empty(1, 256, 512, dtype=torch.bfloat16, device=device) + expected = torch.empty_like(out) + dequantize_and_gather_k_cache(out, *args) + _per_token_nvfp4_gather(expected, *args) + assert torch.equal(out.view(torch.int16), expected.view(torch.int16)) + + @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="The V4.1 records are SM100-only", diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 1ab4d8def091..23aa3d3fe595 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -602,8 +602,12 @@ def _dequantize_and_gather_k_nvfp4_kernel( quant_block: tl.constexpr, # 16 cache_block_size: tl.constexpr, block_stride: tl.constexpr, + TILE: tl.constexpr, # tokens per iteration ): - """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace.""" + """Gather and dequantize V4.1 NVFP4 rows into a bf16 workspace, TILE tokens + per iteration. A nibble placed as fp16 bits (sign at 15, exponent and + mantissa at 11:9) reads as e2m1 * 2^-14; times scale * 2^14 the fp32 product + is exact, so the bf16 result matches a direct decode.""" batch_idx = tl.program_id(0) worker_id = tl.program_id(1) num_workers = tl.num_programs(1) @@ -616,35 +620,52 @@ def _dequantize_and_gather_k_nvfp4_kernel( start_pos = seq_len - gather_len packed_bytes: tl.constexpr = head_dim // 2 + group_bytes: tl.constexpr = quant_block // 2 + # Byte b of a record packs elements 2b (low nibble) and 2b + 1, which share + # scale b // group_bytes: index the bytes as [scale, byte within group]. + byte = ( + tl.arange(0, scale_dim)[:, None] * group_bytes + + tl.arange(0, group_bytes)[None, :] + ) + t = tl.arange(0, TILE) d = tl.arange(0, head_dim) - for i in range(worker_id, gather_len, num_workers): + block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq + out_row_ptr = out_ptr + batch_idx.to(tl.int64) * out_stride0 + for tile_start in range(worker_id * TILE, gather_len, num_workers * TILE): + i = tile_start + t + valid = i < gather_len pos = start_pos + i - block_in_seq = pos // cache_block_size - pos_in_block = pos % cache_block_size - physical_block_idx = tl.load( - block_table_ptr + batch_idx * max_blocks_per_seq + block_in_seq + block_table_row_ptr + pos // cache_block_size, mask=valid, other=0 ) page = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride - token_ptr = page + pos_in_block * (packed_bytes + scale_dim) - - packed = tl.load(token_ptr + tl.arange(0, packed_bytes)) - # Even element in the low nibble, odd in the high nibble. - codes = tl.interleave((packed & 0xF).to(tl.int32), (packed >> 4).to(tl.int32)) - # e2m1 magnitudes: 0, 0.5, 1, 1.5, 2, 3, 4, 6. - mag_code = codes & 7 - e = (mag_code >> 1).to(tl.float32) - m = (mag_code & 1).to(tl.float32) - mag = tl.where(mag_code < 2, m * 0.5, (1.0 + m * 0.5) * tl.exp2(e - 1.0)) - vals = tl.where(codes >= 8, -mag, mag) - - sf = tl.load(token_ptr + packed_bytes + tl.arange(0, scale_dim)) - scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) - tiles = tl.reshape(vals, (scale_dim, quant_block)) - 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)) + token_ptr = page + (pos % cache_block_size) * (packed_bytes + scale_dim) + + packed = tl.load( + token_ptr[:, None, None] + byte[None, :, :], + mask=valid[:, None, None], + other=0, + ).to(tl.int32) + sf = tl.load( + token_ptr[:, None] + packed_bytes + tl.arange(0, scale_dim)[None, :], + mask=valid[:, None], + other=0, + ) + scale = sf.to(tl.float8e4nv, bitcast=True).to(tl.float32) * 16384.0 + scale = scale[:, :, None] + + lo = packed & 0xF + hi = packed >> 4 + lo = ((lo & 7) << 9) | ((lo & 8) << 12) + hi = ((hi & 7) << 9) | ((hi & 8) << 12) + lo = lo.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32) * scale + hi = hi.to(tl.int16).to(tl.float16, bitcast=True).to(tl.float32) * scale + dequant = tl.reshape(tl.join(lo, hi), (TILE, head_dim)) + + out_ptrs = ( + out_row_ptr + (offset + i).to(tl.int64)[:, None] * out_stride1 + d[None, :] + ) + tl.store(out_ptrs, dequant.to(tl.bfloat16), mask=valid[:, None]) # The per-token loop is a chain of dependent loads, so the grid has to hold @@ -653,6 +674,7 @@ def _dequantize_and_gather_k_nvfp4_kernel( _GATHER_STEPS_PER_WORKER = 8 _GATHER_MIN_WORKERS = 128 _GATHER_MAX_WORKERS = 8192 +_NVFP4_GATHER_TILE = 8 def gather_num_workers(max_gather_len: int | None) -> int: @@ -682,9 +704,11 @@ def dequantize_and_gather_k_cache_triton( num_reqs = seq_lens.shape[0] NUM_WORKERS = 128 if k_cache.shape[-1] == V41_NVFP4_BYTES_PER_TOKEN: - # Keep short gathers unchanged; cap total workers for multi-request chunks. + # One tile per worker up to the cap: the free width of ``out`` bounds + # every request's gather, and the kernel loops over any remainder. num_workers = min( - max(128, 2048 // max(1, num_reqs)), max(128, out.shape[1] - offset) + cdiv(max(1, out.shape[1] - offset), _NVFP4_GATHER_TILE), + _GATHER_MAX_WORKERS, ) _dequantize_and_gather_k_nvfp4_kernel[(num_reqs, num_workers)]( out, @@ -701,6 +725,8 @@ def dequantize_and_gather_k_cache_triton( quant_block=V41_NVFP4_QUANT_BLOCK, cache_block_size=block_size, block_stride=k_cache.stride(0), + TILE=_NVFP4_GATHER_TILE, + num_warps=2, ) return From 9fc019a749cbf3235cc06aa5357b17fd37cb648a Mon Sep 17 00:00:00 2001 From: Paolo Perrone <58889145+paolo-perrone@users.noreply.github.com> Date: Wed, 28 Oct 2026 01:31:11 -0600 Subject: [PATCH 048/406] [Frontend] Support /v1/completions in run-batch (#64097) Signed-off-by: paoloap <58889145+paoloap-py@users.noreply.github.com> Co-authored-by: paoloap <58889145+paoloap-py@users.noreply.github.com> Co-authored-by: Cyrus Leung --- examples/features/openai_batch/README.md | 6 +- tests/entrypoints/launchers/test_run_batch.py | 74 +++++++++++++++++++ vllm/entrypoints/launchers/run_batch.py | 39 +++++++--- 3 files changed, 107 insertions(+), 12 deletions(-) diff --git a/examples/features/openai_batch/README.md b/examples/features/openai_batch/README.md index c1bb32937688..fff4acc2e601 100644 --- a/examples/features/openai_batch/README.md +++ b/examples/features/openai_batch/README.md @@ -13,9 +13,9 @@ The OpenAI batch file format consists of a series of json objects on new lines. Each line represents a separate request. See the [OpenAI package reference](https://platform.openai.com/docs/api-reference/batch/requestInput) for more details. ```{note} -We currently support `/v1/chat/completions`, `/v1/embeddings`, `/score`, `/rerank`, -`/v1/audio/transcriptions`, and `/v1/audio/translations` endpoints (`/v1/completions` -coming soon). Score and rerank match any URL ending in `/score` or `/rerank`, such as +We currently support `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`, +`/score`, `/rerank`, `/v1/audio/transcriptions`, and `/v1/audio/translations` +endpoints. Score and rerank match any URL ending in `/score` or `/rerank`, such as `/v1/score`. ``` diff --git a/tests/entrypoints/launchers/test_run_batch.py b/tests/entrypoints/launchers/test_run_batch.py index 04e6e71e70fc..6595e5d0b153 100644 --- a/tests/entrypoints/launchers/test_run_batch.py +++ b/tests/entrypoints/launchers/test_run_batch.py @@ -36,6 +36,8 @@ validate_batch, validate_run_batch_args, ) +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse +from vllm.entrypoints.openai.completion.protocol import CompletionResponse from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.exceptions import VLLMValidationError from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -136,6 +138,32 @@ ] ) +INPUT_COMPLETIONS_BATCH = "\n".join( + json.dumps(req) + for req in [ + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": CHAT_MODEL_NAME, + "prompt": "Hello world!", + "max_tokens": 10, + }, + }, + { + "custom_id": "request-2", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": CHAT_MODEL_NAME, + "prompt": ["The capital of France is", "One plus one is"], + "max_tokens": 10, + }, + }, + ] +) + INVALID_INPUT_BATCH = "\n".join( json.dumps(req) for req in [ @@ -525,6 +553,52 @@ def test_completions(): BatchRequestOutput.model_validate_json(line) +def test_completions_endpoint(): + with ( + tempfile.NamedTemporaryFile("w") as input_file, + tempfile.NamedTemporaryFile("r") as output_file, + ): + # A chat line rides along: /v1/chat/completions and /v1/completions + # share their last URL segment, so each must reach its own handler. + chat_line = json.dumps( + json.loads(INPUT_BATCH.strip().split("\n")[0]) | {"custom_id": "chat-1"} + ) + input_file.write(INPUT_COMPLETIONS_BATCH + "\n" + chat_line) + input_file.flush() + proc = subprocess.Popen( + [ + "vllm", + "run-batch", + "-i", + input_file.name, + "-o", + output_file.name, + "--model", + CHAT_MODEL_NAME, + ], + ) + proc.communicate() + proc.wait() + assert proc.returncode == 0, f"{proc=}" + + contents = output_file.read() + outputs = { + output.custom_id: output + for output in map( + BatchRequestOutput.model_validate_json, contents.strip().split("\n") + ) + } + assert len(outputs) == 3 + for output in outputs.values(): + assert output.error is None + assert isinstance(outputs["request-1"].response.body, CompletionResponse) + # A list prompt returns one choice per prompt. + completion = outputs["request-2"].response.body + assert isinstance(completion, CompletionResponse) + assert len(completion.choices) == 2 + assert isinstance(outputs["chat-1"].response.body, ChatCompletionResponse) + + def test_completions_invalid_input(): """Ensure that we fail when the input doesn't conform to the openai api.""" with ( diff --git a/vllm/entrypoints/launchers/run_batch.py b/vllm/entrypoints/launchers/run_batch.py index 94c543a8ab44..21887661aab8 100644 --- a/vllm/entrypoints/launchers/run_batch.py +++ b/vllm/entrypoints/launchers/run_batch.py @@ -40,6 +40,10 @@ ChatCompletionRequest, ChatCompletionResponse, ) +from vllm.entrypoints.openai.completion.protocol import ( + CompletionRequest, + CompletionResponse, +) from vllm.entrypoints.pooling.embed.protocol import ( EmbeddingRequest, EmbeddingResponse, @@ -174,12 +178,14 @@ class BatchRequestInput(OpenAIBaseModel): # supported. method: str - # The OpenAI API relative URL to be used for the request. Currently - # /v1/chat/completions is supported. + # The OpenAI API relative URL to be used for the request, such as + # /v1/chat/completions or /v1/completions. url: str - # The parameters of the request. - body: BatchRequestInputBody + # The parameters of the request. CompletionRequest is kept out of + # BatchRequestInputBody, which an unknown URL is parsed against: its + # validator raises VLLMValidationError, which pydantic does not catch. + body: BatchRequestInputBody | CompletionRequest @field_validator("body", mode="plain") @classmethod @@ -188,6 +194,8 @@ def check_type_for_url(cls, value: Any, info: ValidationInfo): url: str = info.data["url"] if url == "/v1/chat/completions": return ChatCompletionRequest.model_validate(value) + if url == "/v1/completions": + return CompletionRequest.model_validate(value) if url == "/v1/embeddings": return TypeAdapter(EmbeddingRequest).validate_python(value) if url.endswith("/score"): @@ -203,6 +211,7 @@ def check_type_for_url(cls, value: Any, info: ValidationInfo): AllResponse: TypeAlias = ( ChatCompletionResponse + | CompletionResponse | EmbeddingResponse | ScoreResponse | RerankResponse @@ -797,6 +806,7 @@ async def build_endpoint_registry( # Get serving objects from state (defaulting to None if not set) openai_serving_chat = getattr(state, "openai_serving_chat", None) + openai_serving_completion = getattr(state, "openai_serving_completion", None) openai_serving_transcription = getattr(state, "openai_serving_transcription", None) openai_serving_translation = getattr(state, "openai_serving_translation", None) serving_embedding = getattr(state, "serving_embedding", None) @@ -806,7 +816,7 @@ async def build_endpoint_registry( # Registry of endpoint configurations endpoint_registry: dict[str, EndpointConfig] = { - "completions": { + "chat_completions": { "url": "/v1/chat/completions", "handler_getter": lambda: ( openai_serving_chat.create_chat_completion @@ -815,6 +825,15 @@ async def build_endpoint_registry( ), "wrapper_fn": None, }, + "completions": { + "url": "/v1/completions", + "handler_getter": lambda: ( + openai_serving_completion.create_completion + if openai_serving_completion is not None + else None + ), + "wrapper_fn": None, + }, "embeddings": { "url": "/v1/embeddings", "handler_getter": lambda: ( @@ -941,11 +960,13 @@ async def run_one_request( """Route a single line of the batch to its endpoint handler.""" request = BatchRequestInput.model_validate_json(request_json) - endpoint_key = request.url.split("/")[-1] - + # Match on the whole URL: /v1/chat/completions and /v1/completions share + # their last segment. result = None - if endpoint_key in endpoint_registry: - result = handle_endpoint_request(request, endpoint_registry[endpoint_key]) + for endpoint in endpoint_registry.values(): + result = handle_endpoint_request(request, endpoint) + if result is not None: + break if result is None: result = make_async_error_request_output( From 5118d990fb9d28768b4ec6c7328322f96186f915 Mon Sep 17 00:00:00 2001 From: Yusen Gao Date: Wed, 28 Oct 2026 15:53:31 +0800 Subject: [PATCH 049/406] [Bugfix][Spec Decode] Align adaptive draft prefill graph padding (#62796) Signed-off-by: Yusen Gao --- tests/v1/cudagraph/test_cudagraph_manager.py | 130 ++++++++++++++++++ vllm/benchmarks/forward/speculator.py | 1 + vllm/v1/worker/gpu/dp_utils.py | 23 ++-- vllm/v1/worker/gpu/spec_decode/speculator.py | 1 + .../target_dependent_ar/speculator.py | 3 + 5 files changed, 150 insertions(+), 8 deletions(-) diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index cfe77253e29b..a81586605673 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/tests/v1/cudagraph/test_cudagraph_manager.py @@ -19,6 +19,7 @@ from vllm.distributed.device_communicators import pynccl_allocator from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils +from vllm.v1.worker.gpu import dp_utils from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBuffers @@ -26,6 +27,9 @@ from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import ( cudagraph_utils as spec_cudagraph_utils, ) +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import ( + speculator as spec_module, +) from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.cudagraph_utils import ( SpeculatorCudaGraphManager, ) @@ -454,6 +458,132 @@ def dispatch(max_query_len): assert dispatch(_DECODE_QUERY_LEN + 1).cg_mode == CUDAGraphMode.PIECEWISE +@pytest.mark.parametrize( + "adaptive,num_tokens,has_prefill,decode_eligible,dp_size,peer_query_len", + [ + (False, 16, False, True, 1, None), + (True, 16, False, True, 1, None), + (True, 8, False, True, 1, None), + (True, 16, True, False, 1, None), + pytest.param(True, 16, True, True, 1, None, id="decode-eligible-prefill-tail"), + pytest.param(True, 16, False, True, 2, None, id="dp-peer-prefill"), + pytest.param(True, 8, False, True, 2, 4, id="dp-peer-longer-decode"), + ], +) +def test_draft_prefill_matches_shared_target_metadata_shape( + monkeypatch, + adaptive, + num_tokens, + has_prefill, + decode_eligible, + dp_size, + peer_query_len, +): + """Draft prefill must reuse the target's padded request layout, even after AV.""" + target = _make_spec_decode_manager( + monkeypatch, + decode_query_len=4, + capture_sizes=[1, 2, 4, 8, 16, 32], + varlen_decode=adaptive, + ) + collective_calls = 0 + + def all_reduce(tensor): + nonlocal collective_calls + collective_calls += 1 + assert collective_calls == 1, "Draft prefill must reuse the target's DP sync" + tensor[:, 1] = torch.tensor( + [ + 16, + CUDAGraphMode.PIECEWISE.value + if peer_query_len is None + else CUDAGraphMode.FULL.value, + peer_query_len or 0, + peer_query_len or -1, + 0, + 4, + 0, + ] + ) + + monkeypatch.setattr( + dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_all_reduce_=all_reduce) + ) + max_query_len = num_tokens // 4 + target_desc, target_sync = dispatch_cg_and_sync_dp( + target, + 4, + num_tokens, + max_query_len if decode_eligible else None, + dp_size=dp_size, + dp_rank=0, + max_query_len=None if has_prefill else max_query_len, + ) + hidden_states = torch.zeros(target_desc.num_tokens, 1) + speculator = SimpleNamespace( + vllm_config=target.vllm_config, + device=torch.device("cpu"), + num_speculative_steps=3, + enable_adaptive_verification=adaptive, + max_model_len=4096, + max_num_reqs=8, + dp_size=dp_size, + dp_rank=0, + pcp_manager=None, + hidden_states=torch.empty_like(hidden_states), + last_token_indices=None, + current_draft_step=None, + input_buffers=None, + _copy_request_inputs=lambda *a, **kw: None, + ) + spec_module.TargetDependentARSpeculator.init_cudagraph_manager( + speculator, CUDAGraphMode.FULL_AND_PIECEWISE + ) + speculator.prefill_cudagraph_manager._graphs_captured = True + batch = SimpleNamespace( + num_reqs=4, + num_tokens=num_tokens, + num_tokens_after_padding=target_desc.num_tokens, + num_scheduled_tokens=torch.full((4,), max_query_len).numpy(), + seq_lens_cpu_upper_bound=torch.full((8,), 128), + idx_mapping=None, + has_prefill=has_prefill, + decode_graph_eligible=decode_eligible, + ) + monkeypatch.setattr(spec_module, "prepare_prefill_inputs", lambda *a: None) + dispatch = spec_module.dispatch_cg_and_sync_dp + + class Dispatched(Exception): + pass + + def stop_after_dispatch(*args, **kwargs): + raise Dispatched(dispatch(*args, **kwargs)[0]) + + monkeypatch.setattr(spec_module, "dispatch_cg_and_sync_dp", stop_after_dispatch) + with pytest.raises(Dispatched) as result: + spec_module.TargetDependentARSpeculator.propose( + speculator, + batch, + {}, + {}, + hidden_states, + None, + num_sampled=None, + num_rejected=None, + last_sampled=None, + next_prefill_tokens=None, + temperature=None, + seeds=None, + dp_sync_state=target_sync, + ) + + assert result.value.args[0] == target_desc + if dp_size > 1: + assert collective_calls == 1 + assert target_sync is not None + assert target_sync.max_query_len == peer_query_len + + def test_mixed_batch_at_decode_only_token_count_still_gets_a_graph(monkeypatch): """A mixed batch must not fall to eager where only decode graphs are staged. diff --git a/vllm/benchmarks/forward/speculator.py b/vllm/benchmarks/forward/speculator.py index bbecdce8579d..0e35cc58d2fd 100644 --- a/vllm/benchmarks/forward/speculator.py +++ b/vllm/benchmarks/forward/speculator.py @@ -45,6 +45,7 @@ def prepare(self, input_batch, context): dtype=torch.int32, ), uniform_token_count=None, + max_query_len=None, eager=True, num_reqs=self.batch.num_reqs, ) diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index fb9f4c02d3e4..ab452683ab82 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -32,6 +32,8 @@ class DPSyncState: # Agreed uniform decode length. None means the ranks disagreed, which is an # answer, not "unknown". uniform_token_count: int | None + # Agreed query-length bound. None disables varlen decode graphs. + max_query_len: int | None # Whether the ranks agreed to run eager. A dispatch reusing this must run # eager too; no shape was agreed, so picking a graph per rank would diverge. eager: bool @@ -94,6 +96,11 @@ def sync_cudagraph_and_dp_padding( ): synced_uniform_token_count = None + # Varlen decode graphs need the same query-length bound on every rank. + synced_max_query_len: int | None = None + if bool(torch.all(max_query_lens_across_dp != -1)): + synced_max_query_len = int(max_query_lens_across_dp.max()) + if torch.all(num_tokens_across_dp == 0).item(): synced_desc = BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=0, num_reqs=0 @@ -148,6 +155,7 @@ def sync_cudagraph_and_dp_padding( num_tokens_across_dp, ubatch_num_tokens ), uniform_token_count=synced_uniform_token_count, + max_query_len=synced_max_query_len, eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE, num_reqs=int(num_reqs_across_dp.max()), skip_drafts=skip_drafts, @@ -167,6 +175,7 @@ def sync_cudagraph_and_dp_padding( DPSyncState( num_tokens_across_dp=num_tokens_across_dp, uniform_token_count=synced_uniform_token_count, + max_query_len=synced_max_query_len, eager=True, num_reqs=int(num_reqs_across_dp.max()), skip_drafts=skip_drafts, @@ -179,12 +188,6 @@ def sync_cudagraph_and_dp_padding( ) synced_num_tokens = int(num_tokens_across_dp.max()) - # Varlen decode graphs are selected by the query-length bound, so ranks must agree - # on it or they pad to different token counts below. - synced_max_query_len: int | None = None - if bool(torch.all(max_query_lens_across_dp != -1)): - synced_max_query_len = int(max_query_lens_across_dp.max()) - # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs # so we don't perform request padding for PIECEWISE graphs. # num_active_loras is per-rank and doesn't need cross-rank agreement. @@ -202,6 +205,7 @@ def sync_cudagraph_and_dp_padding( return synced_desc, DPSyncState( num_tokens_across_dp=num_tokens_across_dp, uniform_token_count=synced_uniform_token_count, + max_query_len=synced_max_query_len, eager=False, num_reqs=int(num_reqs_across_dp.max()), skip_drafts=skip_drafts, @@ -243,7 +247,8 @@ def dispatch_cg_and_sync_dp( dp_rank: This rank's index in the DP group. max_query_len: Upper bound on per-request query length, for selecting varlen decode graphs. None, as for a batch with a prefill, keeps - the batch out of graphs that constrain it. + the batch out of graphs that constrain it. Reusing a sync replaces + this with `dp_sync_state.max_query_len`, including when it is None. need_eager: Force `CUDAGraphMode.NONE` instead of dispatching. num_active_loras: Active LoRA count for this rank. Does not need cross-rank agreement; it never changes a bucket's token count. @@ -286,7 +291,9 @@ def dispatch_cg_and_sync_dp( num_tokens, dp_sync_state.uniform_token_count if dp_sync_state else uniform_token_count, num_active_loras=num_active_loras, - max_query_len=max_query_len, + max_query_len=( + dp_sync_state.max_query_len if dp_sync_state else max_query_len + ), ) if dp_size == 1: diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 517103949fee..87fbc3704f10 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -612,6 +612,7 @@ def _build_uniform_batch_dp_sync( target_dp_sync_state.num_tokens_across_dp, num_batch_tokens ), uniform_token_count=num_query_per_req, + max_query_len=num_query_per_req, eager=False, ), num_batch_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 7b70fd0ee0bd..95aa5959972e 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -93,6 +93,8 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: self.device, cudagraph_mode, self.num_speculative_steps + 1, + # Draft prefill shares the target's padded attention metadata. + varlen_decode=self.enable_adaptive_verification, ) # PIECEWISE cudagraphs are not supported for draft decodes. @@ -266,6 +268,7 @@ def propose( uniform_token_count, dp_size=self.dp_size, dp_rank=self.dp_rank, + max_query_len=None if input_batch.has_prefill else int(max_query_len), need_eager=is_profile, dp_sync_state=dp_sync_state, ) From 1cb31be7028457c7eae7b6fabca1941d5d3c4730 Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Wed, 28 Oct 2026 16:27:07 +0800 Subject: [PATCH 050/406] [ROCm][MiniMax-M3.1] Fuse MXFP8 activation quant into the Gemma input_layernorm for qkv_proj (#63637) Signed-off-by: RuibinCheung Co-authored-by: Cursor Co-authored-by: Shanshan Shen <467638484@qq.com> Co-authored-by: vllm-agent --- .../fusion/test_quant_activation_contract.py | 4 + tests/kernels/test_minimax_m3_amd_ops.py | 157 +++++++++++++++++- .../kernels/linear/mxfp8/rocm_native.py | 78 +++++++-- .../quark/schemes/quark_ocp_mx.py | 5 + vllm/models/minimax_m3/amd/model.py | 51 +++++- vllm/models/minimax_m3/amd/ops/__init__.py | 2 + .../minimax_m3/amd/ops/gemma_rmsnorm.py | 102 +++++++++++- 7 files changed, 371 insertions(+), 28 deletions(-) diff --git a/tests/fusion/test_quant_activation_contract.py b/tests/fusion/test_quant_activation_contract.py index da507e5caf50..0931fc819bfd 100644 --- a/tests/fusion/test_quant_activation_contract.py +++ b/tests/fusion/test_quant_activation_contract.py @@ -16,6 +16,9 @@ FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, +) from vllm.model_executor.kernels.linear.nvfp4.base import ( NvFp4LinearKernel, NvFp4LinearLayerConfig, @@ -68,6 +71,7 @@ AiterPerTokenFp8ScaledMMLinearKernel, FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, + RocmDotScaledMxfp8LinearKernel, } diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 868c2fcac7be..4e277470a2d8 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -6,9 +6,13 @@ tests assert the two agree within tolerance: * Gemma RMSNorm (plain + fused-add-residual) -> fp32 PyTorch normalize + * Fused-add Gemma RMSNorm + MXFP8 epilogue -> norm, then mxfp8_e4m3_quantize + (bitwise) * SwiGLU-OAI (split layout) -> fp32 PyTorch elementwise * Fused MXFP8 activation quant (Triton) -> _mxfp8_e4m3_quantize_torch * Native MXFP8 linear (dot_scaled) -> dequant-to-bf16 @ matmul + * Native MXFP8 linear on the fused norm output -> norm, then the linear + (bitwise) * Native MXFP8 MoE (dot_scaled grouped GEMM) -> dequant-to-bf16 MoE math * Strided K/V sparse-PA cache insert -> contiguous-copy insert @@ -39,10 +43,15 @@ _mxfp8_e4m3_quantize_torch, _mxfp8_e4m3_quantize_triton, dequant_mxfp8_to_bf16, + mxfp8_e4m3_quantize, +) +from vllm.models.minimax_m3.amd.model import ( # noqa: E402 + _can_fuse_input_norm_mxfp8, + _kv_insert_operand, ) -from vllm.models.minimax_m3.amd.model import _kv_insert_operand # noqa: E402 from vllm.models.minimax_m3.amd.ops import ( # noqa: E402 gemma_fused_add_rmsnorm, + gemma_fused_add_rmsnorm_mxfp8, gemma_rmsnorm, swiglu_oai_quantize_mxfp8, swiglu_oai_split, @@ -144,6 +153,61 @@ def test_num_warps_monotonic(): assert _num_warps(128) <= _num_warps(2048) <= _num_warps(8192) +def _mxfp8_edge_case_inputs(n: int, dtype: torch.dtype): + """(x, residual, weight) whose normed output covers the MXFP8 block cases. + + Through the weight, every row gets a huge block (scale above 1.0), an + all-zero block, and a tiny block (``1 + w == 2**-8``). Rows: 32 ordinary + rows with a residual, then an all-zero row, a row of tiny values, a row + whose first block dwarfs the rest, and a row of partly-zero blocks. + """ + torch.manual_seed(0) + w = torch.randn(n, device=DEVICE, dtype=dtype) * 0.1 + w[:32] = 300.0 + w[32:64] = -1.0 + 2**-8 + w[64:96] = -1.0 + + ordinary = torch.randn(32, n, device=DEVICE, dtype=dtype) + zero = torch.zeros(1, n, device=DEVICE, dtype=dtype) + tiny = torch.full((1, n), torch.finfo(dtype).tiny * 16, device=DEVICE, dtype=dtype) + dominated = torch.randn(1, n, device=DEVICE, dtype=dtype) * 1e-2 + dominated[:, :32] = 1e4 + partly_zero = torch.randn(1, n, device=DEVICE, dtype=dtype) + partly_zero[:, ::2] = 0.0 + partly_zero[:, 128:192] = 0.0 + x = torch.cat([ordinary, zero, tiny, dominated, partly_zero]) + + res = torch.zeros_like(x) + res[:32] = torch.randn(32, n, device=DEVICE, dtype=dtype) + return x, res, w + + +@pytest.mark.parametrize("n", [2048, 3072, 6144]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_gemma_fused_add_rmsnorm_mxfp8_bitwise_eq_unfused(n, dtype): + """The MXFP8 epilogue writes the same values, E8M0 scales and residual as + ``gemma_fused_add_rmsnorm`` followed by ``mxfp8_e4m3_quantize``.""" + x, res, w = _mxfp8_edge_case_inputs(n, dtype) + m = x.shape[0] + # A leading batch dim is flattened into the quantized rows. + x, res = x.unsqueeze(0), res.unsqueeze(0) + + got_q, got_scale, got_res = gemma_fused_add_rmsnorm_mxfp8(x, res, w, EPS) + normed, ref_res = gemma_fused_add_rmsnorm(x, res, w, EPS) + ref_q, ref_scale = mxfp8_e4m3_quantize(normed.reshape(m, n)) + + assert got_q.shape == (m, n) and got_q.dtype == torch.float8_e4m3fn + assert got_scale.shape == (m, n // 32) and got_scale.dtype == torch.uint8 + assert got_res.shape == x.shape + # Both scale clamps are reached: E8M0 0 (all-zero/tiny) and above 1.0. + assert (ref_scale == 0).any() and (ref_scale > 127).any() + assert not ref_q.float().isnan().any() + assert torch.equal(got_q.view(torch.uint8), ref_q.view(torch.uint8)) + assert torch.equal(got_scale, ref_scale) + assert torch.equal(got_res, ref_res) + + # --------------------------------------------------------------------------- # # SwiGLU-OAI (split layout) # --------------------------------------------------------------------------- # @@ -511,6 +575,97 @@ def _spy(*a, **k): assert _relerr(out.float(), out_ref.float()) < tol +# --------------------------------------------------------------------------- # +# input_layernorm -> qkv_proj with the MXFP8 activation quant fused into the norm +# --------------------------------------------------------------------------- # +def _quark_mxfp8_linear(n: int, k: int, block32: bool = False): + """A Quark MXFP8 linear layer (weights loaded, not yet processed).""" + from vllm.model_executor.layers.quantization.quark.schemes.quark_ocp_mx import ( + QuarkOCP_MX, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp8Dynamic, + kMxfp8Static, + ) + from vllm.model_executor.model_loader.weight_utils import default_weight_loader + + 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) + if block32: + w_scale = w_scale[::32].repeat_interleave(32, dim=0) + scheme = QuarkOCP_MX( + weight_quant_key=kMxfp8Static, activation_quant_key=kMxfp8Dynamic + ) + with torch.device(DEVICE): + layer = torch.nn.Module() + scheme.create_weights( + layer=layer, + output_partition_sizes=[n], + input_size_per_partition=k, + params_dtype=torch.bfloat16, + weight_loader=default_weight_loader, + ) + layer.weight.data.copy_(w_fp8) + layer.weight_scale.data.copy_(w_scale) + # As QuarkLinearMethod.create_weights does. + layer.scheme = scheme + return layer, scheme + + +@requires_gfx950 +@pytest.mark.parametrize( + "n,k,block32", + [ + (2560, 6144, False), # MiniMax-M3.1 qkv_proj, TP4 + (256, 2048, True), # 32x32 block scales: the block32 GEMM + (5120, 576, False), # K % 128 != 0: BF16 weight after load + ], +) +@pytest.mark.parametrize("m", [1, 37]) +@torch.inference_mode() +def test_fused_input_norm_mxfp8_qkv_proj(dist_init, n, k, block32, m): + """The linear on the fused norm's MXFP8 output equals the linear on the + unfused norm output, wherever the producer turns the fusion on.""" + layer, scheme = _quark_mxfp8_linear(n, k, block32) + fused = _can_fuse_input_norm_mxfp8(layer) + assert fused == (k % 128 == 0) + scheme.process_weights_after_loading(layer) + + from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp8Dynamic, + ) + + x = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16) + res = torch.randn(m, k, device=DEVICE, dtype=torch.bfloat16) + norm_w = torch.randn(k, device=DEVICE, dtype=torch.bfloat16) * 0.1 + normed, _ = gemma_fused_add_rmsnorm(x, res, norm_w, EPS) + x_q, x_scale, _ = gemma_fused_add_rmsnorm_mxfp8(x, res, norm_w, EPS) + qa = QuantizedActivation(x_q, x_scale, x.dtype, x.shape, kMxfp8Dynamic) + + got = scheme.apply_weights(layer, qa) + ref = scheme.apply_weights(layer, normed) + assert got.shape == ref.shape == (m, n) and got.dtype == ref.dtype + if fused: + assert torch.equal(got, ref) + else: + # A QuantizedActivation reaching a BF16 weight is dequantized back. + assert _relerr(got, ref) < 5e-2 + + +@requires_gfx950 +def test_fused_input_norm_mxfp8_off_under_batch_invariance(dist_init, monkeypatch): + import vllm.envs as envs + + layer, _ = _quark_mxfp8_linear(2560, 6144) + assert _can_fuse_input_norm_mxfp8(layer) + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) + assert not _can_fuse_input_norm_mxfp8(layer) + + # ── 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 e46b61275c66..134567555cd7 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -3,21 +3,31 @@ """Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``. 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. ``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. +activations are MXFP8-quantized per token, or arrive pre-quantized as a +``QuantizedActivation`` with row-major (not swizzled) E8M0 scales. Uses the +CDNA4 hardware microscaling 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 import vllm.envs as envs +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + as_quantized_activation, +) from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, MXFP8_SCALE_DTYPE, dequant_mxfp8_to_bf16, mxfp8_e4m3_quantize, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp8Dynamic, +) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -87,10 +97,20 @@ def _mxfp8_dot_scaled_linear( w: torch.Tensor, # [N, K] fp8 e4m3 w_scale: torch.Tensor, # [N, K//32] uint8 (E8M0) ) -> torch.Tensor: - M, K = x.shape - N = w.shape[0] x_q, x_scale = mxfp8_e4m3_quantize(x) - out = torch.empty((M, N), dtype=x.dtype, device=x.device) + return _mxfp8_dot_scaled_gemm(x_q, x_scale, w, w_scale, x.dtype) + + +def _mxfp8_dot_scaled_gemm( + x_q: torch.Tensor, # [M, K] fp8 e4m3 + x_scale: torch.Tensor, # [M, K//32] uint8 (E8M0), row-major + w: torch.Tensor, # [N, K] fp8 e4m3 + w_scale: torch.Tensor, # [N, K//32] uint8 (E8M0) + out_dtype: torch.dtype, +) -> torch.Tensor: + M, K = x_q.shape + N = w.shape[0] + out = torch.empty((M, N), dtype=out_dtype, device=x_q.device) BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _select_cfg(M, N, K) grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) _mxfp8_linear_kernel[grid]( @@ -261,10 +281,24 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter(layer, "weight", weight.contiguous()) replace_parameter(layer, "weight_scale", weight_scale) + def input_quant_key(self) -> QuantKey: + # Row-major [M, K // 32] E8M0 scales, unlike FlashInfer's swizzled ones + # under the same key: producers must also check the kernel type. + return kMxfp8Dynamic + + @staticmethod + def supports_quantized_input(layer: torch.nn.Module) -> bool: + """Whether ``apply_weights`` reads a ``QuantizedActivation`` for + ``layer`` as is (fp8 weights), rather than dequantizing it back.""" + return layer.weight.element_size() == 1 and ( + layer.weight_scale.shape[0] != layer.weight.shape[0] + or layer.weight.shape[1] % _DOT_SCALED_K_ALIGN == 0 + ) + def apply_weights( self, layer: torch.nn.Module, - x: torch.Tensor, + x: torch.Tensor | QuantizedActivation, bias: torch.Tensor | None = None, ) -> torch.Tensor: if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE: @@ -272,23 +306,39 @@ def apply_weights( f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got " f"{layer.weight_scale.dtype}." ) - out_shape = (*x.shape[:-1], layer.weight.shape[0]) - x2d = x.reshape(-1, x.shape[-1]) + qa = as_quantized_activation(x, self.input_quant_key()) + if qa is not None and not self.supports_quantized_input(layer): + # BF16 weight (K % 128 != 0): serve the dequantized activation. + x = dequant_mxfp8_to_bf16(qa.data, qa.scale).to(qa.orig_dtype) + x, qa = x.reshape(qa.orig_shape), None + if qa is not None: + out_dtype, in_shape = qa.orig_dtype, qa.orig_shape + x_q, x_scale = qa.data.reshape(-1, in_shape[-1]), qa.scale + else: + assert isinstance(x, torch.Tensor) + out_dtype, in_shape = x.dtype, x.shape + x2d = x.reshape(-1, in_shape[-1]) + out_shape = (*in_shape[:-1], layer.weight.shape[0]) if layer.weight_scale.shape[0] != layer.weight.shape[0]: # One scale row per 32 weight rows (see process_weights_after_loading). - x_q, x_scale = mxfp8_e4m3_quantize(x2d) + if qa is None: + x_q, x_scale = mxfp8_e4m3_quantize(x2d) out = rocm_mxfp8_block32_gemm( - x_q, x_scale, layer.weight, layer.weight_scale, x.dtype + x_q, x_scale, layer.weight, layer.weight_scale, out_dtype + ) + elif qa is not None: + out = _mxfp8_dot_scaled_gemm( + x_q, x_scale, layer.weight, layer.weight_scale, out_dtype ) elif layer.weight.element_size() >= 2: - out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype)) + out = torch.nn.functional.linear(x2d, layer.weight.to(out_dtype)) elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0: out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale) else: # 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 = torch.nn.functional.linear(x2d, w_bf16).to(out_dtype) out = out.reshape(out_shape) if bias is not None: out = out + bias diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index ee089be321bf..2caf3b7ac8f7 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -15,6 +15,9 @@ init_mxfp6_linear_kernel, init_mxfp8_linear_kernel, ) +from vllm.model_executor.layers.fusion.quant_activation import ( + expose_input_quant_key, +) from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_SCALE_DTYPE, MXFP8_VALUE_DTYPE, @@ -246,6 +249,8 @@ def create_weights( self.ocp_mx_linear = init_mxfp8_linear_kernel( weight_shape=layer.weight.shape ) + if current_platform.is_rocm(): + expose_input_quant_key(layer, self.ocp_mx_linear) else: raise NotImplementedError( "No OCP MX linear kernel for weight_quant_key=" diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 3222280b65ab..7b7459770d87 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -36,6 +36,9 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger +from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, +) from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention.attention import set_default_quant_scales from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -50,6 +53,10 @@ from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, ) +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + get_input_quant_key, +) from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, MinimaxM3QKVParallelLinearWithIndexer, @@ -61,6 +68,7 @@ from vllm.model_executor.layers.quantization.utils.config_utils import ( is_shared_expert_quant_fse_compatible, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -97,6 +105,7 @@ ) from vllm.models.minimax_m3.amd.ops import ( gemma_fused_add_rmsnorm, + gemma_fused_add_rmsnorm_mxfp8, gemma_rmsnorm, swiglu_oai_split, ) @@ -569,7 +578,7 @@ def __init__( def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) # Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place (dense @@ -980,7 +989,7 @@ def _insert_aiter_sparse_pa_kv( def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, ) -> torch.Tensor: # Single fused projection emitting [q | k | v | index_q | index_k]. qkv, _ = self.qkv_proj(hidden_states) @@ -1283,6 +1292,24 @@ def _run_attention( return self.impl.forward(self, query, self.kv_cache, output) +def _can_fuse_input_norm_mxfp8(qkv_proj: nn.Module) -> bool: + """Whether ``input_layernorm`` can hand ``qkv_proj`` an MXFP8 activation. + + The quant key does not encode the scale layout (FlashInfer wants swizzled + scales under the same key), so this also requires the ROCm ``dot_scaled`` + kernel, which reads the row-major scales the fused norm writes. + """ + if envs.VLLM_BATCH_INVARIANT or get_input_quant_key(qkv_proj) != kMxfp8Dynamic: + return False + # Quark keeps the kernel on the layer's scheme, ModelOpt on the method. + kernel = getattr(getattr(qkv_proj, "scheme", None), "ocp_mx_linear", None) + if kernel is None: + kernel = getattr(getattr(qkv_proj, "quant_method", None), "kernel", None) + return isinstance( + kernel, RocmDotScaledMxfp8LinearKernel + ) and kernel.supports_quantized_input(qkv_proj) + + class MiniMaxM3DecoderLayer(nn.Module): def __init__( self, @@ -1356,6 +1383,9 @@ def __init__( self.post_attention_layernorm = MiniMAXGemmaRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) + self.fuse_input_norm_mxfp8 = _can_fuse_input_norm_mxfp8(self.self_attn.qkv_proj) + if self.fuse_input_norm_mxfp8: + logger.info_once("Fusing input_layernorm with qkv_proj's MXFP8 input quant") def forward( self, @@ -1364,14 +1394,25 @@ def forward( residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: # Self Attention + attn_input: torch.Tensor | QuantizedActivation if residual is None: residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) + attn_input = self.input_layernorm(hidden_states) + elif self.fuse_input_norm_mxfp8: + x_q, x_scale, residual = gemma_fused_add_rmsnorm_mxfp8( + hidden_states, + residual, + self.input_layernorm.weight, + self.input_layernorm.variance_epsilon, + ) + attn_input = QuantizedActivation( + x_q, x_scale, hidden_states.dtype, hidden_states.shape, kMxfp8Dynamic + ) else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) + attn_input, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn( positions=positions, - hidden_states=hidden_states, + hidden_states=attn_input, ) hidden_states, residual = fused_allreduce_gemma_rms_norm( diff --git a/vllm/models/minimax_m3/amd/ops/__init__.py b/vllm/models/minimax_m3/amd/ops/__init__.py index 22d96f9de979..a9261bdfd236 100644 --- a/vllm/models/minimax_m3/amd/ops/__init__.py +++ b/vllm/models/minimax_m3/amd/ops/__init__.py @@ -9,6 +9,7 @@ from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import ( gemma_fused_add_rmsnorm, + gemma_fused_add_rmsnorm_mxfp8, gemma_rmsnorm, ) from vllm.models.minimax_m3.amd.ops.swiglu_oai import ( @@ -19,6 +20,7 @@ __all__ = [ "gemma_rmsnorm", "gemma_fused_add_rmsnorm", + "gemma_fused_add_rmsnorm_mxfp8", "swiglu_oai_split", "swiglu_oai_quantize_mxfp8", ] diff --git a/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py b/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py index cb74877f6824..208f3432d0fa 100644 --- a/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py +++ b/vllm/models/minimax_m3/amd/ops/gemma_rmsnorm.py @@ -8,9 +8,13 @@ pow, mean, rsqrt, two muls, cast) — each a separate kernel launch materializing fp32 intermediates. These kernels collapse that into a single pass per row. -Two entry points: +Entry points: * ``gemma_rmsnorm(x, w, eps)`` -> normalized tensor * ``gemma_fused_add_rmsnorm(x, res, w, eps)`` -> (normalized, x + res) + * ``gemma_fused_add_rmsnorm_mxfp8(x, res, w, eps)`` + -> (MXFP8 values, E8M0 scales, x + res): the fused add + norm with the + per-32-block MXFP8 activation quant as an epilogue, so a following + MXFP8 linear need not quantize its input again. Both normalize over the last dim and broadcast ``weight`` (shape [N]) over it, so they serve both the full-hidden norms (input/post-attn/final) and the @@ -21,6 +25,11 @@ import torch +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + MXFP8_VALUE_DTYPE, +) from vllm.triton_utils import tl, triton @@ -59,6 +68,7 @@ def _gemma_fused_add_rmsnorm_kernel( w_ptr, out_ptr, res_out_ptr, + scale_ptr, n_cols, stride_xrow, stride_xcol, @@ -66,6 +76,9 @@ def _gemma_fused_add_rmsnorm_kernel( stride_rcol, eps, BLOCK_N: tl.constexpr, + QUANT_MXFP8: tl.constexpr, + FP8_MAX: tl.constexpr, + TINY: tl.constexpr, ): row = tl.program_id(0) cols = tl.arange(0, BLOCK_N) @@ -87,11 +100,35 @@ def _gemma_fused_add_rmsnorm_kernel( rstd = 1.0 / tl.sqrt(var + eps) w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32) out = s * rstd * (1.0 + w) - tl.store( - out_ptr + row * n_cols + cols, - out.to(out_ptr.dtype.element_ty), - mask=mask, - ) + if QUANT_MXFP8: + # Same math as mxfp8_e4m3_quantize on the input-dtype normed row, so + # the result is bit-identical to norm -> bf16 -> quant. + y = tl.reshape( + out.to(x_ptr.dtype.element_ty).to(tl.float32), (BLOCK_N // 32, 32) + ) + amax = tl.maximum(tl.max(tl.abs(y), axis=1), TINY) + sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0 + sb = tl.minimum(tl.maximum(sb, 0.0), 254.0) + # exp2(sb - 127) is subnormal at sb == 0 and flushes to zero on CDNA, so + # scale by the (exact, normal) reciprocal instead of dividing. + q = y * tl.exp2(127.0 - sb)[:, None] + tl.store( + out_ptr + row * n_cols + cols, + tl.reshape(q, (BLOCK_N,)).to(out_ptr.dtype.element_ty), + mask=mask, + ) + blocks = tl.arange(0, BLOCK_N // 32) + tl.store( + scale_ptr + row * (n_cols // 32) + blocks, + sb.to(scale_ptr.dtype.element_ty), + mask=blocks < n_cols // 32, + ) + else: + tl.store( + out_ptr + row * n_cols + cols, + out.to(out_ptr.dtype.element_ty), + mask=mask, + ) def _num_warps(block_n: int) -> int: @@ -129,12 +166,57 @@ def gemma_fused_add_rmsnorm( weight: torch.Tensor, eps: float, ) -> tuple[torch.Tensor, torch.Tensor]: + out, _, res_out = _gemma_fused_add_rmsnorm(x, residual, weight, eps, False) + return out.reshape(x.shape), res_out + + +def gemma_fused_add_rmsnorm_mxfp8( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``gemma_fused_add_rmsnorm`` with the normed output quantized to MXFP8. + + Bit-identical to ``mxfp8_e4m3_quantize(gemma_fused_add_rmsnorm(...)[0])``. + + Returns: + ``(x_q, x_scale, residual_out)``: ``x_q`` float8_e4m3fn ``[M, N]`` and + ``x_scale`` uint8 E8M0 ``[M, N // 32]`` (row-major, not swizzled), with + the leading dims of ``x`` flattened into ``M``; ``residual_out`` has + the shape of ``x``. + + """ + n = x.shape[-1] + assert n % MXFP8_BLOCK_SIZE == 0, ( + f"MXFP8 needs the hidden size to be a multiple of {MXFP8_BLOCK_SIZE}, got {n}" + ) + x_q, x_scale, res_out = _gemma_fused_add_rmsnorm(x, residual, weight, eps, True) + assert x_scale is not None + return x_q, x_scale, res_out + + +def _gemma_fused_add_rmsnorm( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + quant_mxfp8: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: orig_shape = x.shape n = orig_shape[-1] x2 = x.reshape(-1, n) r2 = residual.reshape(-1, n) m = x2.shape[0] - out = torch.empty((m, n), dtype=x.dtype, device=x.device) + out_dtype = MXFP8_VALUE_DTYPE if quant_mxfp8 else x.dtype + out = torch.empty((m, n), dtype=out_dtype, device=x.device) + scale = ( + torch.empty( + (m, n // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device + ) + if quant_mxfp8 + else None + ) res_out = torch.empty((m, n), dtype=x.dtype, device=x.device) block_n = triton.next_power_of_2(n) _gemma_fused_add_rmsnorm_kernel[(m,)]( @@ -143,6 +225,7 @@ def gemma_fused_add_rmsnorm( weight, out, res_out, + scale, n, x2.stride(0), x2.stride(1), @@ -150,6 +233,9 @@ def gemma_fused_add_rmsnorm( r2.stride(1), eps, BLOCK_N=block_n, + QUANT_MXFP8=quant_mxfp8, + FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max), + TINY=float(torch.finfo(torch.float32).tiny), num_warps=_num_warps(block_n), ) - return out.reshape(orig_shape), res_out.reshape(orig_shape) + return out, scale, res_out.reshape(orig_shape) From 066d9799347c88c14e91baa9d849e9595ff8a3d3 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Wed, 28 Oct 2026 01:50:32 -0700 Subject: [PATCH 051/406] [KV Connector][NIXL] Select push mode for NixlConnector via transfer_mode (#64230) Signed-off-by: Juntian Liu Signed-off-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 Co-authored-by: Yifan Qiao Co-authored-by: Codex --- docs/design/nixl_kv_push_connector.md | 11 +++++++++++ tests/v1/kv_connector/unit/test_config.py | 18 ++++++++++++++++++ .../kv_transfer/kv_connector/factory.py | 5 +++++ 3 files changed, 34 insertions(+) diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md index 621b0f7c9dd7..d29a21629e96 100644 --- a/docs/design/nixl_kv_push_connector.md +++ b/docs/design/nixl_kv_push_connector.md @@ -6,6 +6,17 @@ prefill completes. `NixlPushConnector` adds a **push-based** alternative in which P writes the KV blocks directly into D's pre-allocated memory via `NIXL WRITE`. +Enable push with `"kv_connector": "NixlPushConnector"`. Alternatively, use +`"kv_connector": "NixlConnector"` and set +`"kv_connector_extra_config": {"transfer_mode": "push"}`. + +For the built-in `NixlConnector`, `transfer_mode` defaults to `"pull"` and accepts +only `"pull"` or `"push"`; other values raise `ValueError`. +Explicit `NixlPullConnector` and `NixlPushConnector` names ignore this option. +`kv_connector_module_path` takes precedence over built-in selection. +Under `MultiConnector`, set `transfer_mode` in the NIXL child's +`kv_connector_extra_config`. + This document describes the threading, queues, and scheduling interactions specific to the push design. The pull-mode design is unchanged; the push connector reuses the same handshake, NIXL agent diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index b2853195e3a3..2c413708b6aa 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -215,3 +215,21 @@ def test_kv_offloading_size_only_uses_native_default(): assert kv_transfer_config.kv_connector == "OffloadingConnector" assert kv_transfer_config.kv_role == "kv_both" assert kv_connector_extra_config["cpu_bytes_to_use"] == 4.0 * (1 << 30) + + +@pytest.mark.parametrize( + "mode,expected", + [(None, "NixlPullConnector"), ("push", "NixlPushConnector"), ("Push", None)], +) +def test_nixl_connector_transfer_mode(mode, expected): + """Frameworks that only accept "NixlConnector" select push via extra config.""" + ktc = KVTransferConfig( + kv_connector="NixlConnector", + kv_role="kv_both", + kv_connector_extra_config={} if mode is None else {"transfer_mode": mode}, + ) + if expected is None: + with pytest.raises(ValueError, match="transfer_mode"): + KVConnectorFactory.get_connector_class(ktc) + else: + assert KVConnectorFactory.get_connector_class(ktc).__name__ == expected diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index 97064e06cdb8..2eefa15ebeb2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -123,6 +123,11 @@ def get_connector_class( logger.error(msg) raise ValueError(msg) elif connector_name in cls._registry: + if connector_name == "NixlConnector": + mode = kv_transfer_config.get_from_extra_config("transfer_mode", "pull") + if mode not in ("pull", "push"): + raise ValueError(f"Invalid NixlConnector transfer_mode: {mode!r}") + connector_name = f"Nixl{mode.capitalize()}Connector" connector_cls = cls._registry[connector_name]() else: raise ValueError(f"Unsupported connector type: {connector_name}") From 317081fb6dc8b125a3134711a7750707bfd55371 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Wed, 28 Oct 2026 16:51:18 +0800 Subject: [PATCH 052/406] [Bugfix][DSA] Reserve the decode indexer logits in the memory profile (#62860) Signed-off-by: Canlin Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../v1/attention/test_indexer_tp_row_shard.py | 4 ++ .../layers/sparse_attn_indexer.py | 44 ++++++++++++++++--- 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/tests/v1/attention/test_indexer_tp_row_shard.py b/tests/v1/attention/test_indexer_tp_row_shard.py index 4bbbf93eb905..ae683df8a1a3 100644 --- a/tests/v1/attention/test_indexer_tp_row_shard.py +++ b/tests/v1/attention/test_indexer_tp_row_shard.py @@ -210,6 +210,10 @@ def run_dummy(tp_world: int) -> None: decode_context_parallel_size=1, prefill_context_parallel_size=1, ), + scheduler_config=SimpleNamespace( + max_num_seqs=4, max_num_batched_tokens=512 + ), + num_speculative_tokens=0, ), ) m.setattr(sparse_indexer.envs, "VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", 1) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index fb3eb528a856..1c89724c2f60 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Custom Sparse Attention Indexer layers.""" +import math + import torch import torch.distributed as dist @@ -46,6 +48,7 @@ mqa_logits_use_mx_scales, ) from vllm.utils.import_utils import has_cutedsl +from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, @@ -478,6 +481,10 @@ def sparse_attn_indexer( ) assert candidate_block_size > 0 + # DeepSelect ranks SM100's BF16 decode logits directly; other top-k + # backends and the DCP merge read an FP32 copy. + fp32_decode_logits = topk_backend != "deep_select" or dcp_world_size != 1 + # assert isinstance(attn_metadata, dict) if not isinstance(attn_metadata, dict): # Reserve workspace for indexer during profiling run @@ -498,11 +505,34 @@ def sparse_attn_indexer( profile_specs.extend(gather_spec * 2) current_workspace_manager().get_simultaneous(*profile_specs) - # Dummy allocation to simulate for peak logits tensor memory during inference. - # FP8 elements so elements == bytes - max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + # Decode rows per step: each decode is padded to 1 + speculative tokens. + # The profiling run sets the vLLM config; other dummy runs may not. + max_decode_tokens = 0 + if (cfg := get_current_vllm_config_or_none()) is not None: + max_decode_tokens = min( + cfg.scheduler_config.max_num_seqs, + cfg.scheduler_config.max_num_batched_tokens, + ) * (1 + cfg.num_speculative_tokens) + + # Dummy allocation for the peak logits memory during inference: the + # prefill chunks' budget, or the decode logits, which eager and + # piecewise steps allocate outside the CUDA graph pools. DeepGEMM emits + # them in BF16 on SM100 and FP32 elsewhere, each row holding whole KV + # splits at a 1024-byte stride; the FP32 copy keeps that stride. + if current_platform.is_device_capability_family(100): + logits_dtype, split_kv = torch.bfloat16, 384 + else: + logits_dtype, split_kv = torch.float32, 256 + row_align = math.lcm(split_kv, 1024 // logits_dtype.itemsize) + elem_bytes = logits_dtype.itemsize + if logits_dtype != torch.float32 and fp32_decode_logits: + elem_bytes += torch.float32.itemsize + max_logits_bytes = max( + envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024, + max_decode_tokens * round_up(max_model_len, row_align) * elem_bytes, + ) _ = torch.empty( - max_logits_elems, dtype=torch.uint8, device=hidden_states.device + max_logits_bytes, dtype=torch.uint8, device=hidden_states.device ) # The dummy run returns before exchange_topk executes, so reserve its @@ -887,9 +917,9 @@ def sparse_attn_indexer( max_model_len=max_model_len, indices=decode_metadata.indices, ) - # Keeping bf16 halves top-k/candidate reads; other top-k backends and - # DCP need fp32. SM100 decode is varlen: one context length per row. - if topk_backend != "deep_select" or dcp_world_size != 1: + # Keeping bf16 halves top-k/candidate reads. SM100 decode is varlen: + # one context length per row. + if fp32_decode_logits: logits = as_fp32_logits(logits, seq_lens.view(-1)) num_rows = logits.shape[0] if candidate_blocks is not None: From deb1a97eb8303ab412ad3dee11f335fe8db53ebc Mon Sep 17 00:00:00 2001 From: Subhash Polisetti <49185065+subhashpolisetti@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:17:11 -0700 Subject: [PATCH 053/406] [Bugfix][Benchmark] Restore --input-len/--output-len/--prefix-len in vllm bench throughput (#64227) Signed-off-by: subhash polisetti --- tests/benchmarks/test_throughput_cli.py | 64 +++++++++++++++++++++++++ vllm/benchmarks/throughput.py | 56 ++++++++++++++-------- 2 files changed, 101 insertions(+), 19 deletions(-) diff --git a/tests/benchmarks/test_throughput_cli.py b/tests/benchmarks/test_throughput_cli.py index 37ced49838c3..c994d9ed0bf7 100644 --- a/tests/benchmarks/test_throughput_cli.py +++ b/tests/benchmarks/test_throughput_cli.py @@ -9,6 +9,7 @@ """ import subprocess +import warnings from argparse import Namespace from types import SimpleNamespace @@ -22,6 +23,7 @@ add_cli_args, assign_loras, get_requests, + validate_args, ) from vllm.utils.argparse_utils import FlexibleArgumentParser @@ -322,6 +324,68 @@ def test_to_serve_args_keeps_legacy_flag_fallback() -> None: assert not serve_args.enable_multimodal_chat +def _parse_throughput_args(*argv: str) -> Namespace: + parser = FlexibleArgumentParser() + add_cli_args(parser) + return parser.parse_args(["--model", MODEL_NAME, *argv]) + + +@pytest.mark.parametrize( + ("argv", "expected"), + [ + # Neither flag: the RandomDataset defaults. + ([], (1024, 128, 0)), + # Legacy flags alone are honored. + ( + ["--input-len", "2048", "--output-len", "512", "--prefix-len", "64"], + (2048, 512, 64), + ), + # Explicit --random-* wins, including an explicit 0. + ( + [ + "--input-len", + "2048", + "--random-input-len", + "300", + "--prefix-len", + "64", + "--random-prefix-len", + "0", + ], + (300, 128, 0), + ), + ], +) +def test_parsed_lengths_follow_flag_precedence(argv, expected) -> None: + """Lengths from the real parser: --random-* > legacy flags > defaults.""" + serve_args = _to_serve_args(_parse_throughput_args(*argv)) + + assert ( + serve_args.random_input_len, + serve_args.random_output_len, + serve_args.random_prefix_len, + ) == expected + + +@pytest.mark.parametrize( + "argv", + [ + ["--input-len", "32", "--output-len", "1", "--prefix-len", "8"], + ["--random-input-len", "32", "--random-output-len", "1"], + ["--random-prefix-len", "64"], + ["--random-prefix-len", "0"], + ], +) +def test_one_flag_per_length_does_not_warn(argv) -> None: + """Only the legacy or only the --random-* flag: no "Both ... are specified".""" + args = _parse_throughput_args(*argv) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + validate_args(args) + + assert not [str(w.message) for w in caught if "are specified" in str(w.message)] + + @pytest.mark.benchmark def test_get_requests_random_mm_with_mm_processor_args( hf_tokenizer: PreTrainedTokenizerBase, diff --git a/vllm/benchmarks/throughput.py b/vllm/benchmarks/throughput.py index 76f6a72fd98a..b48e5da6895c 100644 --- a/vllm/benchmarks/throughput.py +++ b/vllm/benchmarks/throughput.py @@ -17,6 +17,7 @@ from vllm.benchmarks.datasets import ( BenchmarkDataset, + RandomDataset, SampleRequest, add_random_dataset_base_args, add_random_multimodal_dataset_args, @@ -523,6 +524,14 @@ def save_to_pytorch_benchmark_format( write_to_json(pt_file, pt_records) +# --random-* length -> (legacy option it overrides, value when neither is set). +_RANDOM_LEN_FALLBACKS = { + "random_input_len": ("input_len", RandomDataset.DEFAULT_INPUT_LEN), + "random_output_len": ("output_len", RandomDataset.DEFAULT_OUTPUT_LEN), + "random_prefix_len": ("prefix_len", RandomDataset.DEFAULT_PREFIX_LEN), +} + + def _to_serve_args(args: argparse.Namespace) -> argparse.Namespace: """Translate throughput args into a namespace the shared get_samples reads. @@ -539,16 +548,16 @@ def _to_serve_args(args: argparse.Namespace) -> argparse.Namespace: """ d = vars(args).copy() - # random_*: prefer --random-* over legacy --input/output/prefix-len. - # Callers without the legacy flags (e.g. bench mm-processor) omit them, so - # read them optionally. Identity checks, not truthiness: an explicit 0 - # (the default of --random-prefix-len) must not fall through to them. - input_len = getattr(args, "input_len", None) + # random_*: prefer --random-* over legacy --input/output/prefix-len, then + # the RandomDataset defaults. Callers without the legacy flags (e.g. bench + # mm-processor) omit them, so read them optionally. Identity checks, not + # truthiness: an explicit 0 must not fall through to them. + for name, (legacy, default) in _RANDOM_LEN_FALLBACKS.items(): + value = getattr(args, name, None) + if value is None: + value = getattr(args, legacy, None) + d[name] = default if value is None else value output_len = getattr(args, "output_len", None) - prefix_len = getattr(args, "prefix_len", None) - d["random_input_len"] = getattr(args, "random_input_len", input_len) - d["random_output_len"] = getattr(args, "random_output_len", output_len) - d["random_prefix_len"] = getattr(args, "random_prefix_len", prefix_len) # --output-len maps to the per-dataset output-len entry points get_samples # reads. None passes through (each dataset applies its own default), # matching prior throughput behaviour of omitting output_len when unset. @@ -642,12 +651,6 @@ def validate_args(args): if not args.dataset_path and args.dataset_name != "prefix_repetition": print("When dataset path is not set, it will default to random dataset") args.dataset_name = "random" - random_input_len = getattr(args, "random_input_len", None) - if args.input_len is None and random_input_len is None: - raise ValueError( - "Either --input-len or --random-input-len must be provided " - "for a random dataset" - ) # === Dataset Name Specific Checks === # --hf-subset and --hf-split: only used @@ -820,14 +823,18 @@ def add_cli_args(parser: FlexibleArgumentParser): "--input-len", type=int, default=None, - help="Input prompt length for each request", + help="Input prompt length for each request. For random datasets, " + "--random-input-len takes precedence; if neither is set, " + f"{RandomDataset.DEFAULT_INPUT_LEN}.", ) parser.add_argument( "--output-len", type=int, default=None, help="Output length for each request. Overrides the " - "output length from the dataset.", + "output length from the dataset. For random datasets, " + "--random-output-len takes precedence; if neither is set, " + f"{RandomDataset.DEFAULT_OUTPUT_LEN}.", ) parser.add_argument( "--n", type=int, default=1, help="Number of generated sequences per prompt." @@ -907,9 +914,10 @@ def add_cli_args(parser: FlexibleArgumentParser): parser.add_argument( "--prefix-len", type=int, - default=0, + default=None, help="Number of fixed prefix tokens before the random " - "context in a request (default: 0).", + "context in a request. --random-prefix-len takes precedence; if " + f"neither is set, {RandomDataset.DEFAULT_PREFIX_LEN}.", ) # hf dataset @@ -975,6 +983,16 @@ def add_cli_args(parser: FlexibleArgumentParser): # (random, random-mm, random-rerank) add_random_dataset_base_args(parser) + # Unset unless passed, so the legacy --input-len/--output-len/--prefix-len + # apply when given; _to_serve_args falls back to the RandomDataset defaults. + for action in parser._actions: + if action.dest in _RANDOM_LEN_FALLBACKS: + legacy, default = _RANDOM_LEN_FALLBACKS[action.dest] + action.default = None + action.help = ( + f"{action.help} Takes precedence over " + f"--{legacy.replace('_', '-')}; if neither is set, {default}." + ) add_random_multimodal_dataset_args(parser) # ASR dataset From 76c4ad2f1544f8a9f9dbb9e0136cd4085fc55f50 Mon Sep 17 00:00:00 2001 From: nrenvo Date: Wed, 28 Oct 2026 17:22:33 +0800 Subject: [PATCH 054/406] =?UTF-8?q?[Bugfix][Multimodal]=20Fix=20DeepSeek-O?= =?UTF-8?q?CR=20zero-sum=20patch=20guard=20in=20=5Fencode=E2=80=A6=20(#640?= =?UTF-8?q?91)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: nrenvo Co-authored-by: vllm-agent --- .../generation/test_deepseek_ocr.py | 58 +++++++++++++++++++ vllm/model_executor/models/deepseek_ocr.py | 2 +- vllm/model_executor/models/deepseek_ocr2.py | 2 +- 3 files changed, 60 insertions(+), 2 deletions(-) create mode 100644 tests/models/multimodal/generation/test_deepseek_ocr.py diff --git a/tests/models/multimodal/generation/test_deepseek_ocr.py b/tests/models/multimodal/generation/test_deepseek_ocr.py new file mode 100644 index 000000000000..6dc9ea52434f --- /dev/null +++ b/tests/models/multimodal/generation/test_deepseek_ocr.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Guard tests for DeepSeek-OCR _encode_local_features. + +_encode_local_features must skip processing only for an empty-element +patches tensor, not for a zero-sum one. + +Regression for #64075: the old guard ``torch.sum(patches).item() == 0`` +returned None for a 1536x768 half-white/half-black image whose two +normalised crop patches cancel to sum==0, causing a shape mismatch in +_merge_multimodal_embeddings and a fatal EngineCore crash. + +The correct guard is ``patches.numel() == 0``. +""" + +import pytest +import torch + +from vllm.model_executor.models.deepseek_ocr import DeepseekOCRForCausalLM +from vllm.model_executor.models.deepseek_ocr2 import DeepseekOCR2ForCausalLM + + +class TestDeepseekOCREncodeLocalFeaturesGuard: + def test_ocr2_empty_patches_returns_none(self): + """Zero-element tensor → skip local encoding (non-tiled image).""" + empty = torch.empty(0, 3, 224, 224) + assert DeepseekOCR2ForCausalLM._encode_local_features(None, empty) is None + + def test_ocr2_zero_sum_patches_enters_model_body(self): + """Non-empty patches that happen to sum to zero must NOT return None. + + The method will raise AttributeError when it tries to call + self.sam_model on a None self – that is the expected outcome here + because it proves the guard was passed, not that encoding succeeded. + The real crash was a shape-mismatch ValueError that killed EngineCore. + """ + # Two patches that cancel: white (1.0) + black (-1.0) after Normalize. + wb = torch.cat([torch.ones(1, 3, 4, 4), -torch.ones(1, 3, 4, 4)]) + assert wb.numel() > 0 + assert wb.sum().item() == 0.0 + with pytest.raises(AttributeError): + DeepseekOCR2ForCausalLM._encode_local_features(None, wb) + + def test_ocr1_empty_patches_returns_none(self): + """OCR-1: zero-element tensor → skip local encoding.""" + empty = torch.empty(0, 3, 224, 224) + crop_shape = torch.tensor([0, 0]) + assert ( + DeepseekOCRForCausalLM._encode_local_features(None, empty, crop_shape) + is None + ) + + def test_ocr1_zero_sum_patches_enters_model_body(self): + """OCR-1: non-empty zero-sum patches must pass the guard.""" + wb = torch.cat([torch.ones(1, 3, 4, 4), -torch.ones(1, 3, 4, 4)]) + crop_shape = torch.tensor([1, 2]) + with pytest.raises(AttributeError): + DeepseekOCRForCausalLM._encode_local_features(None, wb, crop_shape) diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 3aef695d0608..d137030a337a 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -499,7 +499,7 @@ def _encode_global_features(self, image_tensor: torch.Tensor) -> torch.Tensor: def _encode_local_features( self, patches: torch.Tensor, crop_shape: torch.Tensor ) -> torch.Tensor | None: - if torch.sum(patches).item() == 0: + if patches.numel() == 0: return None local_features_1 = self.sam_model(patches) diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py index 18de7503c721..9139259909e6 100644 --- a/vllm/model_executor/models/deepseek_ocr2.py +++ b/vllm/model_executor/models/deepseek_ocr2.py @@ -343,7 +343,7 @@ def _encode_global_features(self, image_tensor: torch.Tensor) -> torch.Tensor: return features.view(-1, dim) def _encode_local_features(self, patches: torch.Tensor) -> torch.Tensor | None: - if torch.sum(patches).item() == 0: + if patches.numel() == 0: return None local_features = self.sam_model(patches) From 77f58e530450416d217e4816727a35910bf3d9bf Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Wed, 28 Oct 2026 02:22:50 -0700 Subject: [PATCH 055/406] [Bugfix][KV Offload] Count bounded-replay SWA allocations in the lazy walk target (#64229) Signed-off-by: Juntian Liu Co-authored-by: Claude Opus 5.5 --- tests/v1/simple_kv_offload/test_scheduler.py | 31 ++++++++++++++++++++ vllm/v1/simple_kv_offload/manager.py | 9 ++++-- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index c63cd614e77d..dd2028c6eaca 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -53,6 +53,7 @@ KVCacheSpec, KVCacheTensor, MambaSpec, + SlidingWindowMLASpec, SlidingWindowSpec, ) from vllm.v1.metrics.cache_hit_source import CacheHitSource @@ -1136,6 +1137,36 @@ def test_lazy_duplicate_store_skipped() -> None: assert get_cpu_free_blocks(sched) == cpu_free_after_first +# --------------------------------------------------------------------------- +# Test 2d: Lazy walk covers bounded-replay SWA allocations +# --------------------------------------------------------------------------- +def test_lazy_walk_covers_bounded_replay_swa_groups() -> None: + """Lazy: the walk covers what every group, offloaded or not, takes per step.""" + config = _make_kv_cache_config(num_blocks=32, num_groups=3) + for group in config.kv_cache_groups[1:]: + group.kv_cache_spec = SlidingWindowMLASpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=HEAD_SIZE, + dtype=DTYPE, + sliding_window=BLOCK_SIZE * 4, + bounded_replay=True, + ) + fix = make_scheduler( + num_cpu_blocks=32, num_gpu_blocks=32, lazy=True, kv_cache_config=config + ) + gpu_pool = fix.gpu_block_pool + + # One step takes cdiv(64, 16) blocks per group. + num_blocks = 3 * 4 + blocks = _allocate_gpu_blocks(gpu_pool, make_request(num_blocks), num_blocks) + _flush_old_blocks_to_lru_head(gpu_pool, gpu_pool.get_num_free_blocks()) + gpu_pool.free_blocks(blocks) + + meta = fix.scheduler.build_connector_meta(make_scheduler_output({})) + assert sorted(meta.store_gpu_blocks) == [b.block_id for b in blocks] + + # --------------------------------------------------------------------------- # Test 3: LRU eviction order # --------------------------------------------------------------------------- diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 51e6473edf41..6794412402dc 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -342,10 +342,15 @@ def _estimate_lazy_target_blocks( """GPU blocks to keep available (free/offloaded) per step in lazy mode.""" WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks target = 0 - for g in kv_cache_config.prefix_cacheable_groups: + for g in kv_cache_config.kv_cache_groups: spec = g.kv_cache_spec block_size = resolve_dcp_kv_block_size(spec, cp_world_size) - if isinstance(spec, MambaSpec): + if not spec.prefix_cacheable: + # Bounded-replay SWA (DSV4.1) is never offloaded, but its blocks for + # each scheduled token come from the same free queue. + if isinstance(spec, SlidingWindowSpec): + target += cdiv(max_num_batched_tokens, block_size) + elif isinstance(spec, MambaSpec): target += 2 elif isinstance(spec, SlidingWindowSpec): target += cdiv(spec.sliding_window, block_size) + 1 From 5ea8301982648ac444ea01368ecb1b064f6de7bc Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:39:44 -0400 Subject: [PATCH 056/406] [CI] Move Attention, Quantization and MoE kernel tests and their AMD mirrors from fixed to automatic shards (#62077) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/kernels.yaml | 35 ++++++++++-------------------- 1 file changed, 12 insertions(+), 23 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index f4fd3e7e1795..18c69b08f690 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -345,8 +345,9 @@ steps: --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated - pytest -v -s kernels/test_fused_muse_spark_qknorm_rope_kv_insert.py -k rocm -- label: ":nvidia: (L4) Attention Kernels Shard %N" +- label: ":nvidia: (L4) Attention Kernels" key: kernels-attention-test + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 30 @@ -359,15 +360,14 @@ steps: - vllm/utils/flashinfer.py - tests/kernels/attention commands: - - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 7 + - pytest -v -s kernels/attention mirror: amd: - label: ":amd: (MI300) Attention Kernels Shard %N" + label: ":amd: (MI300) Attention Kernels" + automatic_shard: true dind: false device: mi300_1 timeout_in_minutes: 90 - parallelism: 6 depends_on: - image-build-amd source_file_dependencies: @@ -380,15 +380,7 @@ steps: - vllm/envs.py - vllm/platforms/rocm.py commands: - - >- - pytest -v -s - kernels/attention - --ignore=kernels/attention/test_triton_unified_attention_diffkv.py - --ignore=kernels/attention/test_rocm_aiter_mla_decode.py - --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py - --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py - --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py - --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py - label: ":nvidia: (H100) Attention DiffKV Kernels" key: kernels-attention-diffkv-test-h100 @@ -512,8 +504,9 @@ steps: commands: - pytest -v -s kernels/quantization -k "not test_nvfp4_dense_emulation" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":nvidia: (L4) MoE Kernels Shard %N" +- label: ":nvidia: (L4) MoE Kernels" key: kernels-moe-test + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 50 @@ -526,16 +519,12 @@ steps: - vllm/envs.py - vllm/config commands: - - >- - pytest -v -s - kernels/moe - --ignore=kernels/moe/test_modular_oai_triton_moe.py - --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 5 + - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py mirror: amd: - label: ":amd: (MI300) MoE Kernels Shard %N" + label: ":amd: (MI300) MoE Kernels" + automatic_shard: true dind: false device: mi300_1 timeout_in_minutes: 55 From f4fd1ab68e2ba3a76bf6b903662eb6a8ee8202ab Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Wed, 28 Oct 2026 03:27:52 -0700 Subject: [PATCH 057/406] [Bugfix][KV Offload][P2P] Stop DEALER monitor before closing its reader (#63207) Signed-off-by: Maroon Ayoub --- .../tiering/p2p/test_zmq_transport.py | 99 +++++++++++++++++++ vllm/v1/kv_offload/tiering/p2p/control/zmq.py | 4 + 2 files changed, 103 insertions(+) diff --git a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py index eebb9c01ba1e..d7bf17fb925d 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py @@ -4,6 +4,7 @@ from __future__ import annotations +import select import socket import time @@ -17,6 +18,11 @@ _Sockets, ) +# Transports whose libzmq I/O thread is blocked. Terminating their context +# blocks forever, and garbage collection would terminate it, so they stay +# referenced until the interpreter exits (pyzmq skips the term at exit). +_WEDGED_TRANSPORTS: list[ZmqTransport] = [] + def _free_port() -> int: """Find a free TCP port.""" @@ -353,6 +359,99 @@ def test_inbound_message_survives_dead_registration(self): transport_a.close() transport_b.close() + def test_close_with_unsent_message_keeps_transport_usable(self, monkeypatch): + """A connection closed with an unsent message does not wedge the transport. + + A closed DEALER that still holds a message keeps reconnecting to an + unresponsive peer, and its monitor events then block libzmq's I/O + thread, so no socket in the context can send or receive again. + """ + apply_heartbeat = zmq_module._apply_heartbeat + + def apply_heartbeat_fast_handshake(sock: zmq.Socket) -> None: + apply_heartbeat(sock) + sock.setsockopt(zmq.HANDSHAKE_IVL, 200) + + monkeypatch.setattr( + zmq_module, "_apply_heartbeat", apply_heartbeat_fast_handshake + ) + # Accepts TCP connections but never answers the ZMTP greeting. + silent = socket.socket() + silent.bind(("127.0.0.1", 0)) + silent.listen(8) + transport_a, port_a = _make_transport() + transport_b, _ = _make_transport() + wedged = True + try: + stale = transport_b.connect(f"127.0.0.1:{silent.getsockname()[1]}") + stale.send({"type": "stale"}) + stale.close() + # Several ZMTP handshake timeouts on the closed DEALER. + time.sleep(1.0) + + conn = transport_b.connect(f"127.0.0.1:{port_a}") + conn.send({"seq": 1}) + inbound = _wait_for_inbound(transport_a)[0] + assert _wait_for_messages(transport_a, inbound, 1) == [{"seq": 1}] + wedged = False + finally: + transport_a.close() + silent.close() + if wedged: + _WEDGED_TRANSPORTS.append(transport_b) + else: + transport_b.close() + + def test_close_while_connecting_keeps_transport_usable(self): + """Closing a connection whose TCP connect never completed keeps the + transport usable. + + A peer address with no host behind it (a departed pod, a lost node) + never answers the SYN. Tearing down that connect raises EVENT_CLOSED + on libzmq's I/O thread, and a monitor event sent after its reader is + closed blocks that thread, so no socket in the context can send or + receive again. + """ + # Linux drops SYNs to a listener whose accept queue is full, so a + # connect to it stays in progress like one to an address with no host. + full = socket.socket() + full.bind(("127.0.0.1", 0)) + full.listen(0) + addr = full.getsockname() + filler = socket.create_connection(addr, timeout=1.0) + probe = socket.socket() + probe.setblocking(False) + probe.connect_ex(addr) + _, connected, _ = select.select([], [probe], [], 0.5) + probe.close() + if connected: + filler.close() + full.close() + pytest.skip("kernel completes connects to a full accept queue") + + transport_a, port_a = _make_transport() + transport_b, _ = _make_transport() + wedged = True + try: + pending = transport_b.connect(f"127.0.0.1:{addr[1]}") + # Let the I/O thread start the TCP connect. + time.sleep(0.2) + pending.close() + + conn = transport_b.connect(f"127.0.0.1:{port_a}") + conn.send({"seq": 1}) + inbound = _wait_for_inbound(transport_a)[0] + assert _wait_for_messages(transport_a, inbound, 1) == [{"seq": 1}] + wedged = False + finally: + transport_a.close() + filler.close() + full.close() + if wedged: + _WEDGED_TRANSPORTS.append(transport_b) + else: + transport_b.close() + class TestExpandedMonitorEvents: """The monitor mask includes EVENT_DISCONNECTED and EVENT_CLOSED so diff --git a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py index ee0fa819f3db..a64f0463844f 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py @@ -132,6 +132,10 @@ def close(self) -> None: self._closed = True self._dead = True logger.info("ZmqConnection: closing connection to %s", self.peer_id) + # libzmq sends monitor events with a blocking send from its I/O + # thread, so an event raised after the reader is closed (EVENT_CLOSED + # for a connect still in progress) blocks every socket in the context. + self._sockets.dealer.disable_monitor() self._sockets.monitor.close() self._sockets.dealer.close() From 9b9eeb283a195531790c0b9a2747b543df189a80 Mon Sep 17 00:00:00 2001 From: bantianyue Date: Wed, 28 Oct 2026 19:15:44 +0800 Subject: [PATCH 058/406] [Doc] Fix long-text embedding example pooler config and GPU pinning (#64170) Signed-off-by: bantianyue Co-authored-by: Taneem Ibrahim --- .../pooling/embed/openai_embedding_long_text/README.md | 4 ++-- .../pooling/embed/openai_embedding_long_text/service.sh | 7 +++---- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/examples/pooling/embed/openai_embedding_long_text/README.md b/examples/pooling/embed/openai_embedding_long_text/README.md index 2ed04f1d91f2..cafe4975a58b 100644 --- a/examples/pooling/embed/openai_embedding_long_text/README.md +++ b/examples/pooling/embed/openai_embedding_long_text/README.md @@ -46,7 +46,7 @@ The key parameters for chunked processing are in the `--pooler-config`: ```json { - "pooling_type": "auto", + "pooling_type": "MEAN", "use_activation": true, "enable_chunked_processing": true, "max_embed_len": 3072000 @@ -54,7 +54,7 @@ The key parameters for chunked processing are in the `--pooler-config`: ``` !!! note - `pooling_type` sets the model's own pooling strategy for processing within each chunk. The cross-chunk aggregation automatically uses MEAN strategy when input exceeds the model's native maximum length. + `pooling_type` sets the model's own pooling strategy for processing within each chunk and must be one of `CLS`, `LAST` or `MEAN`. `service.sh` accepts `POOLING_TYPE=auto` and resolves it to one of these before building the config. The cross-chunk aggregation automatically uses MEAN strategy when input exceeds the model's native maximum length. #### Chunked Processing Behavior diff --git a/examples/pooling/embed/openai_embedding_long_text/service.sh b/examples/pooling/embed/openai_embedding_long_text/service.sh index 68950d04ee14..1c68b00f9177 100644 --- a/examples/pooling/embed/openai_embedding_long_text/service.sh +++ b/examples/pooling/embed/openai_embedding_long_text/service.sh @@ -20,8 +20,7 @@ API_KEY=${API_KEY:-"your-api-key"} # Enhanced pooling configuration with model-specific defaults POOLING_TYPE=${POOLING_TYPE:-"auto"} # auto, MEAN, CLS, LAST -export VLLM_ENABLE_CHUNKED_PROCESSING=true -export CUDA_VISIBLE_DEVICES=2,3,4,5 +ENABLE_CHUNKED_PROCESSING=true echo "🚀 Starting vLLM Embedding Server with Enhanced Chunked Processing" echo "==================================================================" @@ -68,7 +67,7 @@ echo "📋 Configuration:" echo " - Model: $MODEL_NAME" echo " - Port: $PORT" echo " - GPU Count: $GPU_COUNT" -echo " - Enhanced Chunked Processing: ${VLLM_ENABLE_CHUNKED_PROCESSING}" +echo " - Enhanced Chunked Processing: ${ENABLE_CHUNKED_PROCESSING}" echo " - Max Embed Length: ${MAX_EMBED_LEN} tokens" echo " - Native Pooling Type: $POOLING_TYPE + Normalization" echo " - Cross-chunk Aggregation: MEAN (automatic)" @@ -96,7 +95,7 @@ echo "" echo "🔧 Starting server with enhanced chunked processing configuration..." # Build pooler config JSON -POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"use_activation\": true, \"enable_chunked_processing\": ${VLLM_ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}" +POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"use_activation\": true, \"enable_chunked_processing\": ${ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}" # Start vLLM server with enhanced chunked processing vllm serve "$MODEL_NAME" \ From ab0e5ef7954190de2a74a6715d1382b08f5accc4 Mon Sep 17 00:00:00 2001 From: "Anvit D." Date: Wed, 28 Oct 2026 16:52:34 +0530 Subject: [PATCH 059/406] [Build] keep CUDA PyTorch wheels aligned with image (#59198) Signed-off-by: AnvitDevadiga Co-authored-by: AnvitDevadiga Co-authored-by: Shengqi Chen --- docker/Dockerfile | 31 ++++++++++++++++++++++++------- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index f3cac42652e9..e35206808333 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -253,8 +253,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \ && uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ else \ + # Keep the CUDA-specific PyTorch index first. Using it as an extra + # index with unsafe-best-match can select a higher, incompatible + # torch build from the base image (for example cu130 in a cu129 image). uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi # Track PyTorch lib versions used during build and match in downstream instances. @@ -372,7 +377,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ echo "Installing build requirements..." \ && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi WORKDIR /workspace @@ -732,7 +739,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ echo "Installing dev requirements..." \ && uv pip install --python /opt/venv/bin/python3 -r requirements/dev.txt \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi #################### DEV IMAGE #################### @@ -941,7 +950,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ uv pip install --system \ -r /tmp/requirements-cuda.txt \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi && \ rm /tmp/requirements-cuda.txt /tmp/common.txt @@ -1017,7 +1028,9 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist else \ echo "Installing vLLM..." \ && uv pip install --system dist/*.whl --verbose \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi # Expose the Rust frontend bundled in the wheel on the container's PATH. @@ -1140,7 +1153,9 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist else \ echo "Installing vLLM..." \ && uv pip install --system dist/*.whl --verbose \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH @@ -1150,7 +1165,9 @@ COPY --from=build /tmp/ep_kernels_workspace/dist/wheels.sha256 /tmp/ep-kernels-w RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \ --mount=type=cache,target=/opt/uv/cache \ uv pip install --system ep_kernels/dist/*.whl --verbose \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match # See the matching comment in the vllm-base stage. ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib64:${LD_LIBRARY_PATH} From 0beab6c251e8046c0c6cf2fef114cd23b58cda74 Mon Sep 17 00:00:00 2001 From: liranschour Date: Wed, 28 Oct 2026 14:41:36 +0200 Subject: [PATCH 060/406] [KV Offload][Tiering] Let a tier fail a promotion and defer releasing its chunks (#63779) Signed-off-by: Liran Schour Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Or Ozeri --- .../tiering/test_tiering_offloading.py | 152 +++++++++++++++++- vllm/v1/kv_offload/tiering/base.py | 14 ++ vllm/v1/kv_offload/tiering/manager.py | 98 ++++++++++- vllm/v1/kv_offload/tiering/metrics.py | 22 +++ vllm/v1/kv_offload/tiering/spec.py | 21 +++ 5 files changed, 305 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 1de7850d7972..d3c212b9d1c0 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -10,6 +10,7 @@ 5. Eviction coordination between tiers """ +import threading from collections.abc import Iterable from unittest.mock import MagicMock, patch @@ -40,7 +41,7 @@ TierMatcher, make_offload_key, ) -from vllm.v1.kv_offload.cpu.common import CPUOffloadingInfo +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec, CPUOffloadingInfo from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy from vllm.v1.kv_offload.tiering.base import ( JobResult, @@ -187,6 +188,7 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, + TieringOffloadingMetrics.DEFERRED_RELEASE_CHUNKS, ): assert isinstance(metrics[metric_name], OffloadingGaugeMetadata) assert metrics[metric_name].labelnames == ("tier",) @@ -656,6 +658,154 @@ def submit_partial(job_metadata: TransferJob) -> None: expected_source_keys ) + def _promote_with_deferred_release( + self, monkeypatch, chunks, successful_indices=(), preset=False + ) -> list[threading.Event]: + """Promote ``chunks`` from tier1 into a failure with a deferred release. + + Returns the events the tier reported; setting one releases its job. + """ + for chunk in chunks: + self.secondary_tier1.chunks[chunk] = True + events: list[threading.Event] = [] + + def submit_deferring(job_metadata: TransferJob) -> None: + event = threading.Event() + if preset: + event.set() + events.append(event) + self.secondary_tier1.completed_jobs.append( + JobResult( + job_id=job_metadata.job_id, + success=False, + successful_keys=tuple(chunks[i] for i in successful_indices) + or None, + deferred_release=event, + ) + ) + + monkeypatch.setattr(self.secondary_tier1, "submit_load", submit_deferring) + for chunk in chunks: + assert self.manager.lookup(chunk, _CTX) is LookupResult.HIT_PENDING + self._simulate_on_schedule_end() + self._simulate_on_schedule_end() + return events + + def _reduced_stats(self) -> dict[str, int | float]: + stats = self.manager.get_stats() + assert stats is not None + return stats.reduce() + + def _release(self, events: list[threading.Event]) -> None: + for event in events: + event.set() + self.manager._process_finished_jobs() + + def test_deferred_release_misses_but_holds_chunks_until_released( + self, manager_setup, monkeypatch + ): + """A deferred release fails the keys at once but keeps their chunks out + of the pool until the tier sets the event.""" + self._start_request() + chunks = to_keys(range(2)) + free_before = self.primary_tier._get_num_free_chunks() + + events = self._promote_with_deferred_release(monkeypatch, chunks) + + assert len(events) == 1 + assert [self.primary_tier.lookup(c, _CTX) for c in chunks] == [ + LookupResult.MISS, + LookupResult.MISS, + ] + assert self.primary_tier.num_deferred_release_chunks == 2 + assert self.primary_tier._get_num_free_chunks() == free_before - 2 + assert self.manager.has_pending_work() + self.manager._process_finished_jobs() + assert self.primary_tier.num_deferred_release_chunks == 2 + reduced = self._reduced_stats() + label = self.manager._metrics.tier_label(0) + assert ( + reduced[f"{TieringOffloadingMetrics.DEFERRED_RELEASE_CHUNKS}:{label}"] == 2 + ) + hold = TieringOffloadingMetrics.DEFERRED_RELEASE_HOLD + assert not has_histogram_count(reduced, hold) + + self._release(events) + + reduced = self._reduced_stats() + assert reduced[histogram_count_key(hold, label)] == 1 + + assert self.primary_tier.num_deferred_release_chunks == 0 + assert self.primary_tier._get_num_free_chunks() == free_before + assert self.primary_tier._num_write_pending_chunks == 0 + assert not self.manager._deferred_releases + + def test_deferred_release_restore_gets_a_fresh_chunk( + self, manager_setup, monkeypatch + ): + """A key whose old chunk awaits release can be stored again, into a + different chunk the late writer cannot touch.""" + self._start_request() + chunk = to_keys([7]) + events = self._promote_with_deferred_release(monkeypatch, chunk) + ((held,),) = (d.chunk_ids for d in self.manager._deferred_releases.values()) + + output = self.primary_tier.prepare_write(chunk, _CTX) + assert output is not None + assert isinstance(output.store_spec, CPULoadStoreSpec) + assert list(output.store_spec.chunk_ids) != [held] + self.primary_tier.complete_write(chunk, _CTX, True) + + self._release(events) + assert self.primary_tier.lookup(chunk[0], _CTX) is LookupResult.HIT + + def test_deferred_release_partial_failure_holds_only_failed_chunks( + self, manager_setup, monkeypatch + ): + self._start_request() + chunks = to_keys(range(3)) + events = self._promote_with_deferred_release( + monkeypatch, chunks, successful_indices=(0,) + ) + + assert self.primary_tier.lookup(chunks[0], _CTX) is LookupResult.HIT + assert self.primary_tier.num_deferred_release_chunks == 2 + self._release(events) + assert self.primary_tier.num_deferred_release_chunks == 0 + + def test_deferred_release_already_set_with_result(self, manager_setup, monkeypatch): + """An event already set when its result is reported releases in the + same poll round.""" + self._start_request() + self._promote_with_deferred_release(monkeypatch, to_keys(range(2)), preset=True) + assert self.primary_tier.num_deferred_release_chunks == 0 + assert not self.manager._deferred_releases + + def test_defer_write_release_asserts_on_non_pending_key(self, manager_setup): + """Every failed key must still be pending: a ready or missing key + means a chunk the tier may still write to would go unheld.""" + self._start_request() + chunk = to_keys([5]) + assert self.primary_tier.prepare_write(chunk, _CTX) is not None + self.primary_tier.complete_write(chunk, _CTX, True) + with pytest.raises(AssertionError, match="is not a pending write"): + self.primary_tier.defer_write_release(chunk) + with pytest.raises(AssertionError, match="is not a pending write"): + self.primary_tier.defer_write_release(to_keys([6])) + + def test_reset_after_deferred_release_succeeds(self, manager_setup, monkeypatch): + """drain_jobs() completing the deferred release lets reset_cache proceed.""" + self._start_request() + events = self._promote_with_deferred_release(monkeypatch, to_keys(range(2))) + + def drain_jobs() -> None: + for event in events: + event.set() + + monkeypatch.setattr(self.secondary_tier1, "drain_jobs", drain_jobs) + self.manager.reset_cache() + assert self.primary_tier.num_deferred_release_chunks == 0 + def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup): """Resolved lookups report one sync delay sample per tier and chunk.""" self._start_request() diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index 3e9fccf4c740..cc7fbf0b4247 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Abstract interfaces and data types for the secondary tiering layer.""" +import threading import time from abc import ABC, abstractmethod from collections.abc import Collection, Iterable @@ -66,6 +67,8 @@ class TieringOffloadingMetrics: BACKPRESSURE_STORES_DROPPED = "vllm:kv_offload_tiering_backpressure_stores_dropped" BACKPRESSURE_BLOCKS_DROPPED = "vllm:kv_offload_tiering_backpressure_blocks_dropped" PROMOTION_LATENCY = "vllm:kv_offload_tiering_promotion_latency_seconds" + DEFERRED_RELEASE_CHUNKS = "vllm:kv_offload_tiering_deferred_release_chunks" + DEFERRED_RELEASE_HOLD = "vllm:kv_offload_tiering_deferred_release_hold_seconds" @dataclass @@ -93,6 +96,13 @@ class JobResult: successful_keys: Collection[OffloadKey] | None = None transfer_time: float | None = None transfer_bytes: int | None = None + # Only applicable to failed promotion jobs. Set when the tier's I/O may + # still write into the primary chunks of the failed keys: the manager fails + # those keys now but defers releasing their chunks to the free pool until + # the tier sets the event. The tier may set it from any thread, once its + # I/O can no longer touch those chunks; the manager applies the release on + # its own polling path. + deferred_release: threading.Event | None = None class ParentManager(ABC): @@ -384,6 +394,10 @@ def drain_jobs(self) -> None: would corrupt either the primary memoryview or the secondary backing store. Queued (not-yet-started) transfers may be cancelled, but their failure result must still appear in `get_finished_jobs()`. + + Likewise, every ``JobResult.deferred_release`` event must be set + before this returns, so the tier must bound how long it can defer a + release. """ pass diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 74b6bd7c6b66..b567b189ba05 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -111,6 +111,13 @@ class JobMetadata(NamedTuple): tier_idx: int +class _DeferredRelease(NamedTuple): + event: threading.Event + tier_idx: int + chunk_ids: list[int] + deferred_at: float + + class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): """CPUOffloadingManager with a primary/secondary transfer interface. @@ -144,6 +151,9 @@ def __init__( self.complete_read = self.complete_load self.prepare_write = self.prepare_store self.complete_write = self.complete_store + # Chunks of failed writes the writer may still be filling; their + # release to the free pool is deferred until release_deferred_write(). + self._deferred_release_chunk_ids: set[int] = set() self._kv_memoryview = mmap_region.create_kv_memoryview() @@ -157,6 +167,54 @@ def prepare_read( """ return self._prepare_load(keys, req_context, record_access=False) + def defer_write_release(self, keys: Collection[OffloadKey]) -> list[int]: + """Fail pending writes like ``complete_write(False)``, deferring the release. + + The keys leave the cache (lookups miss and a later store allocates a + fresh chunk), while the chunks stay counted as write-pending and out + of the free pool. Returns the deferred chunk ids, to be passed to + ``release_deferred_write`` once the writer can no longer touch them. + """ + chunk_ids: list[int] = [] + for key in keys: + chunk = self._policy.get(key) + # Every key must still be pending, or a chunk the tier may still + # write to would go unheld. + assert chunk is not None and not chunk.is_ready, ( + f"defer_write_release: key {key!r} is not a pending write" + ) + self._policy.remove(key) + self._deferred_release_chunk_ids.add(chunk.chunk_id) + chunk_ids.append(chunk.chunk_id) + if chunk_ids and self.events is not None: + self.events.append( + OffloadingEvent( + keys=list(keys), + medium=self.medium, + removed=True, + metadata_only=True, + ) + ) + return chunk_ids + + def release_deferred_write(self, chunk_ids: Iterable[int]) -> None: + """Complete a deferred release: return the chunks to the free pool.""" + for chunk_id in chunk_ids: + self._deferred_release_chunk_ids.remove(chunk_id) + self._num_write_pending_chunks -= 1 + self._free_list.append(chunk_id) + + @property + def num_deferred_release_chunks(self) -> int: + return len(self._deferred_release_chunk_ids) + + @override + def reset_cache(self) -> None: + assert not self._deferred_release_chunk_ids, ( + "reset_cache with deferred-release chunks: a tier may still write to them" + ) + super().reset_cache() + def get_kv_memoryview(self) -> memoryview: """Return the memoryview over the primary tier's KV cache buffer. @@ -250,6 +308,9 @@ def __init__( # True: secondary → primary (promotion) # False: primary → secondary (cascade) self._jobs: dict[JobId, JobMetadata] = {} + # Finished promotions whose failed chunks await release until the tier + # sets their JobResult.deferred_release event. + self._deferred_releases: dict[JobId, _DeferredRelease] = {} # Set by reset_cache() until on_reset_flushed(); see lookup(). self._promotions_held = False primary_view = self.primary_tier.get_kv_memoryview() @@ -491,7 +552,18 @@ def _complete_promotion( transfer_job.req_context, True, ) - if failed_keys: + if not failed_keys: + assert completed_job.deferred_release is None + elif completed_job.deferred_release is not None: + chunk_ids = self.primary_tier.defer_write_release(failed_keys) + self._deferred_releases[completed_job.job_id] = _DeferredRelease( + event=completed_job.deferred_release, + tier_idx=job_metadata.tier_idx, + chunk_ids=chunk_ids, + deferred_at=time.monotonic(), + ) + self._metrics.on_release_deferred(job_metadata.tier_idx, len(chunk_ids)) + else: self.primary_tier.complete_write( failed_keys, transfer_job.req_context, @@ -534,6 +606,9 @@ def _process_finished_jobs( f" but belongs to tier #{job_metadata.tier_idx}" ) transfer_job = job_metadata.transfer_job + assert ( + transfer_job.is_promotion or completed_job.deferred_release is None + ) self._metrics.on_job_finished(job_metadata, completed_job) if transfer_job.is_promotion: @@ -547,6 +622,22 @@ def _process_finished_jobs( transfer_job.keys, transfer_job.req_context ) self._update_backpressure(tier, job_metadata, completed_job) + if self._deferred_releases: + self._complete_deferred_releases(i) + + def _complete_deferred_releases(self, tier_idx: int) -> None: + released = [ + job_id + for job_id, deferred in self._deferred_releases.items() + if deferred.tier_idx == tier_idx and deferred.event.is_set() + ] + now = time.monotonic() + for job_id in released: + deferred = self._deferred_releases.pop(job_id) + self.primary_tier.release_deferred_write(deferred.chunk_ids) + self._metrics.on_deferred_release_done( + tier_idx, len(deferred.chunk_ids), now - deferred.deferred_at + ) def _should_store_to_tier( self, tier: SecondaryTierManager, num_blocks: int @@ -1119,8 +1210,10 @@ def has_pending_work(self) -> bool: # In-flight primary<->secondary transfers (pending promotions are # translated to transfer jobs in on_schedule_end), plus any work the # secondary tiers themselves still have outstanding. + # Deferred releases keep the tiers polled until they complete. return ( bool(self._jobs) + or bool(self._deferred_releases) or any(state.pending_cascade_keys for state in self._req_state.values()) or any(tier.has_pending_work() for tier in self.secondary_tiers) ) @@ -1167,6 +1260,9 @@ def reset_cache(self) -> None: # so manager bookkeeping is consistent before the primary reset. self._process_finished_jobs() assert not self._jobs + assert not self._deferred_releases, ( + "drain_jobs() returned with deferred releases still pending" + ) # Deferred promotion submissions reserve primary slots that the # reset below invalidates; their submit_load() has not yet been diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index 555077ec4305..1b28f67987d2 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -42,6 +42,7 @@ class _TierState: active_cascade_count: int = 0 primary_write_chunk_count: int = 0 primary_read_chunk_count: int = 0 + deferred_release_chunk_count: int = 0 class TieringMetricsTracker: @@ -128,6 +129,21 @@ def on_job_finished( self._observe_promotion_latency(job_metadata, result) self._decrement_tier_state(job_metadata) + def on_release_deferred(self, tier_idx: int, chunk_count: int) -> None: + self._tier_states[tier_idx].deferred_release_chunk_count += chunk_count + + def on_deferred_release_done( + self, tier_idx: int, chunk_count: int, hold_duration: float + ) -> None: + state = self._tier_states[tier_idx] + state.deferred_release_chunk_count -= chunk_count + assert state.deferred_release_chunk_count >= 0 + self._stats.observe_histogram( + TieringOffloadingMetrics.DEFERRED_RELEASE_HOLD, + hold_duration, + self.tier_label(tier_idx), + ) + def on_promotion_allocation_failure(self) -> None: self._stats.increase_counter( TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES @@ -184,6 +200,7 @@ def assert_idle(self) -> None: and state.active_cascade_count == 0 and state.primary_write_chunk_count == 0 and state.primary_read_chunk_count == 0 + and state.deferred_release_chunk_count == 0 for state in self._tier_states ) @@ -290,6 +307,11 @@ def _observe_active_transfer_stats(self, stats: OffloadingConnectorStats) -> Non state.active_cascade_count, labelvalues, ) + stats.set_gauge( + TieringOffloadingMetrics.DEFERRED_RELEASE_CHUNKS, + state.deferred_release_chunk_count, + labelvalues, + ) def _observe_resolved_lookup( self, diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 3aec216885b1..4914d7b8172d 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -247,6 +247,27 @@ def build_metric_definitions( ), labelnames=(TIER_LABEL,), ) + metrics[TieringOffloadingMetrics.DEFERRED_RELEASE_CHUNKS] = ( + OffloadingGaugeMetadata( + documentation=( + "Number of primary-tier chunks of failed promotions whose release " + "to the free pool is deferred until the tier confirms its I/O can " + "no longer write to them, labeled by tier." + ), + labelnames=(TIER_LABEL,), + ) + ) + metrics[TieringOffloadingMetrics.DEFERRED_RELEASE_HOLD] = ( + OffloadingHistogramMetadata( + documentation=( + "Histogram of how long a failed promotion's primary-tier " + "chunks were held before their deferred release completed, " + "observed at release, labeled by tier, in seconds." + ), + labelnames=(TIER_LABEL,), + buckets=(0.01, 0.05, 0.1, 0.5, 1, 2.5, 5, 10, 20, 30, 60, 120), + ) + ) metrics[TieringOffloadingMetrics.PROMOTION_LATENCY] = ( OffloadingHistogramMetadata( documentation=( From ed6b8f4cd2b47e4240e7ff5feb5105444245b0c2 Mon Sep 17 00:00:00 2001 From: Yaoyao Ding Date: Wed, 28 Oct 2026 09:53:34 -0400 Subject: [PATCH 061/406] [Bugfix][DSv4.1] Keep the specialized decode's FlashInfer paths on newer FlashInfer (#63914) Signed-off-by: Yaoyao Ding Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .../nvidia/specialized_decode/test_kernels.py | 2 +- .../specialized_decode/test_numerics.py | 20 +++++------ .../nvidia/specialized_decode/AGENTS.md | 22 +++++++----- .../nvidia/specialized_decode/__init__.py | 20 ++++++----- .../specialized_decode/attention/mxfp8.py | 12 +++---- .../attention/sparse_mla.py | 35 ++++++++++++++++--- 6 files changed, 71 insertions(+), 40 deletions(-) diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py index 12ecc14a1303..c3d1deb741b4 100644 --- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_kernels.py @@ -435,7 +435,7 @@ def test_sparse_mla_persistent_kv_counters(dtype, query_lens): """The decode on the shared KV counters equals the decode on fresh ones bitwise, launch after launch, and leaves the counters zeroed.""" if sparse_mla.kv_counter_buffer(1) is None: - pytest.skip("FlashInfer is not the version the persistent counters need") + pytest.skip("FlashInfer's sparse-MLA decode takes no KV counter buffer") torch.manual_seed(0) heads, head_dim, page, top_k = 64, 512, 64, 1024 tokens = sum(query_lens) diff --git a/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py b/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py index 3442d862e491..cc6445c837f0 100644 --- a/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py +++ b/tests/models/deepseek_v41/nvidia/specialized_decode/test_numerics.py @@ -18,13 +18,12 @@ * one CUDA graph of all of these, replayed ``REPLAYS`` times bitwise. Needs four SM100-family GPUs with NVLink multicast and flashinfer-python -0.7.2; about 5 minutes. For a race stress run, run this file as a module -(``python -m tests.models.deepseek_v41.nvidia.specialized_decode.test_numerics``) +0.7.2 or newer; about 5 minutes. For a race stress run, run this file as a +module (``python -m tests.models.deepseek_v41.nvidia.specialized_decode.test_numerics``) with ``--replays 100000``. """ import argparse -import importlib.metadata import os import tempfile from types import SimpleNamespace @@ -809,14 +808,13 @@ def run(replays: int) -> str | None: def test_specialized_decode_numerics(): if torch.accelerator.device_count() < TP: pytest.skip("needs four GPUs") - from vllm.models.deepseek_v41.nvidia.specialized_decode import FLASHINFER_VERSION - - try: - version = importlib.metadata.version("flashinfer-python") - except importlib.metadata.PackageNotFoundError: - version = None - if version != FLASHINFER_VERSION: - pytest.skip(f"needs flashinfer-python {FLASHINFER_VERSION}, found {version}") + from vllm.models.deepseek_v41.nvidia.specialized_decode import ( + MIN_FLASHINFER_VERSION, + uses_checked_flashinfer, + ) + + if not uses_checked_flashinfer("the numerics test"): + pytest.skip(f"needs flashinfer-python >= {MIN_FLASHINFER_VERSION}") skipped = run(REPLAYS) if skipped: pytest.skip(skipped) diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md index 0986fd485687..724b3aee63c1 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/AGENTS.md @@ -36,22 +36,28 @@ collectives and would hang. Which forward steps use the executor depends only on the batch, which every TP rank receives identically, so that needs no collective. -Four features run only on flashinfer-python 0.7.2 (`FLASHINFER_VERSION`, the -version `requirements/cuda.txt` pins). Three use FlashInfer internals: +Two features use FlashInfer internals and run on flashinfer-python 0.7.2 or +newer (`MIN_FLASHINFER_VERSION`): - the publishing `wo_b` GEMM: the CuTe DSL low-latency block-scaled GEMM kernel, and `mm_mxfp8`'s private runner and tuning config (`flashinfer.gemm.gemm_base`) to look up its tuned tactic; - the fused kernel's MXFP8 layer input: the CuTe DSL MXFP8 quantize helpers (`flashinfer.quantization.quantization_cute_dsl_utils`, - `flashinfer.cute_dsl.fp4_common`); -- the pinned MXFP8 GEMM tactics: the low-latency GEMM's `autotune_tactics`. + `flashinfer.cute_dsl.fp4_common`). -The fourth, the sparse-MLA decode's persistent KV counters, passes the -`multi_ctas_kv_counter_buffer` argument that 0.7.1 added. +On an older version each logs a warning once and the executor goes without +it. `test_numerics.py` compares both with FlashInfer's own GEMM and quantize +bitwise; run it when `requirements/cuda.txt` moves to a new FlashInfer. -On any other version each logs a warning once and the executor goes without -it; the MXFP8 GEMMs then tune over every tactic. +Two more depend on FlashInfer's interface, not its version: + +- the pinned MXFP8 GEMM tactics (GB300) wrap the low-latency GEMM's + `autotune_tactics`; without it, or when it does not list a pinned tactic, + the GEMMs tune over every tactic; +- the sparse-MLA decode's persistent KV counters need the + `multi_ctas_kv_counter_buffer` argument (0.7.1 added it); without it each + decode allocates its own. ## Structure diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py index 1d93d5e64f7d..a08a284cdac5 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/__init__.py @@ -16,6 +16,7 @@ import torch import torch.distributed as dist +from packaging.version import Version import vllm.envs as envs from vllm.config import CompilationMode, CUDAGraphMode, VllmConfig @@ -37,10 +38,11 @@ # Rows of the low-batch GEMMs: one DSpark verification step of up to 8 rows # (6 at one request). LOW_BATCH_MAX_ROWS = 8 -# The FlashInfer (requirements/cuda.txt) whose internals the publishing wo_b -# GEMM, the MXFP8 layer input and the pinned MXFP8 tactics were checked -# against. Its sparse-MLA decode takes caller-owned KV counters (since 0.7.1). -FLASHINFER_VERSION = "0.7.2" +# The oldest FlashInfer whose internals the publishing wo_b GEMM and the MXFP8 +# layer input were checked against. Newer releases run them too: test_numerics.py +# compares both with FlashInfer's own GEMM and quantize bitwise on the release +# requirements/cuda.txt pins. +MIN_FLASHINFER_VERSION = "0.7.2" @dataclass(frozen=True) @@ -54,19 +56,19 @@ class DecodeConfig: @cache def uses_checked_flashinfer(feature: str) -> bool: - """Whether the installed FlashInfer is the version whose internals - ``feature`` uses; logs once per feature when it is not.""" + """Whether the installed FlashInfer is at least the version whose internals + ``feature`` was checked against; logs once per feature when it is not.""" try: version = importlib.metadata.version("flashinfer-python") except importlib.metadata.PackageNotFoundError: version = None - if version == FLASHINFER_VERSION: + if version is not None and Version(version) >= Version(MIN_FLASHINFER_VERSION): return True logger.warning( - "DeepSeek-V4.1 specialized decode: %s needs flashinfer-python %s " + "DeepSeek-V4.1 specialized decode: %s needs flashinfer-python >= %s " "(found %s); using the stock path instead.", feature, - FLASHINFER_VERSION, + MIN_FLASHINFER_VERSION, version, ) return False diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py index 4014147ba0be..18a79fb02758 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/mxfp8.py @@ -33,7 +33,7 @@ from vllm.platforms import current_platform from vllm.utils import flashinfer as vllm_flashinfer -from .. import LOW_BATCH_MAX_ROWS, uses_checked_flashinfer +from .. import LOW_BATCH_MAX_ROWS BACKEND = "cutedsl_low_latency" @@ -134,12 +134,10 @@ def run_zeros(linear: nn.Module, rows: int) -> torch.Tensor: @contextmanager def pinned_tactics() -> Iterator[None]: """Offer FlashInfer's tuner only the pinned tactic for pinned problems. The - tuner's tactic list is a FlashInfer internal: off the checked version, - every tactic is tuned.""" - if not ( - current_platform.is_device_capability(103) - and uses_checked_flashinfer("pinning the MXFP8 GEMM tactics") - ): + tuner's tactic list is a FlashInfer internal: when it has no + ``autotune_tactics`` or does not list a pinned tactic, every tactic is + tuned.""" + if not current_platform.is_device_capability(103): yield return module = importlib.import_module( diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py index 383093c27ec6..a0e68f271032 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/attention/sparse_mla.py @@ -18,11 +18,14 @@ backend. """ +import inspect +from functools import cache from typing import TYPE_CHECKING, cast import torch from vllm.forward_context import get_forward_context +from vllm.logger import init_logger from vllm.models.deepseek_v41.sparse_mla import DeepseekV4FlashMLAMetadata from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 @@ -30,7 +33,8 @@ DeepseekV4FlashInferMLAAttention, _get_flashinfer_dsv4_workspace, ) -from .. import uses_checked_flashinfer + +logger = init_logger(__name__) if TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -41,17 +45,40 @@ _kv_counters_by_device: dict[torch.device, torch.Tensor] = {} +@cache +def _takes_kv_counters() -> bool: + """Whether the installed FlashInfer's sparse-MLA decode takes caller-owned + KV counters (``multi_ctas_kv_counter_buffer``, since 0.7.1); logs once when + it does not.""" + try: + import flashinfer.utils + from flashinfer.decode import trtllm_batch_decode_sparse_mla_dsv4 + except ImportError: + return False + if "multi_ctas_kv_counter_buffer" in inspect.signature( + trtllm_batch_decode_sparse_mla_dsv4 + ).parameters and hasattr( + flashinfer.utils, "get_trtllm_gen_multi_ctas_kv_counter_bytes" + ): + return True + logger.warning( + "DeepSeek-V4.1 specialized decode: the installed FlashInfer's sparse-MLA " + "decode takes no KV counter buffer; each decode allocates its own." + ) + return False + + def kv_counter_buffer(max_rows: int) -> torch.Tensor | None: """Zeroed multi-CTA KV counters on the current device for decode launches - of up to ``max_rows`` rows (requests x heads), or None off the checked - FlashInfer version. + of up to ``max_rows`` rows (requests x heads), or None when FlashInfer's + decode takes no counter buffer. Allocated when the executor is built, outside any CUDA graph capture. Every layer shares the buffer, as it shares the workspace: their decode launches run one after another on one stream. Decode launches that overlap on one device would race on it. """ - if not uses_checked_flashinfer("the sparse-MLA decode's persistent KV counters"): + if not _takes_kv_counters(): return None from flashinfer.utils import ( get_device_sm_count, From df028dffa548710fb77320a5c08b87ae58757933 Mon Sep 17 00:00:00 2001 From: Tarun Kumar Date: Wed, 28 Oct 2026 19:46:22 +0530 Subject: [PATCH 062/406] Ship mamba SSU and Helion tuned kernel configs in wheel package_data (#64063) --- setup.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/setup.py b/setup.py index 8859d0a93671..ae9f72e4cf3a 100644 --- a/setup.py +++ b/setup.py @@ -1624,6 +1624,12 @@ def get_rocm_device_requirements() -> list[str]: "py.typed", "libs/*.so*", "model_executor/layers/fused_moe/configs/*.json", + # Tuned kernel configs must be explicit: the setuptools-scm git + # file-finder is inactive when building from a .git-less tree + # (e.g. tarball/sdist builds), which silently drops them from + # the wheel and forces fallback to default launch configs. + "model_executor/layers/mamba/ops/configs/selective_state_update/*.json", + "kernels/helion/configs/*/*.json", "model_executor/layers/quantization/utils/configs/*.json", "entrypoints/serve/instrumentator/static/*.js", "entrypoints/serve/instrumentator/static/*.css", From 6cfb375fd6a2c7aee76b086c19911646711834a5 Mon Sep 17 00:00:00 2001 From: CharlesWu Date: Wed, 28 Oct 2026 23:01:33 +0800 Subject: [PATCH 063/406] [UMBP][5/10] Add MORI embedded DRAM runtime and transfer handling (#61793) Signed-off-by: CharlesWu Co-authored-by: Codex Co-authored-by: TJian Co-authored-by: vllm-agent --- docs/features/README.md | 4 + docs/features/umbp_store_connector_usage.md | 75 ++ .../kv_connector/unit/test_umbp_embedded.py | 125 +++ .../kv_transfer/kv_connector/factory.py | 5 + .../kv_connector/v1/umbp/runtime/__init__.py | 4 + .../kv_connector/v1/umbp/runtime/embedded.py | 880 ++++++++++++++++++ 6 files changed, 1093 insertions(+) create mode 100644 docs/features/umbp_store_connector_usage.md create mode 100644 tests/v1/kv_connector/unit/test_umbp_embedded.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/embedded.py diff --git a/docs/features/README.md b/docs/features/README.md index 1df8960ee354..1618bee0fe69 100644 --- a/docs/features/README.md +++ b/docs/features/README.md @@ -1,5 +1,9 @@ # Features +## KV Cache Offloading + +- [UMBPStoreConnector](umbp_store_connector_usage.md) + ## Compatibility Matrix The tables below show mutually exclusive features and the support on some hardware. diff --git a/docs/features/umbp_store_connector_usage.md b/docs/features/umbp_store_connector_usage.md new file mode 100644 index 000000000000..fa6e1af69ba9 --- /dev/null +++ b/docs/features/umbp_store_connector_usage.md @@ -0,0 +1,75 @@ +# UMBPStoreConnector Usage Guide + +`UMBPStoreConnector` offloads KV cache blocks to host DRAM managed by UMBP, +the unified memory buffer pool of [MORI](https://github.com/ROCm/mori). It +restores them when a later request shares a prefix that has left GPU memory. + +## Requirements + +- MORI built with UMBP (`BUILD_UMBP=ON`), importable as `mori`. +- Prefix caching enabled. + +## Embedded mode + +Embedded mode is the default. Each vLLM worker owns a private DRAM pool in its +own process. The engine that stored a prefix can restore it; independent +engines do not share these pools. + +```bash +vllm serve \ + --enable-prefix-caching \ + --kv-transfer-config '{ + "kv_connector": "UMBPStoreConnector", + "kv_role": "kv_both", + "kv_load_failure_policy": "recompute", + "kv_connector_extra_config": { + "mode": "embedded", + "total_capacity_bytes": 68719476736 + } + }' +``` + +### Embedded options + +Set these options in `kv_connector_extra_config`: + +| Option | Default | Meaning | +| --- | --- | --- | +| `capacity_bytes` | 64 GiB | Pool size per worker. | +| `total_capacity_bytes` | | Pool size per engine, split evenly across its workers. Exclusive with `capacity_bytes`. | +| `lookup_instance` | | Distinguishes independent engines serving the same model on one host. Data-parallel ranks are already distinguished. | +| `lookup_dir` | `/tmp` | Directory of the per-worker lookup sockets. | +| `dram_high_watermark`, `dram_low_watermark` | MORI defaults | Pool occupancy at which eviction starts and stops. | +| `dram_use_hugepages`, `dram_hugepage_size`, `dram_numa_node`, `dram_prefault` | MORI defaults | Allocation of each worker's pool. | + +## Shared options + +Set these options in `kv_connector_extra_config`: + +| Option | Default | Meaning | +| --- | --- | --- | +| `load_async` | `true` | Load in the background while the request waits. Models with Mamba-style or sparse-attention indexer layers always load asynchronously, since those layers read restored state before a synchronous load could have landed. | +| `lookup_async` | `false` | Look up the pool off the scheduler thread; the request waits until the lookup completes. | +| `kv_event_max_objects` | `65536` | Maximum objects tracked for KV events. When full, new stored-object notifications are suppressed until tracked objects leave the pool. Storage and lookup continue normally. | +| `enable_lookup` | `true` | Restore from the pool. With `false` the connector only stores. | +| `save_decode_cache` | `false` | Also store blocks produced by decoding. | +| `enable_partial_hash_hits` | `false` | Allow hits that end inside a block, at core's prefix-match unit (`cache_config.prefix_match_unit`). | +| `num_workers`, `timeout_ms` | `4`, `30000` | Transfer threads per worker, and the transfer timeout. A synchronous wait that times out fails the step rather than reusing its blocks. | +| `key_namespace` | derived | Overrides the namespace derived from the model, revision, KV layout, cache groups, parallel topology and draft model. | + +Use `kv_load_failure_policy: "recompute"` in `kv_transfer_config`, as in the +example above. An object can be evicted between lookup and load; the default +`fail` policy fails the request instead of recomputing it. + +With `kv_role: "kv_consumer"` the connector only restores and never stores. + +## Limitations + +- Pipeline parallelism is rejected. +- Stores are submitted once per step after the forward pass. + +With KV events enabled, the scheduler verifies residency on every rank before +publishing a stored-object event. It periodically rechecks tracked objects +during engine steps to report removals. These checks do not delay the release +of GPU source blocks after a completed transfer. With events disabled, no +historical object identities are retained for event reporting. diff --git a/tests/v1/kv_connector/unit/test_umbp_embedded.py b/tests/v1/kv_connector/unit/test_umbp_embedded.py new file mode 100644 index 000000000000..3d79e6c9463f --- /dev/null +++ b/tests/v1/kv_connector/unit/test_umbp_embedded.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import tempfile +import threading +from types import SimpleNamespace + +import pytest + +from vllm.distributed.kv_transfer.kv_connector.v1.umbp.data import ( + BlockTransferPlan, + RankTopology, + TransferJobStatus, +) +from vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime import embedded +from vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime.embedded import ( + _lookup_socket_path, + _MoriLookupServer, + _MoriSchedulerHandle, + _MoriWorkerHandle, +) + + +@pytest.fixture +def blocked_native_load(monkeypatch, tmp_path): + started = threading.Event() + release = threading.Event() + finished = threading.Event() + clock = SimpleNamespace(now=100.0) + monkeypatch.setattr(embedded.time, "monotonic", lambda: clock.now) + + def get(keys, *args): + started.set() + assert release.wait(5) + return [True] * len(keys) + + handle = _MoriWorkerHandle( + SimpleNamespace(batch_get_ranges_into_ptr=get, close=lambda: None), + "blocked-load", + RankTopology(), + str(tmp_path), + 1, + 1, + ) + try: + job = handle.load([BlockTransferPlan("running", 0)]) + future = handle._futures[id(job)] + future.add_done_callback(lambda _: finished.set()) + assert started.wait(5) + yield handle, job, future, release, finished, clock + finally: + release.set() + handle.close() + + +def test_async_load_poll_timeout_keeps_buffers_owned(blocked_native_load): + handle, job, future, release, finished, clock = blocked_native_load + assert handle.poll(job) is None + clock.now += 2 + with pytest.raises(TimeoutError, match="buffers still in use"): + handle.poll(job) + assert job.status is TransferJobStatus.RUNNING + assert handle._futures[id(job)] is future + assert not future.done() + + release.set() + assert finished.wait(5) + assert handle.poll(job) is job + assert job.status is TransferJobStatus.COMPLETED + assert id(job) not in handle._futures + + +def test_load_timeout_does_not_count_time_waiting_for_a_worker(blocked_native_load): + handle, _, _, release, _, clock = blocked_native_load + queued = handle.load([BlockTransferPlan("queued", 1)]) + finished = threading.Event() + handle._futures[id(queued)].add_done_callback(lambda _: finished.set()) + clock.now += 2 + assert handle.poll(queued) is None + assert queued.status is TransferJobStatus.RUNNING + + release.set() + assert finished.wait(5) + assert handle.poll(queued) is queued + assert queued.status is TransferJobStatus.COMPLETED + + +@pytest.fixture +def socket_dir(): + """Keep AF_UNIX paths short regardless of CI's TMPDIR.""" + with tempfile.TemporaryDirectory(prefix="umbp-", dir="/tmp") as directory: + yield directory + + +def test_lookup_server_refuses_a_socket_another_engine_serves(socket_dir): + client = SimpleNamespace(batch_exists=lambda keys: [True] * len(keys)) + path = _lookup_socket_path("taken", (0, 0, 0, 0), socket_dir, "dp0") + first = _MoriLookupServer(path, client) + second = _MoriLookupServer(path, client) + first.start() + try: + with pytest.raises(RuntimeError, match="another live engine"): + second.start() + second.close() + scheduler = _MoriSchedulerHandle("taken", RankTopology(), socket_dir, "dp0") + assert scheduler.lookup(["k"]) == [True] + finally: + second.close() + first.close() + + +def test_lookup_server_close_does_not_unlink_a_replacement(socket_dir): + client = SimpleNamespace(batch_exists=lambda keys: [True] * len(keys)) + path = _lookup_socket_path("replaced", (0, 0, 0, 0), socket_dir) + first = _MoriLookupServer(path, client) + first.start() + first.close() + second = _MoriLookupServer(path, client) + second.start() + try: + first.close() + scheduler = _MoriSchedulerHandle("replaced", RankTopology(), socket_dir) + assert scheduler.lookup(["k"]) == [True] + finally: + second.close() diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index 2eefa15ebeb2..c0f0a4e6600f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -237,6 +237,11 @@ def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool: "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.connector", "MooncakeStoreConnector", ) +KVConnectorFactory.register_connector( + "UMBPStoreConnector", + "vllm.distributed.kv_transfer.kv_connector.v1.umbp", + "UMBPStoreConnector", +) KVConnectorFactory.register_connector( "FlexKVConnectorV1", "vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py index 7b2ce33ebf0d..24ac28628bf3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/__init__.py @@ -2,12 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from .base import IUMBPRuntime, UMBPSchedulerHandle, UMBPWorkerHandle +from .embedded import EmbeddedRuntime from .factory import UMBPRuntimeConfig, UMBPRuntimeFactory +UMBPRuntimeFactory.register("embedded", EmbeddedRuntime.from_config) + __all__ = [ "IUMBPRuntime", "UMBPSchedulerHandle", "UMBPWorkerHandle", "UMBPRuntimeConfig", "UMBPRuntimeFactory", + "EmbeddedRuntime", ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/embedded.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/embedded.py new file mode 100644 index 000000000000..818e9ab37cef --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/embedded.py @@ -0,0 +1,880 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MORI-backed embedded DRAM offloading.""" + +from __future__ import annotations + +import contextlib +import hashlib +import os +import socket +import threading +import time +from collections.abc import Iterable, Sequence +from concurrent.futures import Future, ThreadPoolExecutor, TimeoutError +from functools import lru_cache +from pathlib import Path +from typing import Any, NamedTuple, Protocol + +import msgspec +import numpy as np +import regex as re +import torch + +from vllm.logger import init_logger + +from ..data import ( + BlockLoadBatch, + BlockStoreBatch, + BlockTransferPlan, + KVLayoutDescriptor, + RankTopology, + TransferJobState, + TransferJobStatus, +) +from .base import IUMBPRuntime, UMBPSchedulerHandle, UMBPWorkerHandle +from .factory import UMBPRuntimeConfig + +logger = init_logger(__name__) + +# Bounds one scheduler-to-worker lookup round trip on the local socket. +_LOOKUP_TIMEOUT_S = 1.0 + +# MORI evicts asynchronously once the pool crosses its high watermark, so a +# put into a full pool can fail until eviction catches up. +_STORE_RETRIES = 2 +_STORE_RETRY_BACKOFF_S = 0.01 + +_RANK_KEY_PATTERN = re.compile(r":tp(\d+):pcp(\d+):dcp(\d+):pp(\d+):g\d+:") + + +@lru_cache(maxsize=256) +def _rank_namespace_from_key_prefix(prefix: str) -> tuple[int, int, int, int] | None: + match = _RANK_KEY_PATTERN.search(prefix) + if match is None: + return None + tp_rank, pcp_rank, dcp_rank, pp_rank = map(int, match.groups()) + return tp_rank, pcp_rank, dcp_rank, pp_rank + + +class _LookupClient(Protocol): + def batch_exists(self, keys: Sequence[str]) -> Sequence[bool]: ... + + def clear(self) -> bool: ... + + +class _MoriClient(_LookupClient, Protocol): + def register_memory(self, *args: Any) -> bool: ... + + def deregister_memory(self, *args: Any) -> bool: ... + + def batch_get_ranges_into_ptr(self, *args: Any) -> Sequence[bool]: ... + + def batch_put_ranges_from_ptr(self, *args: Any) -> list[bool]: ... + + def flush(self) -> bool: ... + + +def _lookup_socket_path( + namespace: str, + rank_namespace: tuple[int, int, int, int], + lookup_dir: str, + instance: str = "", +) -> str: + # The namespace identifies the model and layout, not the engine: DP ranks + # and independent engines serving one model share it, so the engine's own + # identity must separate their private pools. + identity = f"{namespace}:{rank_namespace}" + if instance: + identity = f"{namespace}:{instance}:{rank_namespace}" + digest = hashlib.sha256(identity.encode()).hexdigest()[:24] + return str(Path(lookup_dir) / f"vllm-umbp-{digest}.sock") + + +def _socket_has_listener(path: str) -> bool: + probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + probe.settimeout(0.5) + try: + probe.connect(path) + except OSError: + return False + finally: + probe.close() + return True + + +class _MoriLookupServer: + """Small worker-local lookup bridge for the scheduler process.""" + + def __init__(self, path: str, client: _LookupClient) -> None: + self.path = path + self.client = client + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._socket: socket.socket | None = None + self._owns_path = False + + def start(self) -> None: + Path(self.path).parent.mkdir(parents=True, exist_ok=True) + if os.path.exists(self.path) and _socket_has_listener(self.path): + raise RuntimeError( + f"UMBP lookup socket {self.path} is served by another live " + "engine; set lookup_instance to a distinct value for each " + "engine that serves this model on this host" + ) + with contextlib.suppress(FileNotFoundError): + os.unlink(self.path) + self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self._socket.bind(self.path) + self._owns_path = True + self._socket.listen(16) + self._socket.settimeout(0.2) + self._thread = threading.Thread( + target=self._serve, + name="umbp-embedded-lookup", + daemon=True, + ) + self._thread.start() + + def _serve(self) -> None: + assert self._socket is not None + while not self._stop.is_set(): + try: + connection, _ = self._socket.accept() + except TimeoutError: + continue + except OSError: + break + with connection: + connection.settimeout(_LOOKUP_TIMEOUT_S) + try: + request = b"" + while not request.endswith(b"\n"): + chunk = connection.recv(65536) + if not chunk: + break + request += chunk + payload = msgspec.json.decode(request) + if isinstance(payload, dict) and payload.get("op") == "clear": + connection.sendall( + msgspec.json.encode(bool(self.client.clear())) + b"\n" + ) + continue + result = self.client.batch_exists(payload) + connection.sendall( + msgspec.json.encode([bool(value) for value in result]) + b"\n" + ) + except Exception: + with contextlib.suppress(OSError): + connection.sendall(b"[]\n") + + def close(self) -> None: + self._stop.set() + if self._socket is not None: + self._socket.close() + if self._thread is not None and self._thread.ident is not None: + self._thread.join(timeout=2) + if self._owns_path: + with contextlib.suppress(FileNotFoundError): + os.unlink(self.path) + self._owns_path = False + + +class _MoriSchedulerHandle(UMBPSchedulerHandle): + def __init__( + self, + namespace: str, + topology: RankTopology, + lookup_dir: str, + instance: str = "", + ) -> None: + self._paths = { + rank: _lookup_socket_path(namespace, rank, lookup_dir, instance) + for rank in topology.all_namespaces() + } + self.last_lookup_diagnostics: dict[str, tuple[Any, ...]] = {} + + @staticmethod + def _request(path: str, payload: Any, timeout: float = _LOOKUP_TIMEOUT_S) -> Any: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + sock.settimeout(timeout) + sock.connect(path) + sock.sendall(msgspec.json.encode(payload) + b"\n") + response = b"" + while not response.endswith(b"\n"): + chunk = sock.recv(65536) + if not chunk: + break + response += chunk + return msgspec.json.decode(response or b"[]") + + def lookup(self, keys: Sequence[str]) -> Sequence[bool]: + unavailable: list[tuple[int, int, int, int]] + if len(self._paths) == 1: + rank, path = next(iter(self._paths.items())) + try: + values = self._request(path, keys) + if len(values) != len(keys): + raise ValueError("lookup returned an invalid result length") + return self._lookup_result(keys, [bool(value) for value in values], []) + except (OSError, ValueError, IndexError): + unavailable = ( + [rank] + if any( + _rank_namespace_from_key_prefix(key.rpartition(":")[0] + ":") + == rank + for key in keys + ) + else [] + ) + return self._lookup_result(keys, [False] * len(keys), unavailable) + result = [False] * len(keys) + keys_by_rank: dict[tuple[int, int, int, int], list[tuple[int, str]]] = {} + unrouted: list[tuple[int, str]] = [] + for index, key in enumerate(keys): + key_rank = _rank_namespace_from_key_prefix(key.rpartition(":")[0] + ":") + if key_rank is None or key_rank not in self._paths: + unrouted.append((index, key)) + else: + keys_by_rank.setdefault(key_rank, []).append((index, key)) + unavailable = [] + for rank, indexed_keys in keys_by_rank.items(): + path = self._paths[rank] + try: + values = self._request(path, [key for _, key in indexed_keys]) + if len(values) != len(indexed_keys): + raise ValueError("lookup returned an invalid result length") + for (index, _), value in zip(indexed_keys, values, strict=True): + result[index] = bool(value) + except (OSError, ValueError, IndexError): + unavailable.append(rank) + for path in self._paths.values(): + if not unrouted: + break + try: + values = self._request(path, [key for _, key in unrouted]) + for (index, _), value in zip(unrouted, values, strict=True): + result[index] = result[index] or bool(value) + except (OSError, ValueError): + continue + return self._lookup_result(keys, result, unavailable) + + def _lookup_result( + self, + keys: Sequence[str], + result: list[bool], + unavailable: list[tuple[int, int, int, int]], + ) -> list[bool]: + missing = tuple( + key for key, exists in zip(keys, result, strict=True) if not exists + ) + self.last_lookup_diagnostics = { + "unavailable_ranks": tuple(unavailable), + "missing_keys": missing, + } + if unavailable: + logger.warning("UMBP lookup unavailable for ranks: %s", unavailable) + elif missing: + logger.debug("UMBP lookup missing %d objects", len(missing)) + return result + + def close(self) -> None: + return + + def clear(self) -> bool: + success = True + contacted = False + for path in self._paths.values(): + try: + cleared = self._request(path, {"op": "clear"}, timeout=2.0) + contacted = True + success = success and bool(cleared) + except (OSError, ValueError): + success = False + return contacted and success + + +class _BlockLoadLayout(NamedTuple): + bases: np.ndarray + strides: np.ndarray + sizes: list[int] + offsets: list[int] + num_blocks: int + num_bytes: int + + +_LoadArgs = tuple[list[str], list[list[int]], list[list[int]], list[list[int]]] +_StoreArgs = tuple[ + list[str], list[int], list[list[int]], list[list[int]], list[list[int]] +] + + +class _MoriWorkerHandle(UMBPWorkerHandle): + def __init__( + self, + client: _MoriClient, + namespace: str, + topology: RankTopology, + lookup_dir: str, + max_workers: int, + timeout_s: float, + layout: KVLayoutDescriptor | None = None, + lookup_instance: str = "", + ) -> None: + self.client = client + self._namespace = namespace + self._topology = topology + self._lookup_dir = lookup_dir + self._lookup_instance = lookup_instance + self._lookup_server: _MoriLookupServer | None = None + self._registered_storages: set[int] = set() + self._gpu_devices: set[int] = set() + self._executor = ThreadPoolExecutor( + max_workers=max_workers, + thread_name_prefix="umbp-embedded-transfer", + ) + self._job_lock = threading.Lock() + self._futures: dict[int, Future[TransferJobState]] = {} + self._store_deadlines: dict[int, float] = {} + self._load_deadlines: dict[int, float] = {} + self._timeout_s = timeout_s + self._layout = layout + self._load_layouts: dict[int, _BlockLoadLayout] = {} + + def register_buffers(self, kv_caches: dict[str, torch.Tensor]) -> None: + try: + from mori.cpp import MemoryLocationType + except ImportError as exc: + raise RuntimeError("MORI UMBP Python bindings are unavailable") from exc + + for cache in kv_caches.values(): + storage = cache.untyped_storage() + storage_ptr = storage.data_ptr() + if storage_ptr in self._registered_storages: + continue + location = ( + MemoryLocationType.GPU + if cache.device.type == "cuda" + else MemoryLocationType.CPU + ) + device = cache.device.index if cache.device.index is not None else -1 + if not self.client.register_memory( + storage_ptr, storage.nbytes(), location, device + ): + raise RuntimeError( + f"MORI UMBP failed to register KV storage 0x{storage_ptr:x}" + ) + self._registered_storages.add(storage_ptr) + if location == MemoryLocationType.GPU and device >= 0: + self._gpu_devices.add(device) + + self._register_load_layouts(kv_caches) + if self._lookup_server is None: + self._lookup_server = _MoriLookupServer( + _lookup_socket_path( + self._namespace, + self._topology.local_namespace, + self._lookup_dir, + self._lookup_instance, + ), + self, + ) + self._lookup_server.start() + + def _register_load_layouts(self, kv_caches: dict[str, torch.Tensor]) -> None: + self._load_layouts = {} + if self._layout is None: + return + for group_id in {r.group_id for r in self._layout.regions}: + bases: list[int] = [] + strides: list[int] = [] + sizes: list[int] = [] + offsets: list[int] = [] + counts: list[int] = [] + for region in self._layout.regions: + if region.group_id != group_id: + continue + cache = kv_caches.get(region.layer_name) + if cache is None or cache.ndim == 0: + break + stride = cache.stride(0) * cache.element_size() + if ( + stride <= 0 + or stride > region.block_stride + or region.block_stride % stride + or region.block_bytes > stride + ): + break + bases.append(cache.data_ptr()) + strides.append(stride) + offsets.append(sum(sizes)) + sizes.append(region.block_bytes) + counts.append(cache.shape[0]) + else: + self._load_layouts[group_id] = _BlockLoadLayout( + np.asarray(bases, dtype=np.uint64), + np.asarray(strides, dtype=np.uint64), + sizes, + offsets, + min(counts), + sum(sizes), + ) + + def _bulk_load_args( + self, plans: Sequence[BlockTransferPlan] + ) -> tuple[_LoadArgs, tuple[int, ...]] | None: + groups: dict[int, list[tuple[int, int]]] = {} + entries: Iterable[tuple[int | None, int]] + if isinstance(plans, (BlockLoadBatch, BlockStoreBatch)): + entries = zip(plans.group_ids, plans.block_ids, strict=True) + keys = plans.keys + else: + if any(p.ranges for p in plans): + return None + entries = ((p.group_id, p.block_id) for p in plans) + keys = [p.key for p in plans] + for index, (group_id, block_id) in enumerate(entries): + if group_id is None: + return None + layout = self._load_layouts.get(group_id) + if layout is None or not 0 <= block_id < layout.num_blocks: + return None + groups.setdefault(group_id, []).append((index, block_id)) + pointers: list[list[int]] = [[]] * len(plans) + sizes: list[list[int]] = [[]] * len(plans) + offsets: list[list[int]] = [[]] * len(plans) + plan_bytes = [0] * len(plans) + for group_id, group_entries in groups.items(): + layout = self._load_layouts[group_id] + ids = np.asarray([block for _, block in group_entries], dtype=np.uint64) + addresses = ( + layout.bases[None, :] + ids[:, None] * layout.strides[None, :] + ).tolist() + for (index, _), row in zip(group_entries, addresses, strict=True): + pointers[index] = row + sizes[index] = layout.sizes + offsets[index] = layout.offsets + plan_bytes[index] = layout.num_bytes + return ( + (keys, pointers, sizes, offsets), + tuple(plan_bytes), + ) + + def load_blocks( + self, plans: Sequence[BlockTransferPlan] + ) -> TransferJobState | None: + prepared = self._bulk_load_args(plans) + if prepared is None: + return None + args, plan_bytes = prepared + return self._submit_load( + plans if isinstance(plans, BlockLoadBatch) else tuple(plans), + args, + plan_bytes, + ) + + def batch_exists(self, keys: Sequence[str]) -> Sequence[bool]: + return [bool(value) for value in self.client.batch_exists(keys)] + + def clear(self) -> bool: + return bool(self.client.clear()) + + @staticmethod + def _range_args(plans: Sequence[BlockTransferPlan], *, for_store: bool = True): + keys, object_sizes, pointers, sizes, offsets = [], [], [], [], [] + for plan in plans: + plan_pointers, plan_sizes, plan_offsets = [], [], [] + object_size = 0 + for item in plan.ranges: + plan_pointers.append(item.base_address) + plan_sizes.append(item.length) + plan_offsets.append(item.object_offset) + if for_store: + end = item.object_offset + item.length + if end > object_size: + object_size = end + keys.append(plan.key) + if for_store: + object_sizes.append(object_size) + pointers.append(plan_pointers) + sizes.append(plan_sizes) + offsets.append(plan_offsets) + return keys, object_sizes, pointers, sizes, offsets + + def load(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState: + return self._submit_load(tuple(plans)) + + def _submit_load( + self, + plans: tuple[BlockTransferPlan, ...] | BlockLoadBatch, + range_args: _LoadArgs | None = None, + plan_bytes: tuple[int, ...] | None = None, + ) -> TransferJobState: + job = TransferJobState(plans, plan_bytes=plan_bytes) + job.start() + future = self._executor.submit(self._load_sync, job, plans, range_args) + with self._job_lock: + self._futures[id(job)] = future + return job + + def _load_sync( + self, + job: TransferJobState, + plans: tuple[BlockTransferPlan, ...] | BlockLoadBatch, + range_args: _LoadArgs | None = None, + ) -> TransferJobState: + with self._job_lock: + self._load_deadlines[id(job)] = time.monotonic() + self._timeout_s + if not plans: + job.complete() + return job + if range_args is None: + keys, _, pointers, sizes, offsets = self._range_args(plans, for_store=False) + else: + keys, pointers, sizes, offsets = range_args + request_id = ( + plans.request_id + if isinstance(plans, BlockLoadBatch) + else plans[0].request_id + ) + started_at = time.monotonic() + logger.debug( + "MORI UMBP range load started request=%s plans=%d ranges=%d bytes=%d", + request_id, + len(plans), + sum(len(plan_sizes) for plan_sizes in sizes), + sum(sum(plan_sizes) for plan_sizes in sizes), + ) + results = self.client.batch_get_ranges_into_ptr(keys, pointers, sizes, offsets) + logger.debug( + "MORI UMBP range load returned request=%s plans=%d elapsed=%.3fs", + request_id, + len(plans), + time.monotonic() - started_at, + ) + completed = [key for key, ok in zip(keys, results, strict=True) if ok] + failed = [key for key, ok in zip(keys, results, strict=True) if not ok] + if completed: + job.complete(completed) + if failed: + job.fail(failed, "MORI UMBP range load failed") + return job + + def store_blocks( + self, plans: Sequence[BlockTransferPlan] + ) -> TransferJobState | None: + """Store whole blocks from registered layouts; None requests ranges.""" + prepared = self._bulk_load_args(plans) + if prepared is None: + return None + (keys, pointers, sizes, offsets), plan_bytes = prepared + return self._submit_store( + plans if isinstance(plans, BlockStoreBatch) else tuple(plans), + (keys, list(plan_bytes), pointers, sizes, offsets), + plan_bytes, + ) + + def store(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState: + return self._submit_store(tuple(plans)) + + def _submit_store( + self, + plans: tuple[BlockTransferPlan, ...] | BlockStoreBatch, + range_args: _StoreArgs | None = None, + plan_bytes: tuple[int, ...] | None = None, + ) -> TransferJobState: + job = TransferJobState(plans, plan_bytes=plan_bytes) + job.start() + ready_events: list[torch.Event] = [] + for device in self._gpu_devices: + with torch.device(f"cuda:{device}"): + event = torch.Event() + event.record(torch.accelerator.current_stream()) + ready_events.append(event) + future = self._executor.submit( + self._store_sync, job, plans, ready_events, range_args + ) + with self._job_lock: + self._futures[id(job)] = future + return job + + def _store_sync( + self, + job: TransferJobState, + plans: tuple[BlockTransferPlan, ...] | BlockStoreBatch, + ready_events: list[torch.Event], + range_args: _StoreArgs | None = None, + ) -> TransferJobState: + # Timed from here: a store queued behind others has not stalled. + with self._job_lock: + self._store_deadlines[id(job)] = time.monotonic() + self._timeout_s + for event in ready_events: + event.synchronize() + if not plans: + job.complete() + return job + keys, object_sizes, pointers, sizes, offsets = ( + self._range_args(plans) if range_args is None else range_args + ) + results = self.client.batch_put_ranges_from_ptr( + keys, object_sizes, pointers, sizes, offsets + ) + failed_indices = [index for index, ok in enumerate(results) if not ok] + for attempt in range(_STORE_RETRIES): + if not failed_indices: + break + self.client.flush() + time.sleep(_STORE_RETRY_BACKOFF_S * (attempt + 1)) + for index in failed_indices.copy(): + retry = self.client.batch_put_ranges_from_ptr( + [keys[index]], + [object_sizes[index]], + [pointers[index]], + [sizes[index]], + [offsets[index]], + ) + if retry and retry[0]: + results[index] = True + failed_indices.remove(index) + if failed_indices: + logger.warning( + "MORI UMBP store retry %d left %d/%d objects pending", + attempt + 1, + len(failed_indices), + len(plans), + ) + if not self.client.flush(): + job.fail([plan.key for plan in plans], "MORI UMBP flush failed") + return job + completed = [plan.key for plan, ok in zip(plans, results, strict=True) if ok] + failed = [plan.key for plan, ok in zip(plans, results, strict=True) if not ok] + if completed: + job.complete(completed) + if failed: + job.fail(failed, "MORI UMBP range store failed") + return job + + def wait(self, job: TransferJobState) -> TransferJobState: + with self._job_lock: + future = self._futures.get(id(job)) + if future is None: + return job + try: + return future.result(timeout=self._timeout_s) + except TimeoutError as exc: + if not future.done(): + # The transfer still owns GPU pointers. Abort this engine step + # rather than reporting completion and allowing block reuse. + raise TimeoutError( + "MORI UMBP transfer timed out with buffers still in use" + ) from exc + job.fail([plan.key for plan in job.plans], str(exc)) + return job + except Exception as exc: + job.fail([plan.key for plan in job.plans], str(exc)) + return job + finally: + if future.done(): + with self._job_lock: + self._futures.pop(id(job), None) + self._store_deadlines.pop(id(job), None) + self._load_deadlines.pop(id(job), None) + + def poll(self, job: TransferJobState) -> TransferJobState | None: + with self._job_lock: + future = self._futures.get(id(job)) + if future is None: + if job.status in ( + TransferJobStatus.COMPLETED, + TransferJobStatus.FAILED, + TransferJobStatus.CANCELLED, + ): + return job + return None + if not future.done(): + with self._job_lock: + deadline = self._store_deadlines.get( + id(job), self._load_deadlines.get(id(job)) + ) + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError( + "MORI UMBP transfer timed out with buffers still in use" + ) + return None + with self._job_lock: + self._futures.pop(id(job), None) + self._store_deadlines.pop(id(job), None) + self._load_deadlines.pop(id(job), None) + try: + return future.result() + except Exception as exc: + job.fail([plan.key for plan in job.plans], str(exc)) + return job + + def publish(self, job: TransferJobState) -> None: + if job.status.value != "completed": + raise RuntimeError("cannot publish an incomplete MORI UMBP job") + + def cancel(self, job: TransferJobState) -> TransferJobState: + with self._job_lock: + future = self._futures.get(id(job)) + if future is None: + return job + if future.cancel(): + with self._job_lock: + self._futures.pop(id(job), None) + self._store_deadlines.pop(id(job), None) + self._load_deadlines.pop(id(job), None) + job.cancel("preempted") + return job + return self.wait(job) + + def take_evicted_keys(self) -> Sequence[str]: + return () + + def close(self) -> None: + self._executor.shutdown(wait=True, cancel_futures=True) + if self._lookup_server is not None: + self._lookup_server.close() + for storage_ptr in self._registered_storages: + self.client.deregister_memory(storage_ptr) + self._registered_storages.clear() + close = getattr(self.client, "close", None) + if callable(close): + close() + else: + self.client.flush() + + +def _configure_dram(client_config: Any, options: dict[str, Any]) -> None: + """Apply embedded-only host-DRAM options to MORI's client config.""" + dram = client_config.dram + dram.capacity_bytes = options.get("capacity_bytes", 64 * 1024**3) + option_fields = { + "dram_use_shared_memory": "use_shared_memory", + "dram_shm_name": "shm_name", + "dram_high_watermark": "high_watermark", + "dram_low_watermark": "low_watermark", + "dram_use_hugepages": "use_hugepages", + "dram_hugepage_size": "hugepage_size", + "dram_numa_node": "numa_node", + "dram_prefault": "prefault", + } + for option, field in option_fields.items(): + if option in options: + setattr(dram, field, options[option]) + + +def _resolve_embedded_options(config: UMBPRuntimeConfig) -> dict[str, Any]: + options = dict(config.options) + if "capacity_bytes" in options and "total_capacity_bytes" in options: + raise ValueError( + "capacity_bytes and total_capacity_bytes are mutually exclusive" + ) + for name in ("capacity_bytes", "total_capacity_bytes", "dram_hugepage_size"): + if name in options and (type(options[name]) is not int or options[name] <= 0): + raise ValueError(f"{name} must be a positive integer") + for name in ("dram_use_shared_memory", "dram_use_hugepages", "dram_prefault"): + if name in options and type(options[name]) is not bool: + raise ValueError(f"{name} must be a boolean") + for name in ("dram_high_watermark", "dram_low_watermark"): + if name in options and ( + type(options[name]) not in (int, float) or not 0 < options[name] <= 1 + ): + raise ValueError(f"{name} must be in (0, 1]") + low, high = options.get("dram_low_watermark"), options.get("dram_high_watermark") + if low is not None and high is not None and low > high: + raise ValueError("dram_low_watermark must not exceed dram_high_watermark") + if "dram_numa_node" in options and ( + type(options["dram_numa_node"]) is not int or options["dram_numa_node"] < -1 + ): + raise ValueError("dram_numa_node must be an integer >= -1") + if "dram_shm_name" in options and ( + not isinstance(options["dram_shm_name"], str) or not options["dram_shm_name"] + ): + raise ValueError("dram_shm_name must be a non-empty string") + if { + "master_address", + "node_address", + "io_engine_host", + "peer_service_port", + } & options.keys(): + raise ValueError("embedded UMBP cannot configure distributed-only options") + total = options.pop("total_capacity_bytes", None) + if total is not None: + if total < config.rank_count: + raise ValueError( + "total_capacity_bytes must provide at least one byte per rank" + ) + options["capacity_bytes"] = total // config.rank_count + options["_configured_total_capacity_bytes"] = total + return options + + +class EmbeddedRuntime(IUMBPRuntime): + """MORI-backed embedded DRAM runtime.""" + + def __init__(self, config: UMBPRuntimeConfig) -> None: + if config.options.get("backend") == "memory": + raise ValueError( + "The memory backend is test-only; use MORI for embedded mode" + ) + self.options = _resolve_embedded_options(config) + self.lookup_dir = self.options.get("lookup_dir", "/tmp") + instance = self.options.get("lookup_instance", "") + if not isinstance(instance, str): + raise ValueError("lookup_instance must be a string") + self.lookup_instance = f"{instance}:dp{config.dp_index}" + + def create_scheduler_handle( + self, + namespace: str, + topology: RankTopology, + layout: KVLayoutDescriptor, + ) -> UMBPSchedulerHandle: + del layout + return _MoriSchedulerHandle( + namespace, topology, self.lookup_dir, self.lookup_instance + ) + + def create_worker_handle( + self, + namespace: str, + topology: RankTopology, + layout: KVLayoutDescriptor, + ) -> UMBPWorkerHandle: + try: + from mori.cpp import UMBPClient, UMBPConfig + except ImportError as exc: + raise RuntimeError( + "Embedded UMBP requires MORI built with BUILD_UMBP=ON" + ) from exc + + client_config = UMBPConfig() + _configure_dram(client_config, self.options) + total = self.options.get("_configured_total_capacity_bytes") + if total is None: + logger.info( + "Embedded UMBP DRAM capacity: %.2f GiB per rank", + client_config.dram.capacity_bytes / 1024**3, + ) + else: + logger.info( + "Embedded UMBP DRAM capacity: %.2f GiB total, %.2f GiB per rank", + total / 1024**3, + client_config.dram.capacity_bytes / 1024**3, + ) + return _MoriWorkerHandle( + UMBPClient(client_config), + namespace, + topology, + self.lookup_dir, + int(self.options.get("num_workers", 4)), + float(self.options.get("timeout_ms", 30000)) / 1000, + layout, + lookup_instance=self.lookup_instance, + ) + + @classmethod + def from_config(cls, config: UMBPRuntimeConfig) -> IUMBPRuntime: + return cls(config) From a1cfd2c1b584c4ac197db9cc9ac84a6d973d39df Mon Sep 17 00:00:00 2001 From: liranschour Date: Wed, 28 Oct 2026 17:09:05 +0200 Subject: [PATCH 064/406] [Bugfix][KV Offload][P2P] Retain unconfirmed source WRITEs across session teardown (#64250) Signed-off-by: Liran Schour --- .../tiering/p2p/test_data_transport.py | 71 ++++++++ .../v1/kv_offload/tiering/p2p/test_manager.py | 171 ++++++++++++++++++ .../kv_offload/tiering/p2p/test_sessions.py | 37 ++++ vllm/v1/kv_offload/tiering/p2p/data/base.py | 27 ++- vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 66 ++++++- vllm/v1/kv_offload/tiering/p2p/manager.py | 109 ++++++++++- .../kv_offload/tiering/p2p/session/server.py | 25 ++- .../kv_offload/tiering/p2p/session/session.py | 11 +- 8 files changed, 500 insertions(+), 17 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py index 46822f369030..f3d60e8079db 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -909,6 +909,77 @@ def test_raising_state_check_warns_once(self, monkeypatch): transport.poll(peer_id="peer:1") assert sum("check_xfer_state failed" in w for w in warnings) == 1 + def test_remove_remote_peer_retires_running_transfer(self): + """A WRITE still running at peer removal keeps its handle and the + peer's registration until it stops, and leaves every poll() so a + reconnected peer under the same id never sees it.""" + transport = self._make_transport() + agent = transport._agent + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + tid = transport.write_blocks("peer:1", [0], [1]) + + transport.remove_remote_peer("peer:1") + agent.release_dlist_handle.assert_not_called() + agent.remove_remote_agent.assert_not_called() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + agent.check_xfer_state.return_value = "DONE" + assert not transport.poll(peer_id="peer:1").done + assert not transport.poll().done + assert tid in transport._inflight + + assert transport.poll_retired() == [tid] + assert tid not in transport._inflight + agent.release_xfer_handle.assert_called_once() + # The old registration only; the new one stays registered. + agent.release_dlist_handle.assert_called_once() + agent.remove_remote_agent.assert_called_once() + assert "peer:1" in transport._remote_dlists + + @pytest.mark.parametrize("state", ["PROC", "PEND"]) + def test_poll_retired_keeps_running_transfer(self, state): + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + tid = transport.write_blocks("peer:1", [0], [1]) + transport.remove_remote_peer("peer:1") + transport._agent.check_xfer_state.return_value = state + + assert transport.poll_retired() == [] + assert tid in transport._inflight + transport._agent.release_xfer_handle.assert_not_called() + + def test_poll_retired_state_check_raise_is_not_terminal(self): + """Only a remote disconnect proves a retired WRITE stopped.""" + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + tid = transport.write_blocks("peer:1", [0], [1]) + transport.remove_remote_peer("peer:1") + transport._agent.check_xfer_state.side_effect = nixlNotFoundError("gone") + assert transport.poll_retired() == [] + transport._agent.check_xfer_state.side_effect = _remote_disconnect() + assert transport.poll_retired() == [tid] + assert "peer:1" not in transport._peer_faults + + def test_force_release_retired_releases_registration(self): + transport = self._make_transport() + agent = transport._agent + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + tid = transport.write_blocks("peer:1", [0], [1]) + assert tid is not None + transport.remove_remote_peer("peer:1") + + transport.force_release_retired([tid]) + + assert tid not in transport._inflight + agent.release_xfer_handle.assert_called_once() + agent.remove_remote_agent.assert_called_once() + + def test_remove_remote_peer_without_inflight_releases_at_once(self): + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + transport.remove_remote_peer("peer:1") + transport._agent.release_dlist_handle.assert_called_once() + transport._agent.remove_remote_agent.assert_called_once() + def test_remove_remote_peer_clears_the_fault(self): transport = self._make_transport() for _ in range(3): diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index c47010b70790..56611497bd4b 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -39,6 +39,7 @@ from vllm.v1.kv_offload.tiering.p2p import manager as manager_module from vllm.v1.kv_offload.tiering.p2p.manager import ( _RECONNECT_MAX_BACKOFF_S, + _RETAINED_HOLD_CAP_S, _UNBOUND_STORE_TIMEOUT_S, P2PSecondaryTierManager, _annotate_req_context, @@ -139,6 +140,8 @@ def _make_manager() -> P2PSecondaryTierManager: mgr._reaped_stores = {} mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S mgr._failed_serve_ctxs = [] + mgr._retained = {} + mgr._retained_force_released = 0 return mgr @@ -727,6 +730,7 @@ def __init__( close_stores: list[int] | None = None, close_failed_serves: list[ReqContext] | None = None, last_activity_at: float | None = None, + close_retained: dict[int, frozenset[int]] | None = None, ) -> None: self.peer_id = peer_id self.alive = alive @@ -741,6 +745,7 @@ def __init__( self._close_req_ids = close_req_ids or [] self._close_stores = close_stores or [] self._close_failed_serves = close_failed_serves or [] + self._close_retained = close_retained or {} self.requests: list[tuple[int, str]] = [] self.stores_added: list[tuple[str, list, object, int]] = [] self.attached: list[object] = [] @@ -788,6 +793,7 @@ def close(self): failed_req_ids=self._close_req_ids, failed_stores=self._close_stores, failed_serves=self._close_failed_serves, + retained=self._close_retained, ) @@ -2325,6 +2331,171 @@ def test_disabled_when_zero(self, monkeypatch): # --------------------------------------------------------------------------- +class _RetiringFakeData: + """Fake DataTransport whose retired transfers stop on demand.""" + + def __init__(self) -> None: + self.stopped: list[int] = [] + self.force_released: list[int] = [] + self.removed: list[str] = [] + + def remove_remote_peer(self, pid) -> None: + self.removed.append(pid) + + def poll_retired(self) -> list[int]: + stopped, self.stopped = self.stopped, [] + return stopped + + def force_release_retired(self, transfer_ids) -> None: + self.force_released.extend(transfer_ids) + + +class TestRetainedTransfers: + """A reaped session's still-running transfers keep their source pins + until the transport reports them stopped, or the hold cap expires.""" + + PEER = "10.0.0.1:8000" + + def _reap(self, retained, close_stores=None): + mgr = _make_manager() + data = _RetiringFakeData() + mgr._data = data # type: ignore[assignment] + mgr._sessions[self.PEER] = _FakeSession( # type: ignore[assignment] + peer_id=self.PEER, + alive=False, + connected=True, + close_stores=close_stores, + close_retained=retained, + ) + mgr._reap_dead_sessions() + return mgr, data + + def test_reap_holds_pins_until_transfer_stops(self): + mgr, data = self._reap({7: frozenset({1})}, close_stores=[2]) + assert mgr._finished_jobs == [JobResult(job_id=2, success=False)] + assert data.removed == [self.PEER] + + mgr._poll_retained() + assert mgr._finished_jobs == [JobResult(job_id=2, success=False)] + + data.stopped = [7] + mgr._poll_retained() + assert mgr._finished_jobs[-1] == JobResult(job_id=1, success=False) + assert mgr._retained == {} + + def test_shared_job_held_until_last_transfer_stops(self): + mgr, data = self._reap({7: frozenset({1}), 8: frozenset({1})}) + data.stopped = [7] + mgr._poll_retained() + assert mgr._finished_jobs == [] + data.stopped = [8] + mgr._poll_retained() + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] + + def test_hold_cap_force_releases(self, monkeypatch): + now = [1000.0] + monkeypatch.setattr(manager_module.time, "monotonic", lambda: now[0]) + mgr, data = self._reap({7: frozenset({1})}) + + now[0] += _RETAINED_HOLD_CAP_S - 1 + mgr._poll_retained() + assert data.force_released == [] + + now[0] += 1 + mgr._poll_retained() + assert data.force_released == [7] + assert mgr._finished_jobs == [JobResult(job_id=1, success=False)] + assert mgr._retained_force_released == 1 + assert mgr._retained == {} + + def test_idle_eviction_retains_too(self, monkeypatch): + monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10") + monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1") + import vllm.envs + + vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None) + vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None) + mgr = _make_manager() + mgr._data = _RetiringFakeData() # type: ignore[assignment] + mgr._sessions[self.PEER] = _FakeSession( # type: ignore[assignment] + peer_id=self.PEER, + last_activity_at=time.monotonic() - 20, + close_retained={7: frozenset({1})}, + ) + mgr._reap_idle_sessions() + assert self.PEER not in mgr._sessions + assert mgr._retained[7][:2] == (self.PEER, frozenset({1})) + assert mgr._finished_jobs == [] + + def test_shutdown_fails_retained_jobs(self): + mgr, _ = self._reap({7: frozenset({1})}) + mgr._data = _ShutdownFakeData() # type: ignore[assignment] + mgr._control = _ShutdownFakeControl() # type: ignore[assignment] + mgr.shutdown() + assert JobResult(job_id=1, success=False) in mgr._finished_jobs + assert mgr._retained == {} + + def test_outbound_waits_for_retained_transfers(self): + """A new session would load the peer's NIXL agent under the same + name while the reaped one's registration still holds it.""" + mgr, data = self._reap({7: frozenset({1})}) + + class NoConnect: + def connect(self, peer_id): + raise AssertionError("connect() while transfers are retained") + + mgr._control = NoConnect() # type: ignore[assignment] + assert mgr._get_or_create_session(self.PEER) is None + + data.stopped = [7] + mgr._poll_retained() + assert not mgr._retained_peer(self.PEER) + + def test_inbound_is_told_to_disconnect_while_retained(self): + mgr, data = self._reap({7: frozenset({1})}) + + class Conn(_RecordingConn): + def __init__(self, peer_id: str) -> None: + super().__init__(peer_id) + self.sent: list[dict] = [] + + def send(self, msg: dict) -> None: + self.sent.append(msg) + + conn = Conn(self.PEER) + mgr._accept_new_peers([conn]) # type: ignore[list-item] + assert conn.sent == [{TYPE_KEY: DisconnectMsg.TYPE}] + assert conn.close_calls == 1 + assert self.PEER not in mgr._sessions + + data.stopped = [7] + mgr._poll_retained() + assert not mgr._retained_peer(self.PEER) + + def test_other_peers_are_not_held(self): + mgr, _ = self._reap({7: frozenset({1})}) + assert not mgr._retained_peer("10.0.0.2:8000") + + def test_drain_jobs_waits_for_retained_transfers(self, monkeypatch): + """reset_cache() asserts every job finished once drain_jobs() + returns, so a retained transfer's store jobs must be settled.""" + mgr, data = self._reap({7: frozenset({1})}) + mgr._control = _DrainCtrl() # type: ignore[assignment] + polls = [0] + + def poll_retired(): + polls[0] += 1 + return [7] if polls[0] == 3 else [] + + data.poll_retired = poll_retired # type: ignore[method-assign] + monkeypatch.setattr(manager_module.time, "sleep", lambda s: None) + mgr.drain_jobs() + + assert polls[0] == 3 + assert mgr._retained == {} + assert JobResult(job_id=1, success=False) in mgr._finished_jobs + + class TestReconnectBackoff: """A peer whose session ended on a data-plane fault reconnects at once the first time; repeats back off from 1 s up to _RECONNECT_MAX_BACKOFF_S.""" diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 7856867d6a96..368039003ce6 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -2984,6 +2984,43 @@ def test_close_returns_pending_loads_and_stores(self): assert set(result.failed_stores) == {10} assert result.failed_serves == [] + def test_close_retains_running_transfer(self): + """close() must not release a WRITE that may still be landing. + + It cancels in wait mode; a transfer still running is reported as + retained with its store job, which stays pinned, and only the jobs + no transfer reads are returned as failed. + """ + session, conn, data = _make_session() + _activate(session, conn) + session.add_stored_blocks("req", [OffloadKey(b"key")], [0], job_id=1) + session.add_stored_blocks("idle", [OffloadKey(b"other")], [1], job_id=2) + _fetch(conn, "req", [b"key"], [5]) + session.poll() + (tid,) = data._transfers + data._cancel_still_inflight.add(tid) + + result = session.close() + + assert data._cancel_calls == [([tid], "wait")] + assert tid in data._transfers + assert result.retained == {tid: frozenset({1})} + assert result.failed_stores == [2] + + def test_close_releases_stopped_transfer(self): + """A transfer that has already stopped is released, not retained.""" + session, conn, data = _make_session() + _activate(session, conn) + session.add_stored_blocks("req", [OffloadKey(b"key")], [0], job_id=1) + _fetch(conn, "req", [b"key"], [5]) + session.poll() + + result = session.close() + + assert data._transfers == {} + assert result.retained == {} + assert result.failed_stores == [1] + def test_send_failure_marks_connection_dead(self): """A raising send must mark the connection dead, not silently drop the message — otherwise the session lingers alive, is never reaped, diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py index 30456ca2e22d..4a8ca7c10101 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py @@ -65,6 +65,14 @@ state check failed stays inflight unless the failure proves it can no longer write, so the peer's teardown is what bounds it. +remove_remote_peer() does not stop a transfer that is still running: the +transport retires it instead. A retired transfer leaves every poll(), so +a reconnected peer under the same id never sees it; poll_retired() +reports it once it has stopped, and only then are its handle and the +removed peer's registration released. The caller keeps the transfer's +source pins until then, or until it gives up and calls +force_release_retired(). + The transport never recovers a peer itself. The session turns a peer fault into the same connection death as a missed heartbeat, so the manager's reap is the only teardown: DisconnectMsg makes the peer tear @@ -261,10 +269,27 @@ def add_remote_peer_async( def remove_remote_peer(self, peer_id: str) -> None: """Unregister a remote peer and release associated resources. - Inflight transfers to this peer should be cancelled first. + Transfers to this peer still inflight are retired, not released + (see "Failure model"); the peer's registration is released once + they have stopped. """ ... + def poll_retired(self) -> list[int]: + """Release retired transfers that have stopped; return their ids. + + Must be called periodically while any transfer is retired. The + default transport retires nothing. + """ + return [] + + def force_release_retired(self, transfer_ids: Iterable[int]) -> None: + """Release retired transfers that have not stopped. + + Unsafe while one may still be writing; for a bounded hold only. + """ + self.cancel(transfer_ids, mode="immediate") + @abstractmethod def write_blocks( self, diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index 382ed7576527..c62dab63eb22 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -105,11 +105,18 @@ class _Inflight(NamedTuple): transport is shared across all peer sessions of the engine. ``struck`` is set once the transfer has counted against its peer, so a state check that keeps raising charges it once, not once per poll. + ``retired`` is set once its peer is removed while it still runs; only + poll_retired() sees it from then on. """ peer_id: str handle: object struck: bool = False + retired: bool = False + + +# A removed peer's (peer_id, dlist, nixl_name, retired transfer_ids). +_DeferredRemote = tuple[str, object | None, str | None, set[int]] class NixlTransport(DataTransport): @@ -144,6 +151,9 @@ def __init__( self._peer_strikes: dict[str, int] = {} # peer_id → fault cause; reported by poll(peer_id) until removal. self._peer_faults: dict[str, str] = {} + # Remote handles of removed peers that retired transfers still + # reference. Released by poll_retired() once those transfers are gone. + self._deferred_remote: list[_DeferredRemote] = [] # Registration is O(num_blocks) NIXL work — add_remote_agent plus a # prep_xfer_dlist over every block of the peer's region — so it runs # off the scheduler thread, mirroring the NIXL connector's @@ -263,6 +273,17 @@ def remove_remote_peer(self, peer_id: str) -> None: dlist = self._remote_dlists.pop(peer_id, None) self._peer_strikes.pop(peer_id, None) self._peer_faults.pop(peer_id, None) + # Transfers still inflight here may still be writing. Retire them so + # a reconnected session under the same peer_id never sees them, and + # keep the remote handles they were prepped against until they end. + retired: set[int] = set() + for tid, entry in self._inflight.items(): + if entry.peer_id == peer_id and not entry.retired: + self._inflight[tid] = entry._replace(retired=True) + retired.add(tid) + if retired: + self._deferred_remote.append((peer_id, dlist, nixl_name, retired)) + return if self._agent is not None: self._release_remote_handles(peer_id, dlist, nixl_name) @@ -446,7 +467,7 @@ def poll(self, peer_id: str | None = None) -> PollResult: failures: list[tuple[str, str]] | None = None for transfer_id, entry in self._inflight.items(): - if peer_id is not None and entry.peer_id != peer_id: + if entry.retired or (peer_id is not None and entry.peer_id != peer_id): continue try: state = self._agent.check_xfer_state(entry.handle) @@ -515,6 +536,48 @@ def _poll_result( peer_fault=fault, ) + def poll_retired(self) -> list[int]: + """Release retired transfers that have stopped; return their ids. + + A transfer stopped once NIXL reports it terminal, by the same rule + as ``cancel(mode="wait")``, but without charging its peer: the + peer is gone, and a reconnected one under the same id is not at + fault. A removed peer's remote handles are released once none of + its retired transfers remain. + """ + stopped: list[int] = [] + for tid, entry in self._inflight.items(): + if not entry.retired: + continue + try: + state = self._agent.check_xfer_state(entry.handle) + except Exception as exc: + if not _is_terminal_state_error(exc): + continue + state = "ERR" + if state not in ("PROC", "PEND"): + stopped.append(tid) + if not stopped: + return [] + self._release_handles([self._inflight.pop(tid).handle for tid in stopped]) + self._release_deferred_remote() + return stopped + + def force_release_retired(self, transfer_ids: Iterable[int]) -> None: + """Release retired transfers that never stopped (bounded hold).""" + self.cancel(transfer_ids, mode="immediate") + self._release_deferred_remote() + + def _release_deferred_remote(self) -> None: + pending: list[_DeferredRemote] = [] + for peer_id, dlist, nixl_name, tids in self._deferred_remote: + tids &= self._inflight.keys() + if tids: + pending.append((peer_id, dlist, nixl_name, tids)) + elif self._agent is not None: + self._release_remote_handles(peer_id, dlist, nixl_name) + self._deferred_remote = pending + def cancel( self, transfer_ids: Iterable[int], @@ -594,6 +657,7 @@ def close(self) -> None: return self._release_handles([entry.handle for entry in self._inflight.values()]) self._inflight.clear() + self._release_deferred_remote() for peer_id in list(self._remote_dlists): self.remove_remote_peer(peer_id) if self._local_dlist is not None: diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 52aab4bfd266..696dde3bf2ee 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -71,6 +71,13 @@ # enough to avoid busy-spinning the scheduler thread. _DRAIN_SLEEP_S = 0.001 +# Longest a reaped session's still-running transfer keeps its source pins. +# NIXL may never report a wedged transfer terminal, and a hold without a +# bound would look like a full cache; past this the handle is released. +# The peer cannot reconnect while it holds retained transfers (see +# _retained_peer), so this also bounds how long it stays unreachable. +_RETAINED_HOLD_CAP_S = 30.0 + # Idle sessions are only reaped when remaining capacity is at most this # many slots. Session setup is expensive; below this headroom, keep idle # peers so a later burst does not pay handshake again. @@ -400,6 +407,10 @@ def __init__( # at the top of the next ``serve_external_requests`` where the # handle is valid. self._failed_serve_ctxs: list[ReqContext] = [] + # transfer_id → (peer_id, store job_ids, retained_at) for transfers + # of reaped sessions that may still be writing; see _poll_retained. + self._retained: dict[int, tuple[str, frozenset[int], float]] = {} + self._retained_force_released = 0 # ------------------------------------------------------------------ # SecondaryTierManager interface @@ -657,9 +668,11 @@ def drain_jobs(self) -> None: """Block until every submitted load/store job has completed or failed. Loops calling ``_poll_once()`` until no session has outstanding - inbound loads or in-flight outbound stores. Mid-flight transfers - are NOT cancelled — the caller (``TieringOffloadingManager.reset_cache``) - needs the primary memoryview to be quiescent, not aborted. Results + inbound loads or in-flight outbound stores, and no reaped session's + transfer is still retained (bounded by its hold cap). Mid-flight + transfers are NOT cancelled — the caller + (``TieringOffloadingManager.reset_cache``) needs the primary + memoryview to be quiescent, not aborted. Results accumulate in ``_finished_jobs`` and are surfaced by the next ``get_finished_jobs()`` call. """ @@ -667,7 +680,9 @@ def drain_jobs(self) -> None: warned = False while True: self._poll_once() - pending = any(s.has_pending_work for s in self._sessions.values()) + pending = bool(self._retained) or any( + s.has_pending_work for s in self._sessions.values() + ) if not pending: return if not warned and time.monotonic() - start > 5.0: @@ -730,7 +745,7 @@ def _get_or_create_session(self, peer_id: str) -> P2PSession | None: session = self._sessions.get(peer_id) if session is not None: return session - if self._reconnect_wait(peer_id) > 0: + if self._reconnect_wait(peer_id) > 0 or self._retained_peer(peer_id): return None max_peers = envs.VLLM_P2P_MAX_PEERS @@ -803,6 +818,17 @@ def _accept_new_peers(self, new_connections: Sequence[ControlConnection]) -> Non conn.send({TYPE_KEY: DisconnectMsg.TYPE}) conn.close() continue + if self._retained_peer(conn.peer_id): + logger.info( + "P2P %s: rejecting inbound peer %s — its reaped " + "session's transfers are still running", + self._local_id, + conn.peer_id, + ) + with contextlib.suppress(Exception): + conn.send({TYPE_KEY: DisconnectMsg.TYPE}) + conn.close() + continue try: existing = self._sessions.get(conn.peer_id) if existing is not None: @@ -852,6 +878,7 @@ def _reap_dead_sessions(self) -> None: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) for job_id in close_result.failed_stores: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + self._retain(pid, close_result.retained) # Fail every client-side request (in-flight loads plus unresolved # symmetric-P2P probes) toward the dead peer so lookup() returns # MISS (local prefill) instead of RETRY forever — even if a fresh @@ -877,6 +904,70 @@ def _reap_dead_sessions(self) -> None: fault_info, ) + def _retain(self, peer_id: str, retained: dict[int, frozenset[int]]) -> None: + """Hold a reaped session's still-running transfers and their pins.""" + if not retained: + return + now = time.monotonic() + for tid, job_ids in retained.items(): + self._retained[tid] = (peer_id, job_ids, now) + logger.warning( + "P2P %s: peer %s reaped with %d transfers still running; holding " + "their source pins until they stop (at most %.0fs)", + self._local_id, + peer_id, + len(retained), + _RETAINED_HOLD_CAP_S, + ) + + def _poll_retained(self) -> None: + """Unpin retained transfers once they stop, or at the hold cap. + + A store job is failed only once no retained transfer reads it. + """ + if not self._retained: + return + released = list(self._data.poll_retired()) + deadline = time.monotonic() - _RETAINED_HOLD_CAP_S + expired = [ + tid + for tid, (_, _, retained_at) in self._retained.items() + if retained_at <= deadline and tid not in released + ] + if expired: + self._data.force_release_retired(expired) + self._retained_force_released += len(expired) + logger.error( + "P2P %s: %d transfers of reaped peers still running after " + "%.0fs; releasing their handles and source pins anyway " + "(%d so far)", + self._local_id, + len(expired), + _RETAINED_HOLD_CAP_S, + self._retained_force_released, + ) + jobs: set[int] = set() + for tid in (*released, *expired): + entry = self._retained.pop(tid, None) + if entry is not None: + jobs |= entry[1] + for job_id in jobs - self._retained_job_ids(): + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + + def _retained_job_ids(self) -> set[int]: + return {jid for _, job_ids, _ in self._retained.values() for jid in job_ids} + + def _retained_peer(self, peer_id: str) -> bool: + """Whether a reaped session of *peer_id* still has running transfers. + + Its registration stays until they stop. A new session would load + the peer's NIXL agent again under the same name: NIXL keeps the old + UCX endpoints, possibly failed, and the old registration's release + removes the agent under the new one. So the peer stays out until + the transfers stop, at most _RETAINED_HOLD_CAP_S. + """ + return any(pid == peer_id for pid, _, _ in self._retained.values()) + def _schedule_reconnect_backoff(self, peer_id: str) -> float: """Record a data-plane fault for *peer_id*; return its reconnect wait. @@ -988,6 +1079,7 @@ def _reap_idle_sessions(self) -> None: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) for job_id in close_result.failed_stores: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + self._retain(pid, close_result.retained) self._failed_req_ids.update(close_result.failed_req_ids) self._failed_serve_ctxs.extend(close_result.failed_serves) self._data.remove_remote_peer(pid) @@ -1066,6 +1158,7 @@ def _poll_once(self) -> None: self._reap_dead_sessions() self._reap_unbound_stores() self._reap_idle_sessions() + self._poll_retained() # ------------------------------------------------------------------ # Lifecycle @@ -1090,6 +1183,11 @@ def shutdown(self) -> None: ) self._unbound_stores.clear() self._reaped_stores.clear() + # The drain above already waited on retained transfers too, and + # _data.close() releases whatever is left. + for job_id in self._retained_job_ids(): + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + self._retained.clear() self._control.close() self._data.close() @@ -1104,6 +1202,7 @@ def _drain_inflight_for_shutdown(self) -> None: wedged peer must not hang shutdown. """ ids = [tid for s in self._sessions.values() for tid in s._server._inflight] + ids.extend(self._retained) if not ids: return deadline = time.monotonic() + _SHUTDOWN_DRAIN_TIMEOUT_S diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index eff19222ea29..0f057f080f11 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -860,19 +860,28 @@ def drain_pending_aborts(self) -> None: for kv_request_id, round_seq in list(self._pending_aborts): self._drain_abort(kv_request_id, round_seq) - def close(self) -> tuple[list[int], list[ReqContext]]: - """Tear down. Cancels inflight. + def close( + self, + ) -> tuple[list[int], list[ReqContext], dict[int, frozenset[int]]]: + """Tear down. Cancels inflight, retaining what may still write. - Returns ``(failed_store_job_ids, failed_serves)`` where + Returns ``(failed_store_job_ids, failed_serves, retained)``. ``failed_serves`` are synthetic lookup ctxs still owing a ``parent.on_request_finished``. The session is going away with no parent handle in hand, so the manager flushes these in its next - ``serve_external_requests``. + ``serve_external_requests``. ``retained`` maps each transfer NIXL + does not report terminal to the store jobs it still reads: like the + abort drain, close never releases such a transfer, so those jobs + stay pinned and are left out of ``failed_store_job_ids`` until the + transport reports the transfer stopped. """ - failed_stores = list(self._store_jobs.keys()) - self._store_jobs.clear() + retained: dict[int, frozenset[int]] = {} if self._inflight: - self._transport.cancel(list(self._inflight.keys())) + still = self._transport.cancel(list(self._inflight.keys()), mode="wait") + retained = {tid: frozenset(self._inflight[tid].job_ids) for tid in still} + held = set().union(*retained.values()) + failed_stores = [jid for jid in self._store_jobs if jid not in held] + self._store_jobs.clear() self._inflight.clear() self._pending_store_results.clear() # Surface every synthetic ctx still owing on_request_finished so @@ -888,7 +897,7 @@ def close(self) -> tuple[list[int], list[ReqContext]]: self._pending_aborts.clear() self._slow_aborts.clear() self._finished_lookup_ctxs.clear() - return failed_stores, failed_serves + return failed_stores, failed_serves, retained @property def has_inflight_transfers(self) -> bool: diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index 3c435bb7d334..d23c306d3b26 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -91,13 +91,17 @@ class SessionCloseResult(NamedTuple): server stores) the manager must mark failed. `failed_req_ids` is every client-side kv_request_id whose lookup() must fail (in-flight loads plus unresolved probes); `failed_serves` is the server-side lookup state the - dead peer can no longer resolve. + dead peer can no longer resolve. `retained` maps each server transfer + that may still be writing to the store jobs it reads: they are not in + `failed_stores`, and stay pinned until the transport reports it stopped. """ failed_jobs: list[int] # client load job_ids failed_req_ids: list[str] # client kv_request_ids (loads + probes) failed_stores: list[int] # server store job_ids failed_serves: list[ReqContext] # server-side lookup ctxs needing release + # transfer_id → store job_ids it still reads + retained: dict[int, frozenset[int]] class P2PSession: @@ -384,6 +388,8 @@ def close(self) -> SessionCloseResult: failed_serves: synthetic lookup ctxs still owing ``parent.on_request_finished`` (the manager flushes these on its next ``serve_external_requests``). + retained: server transfers that may still be writing, with the + store jobs they read (held until they stop). """ # Abandon any in-flight registration: nothing will ack this peer # now. The manager's remove_remote_peer bumps the transport's @@ -394,7 +400,7 @@ def close(self) -> SessionCloseResult: self._pending_registration = None client_result = self._client.close() - failed_stores, failed_serves = self._server.close() + failed_stores, failed_serves, retained = self._server.close() if self._conn is not None: with contextlib.suppress(Exception): @@ -407,6 +413,7 @@ def close(self) -> SessionCloseResult: failed_req_ids=client_result.failed_req_ids, failed_stores=failed_stores, failed_serves=failed_serves, + retained=retained, ) # ------------------------------------------------------------------ From c91ab83dced6f9a5b5227c6f83595abbf4868d6e Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:39:26 +0000 Subject: [PATCH 065/406] [Model] Remove Keye-VL and Keye-VL-1.5 (#64287) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- docs/models/supported_models.md | 2 - .../vision_language_multi_image_offline.py | 76 - .../multimodal/vision_language_offline.py | 76 - tests/model_executor/test_keye_mrope.py | 147 -- tests/model_executor/test_keye_vl1_5_mrope.py | 147 -- .../models/multimodal/generation/test_keye.py | 76 - tests/models/registry.py | 8 - vllm/model_executor/models/keye.py | 1684 ----------------- vllm/model_executor/models/keye_vl1_5.py | 707 ------- vllm/model_executor/models/registry.py | 7 +- 10 files changed, 2 insertions(+), 2928 deletions(-) delete mode 100644 tests/model_executor/test_keye_mrope.py delete mode 100644 tests/model_executor/test_keye_vl1_5_mrope.py delete mode 100644 tests/models/multimodal/generation/test_keye.py delete mode 100644 vllm/model_executor/models/keye.py delete mode 100644 vllm/model_executor/models/keye_vl1_5.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 8520a3347fda..3d6172d8ea9e 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -581,8 +581,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, InternVL 2.0 | T + IE+ + (VE+) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ | | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + IE+ + VE+ | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ | | `KananaVForConditionalGeneration` | Kanana-V | T + I+ | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ | -| `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ | -| `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ | | `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A+ | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ | | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I+ | `moonshotai/Kimi-K2.5` | | ✅︎ | | `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I+ | `moonshotai/Kimi-K3` | | ✅︎ | diff --git a/examples/generate/multimodal/vision_language_multi_image_offline.py b/examples/generate/multimodal/vision_language_multi_image_offline.py index 7565f1e42b51..d829503750dc 100644 --- a/examples/generate/multimodal/vision_language_multi_image_offline.py +++ b/examples/generate/multimodal/vision_language_multi_image_offline.py @@ -454,80 +454,6 @@ def load_internvl(question: str, image_urls: list[str]) -> ModelRequestData: ) -def load_keye_vl(question: str, image_urls: list[str]) -> ModelRequestData: - model_name = "Kwai-Keye/Keye-VL-8B-Preview" - - engine_args = EngineArgs( - model=model_name, - trust_remote_code=True, - max_model_len=8192, - max_num_seqs=5, - limit_mm_per_prompt={"image": len(image_urls)}, - ) - - placeholders = [{"type": "image", "image": url} for url in image_urls] - messages = [ - { - "role": "user", - "content": [ - *placeholders, - {"type": "text", "text": question}, - ], - }, - ] - - processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True) - - prompt = processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - - image_data = [fetch_image(url) for url in image_urls] - - return ModelRequestData( - engine_args=engine_args, - prompt=prompt, - image_data=image_data, - ) - - -def load_keye_vl1_5(question: str, image_urls: list[str]) -> ModelRequestData: - model_name = "Kwai-Keye/Keye-VL-1_5-8B" - - engine_args = EngineArgs( - model=model_name, - trust_remote_code=True, - max_model_len=32768, - max_num_seqs=5, - limit_mm_per_prompt={"image": len(image_urls)}, - ) - - placeholders = [{"type": "image", "image": url} for url in image_urls] - messages = [ - { - "role": "user", - "content": [ - *placeholders, - {"type": "text", "text": question}, - ], - }, - ] - - processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True) - - prompt = processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - - image_data = [fetch_image(url) for url in image_urls] - - return ModelRequestData( - engine_args=engine_args, - prompt=prompt, - image_data=image_data, - ) - - def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData: model_name = "moonshotai/Kimi-VL-A3B-Instruct" @@ -1349,8 +1275,6 @@ def load_molmo2(question: str, image_urls: list[str]) -> ModelRequestData: "idefics3": load_idefics3, "interns1": load_interns1, "internvl_chat": load_internvl, - "keye_vl": load_keye_vl, - "keye_vl1_5": load_keye_vl1_5, "kimi_vl": load_kimi_vl, "llama4": load_llama4, "llava": load_llava, diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index c1448f1cbe87..d6d0a3f3a555 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -935,80 +935,6 @@ def run_kanana_v(questions: list[str], modality: str) -> ModelRequestData: ) -# Keye-VL -def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData: - model_name = "Kwai-Keye/Keye-VL-8B-Preview" - - mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} - engine_args = EngineArgs( - model=model_name, - max_model_len=8192, - trust_remote_code=True, - limit_mm_per_prompt=mm_limit, - ) - - image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>" - video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>" - - if modality == "image": - placeholder = image_placeholder - elif modality == "video": - placeholder = video_placeholder - elif modality == "image+video": - placeholder = image_placeholder + video_placeholder - - prompts = [ - ( - f"<|im_start|>user\n{placeholder}" - f"{question}<|im_end|>\n" - "<|im_start|>assistant\n" - ) - for question in questions - ] - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - -# Keye-VL-1.5 -def run_keye_vl1_5(questions: list[str], modality: str) -> ModelRequestData: - model_name = "Kwai-Keye/Keye-VL-1.5-8B" - - mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1} - engine_args = EngineArgs( - model=model_name, - max_model_len=8192, - trust_remote_code=True, - limit_mm_per_prompt=mm_limit, - ) - - image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>" - video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>" - - if modality == "image": - placeholder = image_placeholder - elif modality == "video": - placeholder = video_placeholder - elif modality == "image+video": - placeholder = image_placeholder + video_placeholder - - prompts = [ - ( - f"<|im_start|>user\n{placeholder}" - f"{question}<|im_end|>\n" - "<|im_start|>assistant\n" - ) - for question in questions - ] - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - # Kimi-VL def run_kimi_vl(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -2235,8 +2161,6 @@ def run_step_vl(questions: list[str], modality: str) -> ModelRequestData: "interns1_pro": run_interns1_pro, "internvl_chat": run_internvl, "kanana_v": run_kanana_v, - "keye_vl": run_keye_vl, - "keye_vl1_5": run_keye_vl1_5, "kimi_vl": run_kimi_vl, "kimi_k25": run_kimi_k25, "lightonocr": run_lightonocr, diff --git a/tests/model_executor/test_keye_mrope.py b/tests/model_executor/test_keye_mrope.py deleted file mode 100644 index 23f497468ed4..000000000000 --- a/tests/model_executor/test_keye_mrope.py +++ /dev/null @@ -1,147 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass, field - -import pytest -import torch - -from vllm.model_executor.models.keye import KeyeForConditionalGeneration -from vllm.multimodal.inputs import ( - MultiModalFeatureSpec, - MultiModalFieldElem, - MultiModalKwargsItem, - PlaceholderRange, -) - -pytestmark = pytest.mark.skip_global_cleanup - - -@pytest.fixture(autouse=True, scope="module") -def _force_cpu_default_device(): - original = torch.get_default_device() - torch.set_default_device("cpu") - yield - torch.set_default_device(original) - - -@dataclass -class DummyVisionConfig: - spatial_merge_size: int = 2 - - -@dataclass -class DummyConfig: - vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig) - - -def make_model(config: DummyConfig) -> KeyeForConditionalGeneration: - model = object.__new__( - KeyeForConditionalGeneration # type: ignore[type-abstract] # protocol attrs unset - ) - model.config = config - return model - - -def make_mm_feature( - *, - modality: str, - offset: int, - length: int, - grid_thw: tuple[int, int, int] | list[tuple[int, int, int]], -) -> MultiModalFeatureSpec: - field_name = "image_grid_thw" if modality == "image" else "video_grid_thw" - return MultiModalFeatureSpec( - data=MultiModalKwargsItem( - { - field_name: MultiModalFieldElem( - data=torch.tensor(grid_thw), - field=None, # HACK. - ), - } - ), - modality=modality, - identifier="DUMMY", - mm_position=PlaceholderRange(offset=offset, length=length), - ) - - -def test_get_mrope_input_positions_text_only(): - model = make_model(DummyConfig()) - - positions, delta = model.get_mrope_input_positions( - input_tokens=[11, 12, 13, 14, 15], - mm_features=[], - ) - - expected = torch.tensor( - [ - [0, 1, 2, 3, 4], - [0, 1, 2, 3, 4], - [0, 1, 2, 3, 4], - ] - ) - - assert torch.equal(positions, expected) - assert delta == 0 - - -def test_get_mrope_input_positions_single_image(): - model = make_model(DummyConfig()) - mm_features = [ - make_mm_feature( - modality="image", - offset=1, - length=4, - grid_thw=(1, 4, 4), - ) - ] - - positions, delta = model.get_mrope_input_positions( - input_tokens=[10, 20, 21, 22, 23, 30, 31], - mm_features=mm_features, - ) - - expected = torch.tensor( - [ - [0, 1, 1, 1, 1, 3, 4], - [0, 1, 1, 2, 2, 3, 4], - [0, 1, 2, 1, 2, 3, 4], - ] - ) - - assert torch.equal(positions, expected) - assert delta == -2 - - -def test_get_mrope_input_positions_interleaved_image_and_video(): - model = make_model(DummyConfig()) - mm_features = [ - make_mm_feature( - modality="image", - offset=1, - length=4, - grid_thw=(1, 4, 4), - ), - make_mm_feature( - modality="video", - offset=7, - length=4, - grid_thw=[(2, 4, 2)], - ), - ] - - positions, delta = model.get_mrope_input_positions( - input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 42, 43, 50, 51], - mm_features=mm_features, - ) - - expected = torch.tensor( - [ - [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 7, 9, 10], - [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 7, 9, 10], - ] - ) - - assert torch.equal(positions, expected) - assert delta == -2 diff --git a/tests/model_executor/test_keye_vl1_5_mrope.py b/tests/model_executor/test_keye_vl1_5_mrope.py deleted file mode 100644 index 8942bee25585..000000000000 --- a/tests/model_executor/test_keye_vl1_5_mrope.py +++ /dev/null @@ -1,147 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass, field - -import pytest -import torch - -from vllm.model_executor.models.keye_vl1_5 import KeyeVL1_5ForConditionalGeneration -from vllm.multimodal.inputs import ( - MultiModalFeatureSpec, - MultiModalFieldElem, - MultiModalKwargsItem, - PlaceholderRange, -) - -pytestmark = pytest.mark.skip_global_cleanup - - -@pytest.fixture(autouse=True, scope="module") -def _force_cpu_default_device(): - original = torch.get_default_device() - torch.set_default_device("cpu") - yield - torch.set_default_device(original) - - -@dataclass -class DummyVisionConfig: - spatial_merge_size: int = 2 - - -@dataclass -class DummyConfig: - vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig) - - -def make_model(config: DummyConfig) -> KeyeVL1_5ForConditionalGeneration: - model = object.__new__( - KeyeVL1_5ForConditionalGeneration # type: ignore[type-abstract] # protocol attrs unset - ) - model.config = config - return model - - -def make_mm_feature( - *, - modality: str, - offset: int, - length: int, - grid_thw: tuple[int, int, int] | list[tuple[int, int, int]], - is_embed: list[bool] | None = None, -) -> MultiModalFeatureSpec: - field_name = "image_grid_thw" if modality == "image" else "video_grid_thw" - return MultiModalFeatureSpec( - data=MultiModalKwargsItem( - { - field_name: MultiModalFieldElem( - data=torch.tensor(grid_thw), - field=None, # HACK. - ), - } - ), - modality=modality, - identifier="DUMMY", - mm_position=PlaceholderRange( - offset=offset, - length=length, - is_embed=None if is_embed is None else torch.tensor(is_embed), - ), - ) - - -def test_get_mrope_input_positions_text_only(): - model = make_model(DummyConfig()) - - positions, delta = model.get_mrope_input_positions( - input_tokens=[11, 12, 13, 14, 15], - mm_features=[], - ) - - expected = torch.tensor( - [ - [0, 1, 2, 3, 4], - [0, 1, 2, 3, 4], - [0, 1, 2, 3, 4], - ] - ) - - assert torch.equal(positions, expected) - assert delta == 0 - - -def test_get_mrope_input_positions_single_image(): - model = make_model(DummyConfig()) - mm_features = [ - make_mm_feature( - modality="image", - offset=1, - length=4, - grid_thw=(1, 4, 4), - ) - ] - - positions, delta = model.get_mrope_input_positions( - input_tokens=[10, 20, 21, 22, 23, 30, 31], - mm_features=mm_features, - ) - - expected = torch.tensor( - [ - [0, 1, 1, 1, 1, 3, 4], - [0, 1, 1, 2, 2, 3, 4], - [0, 1, 2, 1, 2, 3, 4], - ] - ) - - assert torch.equal(positions, expected) - assert delta == -2 - - -def test_get_mrope_input_positions_video_uses_embed_ranges(): - model = make_model(DummyConfig()) - mm_features = [ - make_mm_feature( - modality="video", - offset=1, - length=8, - grid_thw=[(2, 4, 2)], - is_embed=[False, False, True, True, False, False, True, True], - ) - ] - - positions, delta = model.get_mrope_input_positions( - input_tokens=[10, 101, 102, 20, 21, 103, 104, 30, 31, 40, 41], - mm_features=mm_features, - ) - - expected = torch.tensor( - [ - [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10], - [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10], - ] - ) - - assert torch.equal(positions, expected) - assert delta == 0 diff --git a/tests/models/multimodal/generation/test_keye.py b/tests/models/multimodal/generation/test_keye.py deleted file mode 100644 index e9d6626a20d6..000000000000 --- a/tests/models/multimodal/generation/test_keye.py +++ /dev/null @@ -1,76 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import NamedTuple - -import pytest -from PIL.Image import Image -from transformers import AutoProcessor - -from vllm import LLM, EngineArgs, SamplingParams -from vllm.multimodal.utils import encode_image_url - -MODEL_NAME = "Kwai-Keye/Keye-VL-8B-Preview" - -QUESTION = "What is the content of each image?" - - -class ModelRequestData(NamedTuple): - engine_args: EngineArgs - prompt: str - image_data: list[Image] - stop_token_ids: list[int] | None = None - chat_template: str | None = None - sampling_params: SamplingParams | None = None - - -@pytest.mark.parametrize("question", [QUESTION]) -def test_keye_vl(image_assets, question: str): - images = [asset.pil_image for asset in image_assets] - image_urls = [encode_image_url(image) for image in images] - - placeholders = [{"type": "image", "image": url} for url in image_urls] - messages = [ - { - "role": "user", - "content": [ - *placeholders, - {"type": "text", "text": question}, - ], - }, - ] - - processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True) - - prompt = processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - - llm = LLM( - model=MODEL_NAME, - trust_remote_code=True, - max_model_len=8192, - max_num_seqs=5, - limit_mm_per_prompt={"image": len(image_urls)}, - seed=42, - ) - - sampling_params = SamplingParams( - temperature=0.0, max_tokens=256, stop_token_ids=None - ) - - outputs = llm.generate( - { - "prompt": prompt, - "multi_modal_data": {"image": images}, - }, - sampling_params=sampling_params, - ) - - print("-" * 50) - for o in outputs: - generated_text = o.outputs[0].text - print(generated_text) - assert len(generated_text) > 10, ( - f"Generated text is too short: {generated_text}" - ) - print("-" * 50) diff --git a/tests/models/registry.py b/tests/models/registry.py index 8ff1654ed0c9..58e6208ec7f5 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1046,14 +1046,6 @@ def check_available_online( "kakaocorp/kanana-1.5-v-3b-instruct", trust_remote_code=True, ), - "KeyeForConditionalGeneration": _HfExamplesInfo( - "Kwai-Keye/Keye-VL-8B-Preview", - trust_remote_code=True, - ), - "KeyeVL1_5ForConditionalGeneration": _HfExamplesInfo( - "Kwai-Keye/Keye-VL-1_5-8B", - trust_remote_code=True, - ), "MoonshotKimiaForCausalLM": _HfExamplesInfo( "moonshotai/Kimi-Audio-7B-Instruct", tokenizer_mode="kimi_audio", diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py deleted file mode 100644 index 83ff14dcf948..000000000000 --- a/vllm/model_executor/models/keye.py +++ /dev/null @@ -1,1684 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import math -from abc import abstractmethod -from collections.abc import Iterable, Iterator, Mapping, Sequence -from functools import partial -from typing import Annotated, Any, Literal, TypeAlias, TypeVar - -import numpy as np -import torch -import torch.nn as nn -from einops import rearrange -from transformers import BaseImageProcessor, PreTrainedConfig -from transformers.activations import GELUActivation -from transformers.feature_extraction_utils import BatchFeature -from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling -from transformers.utils import torch_int - -from vllm.config import VllmConfig -from vllm.config.multimodal import MultiModalDummyOptions -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.inputs import ModalityData, MultiModalDataDict -from vllm.logger import init_logger -from vllm.model_executor.layers.attention import ( - MMEncoderAttention, -) -from vllm.model_executor.layers.conv import Conv2dLayer -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding.common import ( - ApplyRotaryEmb, -) -from vllm.model_executor.models.module_mapping import MultiModelKeys -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import ( - ImageItem, - MultiModalFeatureSpec, - MultiModalFieldConfig, - MultiModalKwargsItem, - MultiModalKwargsItems, - VideoItem, -) -from vllm.multimodal.parse import ( - DictEmbeddingItems, - ImageSize, - ModalityDataItems, - MultiModalDataItems, - MultiModalDataParser, -) -from vllm.multimodal.processing import ( - BaseDummyInputsBuilder, - BaseMultiModalProcessor, - BaseProcessingInfo, - PromptReplacement, - PromptUpdate, -) -from vllm.sequence import IntermediateTensors -from vllm.utils.tensor_schema import TensorSchema, TensorShape -from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d - -from .interfaces import ( - MultiModalEmbeddings, - SupportsLoRA, - SupportsMRoPE, - SupportsMultiModal, - SupportsPP, -) -from .siglip import SiglipMLP -from .utils import ( - AutoWeightsLoader, - WeightsMapper, - init_vllm_registered_model, - maybe_prefix, -) -from .vision import is_vit_use_data_parallel - -logger = init_logger(__name__) - - -def smart_resize( - height: int, - width: int, - factor: int, - min_pixels: int, - max_pixels: int, -): - if height < factor: - logger.warning( - "smart_resize: height=%s < factor=%s, reset height=factor", - height, - factor, - ) - width = round((width * factor) / height) - height = factor - - if width < factor: - logger.warning( - "smart_resize: width=%s < factor=%s, reset width=factor", - width, - factor, - ) - height = round((height * factor) / width) - width = factor - - if max(height, width) / min(height, width) > 200: - raise ValueError( - "absolute aspect ratio must be smaller than 200, got " - "{max(height, width) / min(height, width)}" - ) - h_bar = round(height / factor) * factor - w_bar = round(width / factor) * factor - if h_bar * w_bar > max_pixels: - beta = math.sqrt((height * width) / max_pixels) - h_bar = math.floor(height / beta / factor) * factor - w_bar = math.floor(width / beta / factor) * factor - elif h_bar * w_bar < min_pixels: - beta = math.sqrt(min_pixels / (height * width)) - h_bar = math.ceil(height * beta / factor) * factor - w_bar = math.ceil(width * beta / factor) * factor - return h_bar, w_bar - - -class KeyeImagePixelInputs(TensorSchema): - """Dimensions: - - bnp: Batch size * Number of patches - - c: Number of channels - - ps: Patch size - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) - """ - - type: Literal["pixel_values"] - pixel_values: Annotated[ - torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"}) - ] - image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)] - - -class KeyeImageEmbeddingInputs(TensorSchema): - """Dimensions: - - nf: Number of image features - - hs: Hidden size (must match the hidden size of language model - backbone) - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) - """ - - type: Literal["image_embeds"] - image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")] - image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)] - - -KeyeImageInputs: TypeAlias = KeyeImagePixelInputs | KeyeImageEmbeddingInputs - - -class KeyeVideoPixelInputs(TensorSchema): - """Dimensions: - - bnp: Batch size * Number of patches - - c: Number of channels - - ps: Patch size - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) - """ - - type: Literal["pixel_values_videos"] - pixel_values_videos: Annotated[ - torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"}) - ] - video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)] - - -class KeyeVideoEmbeddingInputs(TensorSchema): - """Dimensions: - - nf: Number of video features - - hs: Hidden size (must match the hidden size of language model - backbone) - - nv: Number of videos - - g: Grid dimensions (3 for t, h, w) - """ - - type: Literal["video_embeds"] - video_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")] - video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)] - - -KeyeVideoInputs: TypeAlias = KeyeVideoPixelInputs | KeyeVideoEmbeddingInputs - - -class KeyeVisionEmbeddings(nn.Module): - def __init__(self, config: PreTrainedConfig): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.image_size = config.image_size - self.patch_size = config.patch_size - - self.patch_embedding = Conv2dLayer( - in_channels=config.num_channels, - out_channels=self.embed_dim, - kernel_size=self.patch_size, - stride=self.patch_size, - padding="valid", - ) - - self.num_patches = (self.image_size // self.patch_size) ** 2 - self.num_positions = self.num_patches - self.cache_position_embedding: dict[tuple[int, int], torch.Tensor] = {} - self.cache_position_count: dict[tuple[int, int], int] = {} - self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) - self.packing_position_embedding = nn.Embedding(32768, self.embed_dim) - - self.register_buffer( - "position_ids", - torch.arange(self.num_positions).expand((1, -1)), - persistent=False, - ) - - def interpolate_pos_encoding( - self, - embeddings: torch.Tensor, - height: int, - width: int, - is_after_patchify: bool = False, - ) -> torch.Tensor: - num_positions = self.position_embedding.weight.shape[0] - - patch_pos_embed = self.position_embedding.weight.unsqueeze(0) - - dim = embeddings.shape[-1] - - if is_after_patchify: - new_height = height - new_width = width - else: - new_height = height // self.patch_size - new_width = width // self.patch_size - - sqrt_num_positions = torch_int(num_positions**0.5) - patch_pos_embed = patch_pos_embed.reshape( - 1, sqrt_num_positions, sqrt_num_positions, dim - ) - patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) - - patch_pos_embed = nn.functional.interpolate( - patch_pos_embed, - size=(new_height, new_width), - mode="bilinear", - align_corners=False, - ) - - patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) - return patch_pos_embed - - def forward( - self, - pixel_values: torch.FloatTensor, - position_ids: torch.Tensor | None = None, - image_grid_thw: list[tuple[int, int, int]] | None = None, - interpolate_pos_encoding=False, - ) -> torch.Tensor: - if pixel_values.dim() == 4: - pixel_values = pixel_values.unsqueeze(0) - if pixel_values.dim() == 5: - if position_ids is None: - raise ValueError( - "position_ids cannot be None when pixel_values.dim() is 5." - ) - ( - batch_size, - sequence_len, - channel, - height, - width, - ) = pixel_values.shape - target_dtype = self.patch_embedding.weight.dtype - pixel_values = rearrange(pixel_values, "b l c h w -> (b l) c h w") - patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) - embeddings = patch_embeds.flatten(-2).squeeze(-1) - - if interpolate_pos_encoding and image_grid_thw is not None: - start = 0 - tmp_embeddings = list() - for image_grid in image_grid_thw: - t, h, w = image_grid - end = start + t * h * w - image_embeddings = embeddings[start:end, :] - position_embedding = ( - self.interpolate_pos_encoding(image_embeddings, h, w, True) - .squeeze(0) - .repeat(t, 1) - ) - image_embeddings = image_embeddings + position_embedding - tmp_embeddings.append(image_embeddings) - start = end - embeddings = torch.concat(tmp_embeddings, dim=0).unsqueeze(0) - else: - embeddings = embeddings + self.packing_position_embedding(position_ids) - return embeddings - else: - raise ValueError( - "Unsupported pixel_values dimension:" - f" {pixel_values.dim()}. Expected 4 or 5." - ) - - -def apply_rotary_pos_emb_flashatt( - q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - apply_rotary_emb: ApplyRotaryEmb, -) -> tuple[torch.Tensor, torch.Tensor]: - cos = cos.chunk(2, dim=-1)[0].contiguous() - sin = sin.chunk(2, dim=-1)[0].contiguous() - - q_embed = apply_rotary_emb(q, cos, sin) - k_embed = apply_rotary_emb(k, cos, sin) - - return q_embed, k_embed - - -class KeyeSiglipAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You - Need' paper.""" - - def __init__( - self, - config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.config = config - - hidden_size = config.hidden_size - self.hidden_size = config.hidden_size - use_data_parallel = is_vit_use_data_parallel() - tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size() - self.total_num_heads = config.num_attention_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = config.num_attention_heads - if self.total_num_kv_heads >= tp_size: - assert self.total_num_kv_heads % tp_size == 0 - else: - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = config.hidden_size // self.total_num_heads - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scale = self.head_dim**-0.5 - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.out_proj = RowParallelLinear( - input_size=hidden_size, - output_size=hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.out_proj", - ) - - self.attn = MMEncoderAttention( - num_heads=self.num_heads, - head_size=self.head_dim, - scale=self.scale, - num_kv_heads=self.num_kv_heads, - prefix=f"{prefix}.attn", - ) - - self.apply_rotary_emb = ApplyRotaryEmb( - enforce_enable=True, - enable_fp32_compute=True, - ) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - output_attentions: bool | None = False, - cu_seqlens: torch.Tensor | None = None, - rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split( - [self.q_size, self.kv_size, self.kv_size], - dim=-1, - ) - - if cu_seqlens is None: - raise ValueError("cu_seqlens cannot be None.") - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() - - if rope_emb is None: - q = q.view(*q.shape[:-1], self.num_heads, self.head_dim) - k = k.view( - *k.shape[:-1], - self.num_kv_heads, - self.head_dim, - ) - v = v.view( - *v.shape[:-1], - self.num_kv_heads, - self.head_dim, - ) - else: - cos, sin = rope_emb - q = q.view(*q.shape[:-1], self.num_heads, self.head_dim) - k = k.view( - *k.shape[:-1], - self.num_kv_heads, - self.head_dim, - ) - q, k = apply_rotary_pos_emb_flashatt(q, k, cos, sin, self.apply_rotary_emb) - v = v.view( - *v.shape[:-1], - self.num_kv_heads, - self.head_dim, - ) - - context_layer = self.attn( - query=q, - key=k, - value=v, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - ) - context_layer = rearrange(context_layer, "b s h d -> b s (h d)") - - output, _ = self.out_proj(context_layer) - return output - - -class SigLIPRotaryEmbedding(nn.Module): - def __init__(self, dim: int, theta: float = 10000.0) -> None: - super().__init__() - self.dim = dim - self.theta = theta - self.rope_init() - - def rope_init(self): - inv_freq = 1.0 / ( - self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim) - ) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - def forward(self, seqlen: int) -> torch.Tensor: - seq = torch.arange( - seqlen, - device=self.inv_freq.device, - dtype=self.inv_freq.dtype, - ) - freqs = torch.outer(seq, self.inv_freq) - return freqs - - -class KeyeSiglipEncoderLayer(nn.Module): - def __init__( - self, - config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.embed_dim = config.hidden_size - self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) - self.self_attn = KeyeSiglipAttention( - config, - quant_config=quant_config, - prefix=f"{prefix}.self_attn", - ) - self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) - self.mlp = SiglipMLP( - config, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - output_attentions: bool | None = False, - cu_seqlens: torch.Tensor | None = None, - rope_emb: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> tuple[torch.FloatTensor]: - residual = hidden_states - - hidden_states = self.layer_norm1(hidden_states) - hidden_states = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - output_attentions=output_attentions, - cu_seqlens=cu_seqlens, - rope_emb=rope_emb, - ) - - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.layer_norm2(hidden_states) - hidden_states = self.mlp(hidden_states) - - hidden_states = residual + hidden_states - - return hidden_states - - -class KeyeSiglipEncoder(nn.Module): - def __init__( - self, - config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.config = config - embed_dim = config.hidden_size - num_heads = config.num_attention_heads - head_dim = embed_dim // num_heads - self.layers = nn.ModuleList( - [ - KeyeSiglipEncoderLayer( - config, - quant_config=quant_config, - prefix=f"{prefix}.layers.{layer_idx}", - ) - for layer_idx in range(config.num_hidden_layers) - ] - ) - self.rotary_pos_emb = SigLIPRotaryEmbedding(head_dim // 2) - - @staticmethod - def flatten_list(image_grid_thw): - tmp_image_grid_thw = list() - for image_grid in image_grid_thw: - if isinstance(image_grid, list): - tmp_image_grid_thw.extend(image_grid) - else: - tmp_image_grid_thw.append(image_grid) - return tmp_image_grid_thw - - def forward( - self, - inputs_embeds, - attention_mask: torch.Tensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - cu_seqlens: torch.Tensor | None = None, - image_grid_thw: list[tuple[int, int, int]] | None = None, - height_position_ids: torch.Tensor | None = None, - width_position_ids: torch.Tensor | None = None, - use_rope: bool | None = False, - window_size: int = -1, - vision_or_text: str = "vision", - ) -> BaseModelOutput: - device = inputs_embeds.device - hidden_states = inputs_embeds - if use_rope is True: - flatten_image_grid_thw = self.flatten_list(image_grid_thw) - - if width_position_ids is None or height_position_ids is None: - split_hids = list() - split_wids = list() - for t, h, w in flatten_image_grid_thw: - image_pids = torch.arange(t * h * w, device=device) % (h * w) - sample_hids = image_pids // w - sample_wids = image_pids % w - split_hids.append(sample_hids) - split_wids.append(sample_wids) - width_position_ids = torch.concat(split_wids, dim=0) - height_position_ids = torch.concat(split_hids, dim=0) - - pids = torch.stack( - [height_position_ids, width_position_ids], - dim=-1, - ) - # The ids are built from the grids above, so `h`/`w` bound them - # and the table size is known on the host. - max_grid_size = max(max(h, w) for _, h, w in flatten_image_grid_thw) - rope_emb_max_grid = self.rotary_pos_emb(max_grid_size) - rope_emb = rope_emb_max_grid[pids].flatten(1) - rope_emb = rope_emb.repeat(1, 2) - rope_emb = (rope_emb.cos(), rope_emb.sin()) - else: - rope_emb = None - - attn_cu_seqlens = cu_seqlens - hidden_states = inputs_embeds - assert attention_mask is None - - for encoder_layer in self.layers: - hidden_states = encoder_layer( - hidden_states, - attention_mask, - output_attentions=output_attentions, - cu_seqlens=attn_cu_seqlens, - rope_emb=rope_emb, - ) - return hidden_states - - -class KeyeSiglipVisionTransformer(nn.Module): - def __init__( - self, - config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.config = config - embed_dim = config.hidden_size - - self.embeddings = KeyeVisionEmbeddings(config) - self.encoder = KeyeSiglipEncoder( - config, - quant_config=quant_config, - prefix=f"{prefix}.encoder", - ) - self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) - - def forward( - self, - pixel_values, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - interpolate_pos_encoding: bool | None = False, - attention_mask: torch.Tensor | None = None, - sample_indices: torch.Tensor | None = None, - image_indices: torch.Tensor | None = None, - position_ids: torch.Tensor | None = None, - height_position_ids: torch.Tensor | None = None, - width_position_ids: torch.Tensor | None = None, - cu_seqlens: torch.Tensor | None = None, - padding_mask: torch.Tensor | None = None, - vision_return_embed_list: bool | None = False, - image_grid_thw: list[tuple[int, int, int]] | None = None, - return_pooler_output: bool | None = True, - use_rope: bool | None = False, - window_size: int = -1, - ) -> BaseModelOutputWithPooling: - hidden_states = self.embeddings( - pixel_values, - interpolate_pos_encoding=interpolate_pos_encoding, - position_ids=position_ids, - image_grid_thw=image_grid_thw, - ) - - last_hidden_state = self.encoder( - inputs_embeds=hidden_states, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - attention_mask=attention_mask, - cu_seqlens=cu_seqlens, - image_grid_thw=image_grid_thw, - use_rope=use_rope, - height_position_ids=height_position_ids, - width_position_ids=width_position_ids, - window_size=window_size, - vision_or_text="vision", - ) - - last_hidden_state = self.post_layernorm(last_hidden_state) - - if cu_seqlens is None: - raise ValueError( - "cu_seqlens cannot be None for " - "SiglipVisionTransformer output processing." - ) - # `cu_seqlens` is the running sum of the per-image `t * h * w`, so the - # split sizes are known on the host and slicing needs no device read. - split_sizes = [ - int(np.prod(thw)) for thw in self.encoder.flatten_list(image_grid_thw) - ] - return [ - tensor.squeeze(0) for tensor in last_hidden_state.split(split_sizes, dim=1) - ] - - -class KeyeSiglipVisionModel(nn.Module): - config_class = PreTrainedConfig - main_input_name = "pixel_values" - - 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={"vision_model.head.": None}, - ) - - def __init__( - self, - config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - - self.vision_model = KeyeSiglipVisionTransformer( - config, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "vision_model"), - ) - self.quant_config = quant_config - - @property - def dtype(self) -> torch.dtype: - return self.vision_model.embeddings.patch_embedding.weight.dtype - - @property - def device(self) -> torch.device: - return self.vision_model.embeddings.patch_embedding.weight.device - - def get_input_embeddings(self) -> nn.Module: - return self.vision_model.embeddings.patch_embedding - - def forward( - self, - pixel_values, - sample_indices: torch.Tensor | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - interpolate_pos_encoding: bool = False, - position_ids: torch.Tensor | None = None, - vision_return_embed_list: bool | None = False, - image_grid_thw: list[tuple[int, int, int]] | None = None, - cu_seqlens: torch.Tensor | None = None, - return_pooler_output: bool | None = True, - use_rope: bool | None = False, - window_size: int = -1, - ) -> BaseModelOutputWithPooling: - return self.vision_model( - pixel_values=pixel_values, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - interpolate_pos_encoding=interpolate_pos_encoding, - position_ids=position_ids, - vision_return_embed_list=vision_return_embed_list, - image_grid_thw=image_grid_thw, - sample_indices=sample_indices, - cu_seqlens=cu_seqlens, - return_pooler_output=return_pooler_output, - use_rope=use_rope, - window_size=window_size, - ) - - 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 Projector(nn.Module): - def __init__( - self, - text_config: PreTrainedConfig, - vision_config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.text_config = text_config - self.vision_config = vision_config - self.merge_kernel_size = (2, 2) - - self.hidden_size = ( - self.vision_config.hidden_size - * self.merge_kernel_size[0] - * self.merge_kernel_size[1] - ) - - self.pre_norm = torch.nn.LayerNorm(self.vision_config.hidden_size, eps=1e-05) - self.act = GELUActivation() - - self.linear_1 = ColumnParallelLinear( - self.hidden_size, - self.hidden_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.linear_1", - ) - self.linear_2 = RowParallelLinear( - self.hidden_size, - self.text_config.hidden_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.linear_2", - ) - - def forward( - self, - image_features: torch.Tensor | list[torch.Tensor], - image_grid_thw: list[tuple[int, int, int]], - ) -> torch.Tensor | list[torch.Tensor]: - m1, m2 = self.merge_kernel_size - if isinstance(image_features, (list, tuple)): - processed_features = list() - for image_feature, image_grid in zip(image_features, image_grid_thw): - image_feature = self.pre_norm(image_feature) - t, h, w = image_grid - - image_feature = rearrange( - image_feature, - "(t h p1 w p2) d -> (t h w) (p1 p2 d)", - t=t, - h=h // m1, - p1=m1, - w=w // m2, - p2=m2, - ) - hidden_states, _ = self.linear_1(image_feature) - hidden_states = self.act(hidden_states) - hidden_states, _ = self.linear_2(hidden_states) - processed_features.append(hidden_states) - - return processed_features - - dims = image_features.shape[:-1] - dim = image_features.shape[-1] - image_features = image_features.view(np.prod(dims), dim) - hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size) - hidden_states = self.linear_1(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.linear_2(hidden_states) - - return hidden_states.view(*dims, -1) - - -def _keye_field_config( - hf_inputs: Mapping[str, torch.Tensor], -): - image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) - image_grid_sizes = image_grid_thw.prod(-1) - - video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3))) - video_grid_sizes = video_grid_thw.prod(-1) - - return dict( - pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), - image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), - image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), - pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( - "video", video_grid_sizes - ), - video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes), - video_grid_thw=MultiModalFieldConfig.batched("video", keep_on_cpu=True), - ) - - -class KeyeMultiModalDataParser(MultiModalDataParser): - # The patch grid is what sizes the placeholder range. - embedding_fields = { - "image": {"image_embeds": "values", "image_grid_thw": "metadata"}, - "video": {"video_embeds": "values", "video_grid_thw": "metadata"}, - } - - def _parse_image_data( - self, - data: dict[str, torch.Tensor] | ModalityData[ImageItem], - ) -> ModalityDataItems[Any, Any] | None: - if isinstance(data, dict): - required, optional = self.embedding_field_sets("image") - return DictEmbeddingItems( - data, - modality="image", - required_fields=required, - optional_fields=optional, - fields_factory=_keye_field_config, - ) - - return super()._parse_image_data(data) - - def _parse_video_data( - self, - data: dict[str, torch.Tensor] | ModalityData[VideoItem], - ) -> ModalityDataItems[Any, Any] | None: - if isinstance(data, dict): - required, optional = self.embedding_field_sets("video") - return DictEmbeddingItems( - data, - modality="video", - required_fields=required, - optional_fields=optional, - fields_factory=_keye_field_config, - ) - - return super()._parse_video_data(data) - - -class KeyeProcessingInfo(BaseProcessingInfo): - def get_max_image_size(self) -> int: - return 9999999 # _MAX_IMAGE_SIZE - - def get_max_frame_per_video(self) -> int: - return 16 # _MAX_FRAMES_PER_VIDEO - - def get_image_processor(self, **kwargs: object): - return self.get_hf_processor(**kwargs).image_processor - - def get_data_parser(self): - return KeyeMultiModalDataParser( - expected_hidden_size=self._get_expected_hidden_size(), - allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, - ) - - def get_supported_mm_limits( - self, - ) -> Mapping[str, int | None]: - return {"image": None, "video": None} - - def get_mm_max_tokens_per_item( - self, - seq_len: int, - mm_counts: Mapping[str, int], - ) -> Mapping[str, int]: - return { - "image": self.get_max_image_tokens(), - "video": self.get_max_video_tokens(seq_len), - } - - def _get_vision_info( - self, - *, - image_width: int, - image_height: int, - num_frames: int = 1, - do_resize: bool = True, - image_processor: BaseImageProcessor, - mm_kwargs: Mapping[str, object], - ) -> tuple[ImageSize, int]: - hf_config = self.get_hf_config() - vision_config = hf_config.vision_config - patch_size = vision_config.patch_size - merge_size = vision_config.spatial_merge_size - temporal_patch_size = 1 - - mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs) - size = image_processor.size - if override_size := mm_kwargs.get("size"): - size = size | override_size - if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None: - size = size | {"min_pixels": override_min_pixels} - if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None: - size = size | {"max_pixels": override_max_pixels} - - if do_resize: - resized_height, resized_width = smart_resize( - height=image_height, - width=image_width, - factor=patch_size * merge_size, - min_pixels=size["min_pixels"], - max_pixels=size["max_pixels"], - ) - preprocessed_size = ImageSize(width=resized_width, height=resized_height) - else: - preprocessed_size = ImageSize(width=image_width, height=image_height) - - padded_num_frames = num_frames + (-num_frames % temporal_patch_size) - - grid_t = max(padded_num_frames // temporal_patch_size, 1) - grid_h = preprocessed_size.height // patch_size - grid_w = preprocessed_size.width // patch_size - - num_patches = grid_t * grid_h * grid_w - num_vision_tokens = num_patches // (merge_size**2) - - return preprocessed_size, num_vision_tokens - - def get_num_image_tokens( - self, - *, - image_width: int, - image_height: int, - image_processor: BaseImageProcessor, - mm_kwargs: Mapping[str, object], - ) -> int: - _, num_image_tokens = self._get_vision_info( - image_width=image_width, - image_height=image_height, - image_processor=image_processor, - mm_kwargs=mm_kwargs, - ) - return num_image_tokens - - def get_num_video_tokens( - self, - *, - image_width: int, - image_height: int, - num_frames: int, - image_processor: BaseImageProcessor, - mm_kwargs: Mapping[str, object], - ) -> int: - _, num_video_tokens = self._get_vision_info( - image_width=image_width, - image_height=image_height, - num_frames=num_frames, - image_processor=image_processor, - mm_kwargs=mm_kwargs, - ) - return num_video_tokens - - def get_image_size_with_most_features(self) -> ImageSize: - image_processor = self.get_image_processor() - - max_image_size, _ = self._get_vision_info( - image_width=self.get_max_image_size(), - image_height=self.get_max_image_size(), - image_processor=image_processor, - mm_kwargs={}, - ) - return max_image_size - - def get_max_image_tokens(self) -> int: - image_processor = self.get_image_processor() - target_width, target_height = self.get_image_size_with_most_features() - - return self.get_num_image_tokens( - image_width=target_width, - image_height=target_height, - image_processor=image_processor, - mm_kwargs={}, - ) - - def _get_max_video_frames(self, max_tokens: int) -> int: - image_processor = self.get_image_processor() - target_width, target_height = self.get_image_size_with_most_features() - - num_frames = 0 - - while True: - next_num_frames = num_frames + 1 - next_max_tokens = self.get_num_video_tokens( - image_width=target_width, - image_height=target_height, - num_frames=next_num_frames, - image_processor=image_processor, - mm_kwargs={}, - ) - - if next_max_tokens > max_tokens: - break - - num_frames = next_num_frames - - return num_frames - - def get_num_frames_with_most_features(self, seq_len: int) -> int: - mm_config = self.ctx.get_mm_config() - max_images = mm_config.get_limit_per_prompt("image") - max_videos = mm_config.get_limit_per_prompt("video") - - max_image_tokens = self.get_max_image_tokens() * max_images - max_total_frames = self._get_max_video_frames(seq_len - max_image_tokens) - max_frames_per_video = min( - max_total_frames // max(max_videos, 1), - self.get_max_frame_per_video(), - ) - - return max(max_frames_per_video, 1) - - def get_max_video_tokens(self, seq_len: int) -> int: - image_processor = self.get_image_processor() - target_width, target_height = self.get_image_size_with_most_features() - - return self.get_num_video_tokens( - image_width=target_width, - image_height=target_height, - num_frames=self.get_num_frames_with_most_features(seq_len), - image_processor=image_processor, - mm_kwargs={}, - ) - - -_I = TypeVar("_I", bound=KeyeProcessingInfo) - - -class KeyeBaseDummyInputsBuilder(BaseDummyInputsBuilder[_I]): - def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - - hf_processor = self.info.get_hf_processor() - image_token: str = hf_processor.image_token - video_token: str = hf_processor.video_token - - return image_token * num_images + video_token * num_videos - - def get_dummy_mm_data( - self, - seq_len: int, - mm_counts: Mapping[str, int], - mm_options: MultiModalDummyOptions, - ) -> MultiModalDataDict: - target_width, target_height = self.info.get_image_size_with_most_features() - target_num_frames = self.info.get_num_frames_with_most_features(seq_len) - - mm_data = { - "image": self._get_dummy_images( - width=target_width, - height=target_height, - num_images=mm_counts.get("image", 0), - overrides=mm_options.get("image"), - ), - "video": self._get_dummy_videos( - width=target_width, - height=target_height, - num_frames=target_num_frames, - num_videos=mm_counts.get("video", 0), - overrides=mm_options.get("video"), - ), - } - - return mm_data - - -class KeyeDummyInputsBuilder(KeyeBaseDummyInputsBuilder[KeyeProcessingInfo]): - pass - - -class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]): - def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: - return self.dummy_inputs.get_dummy_text(mm_counts) - - def _get_prompt_updates( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, Any], - out_mm_kwargs: MultiModalKwargsItems, - ) -> Sequence[PromptUpdate]: - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) - tokenizer = self.info.get_tokenizer() - vocab = tokenizer.get_vocab() - - placeholder = { - "image": vocab[hf_processor.image_token], - "video": vocab[hf_processor.video_token], - } - - merge_length = image_processor.merge_size**2 - - def get_replacement_keye(item_idx: int, modality: str): - out_item = out_mm_kwargs[modality][item_idx] - grid_thw = out_item[f"{modality}_grid_thw"].data - assert isinstance(grid_thw, torch.Tensor) - - num_tokens = int(grid_thw.prod()) // merge_length - return [placeholder[modality]] * num_tokens - - return [ - PromptReplacement( - modality=modality, - target=[placeholder[modality]], - replacement=partial(get_replacement_keye, modality=modality), - ) - for modality in ("image", "video") - ] - - def _get_mm_fields_config( - self, - hf_inputs: BatchFeature, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> Mapping[str, MultiModalFieldConfig]: - return _keye_field_config(hf_inputs) - - -class BaseKeyeModule(nn.Module, SupportsMultiModal): - packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], - } - - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_prefix={ - "lm_head.": "language_model.lm_head.", - "model.": "language_model.model.", - } - ) - - supports_tower_connector_lora = True - - @classmethod - def get_placeholder_str(cls, modality: str, i: int) -> str | None: - if modality.startswith("image"): - return "<|vision_start|><|image_pad|><|vision_end|>" - if modality.startswith("video"): - return "<|vision_start|><|video_pad|><|vision_end|>" - - raise ValueError("Only image or video modality is supported") - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config: PreTrainedConfig = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.config = config - - with self._mark_tower_model(vllm_config, {"image", "video"}): - self.visual = KeyeSiglipVisionModel( - config.vision_config, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "visual"), - ) - self.mlp_AR = self._build_projector( - config, - config.vision_config, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "mlp_AR"), - ) - - with self._mark_language_model(vllm_config): - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - prefix=maybe_prefix(prefix, "language_model"), - architectures=["Qwen3ForCausalLM"], - ) - - self.make_empty_intermediate_tensors = ( - self.language_model.make_empty_intermediate_tensors - ) - - @abstractmethod - def _build_projector( - self, - text_config: PreTrainedConfig, - vision_config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> nn.Module: - raise NotImplementedError("Need projector") - - def _process_image_input(self, image_input: Any) -> tuple[torch.Tensor, ...]: - siglip_position_ids = list() - image_grid_hws = list() - sample_indices = list() - cu_seqlens = [0] - - image_grid_thw = image_input["image_grid_thw"] - assert image_grid_thw.ndim == 2 - - image_grid_thw_list = image_grid_thw.tolist() - for idx, thw in enumerate(image_grid_thw_list): - thw_tuple = tuple(thw) - numel = np.prod(thw_tuple) - image_grid_hws.append(thw_tuple) - image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:]) - siglip_position_ids.append(image_position_ids) - sample_indices.append(torch.full((numel,), idx, dtype=torch.int64)) - cu_seqlens.append(cu_seqlens[-1] + numel) - - if image_input["type"] == "image_embeds": - raise ValueError( - "Image embeddings are not supported for this processing path." - ) - else: - pixel_values = image_input["pixel_values"].type(self.visual.dtype) - # These are all built on the host; concat straight into pinned - # buffers so the H2D copies stay non-blocking. - siglip_position_ids = torch.concat( - siglip_position_ids, - dim=0, - out=torch.empty( - sum(t.numel() for t in siglip_position_ids), - dtype=torch.int64, - pin_memory=PIN_MEMORY, - ), - ).to(pixel_values.device, non_blocking=True) - cu_seqlens = async_tensor_h2d( - cu_seqlens, dtype=torch.int32, device=pixel_values.device - ) - sample_indices = torch.concat( - sample_indices, - dim=0, - out=torch.empty( - sum(t.numel() for t in sample_indices), - dtype=torch.int64, - pin_memory=PIN_MEMORY, - ), - ).to(pixel_values.device, non_blocking=True) - - image_embeds = self.visual( - pixel_values=pixel_values, - image_grid_thw=image_grid_hws, - position_ids=siglip_position_ids, - vision_return_embed_list=False, - interpolate_pos_encoding=True, - sample_indices=sample_indices, - cu_seqlens=cu_seqlens, - use_rope=True, - window_size=-1, - ) - image_embeds = tuple(self.mlp_AR(image_embeds, image_grid_thw)) - return image_embeds - - def _process_video_embeds( - self, - video_type: Literal["video_embeds", "pixel_values_videos"], - video_grid_thw: torch.Tensor, - pixel_values_videos: torch.Tensor | None = None, - ) -> torch.Tensor | list[torch.Tensor]: - siglip_position_ids = list() - video_grid_hws = list() - sample_indices = list() - cu_seqlens = [0] - - assert video_grid_thw.ndim == 2 - for idx, sub_thw in enumerate(video_grid_thw): - thw_tuple = tuple(sub_thw.detach().cpu().numpy().tolist()) - numel = np.prod(thw_tuple) - - video_grid_hws.append(thw_tuple) - video_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:]) - siglip_position_ids.append(video_position_ids) - sample_indices.append(torch.full((numel,), idx, dtype=torch.int64)) - cu_seqlens.append(cu_seqlens[-1] + numel) - - if video_type == "video_embeds": - raise ValueError( - "Video embeddings are not supported for this processing path." - ) - else: - assert pixel_values_videos is not None - pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - # These are host-built; concat straight into pinned buffers so - # the H2D copies stay non-blocking. - siglip_position_ids = torch.concat( - siglip_position_ids, - dim=0, - out=torch.empty( - sum(t.numel() for t in siglip_position_ids), - dtype=torch.int64, - pin_memory=PIN_MEMORY, - ), - ).to(pixel_values_videos.device, non_blocking=True) - cu_seqlens = async_tensor_h2d( - cu_seqlens, dtype=torch.int32, device=pixel_values_videos.device - ) - sample_indices = torch.concat( - sample_indices, - dim=0, - out=torch.empty( - sum(t.numel() for t in sample_indices), - dtype=torch.int64, - pin_memory=PIN_MEMORY, - ), - ).to(pixel_values_videos.device, non_blocking=True) - - video_embeds = self.visual( - pixel_values=pixel_values_videos, - image_grid_thw=video_grid_hws, - position_ids=siglip_position_ids, - vision_return_embed_list=True, - interpolate_pos_encoding=True, - sample_indices=sample_indices, - cu_seqlens=cu_seqlens, - use_rope=True, - window_size=-1, - ) - video_embeds = self.mlp_AR(video_embeds, video_grid_thw) - return video_embeds - - def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: - modalities = {} - - for input_key in kwargs: - if ( - input_key in ("pixel_values", "image_embeds") - and "images" not in modalities - ): - modalities["images"] = self._parse_and_validate_image_input(**kwargs) - if ( - input_key in ("pixel_values_videos", "video_embeds") - and "videos" not in modalities - ): - modalities["videos"] = self._parse_and_validate_video_input(**kwargs) - - return modalities - - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: - modalities = self._parse_and_validate_multimodal_inputs(**kwargs) - if not modalities: - return None - - multimodal_embeddings: tuple[torch.Tensor, ...] = () - - for modality in modalities: - if modality == "images": - image_input = modalities["images"] - image_embeddings = self._process_image_input(image_input) - multimodal_embeddings += tuple(image_embeddings) - if modality == "videos": - video_input = modalities["videos"] - video_embeddings = self._process_video_input(video_input) - multimodal_embeddings += tuple(video_embeddings) - return multimodal_embeddings - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: object, - ) -> torch.Tensor | IntermediateTensors: - """Run forward pass for Keye-VL. - - Args: - input_ids: Flattened (concatenated) input_ids corresponding to a - batch. - positions: Flattened (concatenated) position ids corresponding to a - batch. - **NOTE**: If mrope is enabled (default setting for Qwen2-VL - opensource models), the shape will be `(3, seq_len)`, - otherwise it will be `(seq_len,)`. - intermediate_tensors: Intermediate tensors from prior forward pass. - inputs_embeds: Optional tensor of input embeddings. - **kwargs: Multimodal inputs for this batch, forwarded to the - multimodal embedding path. - - """ - if intermediate_tensors is not None: - inputs_embeds = None - - hidden_states = self.language_model.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - ) - - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - return self.language_model.compute_logits(hidden_states) - - 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) - - def get_mm_mapping(self) -> MultiModelKeys: - """Get the module prefix in multimodal models.""" - return MultiModelKeys.from_string_field( - language_model="language_model", - connector="mlp_AR.", - tower_model="visual.", - ) - - 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_mm_embeds * merge_size**2, num_mm_embeds - - -@MULTIMODAL_REGISTRY.register_processor( - KeyeMultiModalProcessor, - info=KeyeProcessingInfo, - dummy_inputs=KeyeDummyInputsBuilder, -) -class KeyeForConditionalGeneration( - BaseKeyeModule, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE -): - def _build_projector( - self, - text_config: PreTrainedConfig, - vision_config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> nn.Module: - return Projector(text_config, vision_config, quant_config, prefix) - - def _parse_and_validate_image_input( - self, **kwargs: object - ) -> KeyeImageInputs | None: - pixel_values = kwargs.pop("pixel_values", None) - image_embeds = kwargs.pop("image_embeds", None) - image_grid_thw = kwargs.pop("image_grid_thw", None) - - if pixel_values is None and image_embeds is None: - return None - - if pixel_values is not None: - return KeyeImagePixelInputs( - type="pixel_values", - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - ) - - return KeyeImageEmbeddingInputs( - type="image_embeds", - image_embeds=image_embeds, - image_grid_thw=image_grid_thw, - ) - - def _parse_and_validate_video_input( - self, **kwargs: object - ) -> KeyeVideoInputs | None: - pixel_values_videos = kwargs.pop("pixel_values_videos", None) - video_embeds = kwargs.pop("video_embeds", None) - video_grid_thw = kwargs.pop("video_grid_thw", None) - - if pixel_values_videos is None and video_embeds is None: - return None - - if pixel_values_videos is not None: - return KeyeVideoPixelInputs( - type="pixel_values_videos", - pixel_values_videos=pixel_values_videos, - video_grid_thw=video_grid_thw, - ) - - return KeyeVideoEmbeddingInputs( - type="video_embeds", - video_embeds=video_embeds, - video_grid_thw=video_grid_thw, - ) - - def _process_video_input( - self, video_input: KeyeVideoInputs - ) -> tuple[torch.Tensor, ...]: - video_type = video_input["type"] - video_grid_thw = video_input["video_grid_thw"] - pixel_values_videos = video_input.get("pixel_values_videos", None) - - return tuple( - self._process_video_embeds(video_type, video_grid_thw, pixel_values_videos) - ) - - @staticmethod - def _split_video_grid_thw( - grid_thw: torch.Tensor | list[list[int]] | list[int], - ) -> list[list[int]]: - """Split video grid_thw along the t dimension into per-frame rows. - - This preserves Keye's current M-RoPE behavior, where a video is emitted - as consecutive frame-level multimodal blocks rather than a single block - spanning the whole video. - """ - if isinstance(grid_thw, list): - if len(grid_thw) == 0: - return [] - if isinstance(grid_thw[0], int): - grid_thw = torch.tensor([grid_thw], dtype=torch.long) - else: - grid_thw = torch.tensor(grid_thw, dtype=torch.long) - elif grid_thw.ndim == 1: - grid_thw = grid_thw.unsqueeze(0) - - if grid_thw.numel() == 0: - return [] - - t, hw = grid_thw[:, 0], grid_thw[:, 1:] - ones = torch.ones_like(hw[:, :1]) - out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0) - return out.tolist() - - def iter_mm_grid_thw( - self, mm_features: list[MultiModalFeatureSpec] - ) -> Iterator[tuple[int, int, int, int]]: - spatial_merge_size = self.config.vision_config.spatial_merge_size - - for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): - if mm_feature.data is None: - raise ValueError("M-RoPE calculation requires multimodal feature data") - - if mm_feature.modality == "image": - grid_thw = mm_feature.data["image_grid_thw"].data - if isinstance(grid_thw, torch.Tensor): - if grid_thw.ndim == 2: - assert grid_thw.shape[0] == 1 - t, h, w = grid_thw[0].tolist() - else: - t, h, w = grid_thw.tolist() - else: - if isinstance(grid_thw[0], list): - assert len(grid_thw) == 1 - t, h, w = grid_thw[0] - else: - t, h, w = grid_thw - - yield ( - mm_feature.mm_position.offset, - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - elif mm_feature.modality == "video": - current_offset = mm_feature.mm_position.offset - for t, h, w in self._split_video_grid_thw( - mm_feature.data["video_grid_thw"].data - ): - llm_grid_h = h // spatial_merge_size - llm_grid_w = w // spatial_merge_size - yield (current_offset, t, llm_grid_h, llm_grid_w) - current_offset += t * llm_grid_h * llm_grid_w - else: - raise ValueError(f"Unsupported modality: {mm_feature.modality}") - - def get_mrope_input_positions( - self, - input_tokens: list[int], - mm_features: list[MultiModalFeatureSpec], - ) -> tuple[torch.Tensor, int]: - llm_pos_ids_list: list = [] - st = 0 - - for ( - offset, - llm_grid_t, - llm_grid_h, - llm_grid_w, - ) in self.iter_mm_grid_thw(mm_features): - text_len = offset - st - - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx - ) - - t_index = ( - ( - torch.arange(llm_grid_t) - .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - ) - .long() - .flatten() - ) - - h_index = ( - torch.arange(llm_grid_h) - .view(1, -1, 1) - .expand(llm_grid_t, -1, llm_grid_w) - .flatten() - ) - w_index = ( - torch.arange(llm_grid_w) - .view(1, 1, -1) - .expand(llm_grid_t, llm_grid_h, -1) - .flatten() - ) - llm_pos_ids_list.append( - torch.stack([t_index, h_index, w_index]) + text_len + st_idx - ) - st = offset + llm_grid_t * llm_grid_h * llm_grid_w - - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx - ) - - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() - - return llm_positions, mrope_position_delta diff --git a/vllm/model_executor/models/keye_vl1_5.py b/vllm/model_executor/models/keye_vl1_5.py deleted file mode 100644 index fc08035fa7dc..000000000000 --- a/vllm/model_executor/models/keye_vl1_5.py +++ /dev/null @@ -1,707 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import itertools -from collections.abc import Iterator, Mapping, Sequence -from functools import partial -from typing import Annotated, Any, Literal, TypeAlias - -import numpy as np -import torch -import torch.nn as nn -from einops import rearrange -from transformers import PreTrainedConfig -from transformers.activations import GELUActivation -from transformers.feature_extraction_utils import BatchFeature - -from vllm.config import VllmConfig -from vllm.inputs import ModalityData -from vllm.logger import init_logger -from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import ( - ImageItem, - MultiModalFeatureSpec, - MultiModalFieldConfig, - MultiModalKwargsItems, - VideoItem, -) -from vllm.multimodal.parse import ( - DictEmbeddingItems, - ModalityDataItems, - MultiModalDataItems, - MultiModalDataParser, -) -from vllm.multimodal.processing import ( - PromptReplacement, - PromptUpdate, - PromptUpdateDetails, -) -from vllm.utils.tensor_schema import TensorSchema, TensorShape - -from .interfaces import SupportsLoRA, SupportsMRoPE, SupportsMultiModal, SupportsPP -from .keye import ( - BaseKeyeModule, - BaseMultiModalProcessor, - KeyeBaseDummyInputsBuilder, - KeyeProcessingInfo, -) - -logger = init_logger(__name__) - - -def split_thw(grid_thw: torch.Tensor) -> torch.Tensor: - """Split grid_thw in t dimension. - - Args: - grid_thw: [N, 3] tensor of [t, h, w] - - Returns: - [Σt, 3] tensor where each row is [1, h, w] - - Example: - >>> grid_thw = torch.tensor([[2, 3, 4], [1, 5, 6]]) - >>> split_thw(grid_thw) - tensor([[1, 3, 4], - [1, 3, 4], - [1, 5, 6]]) - - """ - t = grid_thw[:, 0] - h_w = grid_thw[:, 1:] - ones = torch.ones_like(h_w[:, :1]) - return torch.cat([ones, h_w], dim=1).repeat_interleave(t, dim=0) - - -def get_num_patches( - grid_thw: torch.Tensor, num_frames: list[int] | torch.Tensor -) -> torch.Tensor: - """Return num_patches per video. - - Args: - grid_thw: Tensor with shape [N, 3] containing temporal, height, width - dimensions - num_frames: List or tensor indicating the number of frames per video - - Returns: - Tensor of the number of patches for each video - - Examples: - >>> # Suppose there are 2 videos with a total of 3 grids - >>> grid_thw = torch.tensor( - ... [ - ... [2, 2, 2], # grid 0: 2*2*2=8 patches - ... [2, 2, 2], # grid 1: 2*2*2=8 patches - ... [1, 1, 1], - ... ] - ... ) # grid 2: 1*1*1=1 patches - >>> num_frames = [2, 1] # The first video contains 2 grids, - the second contains 1 grid. - >>> get_num_patches(grid_thw, num_frames) - tensor([16, 1]) # Total patches for first video: 8+8=16, - second video: 1. - - """ - assert len(grid_thw.shape) == 2 - if isinstance(num_frames, torch.Tensor): - num_frames = num_frames.clone().tolist() - - num_grids_per_frame = grid_thw.prod(dim=1) - start_idx_per_video = [0, *itertools.accumulate(num_frames)] - num_patches = [ - num_grids_per_frame[start_idx_per_video[i] : start_idx_per_video[i + 1]].sum() - for i in range(len(num_frames)) - ] - return ( - torch.stack(num_patches) - if num_patches - else torch.zeros(0, dtype=grid_thw.dtype, device=grid_thw.device) - ) - - -class KeyeVL1_5ImagePixelInputs(TensorSchema): - """Dimensions: - - bnp: Batch size * Number of patches - - c: Number of channels - - ps: Patch size - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) - """ - - type: Literal["pixel_values"] - - pixel_values: Annotated[ - torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"}) - ] - - image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)] - - -class KeyeVL1_5ImageEmbeddingInputs(TensorSchema): - """Dimensions: - - nf: Number of image features - - hs: Hidden size (must match the hidden size of language model - backbone) - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) - """ - - type: Literal["image_embeds"] - image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")] - image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)] - - -KeyeVL1_5ImageInputs: TypeAlias = ( - KeyeVL1_5ImagePixelInputs | KeyeVL1_5ImageEmbeddingInputs -) - - -class KeyeVL1_5VideoPixelInputs(TensorSchema): - """Dimensions: - - bnp: Batch size * Number of patches - - c: Number of channels - - ps: Patch size - - ni: Number of images - - g: Grid dimensions (3 for t, h, w) - """ - - type: Literal["pixel_values_videos"] - pixel_values_videos: Annotated[ - torch.Tensor, TensorShape("bnp", 3, "ps", "ps", dynamic_dims={"bnp"}) - ] - video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)] - - num_frames: torch.Tensor - - -class KeyeVL1_5VideoEmbeddingInputs(TensorSchema): - """Dimensions: - - nf: Number of video features - - hs: Hidden size (must match the hidden size of language model - backbone) - - nv: Number of videos - - g: Grid dimensions (3 for t, h, w) - """ - - type: Literal["video_embeds"] - video_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")] - video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)] - num_frames: torch.Tensor - - -KeyeVL1_5VideoInputs: TypeAlias = ( - KeyeVL1_5VideoPixelInputs | KeyeVL1_5VideoEmbeddingInputs -) - - -class KeyeVL1_5Projector(nn.Module): - def __init__( - self, - text_config: PreTrainedConfig, - vision_config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.text_config = text_config - self.vision_config = vision_config - self.merge_kernel_size = (2, 2) - - self.hidden_size = ( - self.vision_config.hidden_size - * self.merge_kernel_size[0] - * self.merge_kernel_size[1] - ) - - self.pre_norm = torch.nn.LayerNorm(self.hidden_size, eps=1e-05) - self.act = GELUActivation() - - self.linear_1 = ColumnParallelLinear( - self.hidden_size, - self.hidden_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.linear_1", - ) - self.linear_2 = RowParallelLinear( - self.hidden_size, - self.text_config.hidden_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.linear_2", - ) - - def forward( - self, - image_features: torch.Tensor | tuple[torch.Tensor] | list[torch.Tensor], - image_grid_thw: list[tuple[int, int, int]], - ) -> torch.Tensor | list[torch.Tensor]: - m1, m2 = self.merge_kernel_size - if isinstance(image_features, (list, tuple)): - processed_features = list() - for image_feature, image_grid in zip(image_features, image_grid_thw): - t, h, w = image_grid - image_feature = rearrange( - image_feature, - "(t h p1 w p2) d -> (t h w) (p1 p2 d)", - t=t, - h=h // m1, - p1=m1, - w=w // m2, - p2=m2, - ) - image_feature = self.pre_norm(image_feature) - hidden_states, _ = self.linear_1(image_feature) - hidden_states = self.act(hidden_states) - hidden_states, _ = self.linear_2(hidden_states) - processed_features.append(hidden_states) - - return processed_features - - dims = image_features.shape[:-1] - dim = image_features.shape[-1] - image_features = image_features.view(np.prod(dims), dim) - hidden_states = self.pre_norm(image_features.view(-1, self.hidden_size)) - hidden_states = self.linear_1(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.linear_2(hidden_states) - - return hidden_states.view(*dims, -1) - - -def _keye_field_config( - hf_inputs: Mapping[str, torch.Tensor], -): - image_grid_thw = hf_inputs.get( - "image_grid_thw", torch.empty((0, 3), dtype=torch.int64) - ) - image_grid_sizes = image_grid_thw.prod(-1) - - video_grid_thw = hf_inputs.get( - "video_grid_thw", torch.empty((0, 3), dtype=torch.int64) - ) - video_grid_thw = split_thw(video_grid_thw) - num_frames = hf_inputs.get("num_frames", video_grid_thw[:, 0]).clone().tolist() - - video_num_patches = get_num_patches(video_grid_thw, num_frames) - - video_num_grids = [] - if len(num_frames) > 0: - i = 0 - j = 1 - cur_frames = num_frames[i] - for t, _, _ in video_grid_thw.tolist(): - cur_frames -= t - if cur_frames == 0: - video_num_grids.append(j) - i += 1 - if i < len(num_frames): - cur_frames = num_frames[i] - j = 1 - else: - j += 1 - video_num_grids = torch.tensor(video_num_grids) - return dict( - pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), - image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), - image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True), - pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( - "video", video_num_patches - ), - video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_num_patches), - video_grid_thw=MultiModalFieldConfig.flat_from_sizes( - "video", video_num_grids, keep_on_cpu=True - ), - num_frames=MultiModalFieldConfig.batched("video", keep_on_cpu=True), - ) - - -class KeyeVL1_5MultiModalDataParser(MultiModalDataParser): - # The patch grid is what sizes the placeholder range. - embedding_fields = { - "image": {"image_embeds": "values", "image_grid_thw": "metadata"}, - "video": {"video_embeds": "values", "video_grid_thw": "metadata"}, - } - - def _parse_image_data( - self, - data: dict[str, torch.Tensor] | ModalityData[ImageItem], - ) -> ModalityDataItems[Any, Any] | None: - if isinstance(data, dict): - required, optional = self.embedding_field_sets("image") - return DictEmbeddingItems( - data, - modality="image", - required_fields=required, - optional_fields=optional, - fields_factory=_keye_field_config, - ) - - return super()._parse_image_data(data) - - def _parse_video_data( - self, - data: dict[str, torch.Tensor] | ModalityData[VideoItem], - ) -> ModalityDataItems[Any, Any] | None: - if isinstance(data, dict): - required, optional = self.embedding_field_sets("video") - return DictEmbeddingItems( - data, - modality="video", - required_fields=required, - optional_fields=optional, - fields_factory=_keye_field_config, - ) - - return super()._parse_video_data(data) - - -class KeyeVL1_5ProcessingInfo(KeyeProcessingInfo): - def get_data_parser(self): - return KeyeVL1_5MultiModalDataParser( - expected_hidden_size=self._get_expected_hidden_size(), - allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, - ) - - def get_max_frame_per_video(self) -> int: - return 2048 - - def get_supported_mm_limits( - self, - ) -> Mapping[str, int | None]: - return {"image": None, "video": 1} - - -class KeyeVL1_5MultiModalProcessor(BaseMultiModalProcessor[KeyeVL1_5ProcessingInfo]): - def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: - return self.dummy_inputs.get_dummy_text(mm_counts) - - def _get_prompt_updates( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, Any], - out_mm_kwargs: MultiModalKwargsItems, - ) -> Sequence[PromptUpdate]: - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) - tokenizer = self.info.get_tokenizer() - vocab = tokenizer.get_vocab() - image_token_id = vocab[hf_processor.image_token] - video_token_id = vocab[hf_processor.video_token] - placeholder = {"image": image_token_id, "video": video_token_id} - merge_length = image_processor.merge_size**2 - - out_mm_kwargs_data = out_mm_kwargs.get_data() - frame_types: list[torch.Tensor] = hf_processor_mm_kwargs.get( - "frame_types", None - ) - timestamps: list[torch.Tensor] = hf_processor_mm_kwargs.get("timestamps", None) - num_videos = mm_items.get_count("video", strict=False) - - if frame_types is None: - frame_types = [None] * num_videos - assert len(frame_types) == num_videos, ( - f"Number of frame_types={len(frame_types)} " - f"doesn't equal to number of videos={num_videos}" - ) - if timestamps is None: - timestamps = [None] * num_videos - assert len(timestamps) == num_videos, ( - f"Number of timestamps={len(timestamps)} " - f"doesn't equal to number of videos={num_videos}" - ) - - video_grid_thw = out_mm_kwargs_data.get( - "video_grid_thw", torch.empty((0, 3), dtype=torch.int64) - ) - num_frames = out_mm_kwargs_data.get( - "num_frames", torch.tensor([], dtype=torch.int64) - ) - - assert len(num_frames) == num_videos, ( - f"Size of num_frames={len(num_frames)} " - f"doesn't equal to number of videos={num_videos}" - ) - - video_grid_hws = split_thw(video_grid_thw) - assert int(num_frames.sum().tolist()) == video_grid_hws.shape[0], ( - f"The first dimension of `video_grid_hws`={video_grid_hws.shape[0]}" - f"doesn't equal to num of frames." - ) - - cu_seqlens = torch.cumsum(torch.tensor([0] + num_frames.tolist()), dim=-1) - - def get_replacement_keye(item_idx: int, modality: str): - """Args: - item_idx(int): The item index of modality to replace - modality(str): The modality - - """ - if modality == "image": - out_item = out_mm_kwargs[modality][item_idx] - grid_thw = out_item[f"{modality}_grid_thw"].data - assert isinstance(grid_thw, torch.Tensor) - - num_tokens = int(grid_thw.prod()) // merge_length - return [image_token_id] * num_tokens - elif modality == "video": - placeholders = [] - video_timestamps = timestamps[item_idx] - video_frame_types = frame_types[item_idx] - grid_thw = video_grid_hws[ - cu_seqlens[item_idx] : cu_seqlens[item_idx + 1] - ] - - nframes = grid_thw.shape[0] - - if video_timestamps is None: - video_timestamps = [""] * nframes - else: - video_timestamps = [format(ts, ".1f") for ts in video_timestamps] - - if video_frame_types is None: - video_frame_types = [0] * nframes - for i, sub_thw in enumerate(grid_thw): - s = f"{hf_processor.frame_token}{video_timestamps[i]}" - if video_frame_types[i] == 1: - s += hf_processor.fast_start - placeholders.extend(tokenizer.encode(s)) - num_frame_tokens = int(sub_thw.prod()) // merge_length - placeholders.extend([video_token_id] * num_frame_tokens) - if video_frame_types[i] == 1: - placeholders.append(vocab[hf_processor.fast_end]) - - return PromptUpdateDetails.select_token_id( - placeholders, embed_token_id=video_token_id - ) - else: - raise ValueError(f"Unsupported modality {modality}") - - return [ - PromptReplacement( - modality=modality, - target=[placeholder[modality]], - replacement=partial(get_replacement_keye, modality=modality), - ) - for modality in ("image", "video") - ] - - def _get_mm_fields_config( - self, - hf_inputs: BatchFeature, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> Mapping[str, MultiModalFieldConfig]: - return _keye_field_config(hf_inputs) - - -class KeyeVL1_5DummyInputsBuilder( - KeyeBaseDummyInputsBuilder[KeyeVL1_5ProcessingInfo] -): ... - - -@MULTIMODAL_REGISTRY.register_processor( - KeyeVL1_5MultiModalProcessor, - info=KeyeVL1_5ProcessingInfo, - dummy_inputs=KeyeVL1_5DummyInputsBuilder, -) -class KeyeVL1_5ForConditionalGeneration( - BaseKeyeModule, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE -): - def _build_projector( - self, - text_config: PreTrainedConfig, - vision_config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> nn.Module: - return KeyeVL1_5Projector(text_config, vision_config, quant_config, prefix) - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - config: PreTrainedConfig = vllm_config.model_config.hf_config - self.merge_size = config.vision_config.spatial_merge_size - super().__init__(vllm_config=vllm_config, prefix=prefix) - - def _parse_and_validate_image_input( - self, **kwargs: object - ) -> KeyeVL1_5ImageInputs | None: - pixel_values = kwargs.pop("pixel_values", None) - image_embeds = kwargs.pop("image_embeds", None) - image_grid_thw = kwargs.pop("image_grid_thw", None) - - if pixel_values is None and image_embeds is None: - return None - - if pixel_values is not None: - return KeyeVL1_5ImagePixelInputs( - type="pixel_values", - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - ) - - return KeyeVL1_5ImageEmbeddingInputs( - type="image_embeds", - image_embeds=image_embeds, - image_grid_thw=image_grid_thw, - ) - - def _parse_and_validate_video_input( - self, **kwargs: object - ) -> KeyeVL1_5VideoInputs | None: - pixel_values_videos = kwargs.pop("pixel_values_videos", None) - video_embeds = kwargs.pop("video_embeds", None) - video_grid_thw = kwargs.pop("video_grid_thw", None) - num_frames = kwargs.pop("num_frames", None) - - if pixel_values_videos is None and video_embeds is None: - return None - - if pixel_values_videos is not None: - return KeyeVL1_5VideoPixelInputs( - type="pixel_values_videos", - pixel_values_videos=pixel_values_videos, - video_grid_thw=video_grid_thw, - num_frames=num_frames, - ) - - return KeyeVL1_5VideoEmbeddingInputs( - type="video_embeds", - video_embeds=video_embeds, - video_grid_thw=video_grid_thw, - num_frames=num_frames, - ) - - def _process_video_input( - self, video_input: KeyeVL1_5VideoInputs - ) -> tuple[torch.Tensor, ...]: - video_type = video_input["type"] - video_grid_thw = split_thw(video_input["video_grid_thw"]) - pixel_values_videos = video_input.get("pixel_values_videos", None) - - video_embeds = self._process_video_embeds( - video_type, video_grid_thw, pixel_values_videos - ) - video_embeds = torch.concat(video_embeds, dim=0) - - num_frames = video_input["num_frames"].clone().tolist() - - num_patches = get_num_patches(video_grid_thw, num_frames).tolist() - - patch_cu_seqlens = torch.cumsum( - torch.tensor([0] + num_patches).detach().clone(), dim=-1 - ) - patch_cu_seqlens = torch.div( - patch_cu_seqlens, self.merge_size**2, rounding_mode="floor" - ) - - new_video_embeds = [] - for idx in range(patch_cu_seqlens.shape[0] - 1): - start = patch_cu_seqlens[idx] - end = patch_cu_seqlens[idx + 1] - new_video_embeds.append(video_embeds[start:end]) - return tuple(new_video_embeds) - - def iter_mm_grid_thw( - self, mm_features: list[MultiModalFeatureSpec] - ) -> Iterator[tuple[int, int, int, int]]: - spatial_merge_size = self.config.vision_config.spatial_merge_size - - for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): - if mm_feature.data is None: - raise ValueError("M-RoPE calculation requires multimodal feature data") - - embed_ranges = mm_feature.mm_position.extract_embeds_range() - if mm_feature.modality == "image": - assert len(embed_ranges) == 1 - grid_thw = mm_feature.data["image_grid_thw"].data - if isinstance(grid_thw, torch.Tensor): - if grid_thw.ndim == 2: - assert grid_thw.shape[0] == 1 - t, h, w = grid_thw[0].tolist() - else: - t, h, w = grid_thw.tolist() - else: - if isinstance(grid_thw[0], list): - assert len(grid_thw) == 1 - t, h, w = grid_thw[0] - else: - t, h, w = grid_thw - - yield ( - embed_ranges[0][0], - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - elif mm_feature.modality == "video": - split_video_grids = split_thw(mm_feature.data["video_grid_thw"].data) - assert len(embed_ranges) == split_video_grids.shape[0] - for (start_idx, end_idx), (t, h, w) in zip( - embed_ranges, split_video_grids.tolist() - ): - llm_grid_h = h // spatial_merge_size - llm_grid_w = w // spatial_merge_size - num_mm_tokens = t * llm_grid_h * llm_grid_w - assert end_idx - start_idx + 1 == num_mm_tokens - yield (start_idx, t, llm_grid_h, llm_grid_w) - else: - raise ValueError(f"Unsupported modality: {mm_feature.modality}") - - def get_mrope_input_positions( - self, - input_tokens: list[int], - mm_features: list[MultiModalFeatureSpec], - ) -> tuple[torch.Tensor, int]: - llm_pos_ids_list: list = [] - st = 0 - - for ( - offset, - llm_grid_t, - llm_grid_h, - llm_grid_w, - ) in self.iter_mm_grid_thw(mm_features): - text_len = offset - st - - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx - ) - - t_index = ( - ( - torch.arange(llm_grid_t) - .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - ) - .long() - .flatten() - ) - - h_index = ( - torch.arange(llm_grid_h) - .view(1, -1, 1) - .expand(llm_grid_t, -1, llm_grid_w) - .flatten() - ) - w_index = ( - torch.arange(llm_grid_w) - .view(1, 1, -1) - .expand(llm_grid_t, llm_grid_h, -1) - .flatten() - ) - llm_pos_ids_list.append( - torch.stack([t_index, h_index, w_index]) + text_len + st_idx - ) - st = offset + llm_grid_t * llm_grid_h * llm_grid_w - - if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - text_len = len(input_tokens) - st - llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx - ) - - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() - - return llm_positions, mrope_position_delta diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index d7f3c3bbf4a2..4050e8c73a58 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -478,11 +478,6 @@ ), "IsaacForConditionalGeneration": ("isaac", "IsaacForConditionalGeneration"), "KananaVForConditionalGeneration": ("kanana_v", "KananaVForConditionalGeneration"), - "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"), - "KeyeVL1_5ForConditionalGeneration": ( - "keye_vl1_5", - "KeyeVL1_5ForConditionalGeneration", - ), "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), "KimiK25ForConditionalGeneration": ("kimi_k25", "KimiK25ForConditionalGeneration"), "KimiK3ForConditionalGeneration": ( @@ -889,6 +884,8 @@ "StableLMEpochForCausalLM": "0.28.0", # Superseded by the generic "Terratorch" architecture "PrithviGeoSpatialMAE": "0.28.0", + "KeyeForConditionalGeneration": "0.33.0", + "KeyeVL1_5ForConditionalGeneration": "0.33.0", } _OOT_SUPPORTED_MODELS = { From ee9ba37199338cf31306a2ec27ec9afa41efc09e Mon Sep 17 00:00:00 2001 From: Dongji Gao Date: Wed, 28 Oct 2026 08:49:14 -0700 Subject: [PATCH 066/406] [Model] Add Nemotron-3.5-Transcribe (#64282) Signed-off-by: Dongji Gao Co-authored-by: Cursor Agent Co-authored-by: vllm-agent --- docs/models/supported_models.md | 14 + .../processing/test_nemotron3_5_transcribe.py | 460 ++++++++ tests/models/registry.py | 8 + tests/v1/spec_decode/test_dflash2.py | 57 + tests/v1/spec_decode/test_lilicorr.py | 6 +- vllm/config/speculative.py | 20 + .../model_executor/layers/logits_processor.py | 25 +- vllm/model_executor/models/config.py | 12 + .../models/nemotron3_5_transcribe.py | 670 +++++++++++ .../models/nemotron3_5_transcribe_audio.py | 1006 +++++++++++++++++ vllm/model_executor/models/qwen3_dflash2.py | 8 +- vllm/model_executor/models/registry.py | 8 + vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 4 + .../configs/nemotron3_5_transcribe.py | 334 ++++++ 15 files changed, 2623 insertions(+), 10 deletions(-) create mode 100644 tests/models/multimodal/processing/test_nemotron3_5_transcribe.py create mode 100644 vllm/model_executor/models/nemotron3_5_transcribe.py create mode 100644 vllm/model_executor/models/nemotron3_5_transcribe_audio.py create mode 100644 vllm/transformers_utils/configs/nemotron3_5_transcribe.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 3d6172d8ea9e..8cb133a66647 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -608,6 +608,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ | | `MuseGlimmerForCausalLM`, `MuseGlimmerForConditionalGeneration` | Muse Glimmer | T + I+ + V+ | `meta-models/Muse-Glimmer-30B` | ✅︎ | ✅︎ | | `MuseSparkForConditionalGeneration` | Muse Spark Open | T + I+ + V+ + A+ | `meta-models/Muse-Spark-Open-816B-A42B-NVFP4`, etc. | ✅︎ | ✅︎ | +| `Nemotron3_5TranscribeForConditionalGeneration` | Nemotron-3.5-Transcribe | T + A | `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-FP8`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-NVFP4` | | ✅︎ | | `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` | ✅︎ | | | `RadioMBartPrefixForCausalLM` | Nemotron Parse 3.5 | T + I | `nvidia/Nemotron-3.5-Parse` | | | | `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ | @@ -744,11 +745,24 @@ Speech2Text models trained specifically for Automatic Speech Recognition. | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ | | `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ | | `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | +| `Nemotron3_5TranscribeForConditionalGeneration` | Nemotron-3.5-Transcribe | `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-FP8`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-NVFP4` | | ✅︎ | | `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ | | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ | | `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ | | `WhisperForConditionalGeneration` | Whisper | `openai/whisper-small`, `openai/whisper-large-v3-turbo`, etc. | | | +!!! note + `Nemotron3_5TranscribeForConditionalGeneration` transcribes with speaker tags. Its chat template turns + thinking on by default, so transcription servers should set + `--default-chat-template-kwargs '{"enable_thinking": false}'`. It supports MTP and DFlash2 speculative + decoding. + + Long audio: by default vLLM refuses audio longer than 600 s and transcription uploads larger than 25 MB + (`VLLM_MAX_AUDIO_DECODE_DURATION_S`, `VLLM_MAX_AUDIO_CLIP_FILESIZE_MB`). A chat request takes one clip of + up to one hour, the model's training limit (`--limit-mm-per-prompt '{"audio": {"length": }}'` + changes it), and `/v1/audio/transcriptions` splits longer audio into pieces of at most one hour, shorter + when `--max-model-len` is below 65,536. + !!! note `VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed. diff --git a/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py b/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py new file mode 100644 index 000000000000..cba7d95ac931 --- /dev/null +++ b/tests/models/multimodal/processing/test_nemotron3_5_transcribe.py @@ -0,0 +1,460 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Nemotron-3.5-Transcribe config checks, audio limits and the audio tower.""" + +import json +from typing import Any + +import pytest +import torch + +from vllm.transformers_utils.configs.nemotron3_5_transcribe import ( + SPEECHLM_EMBED_EXTRA_ROWS, + Nemotron3_5TranscribeConfig, + _checked_online, + _preprocessor_config, +) + +pytestmark = pytest.mark.cpu_test + +ONLINE = { + "window_frames": 500, + "left_context_frames": 50, + "right_context_frames": 50, + "fifo_len": 40, + "spkcache_update_period": 300, + "spkcache_len": 188, +} + +AUDIO_CONFIG = { + "preprocessor": {"sample_rate": 16000, "n_fft": 512, "hop_length": 160}, + "asr_encoder": {"subsampling_factor": 8}, +} + + +def test_default_preprocessor_resolves(): + preprocessor = _preprocessor_config({"normalize": None}) + assert (preprocessor["n_mels"], preprocessor["win_length"]) == (64, 320) + # n_fft defaults to the next power of two at or above the window length. + assert preprocessor["n_fft"] == 512 + longer_window = _preprocessor_config({"normalize": None, "window_size": 0.05}) + assert longer_window["n_fft"] == 1024 + + +@pytest.mark.parametrize( + ("setting", "message"), + [ + ({}, "normalize"), # NeMo normalizes per feature when the key is missing. + ({"normalize": "per_feature"}, "normalize"), + ({"normalize": None, "log": False}, "log-mel"), + ({"normalize": None, "log_zero_guard_type": "clamp"}, "log_zero_guard_type"), + ], +) +def test_unsupported_mel_settings_are_refused(setting, message): + with pytest.raises(ValueError, match=message): + _preprocessor_config(setting) + + +def test_window_settings(): + assert _checked_online(dict(ONLINE)) == ONLINE + for key, value in ( + ("window_frames", 0), + ("left_context_frames", -1), + ("right_context_frames", -1), + ("fifo_len", -1), + ("spkcache_update_period", 0), + ("spkcache_len", 0), + ): + with pytest.raises(ValueError, match=key.split("_")[0]): + _checked_online(ONLINE | {key: value}) + + +def test_transcription_pieces_fit_the_context(): + from vllm.model_executor.models.nemotron3_5_transcribe import ( + _transcription_clip_s, + ) + + assert _transcription_clip_s(AUDIO_CONFIG, 65536) == 3600 # one hour at most + # 12.5 audio tokens and 4 transcript tokens per second in 16,384 - 256 tokens. + assert _transcription_clip_s(AUDIO_CONFIG, 16384) == 977 + assert _transcription_clip_s(AUDIO_CONFIG, 300) == 2 + + +def test_audio_length_sets_the_largest_clip(): + from vllm.model_executor.models.nemotron3_5_transcribe import ( + _max_audio_samples, + _num_audio_embeddings, + ) + + def max_tokens(length, seq_len=65536): + samples = _max_audio_samples(AUDIO_CONFIG, seq_len, length) + return _num_audio_embeddings(AUDIO_CONFIG, samples) + + assert max_tokens(None) == 45_000 # one hour by default + # 69 minutes, as `--limit-mm-per-prompt '{"audio": {"length": 66240000}}'` + # sets it, so a 65-minute clip (48,750 embeddings) fits the encoder cache. + assert max_tokens(66_240_000) == 51_750 + # Never more than the context holds beside a short prompt. + assert max_tokens(115_200_000) == 65_536 - 256 + + +ENCODER = { + "d_model": 64, + "n_heads": 2, + "n_layers": 1, + "self_attention_model": "rope", + "subsampling_factor": 8, +} + +# The smallest config.json in the released checkpoint's NeMo SpeechLM layout. +CHECKPOINT_CONFIG: dict[str, Any] = { + "architectures": ["Nemotron3_5TranscribeForConditionalGeneration"], + "model_type": "nemo_speechlm", + "llm_config": {"model_type": "nemotron_h", "hidden_size": 64, "vocab_size": 1000}, + "perception": { + "preprocessor": {"features": 128, "normalize": None, "window_size": 0.025}, + "modality_adapter": {"_target_": "IdentityConnector", "d_model": 64}, + "output_dim": 64, + }, + "pe_encoder_config": { + "asr_encoder_type": "transformer", + "asr_encoder_cfg": ENCODER, + "diarization_model_cfg": { + "encoder": ENCODER, + "high_resolution": True, + "sortformer_modules": {"use_learnable_sil_emb": True}, + }, + "speaker_feature_mode": "thresholded", + "speaker_activity_threshold": 0.5, + }, +} + + +def test_released_checkpoint_layout_loads(tmp_path): + from vllm.config import ModelConfig + + (tmp_path / "config.json").write_text(json.dumps(CHECKPOINT_CONFIG)) + config = ModelConfig(str(tmp_path)).hf_config + assert isinstance(config, Nemotron3_5TranscribeConfig) + preprocessor = config.audio_config["preprocessor"] + assert (preprocessor["win_length"], preprocessor["hop_length"]) == (400, 160) + assert config.audio_config["online"] == ONLINE + # The audio placeholder sits in embedding rows past the trained vocabulary. + assert config.output_vocab_size == 1000 + assert config.text_config.vocab_size == 1000 + SPEECHLM_EMBED_EXTRA_ROWS + + +@pytest.mark.parametrize( + ("changes", "hint"), + [ + # This model's structure under NeMo's generic name, as older exports have. + ( + {}, + 'set "architectures": ["Nemotron3_5TranscribeForConditionalGeneration"]', + ), + # Another NeMo SpeechLM: no Parallel Expert Encoder. + ({"pe_encoder_config": None}, "NeMo Speech vLLM plugin"), + ], +) +def test_other_nemo_speechlm_architectures_are_refused(tmp_path, changes, hint): + from vllm.transformers_utils.config import get_config + + other = ( + CHECKPOINT_CONFIG + | changes + | {"architectures": ["NeMoSpeechLMForConditionalGeneration"]} + ) + (tmp_path / "config.json").write_text(json.dumps(other)) + with pytest.raises(ValueError) as excinfo: + get_config(str(tmp_path), trust_remote_code=False) + assert hint in str(excinfo.value) + + +def _tiny_encoder(d_model: int, rotary_fraction: float, qk_norm: bool) -> dict: + return { + "feat_in": 16, + "d_model": d_model, + "n_heads": 4, + "n_layers": 2, + "ff_expansion": 4.0, + "qkv_bias": False, + "qk_norm": qk_norm, + "pre_block_norm": True, + "subsampling_factor": 8, + "rope_base": 10000.0, + "rotary_fraction": rotary_fraction, + "xscaling": False, + } + + +# 0.8 s windows, so a few seconds of audio spans several windows and fills and +# compresses the speaker cache. Continuous speaker features carry the speaker +# probabilities into the embeddings unthresholded. +TINY_AUDIO_CONFIG = { + "preprocessor": { + "sample_rate": 16000, + "win_length": 400, + "hop_length": 160, + "n_fft": 512, + "n_mels": 16, + "preemph": 0.97, + "log_zero_guard_value": 2**-24, + "mag_power": 2.0, + "pad_value": 0.0, + }, + "asr_encoder": _tiny_encoder(32, 0.5, True), + "speaker_encoder": _tiny_encoder(24, 1.0, False), + "sortformer": { + "num_speakers": 4, + "fc_d_model": 24, + "tf_d_model": 12, + "spkcache_sil_frames_per_spk": 1, + "pred_score_threshold": 0.25, + "max_index": 99999, + "scores_boost_latest": 0.05, + "strong_boost_rate": 0.75, + "weak_boost_rate": 1.5, + "min_pos_scores_rate": 0.5, + }, + "online": { + "window_frames": 10, + "left_context_frames": 2, + "right_context_frames": 2, + "fifo_len": 3, + "spkcache_update_period": 6, + "spkcache_len": 12, + }, + "fusion": { + "speaker_feature_mode": "continuous", + "speaker_activity_threshold": None, + "spk_kernel_scale": 1.0, + }, + "projection": {"input_dim": 32, "output_dim": 40}, +} + + +@pytest.fixture +def tiny_tower(dist_init): + from vllm.model_executor.layers.linear import ReplicatedLinear + from vllm.model_executor.models.nemotron3_5_transcribe_audio import ( + Nemotron3_5TranscribeAudioTower, + ) + + torch.manual_seed(0) + tower = Nemotron3_5TranscribeAudioTower(TINY_AUDIO_CONFIG).eval() + with torch.no_grad(): + # vLLM's linear layers allocate their weights uninitialized. + for module in tower.modules(): + if isinstance(module, ReplicatedLinear): + module.weight.normal_(0.0, module.input_size**-0.5) + if module.bias is not None: + module.bias.normal_(0.0, 0.1) + featurizer = tower.preprocessor.featurizer + featurizer.fb.uniform_() + window = torch.hann_window(featurizer.win_length, periodic=False) + featurizer.window.copy_(window) + sortformer = tower.encoder.diarization_model.sortformer_modules + sortformer.learnable_sil_emb.normal_() + for module in tower.modules(): + if isinstance(module, torch.nn.LayerNorm): + module.weight.normal_(1.0, 0.1) + module.bias.normal_(0.0, 0.1) + # Weight loading ends with this step; on CPU it selects the GEMM kernel. + for module in tower.modules(): + if isinstance(module, ReplicatedLinear): + module.quant_method.process_weights_after_loading(module) + return tower + + +@pytest.mark.parametrize("num_samples", [160, 161, 1_281, 12_345, 12_960, 48_001]) +def test_audio_tower_returns_one_embedding_per_audio_token(tiny_tower, num_samples): + from vllm.model_executor.models.nemotron3_5_transcribe import ( + _num_audio_embeddings, + ) + + clip = torch.randn(num_samples) * 0.1 + with torch.inference_mode(): + (embeddings,) = tiny_tower.encode([clip], torch.tensor([num_samples])) + num_tokens = _num_audio_embeddings(TINY_AUDIO_CONFIG, num_samples) + assert embeddings.shape == (num_tokens, 40) + + +def test_decoder_weights_read_both_checkpoint_layouts(): + from vllm.model_executor.models.nemotron3_5_transcribe import ( + Nemotron3_5TranscribeForConditionalGeneration as Model, + ) + from vllm.model_executor.models.nemotron3_5_transcribe import _decoder_weights + + nemo = { + "llm.model.embed_tokens.weight": torch.ones(8, 4), + "llm.model.layers.0.mixer._fp32_params.A_log": torch.ones(2), + "llm.model.layers.1.mixer.experts.down_projs": torch.arange(24.0).view(2, 3, 4), + "llm.model.norm.weight": torch.ones(4), + } + # The FP8 and NVFP4 checkpoints: Hugging Face names, one tensor per expert. + hf = { + "backbone.embeddings.weight": torch.ones(8, 4), + "backbone.layers.1.mixer.experts.0.up_proj.weight_scale": torch.ones(()), + "lm_head.weight": torch.ones(8, 4), + } + mapped = Model.hf_to_vllm_mapper.apply((nemo | hf).items()) + decoder = ((name.removeprefix("language_model."), t) for name, t in mapped) + out = dict(_decoder_weights(decoder, vocab_size=10)) + assert sorted(out) == [ + "lm_head.weight", + "model.embed_tokens.weight", + "model.embeddings.weight", + "model.layers.0.mixer.A_log", + "model.layers.1.mixer.experts.0.down_proj.weight", + "model.layers.1.mixer.experts.0.up_proj.weight_scale", + "model.layers.1.mixer.experts.1.down_proj.weight", + "model.norm_f.weight", + ] + for name in ( + "model.embed_tokens.weight", + "model.embeddings.weight", + "lm_head.weight", + ): + assert out[name].shape == (10, 4) and out[name][8:].eq(0).all() + assert torch.equal( + out["model.layers.1.mixer.experts.1.down_proj.weight"], + nemo["llm.model.layers.1.mixer.experts.down_projs"][1].t(), + ) + # Names outside both layouts stay unmapped, so load_weights refuses them. + unknown = [("llm.vision.weight", torch.ones(1))] + assert [name for name, _ in Model.hf_to_vllm_mapper.apply(unknown)] == [ + "llm.vision.weight" + ] + + +def test_only_audio_has_a_placeholder(): + from vllm.model_executor.models.nemotron3_5_transcribe import ( + AUDIO_PLACEHOLDER, + ) + from vllm.model_executor.models.nemotron3_5_transcribe import ( + Nemotron3_5TranscribeForConditionalGeneration as Model, + ) + + for modality in ("audio", "audio_embeds"): + assert Model.get_placeholder_str(modality, 1) == AUDIO_PLACEHOLDER + with pytest.raises(ValueError, match="Only audio modality is supported"): + Model.get_placeholder_str("image", 1) + + +def test_quantization_layer_names_map_onto_the_model(): + from vllm.model_executor.models.nemotron3_5_transcribe import ( + Nemotron3_5TranscribeForConditionalGeneration as Model, + ) + + mapper = Model.hf_to_vllm_mapper.get_rename_mapper() + assert mapper.apply_list( + [ + "backbone.layers.1.mixer.gate", + "lm_head", + "perception.encoder.asr_encoder.pre_encode", + # Names the NVFP4 recipe already gives in vLLM's form stay as they are. + "language_model.model.layers.0.mixer.in_proj", + "mtp.layers.1.mixer.gate", + ] + ) == [ + "language_model.model.layers.1.mixer.gate", + "language_model.lm_head", + "audio_tower.encoder.asr_encoder.pre_encode", + "language_model.model.layers.0.mixer.in_proj", + "mtp.layers.1.mixer.gate", + ] + + +@pytest.mark.usefixtures("dist_init") +def test_fp8_recipe_quantizes_the_asr_encoder_linear_layers(): + from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config + from vllm.model_executor.model_loader.utils import configure_quant_config + from vllm.model_executor.models.nemotron3_5_transcribe import ( + Nemotron3_5TranscribeForConditionalGeneration as Model, + ) + from vllm.model_executor.models.nemotron3_5_transcribe_audio import ( + Nemotron3_5TranscribeAudioTower, + ) + + # As in the FP8 checkpoint's config.json. + quant_config = ModelOptFp8Config.from_config( + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"dynamic": False, "num_bits": 8, "type": "float"}, + "input_activations": { + "dynamic": False, + "num_bits": 8, + "type": "float", + }, + } + }, + "ignore": [ + "perception.encoder.diarization_model*", + "perception.encoder.asr_encoder.pre_encode*", + "perception.proj", + ], + } + ) + configure_quant_config(quant_config, Model) + tower = Nemotron3_5TranscribeAudioTower( + TINY_AUDIO_CONFIG, quant_config=quant_config, prefix="audio_tower" + ) + fp8_layers = sorted( + name + for name, module in tower.named_modules() + if isinstance(getattr(module, "weight", None), torch.Tensor) + and module.weight.dtype == torch.float8_e4m3fn + ) + assert fp8_layers == sorted( + f"encoder.asr_encoder.layers.{i}.{layer}" + for i in range(2) + for layer in ("attn.w_qkv", "attn.out_proj", "ffn.net.0", "ffn.net.3") + ) + + # The checkpoint stores each scale as a scalar. + weights = [] + for name, target in [*tower.named_parameters(), *tower.named_buffers()]: + if name == "encoder.diar_kernel": + continue + tensor = torch.full(target.shape, 0.5).to(target.dtype) + if name.endswith("_scale"): + tensor = torch.tensor(0.25) + weights.append((name, tensor)) + tower.load_weights(weights) + w_qkv = tower.get_submodule("encoder.asr_encoder.layers.0.attn.w_qkv") + assert w_qkv.weight_scale.item() == w_qkv.input_scale.item() == 0.25 + + +def test_mtp_draft_config_keeps_the_quantization_recipe(): + from vllm.config.speculative import SpeculativeConfig + + recipe = {"quant_method": "modelopt", "quant_algo": "FP8", "ignore": ["lm_head"]} + llm_config = CHECKPOINT_CONFIG["llm_config"] | {"num_nextn_predict_layers": 1} + config = Nemotron3_5TranscribeConfig.from_dict( + CHECKPOINT_CONFIG | {"llm_config": llm_config, "quantization_config": recipe} + ) + draft = SpeculativeConfig.hf_config_override(config) + assert draft.architectures == ["Nemotron3_5TranscribeMTPModel"] + # The recipe lives on the SpeechLM config; the MTP head is quantized too. + assert draft.quantization_config == recipe + + +def test_audio_tower_encodes_each_clip_as_if_alone(tiny_tower): + torch.manual_seed(1) + # The 1.3 s and 1.5 s clips span two windows and share a padded batch; the + # 3 s clip spans four windows and is encoded apart, between them in order. + clips = [torch.randn(num_samples) * 0.1 for num_samples in (20_800, 48_000, 24_000)] + lengths = torch.tensor([clip.numel() for clip in clips]) + with torch.inference_mode(): + together = tiny_tower.encode(clips, lengths) + alone = [ + tiny_tower.encode([clip], lengths[index : index + 1])[0] + for index, clip in enumerate(clips) + ] + for got, expected in zip(together, alone): + torch.testing.assert_close(got, expected, rtol=1e-4, atol=1e-4) diff --git a/tests/models/registry.py b/tests/models/registry.py index 58e6208ec7f5..8743775a3368 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1257,6 +1257,9 @@ def check_available_online( "NemotronH_Omni_Reasoning_V3": _HfExamplesInfo( "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", is_available_online=False ), + "Nemotron3_5TranscribeForConditionalGeneration": _HfExamplesInfo( + "nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B", is_available_online=False + ), "OpenCUAForConditionalGeneration": _HfExamplesInfo( "xlangai/OpenCUA-7B", trust_remote_code=True, @@ -1854,6 +1857,11 @@ def check_available_online( speculative_model="nvidia/Nemotron-Super-Placeholder", is_available_online=False, ), + "Nemotron3_5TranscribeMTPModel": _HfExamplesInfo( + "nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B", + speculative_model="nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B", + is_available_online=False, + ), "OpenPanguMTPModel": _HfExamplesInfo( "FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1", trust_remote_code=True, diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 94de5f90a369..53641f32f249 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from unittest import mock import pytest import torch +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.qwen3_dflash2 import _grouped_conv, _score_edges from vllm.platforms import current_platform from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator @@ -797,3 +800,57 @@ def forward(self, hidden_states): torch.testing.assert_close(actual_k, expected_k) torch.testing.assert_close(actual_v, expected_v) assert [projection.calls for projection in projections] == [1, 1] + + +def test_score_edges_keeps_out_of_vocabulary_anchors_in_range(): + """A request in chunked prefill anchors on its next prompt token, which can + lie past the draft vocabulary (a multimodal placeholder).""" + generator = torch.Generator().manual_seed(0) + vocab_size, rank, top_k = 8, 4, 3 + predecessor_table = torch.randn(vocab_size, rank, generator=generator) + successor_table = torch.randn(vocab_size, rank, generator=generator) + candidate_ids = torch.randint(0, vocab_size, (2, 5, top_k), generator=generator) + unary_logits = torch.randn(2, 5, top_k, generator=generator) + hidden = torch.randn(2, 5, rank, generator=generator) + + def score(anchor_token_ids: list[int]) -> torch.Tensor: + return _score_edges( + predecessor_table, + successor_table, + candidate_ids, + unary_logits, + hidden, + torch.tensor(anchor_token_ids), + top_k, + ) + + torch.testing.assert_close(score([vocab_size, 2]), score([vocab_size - 1, 2])) + + +_EMBEDDING = "vllm.model_executor.layers.vocab_parallel_embedding" + + +def test_top_k_tokens_stay_in_the_processor_vocabulary(default_vllm_config): + """A DFlash2 drafter takes candidates from the target's head, which can have + rows past the draft vocabulary (a padded multimodal vocabulary). The + candidates index the draft's codebooks, so top-k must not return those rows. + """ + head_rows, vocab_size, hidden_size = 140, 128, 16 + # One vocabulary shard, so the reductions need no process group. + with ( + mock.patch(f"{_EMBEDDING}.get_tensor_model_parallel_rank", return_value=0), + mock.patch( + f"{_EMBEDDING}.get_tensor_model_parallel_world_size", return_value=1 + ), + ): + head = ParallelLMHead(head_rows, hidden_size) + with torch.no_grad(): + head.weight.zero_() + head.weight[vocab_size:head_rows] = 1.0 + head.quant_method.process_weights_after_loading(head) + hidden_states = torch.ones(4, hidden_size, dtype=head.weight.dtype) + processor = LogitsProcessor(vocab_size) + + _, ids = processor.get_top_k_tokens(head, hidden_states, 4) + assert (ids < vocab_size).all() + assert (processor.get_top_tokens(head, hidden_states) < vocab_size).all() diff --git a/tests/v1/spec_decode/test_lilicorr.py b/tests/v1/spec_decode/test_lilicorr.py index 7ded16e68e88..f6d7834ebe6e 100644 --- a/tests/v1/spec_decode/test_lilicorr.py +++ b/tests/v1/spec_decode/test_lilicorr.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from typing import Any import pytest @@ -228,8 +228,12 @@ def __call__(self, value, dim=-1): fake = SimpleNamespace( scale=1.0, soft_cap=None, + org_vocab_size=3 * tp_size, _apply_head=lambda *args, rank=rank: shards[rank].clone(), ) + fake._mask_outside_vocab = MethodType( + logits_processor.LogitsProcessor._mask_outside_vocab, fake + ) lm_head = SimpleNamespace( tp_size=tp_size, shard_indices=SimpleNamespace( diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 9a6f1a8afc7b..24b519af3d13 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -45,6 +45,7 @@ "glm_ocr_mtp", "ernie_mtp", "nemotron_h_mtp", + "nemotron3_5_transcribe_mtp", "exaone_moe_mtp", "exaone4_5_mtp", "qwen3_next_mtp", @@ -884,6 +885,25 @@ def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig: # Promote VLM's text_config so MTP detection below fires correctly hf_config = hf_config.text_config + if ( + hf_config.architectures[0] + == "Nemotron3_5TranscribeForConditionalGeneration" + ): + # Its MTP head reads NeMo checkpoint names and takes the audio + # embeddings, so it has its own draft class. + quantization_config = getattr(hf_config, "quantization_config", None) + hf_config = hf_config.text_config + if quantization_config is not None: + # The FP8 and NVFP4 checkpoints quantize the MTP head too. + hf_config.quantization_config = copy.deepcopy(quantization_config) + hf_config.model_type = "nemotron3_5_transcribe_mtp" + hf_config.update( + { + "n_predict": hf_config.num_nextn_predict_layers, + "architectures": ["Nemotron3_5TranscribeMTPModel"], + } + ) + if ( hf_config.model_type in {"nemotron_h", "nemotron_h_puzzle"} and hasattr(hf_config, "num_nextn_predict_layers") diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index 04541212fe0e..1ca1d71381bc 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -198,6 +198,21 @@ def _get_logits( logits = logits[..., : self.org_vocab_size] return logits + def _mask_outside_vocab( + self, logits: torch.Tensor, lm_head: VocabParallelEmbedding + ) -> None: + """Mask this shard's padding and any head rows past ``org_vocab_size``. + + A head shared with another model can have rows past this processor's + vocabulary, as ``_get_logits`` drops them. + """ + shard_indices = lm_head.shard_indices + if shard_indices.num_org_vocab_padding > 0: + logits[..., -shard_indices.num_org_vocab_padding :] = -float("inf") + cutoff = self.org_vocab_size - shard_indices.org_vocab_start_index + if cutoff < logits.shape[-1]: + logits[..., max(cutoff, 0) :] = -float("inf") + def get_top_tokens( self, lm_head: VocabParallelEmbedding, @@ -223,10 +238,7 @@ def get_top_tokens( if self.scale != 1.0: logits = logits * self.scale - # Mask out padding entries beyond org_vocab_size on this shard. - num_pad = lm_head.shard_indices.num_org_vocab_padding - if num_pad > 0: - logits[..., -num_pad:] = -float("inf") + self._mask_outside_vocab(logits, lm_head) local_max_vals, local_max_indices = logits.max(dim=-1) @@ -284,10 +296,7 @@ def get_top_k_tokens( if self.scale != 1.0: logits = logits * self.scale - # Mask out padding entries beyond org_vocab_size on this shard. - num_pad = lm_head.shard_indices.num_org_vocab_padding - if num_pad > 0: - logits[..., -num_pad:] = -float("inf") + self._mask_outside_vocab(logits, lm_head) if return_log_probs: log_partition = torch.logsumexp(logits.float(), dim=-1, keepdim=True) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index e40171abcbaa..99d3b80eef3a 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -904,6 +904,15 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: video_kwargs.setdefault("video_backend", "nemotron_vl") +class Nemotron3_5TranscribeForConditionalGenerationConfig(VerifyAndUpdateConfig): + @classmethod + def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: + NemotronHForCausalLMConfig.update_mamba_ssm_cache_dtype( + cache_config=vllm_config.cache_config, + hf_config=vllm_config.model_config.hf_config.text_config, + ) + + class NomicBertModelConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -1292,6 +1301,9 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "NemotronHForCausalLM": NemotronHForCausalLMConfig, "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig, "NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config, + "Nemotron3_5TranscribeForConditionalGeneration": ( + Nemotron3_5TranscribeForConditionalGenerationConfig + ), "NomicBertModel": NomicBertModelConfig, "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig, "Qwen2ForRewardModel": Qwen2ForRewardModelConfig, diff --git a/vllm/model_executor/models/nemotron3_5_transcribe.py b/vllm/model_executor/models/nemotron3_5_transcribe.py new file mode 100644 index 000000000000..6ec2147418ec --- /dev/null +++ b/vllm/model_executor/models/nemotron3_5_transcribe.py @@ -0,0 +1,670 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Inference-only NVIDIA Nemotron-3.5-Transcribe. + +A NeMo speech-language model. The audio tower (log-mel, Parallel Expert +Encoder, projection; see ``nemotron3_5_transcribe_audio.py``) turns each +80 ms of audio into one embedding, and those embeddings replace the +``<|audio|>`` placeholders in the prompt of a NemotronH decoder (hybrid +Mamba, MoE and attention layers). + +The checkpoint uses NeMo names: the audio tower under ``perception.*``, the +decoder under ``llm.*``, the multi-token-prediction head under ``llm.mtp.*`` +and an optional CTC timestamp head, which is not loaded, under +``ctc_timestamp.*``. +""" + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Annotated, Any, Literal, cast + +import numpy as np +import torch +from torch import nn +from transformers import BatchFeature + +from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig +from vllm.config.multimodal import MultiModalDummyOptions +from vllm.config.speech_to_text import SpeechToTextParams +from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt +from vllm.logger import init_logger +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.sequence import IntermediateTensors +from vllm.tokenizers import cached_tokenizer_from_config +from vllm.transformers_utils.configs.nemotron3_5_transcribe import ( + Nemotron3_5TranscribeConfig, +) +from vllm.utils.tensor_schema import TensorSchema, TensorShape + +from .interfaces import ( + IsHybrid, + MultiModalEmbeddings, + SupportsEagle3, + SupportsMambaPrefixCaching, + SupportsMultiModal, + SupportsMultiModalEmbeddings, + SupportsPP, + SupportsTranscription, +) +from .module_mapping import MultiModelKeys +from .nemotron3_5_transcribe_audio import Nemotron3_5TranscribeAudioTower +from .nemotron_h import NemotronHForCausalLM +from .nemotron_h_mtp import NemotronHMTP +from .utils import ( + WeightsMapper, + _merge_multimodal_embeddings, + init_vllm_registered_model, + maybe_prefix, +) + +logger = init_logger(__name__) + +AUDIO_PLACEHOLDER = "<|audio|>" +# Default longest audio used to profile memory, unless --limit-mm-per-prompt +# sets an audio length. +_DEFAULT_MAX_AUDIO_DURATION_S = 3600.0 +# Prompt tokens kept free beside the longest dummy audio. +_DUMMY_TEXT_TOKEN_RESERVE = 256 + +# The transcription instruction of the model card's evaluation prompts. The +# model was trained with the audio after the instruction. +_TRANSCRIPTION_PROMPT = ( + "Audio transcription. Language: as spoken. " + "Wording: verbatim, including fillers and repetitions. " + "Formatting: punctuation and capitalization. Output: transcript only." +) +# Training audio was at most one hour, so transcription sends pieces of at most an +# hour; vLLM splits longer audio at quiet points. +_MAX_TRANSCRIPTION_CLIP_S = 60 * 60 +# Context kept free for the transcript of each second of audio: a fast speaker, +# with speaker tags. +_TRANSCRIPT_TOKENS_PER_S = 4.0 + +# The languages of the model card's FLEURS evaluation. +ISO639_1_SUPPORTED_LANGS = { + "ar": "Arabic", + "bg": "Bulgarian", + "cs": "Czech", + "da": "Danish", + "de": "German", + "el": "Greek", + "en": "English", + "es": "Spanish", + "et": "Estonian", + "fi": "Finnish", + "fr": "French", + "he": "Hebrew", + "hi": "Hindi", + "hr": "Croatian", + "hu": "Hungarian", + "id": "Indonesian", + "it": "Italian", + "ja": "Japanese", + "ko": "Korean", + "lt": "Lithuanian", + "lv": "Latvian", + "mt": "Maltese", + "nl": "Dutch", + "no": "Norwegian", + "pl": "Polish", + "pt": "Portuguese", + "ro": "Romanian", + "ru": "Russian", + "sk": "Slovak", + "sl": "Slovenian", + "sv": "Swedish", + "th": "Thai", + "tl": "Filipino", + "tr": "Turkish", + "uk": "Ukrainian", + "vi": "Vietnamese", + "zh": "Chinese", +} + + +def _num_audio_embeddings(audio_config: Mapping[str, Mapping], num_samples: int) -> int: + """Embeddings the audio tower produces for ``num_samples`` samples.""" + preprocessor = audio_config["preprocessor"] + factor = audio_config["asr_encoder"]["subsampling_factor"] + n_fft = preprocessor["n_fft"] + num_frames = (num_samples + n_fft // 2 * 2 - n_fft) // preprocessor["hop_length"] + return (num_frames + factor - 1) // factor + + +def _transcription_clip_s( + audio_config: Mapping[str, Mapping], max_model_len: int +) -> int: + """Longest piece ``/v1/audio/transcriptions`` sends: at most an hour, and short + enough that its embeddings and transcript fit in the context.""" + sample_rate = audio_config["preprocessor"]["sample_rate"] + hour = 3600 + tokens_per_s = ( + _num_audio_embeddings(audio_config, hour * sample_rate) / hour + + _TRANSCRIPT_TOKENS_PER_S + ) + fits_s = int((max_model_len - _DUMMY_TEXT_TOKEN_RESERVE) / tokens_per_s) + return max(1, min(_MAX_TRANSCRIPTION_CLIP_S, fits_s)) + + +def _max_audio_samples( + audio_config: Mapping[str, Mapping], seq_len: int, override_samples: int | None +) -> int: + """Longest audio whose embeddings fit in ``seq_len`` beside a short prompt.""" + preprocessor = audio_config["preprocessor"] + min_samples = preprocessor["hop_length"] + limit = int( + override_samples or _DEFAULT_MAX_AUDIO_DURATION_S * preprocessor["sample_rate"] + ) + + def num_tokens(num_samples: int) -> int: + return _num_audio_embeddings(audio_config, max(num_samples, min_samples)) + + max_tokens = max(1, seq_len - _DUMMY_TEXT_TOKEN_RESERVE) + if num_tokens(limit) <= max_tokens: + return limit + low, high = min_samples, limit + while low < high: + mid = (low + high + 1) // 2 + if num_tokens(mid) <= max_tokens: + low = mid + else: + high = mid - 1 + return low + + +class Nemotron3_5TranscribeAudioInputs(TensorSchema): + """Dimensions: + - b: Number of audio clips + - t: Number of samples (16 kHz) + """ + + type: Literal["audio_signal"] = "audio_signal" + audio_signal: Annotated[ + torch.Tensor | list[torch.Tensor], + TensorShape("b", "t", dynamic_dims={"t"}), + ] + audio_signal_length: Annotated[torch.Tensor, TensorShape("b")] + + +class Nemotron3_5TranscribeProcessingInfo(BaseProcessingInfo): + def get_hf_config(self) -> Nemotron3_5TranscribeConfig: + return self.ctx.get_hf_config(Nemotron3_5TranscribeConfig) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": None} + + def get_data_parser(self) -> MultiModalDataParser: + return MultiModalDataParser( + target_sr=self.get_sampling_rate(), target_channels=1 + ) + + def get_sampling_rate(self) -> int: + return self.get_hf_config().audio_config["preprocessor"]["sample_rate"] + + def get_min_num_samples(self) -> int: + """Shorter clips are zero-padded so that they yield one embedding.""" + return self.get_hf_config().audio_config["preprocessor"]["hop_length"] + + def get_num_audio_tokens(self, num_samples: int) -> int: + num_samples = max(num_samples, self.get_min_num_samples()) + return _num_audio_embeddings(self.get_hf_config().audio_config, num_samples) + + def get_audio_token_id(self) -> int: + return self.get_tokenizer().convert_tokens_to_ids(AUDIO_PLACEHOLDER) + + def get_max_audio_samples( + self, seq_len: int, override_samples: int | None = None + ) -> int: + """Longest audio whose embeddings fit in ``seq_len`` beside a short prompt.""" + return _max_audio_samples( + self.get_hf_config().audio_config, seq_len, override_samples + ) + + def get_audio_length_override(self) -> int | None: + """The audio length set with ``--limit-mm-per-prompt``, in samples.""" + options = self.ctx.get_mm_config().limit_per_prompt.get("audio") + return getattr(options, "length", None) + + def get_mm_max_tokens_per_item( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> Mapping[str, int]: + # The same longest clip as the dummy inputs, so the encoder cache that + # vLLM sizes from this number admits every clip the limits allow. + samples = self.get_max_audio_samples(seq_len, self.get_audio_length_override()) + return {"audio": self.get_num_audio_tokens(samples)} + + +class Nemotron3_5TranscribeDummyInputsBuilder( + BaseDummyInputsBuilder[Nemotron3_5TranscribeProcessingInfo] +): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return AUDIO_PLACEHOLDER * mm_counts.get("audio", 0) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: MultiModalDummyOptions, + ) -> MultiModalDataDict: + audio_options = mm_options.get("audio") if mm_options else None + override = getattr(audio_options, "length", None) + if override is None: + override = self.info.get_audio_length_override() + return { + "audio": self._get_dummy_audios( + length=self.info.get_max_audio_samples(seq_len, override), + num_audios=mm_counts.get("audio", 0), + ) + } + + +class Nemotron3_5TranscribeMultiModalProcessor( + BaseMultiModalProcessor[Nemotron3_5TranscribeProcessingInfo] +): + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + # There is no Hugging Face processor: the model computes features + # from the waveform, so pass each clip through as float32 samples. + min_samples = self.info.get_min_num_samples() + signals, lengths = [], [] + for audio in cast(Iterable[Any], hf_data.get("audio", [])): + signal = torch.as_tensor(np.asarray(audio, dtype=np.float32)).flatten() + if signal.shape[0] < min_samples: + signal = torch.nn.functional.pad( + signal, (0, min_samples - signal.shape[0]) + ) + signals.append(signal) + lengths.append(signal.shape[0]) + return BatchFeature( + { + "audio_signal": signals, + "audio_signal_length": torch.tensor(lengths, dtype=torch.long), + } + ) + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return dict( + audio_signal=MultiModalFieldConfig.batched("audio"), + audio_signal_length=MultiModalFieldConfig.batched( + "audio", keep_on_cpu=True + ), + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + audio_token_id = self.info.get_audio_token_id() + lengths = out_mm_kwargs.get_data().get("audio_signal_length") + num_tokens = ( + [] + if lengths is None + else [ + self.info.get_num_audio_tokens(int(n)) + for n in torch.as_tensor(lengths).tolist() + ] + ) + + def get_replacement(item_idx: int) -> PromptUpdateDetails: + return PromptUpdateDetails.select_token_id( + [audio_token_id] * num_tokens[item_idx], embed_token_id=audio_token_id + ) + + return [ + PromptReplacement( + modality="audio", + target=[audio_token_id], + replacement=get_replacement, + ) + ] + + +def _unpack_experts( + name: str, tensor: torch.Tensor +) -> Iterator[tuple[str, torch.Tensor]]: + """NeMo packs the routed experts of a MoE layer into one rank-3 tensor per + projection; vLLM loads one transposed matrix per expert.""" + for suffix, projection in ( + (".experts.down_projs", "down_proj"), + (".experts.gate_and_up_projs", "up_proj"), + ): + if name.endswith(suffix): + prefix = name[: -len(suffix)] + for expert in range(tensor.shape[0]): + yield ( + f"{prefix}.experts.{expert}.{projection}.weight", + tensor[expert].t(), + ) + return + yield name, tensor + + +def _pad_vocab_rows(tensor: torch.Tensor, vocab_size: int) -> torch.Tensor: + if tensor.shape[0] >= vocab_size: + return tensor + return torch.cat( + [tensor, tensor.new_zeros(vocab_size - tensor.shape[0], tensor.shape[1])], dim=0 + ) + + +def _decoder_weights( + weights: Iterable[tuple[str, torch.Tensor]], vocab_size: int +) -> Iterator[tuple[str, torch.Tensor]]: + """Shape the mapped decoder tensors for ``NemotronHForCausalLM``. + + NeMo packs the routed experts of each MoE layer into two tensors; the FP8 and + NVFP4 checkpoints have one tensor per expert, with its quantization scales. + The embedding and LM head are zero-padded to the decoder's padded vocabulary. + """ + for name, tensor in weights: + if name in ( + "model.embed_tokens.weight", + "model.embeddings.weight", + "lm_head.weight", + ): + tensor = _pad_vocab_rows(tensor, vocab_size) + yield from _unpack_experts(name, tensor) + + +def _mtp_weights( + weights: Iterable[tuple[str, torch.Tensor]], +) -> Iterator[tuple[str, torch.Tensor]]: + """Select the MTP head's NeMo tensors and rename them for ``NemotronHMTP``. + + The head is ``llm.mtp.*``; the drafter shares the target's embedding and LM + head. + """ + for name, tensor in weights: + if name.startswith("llm.mtp.") and not name.endswith("._extra_state"): + yield from _unpack_experts(name[len("llm.") :], tensor) + + +@MULTIMODAL_REGISTRY.register_processor( + Nemotron3_5TranscribeMultiModalProcessor, + info=Nemotron3_5TranscribeProcessingInfo, + dummy_inputs=Nemotron3_5TranscribeDummyInputsBuilder, +) +class Nemotron3_5TranscribeForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + IsHybrid, + SupportsMambaPrefixCaching, + SupportsEagle3, + SupportsTranscription, +): + packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping + # The BF16 checkpoint names the decoder in NeMo's layout (``llm.*``; Mamba + # ``A_log``/``D``/``dt_bias`` in a float32 holder; final norm ``norm``). The FP8 + # and NVFP4 checkpoints, and their quantization configs, use Hugging Face names. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={".mixer._fp32_params.": ".mixer."}, + orig_to_new_prefix={ + "llm.model.norm.": "language_model.model.norm_f.", + "llm.model.": "language_model.model.", + "llm.lm_head.": "language_model.lm_head.", + "backbone.": "language_model.model.", + "lm_head": "language_model.lm_head", + "perception.": "audio_tower.", + }, + ) + supported_languages = ISO639_1_SUPPORTED_LANGS + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("audio"): + return AUDIO_PLACEHOLDER + raise ValueError("Only audio modality is supported") + + @classmethod + def get_speech_to_text_config( + cls, model_config: ModelConfig, task_type: str + ) -> SpeechToTextConfig: + audio_config = model_config.hf_config.audio_config + return SpeechToTextConfig( + sample_rate=audio_config["preprocessor"]["sample_rate"], + max_audio_clip_s=_transcription_clip_s( + audio_config, model_config.max_model_len + ), + ) + + @classmethod + def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: + """The model card's transcription prompt, whatever the request's language. + + The prompt says "Language: as spoken" and the model detects the + language. No published prompt names it, so the request's language is + only validated. + """ + if stt_params.task_type != "transcribe": + raise ValueError( + "Nemotron-3.5-Transcribe does not support task " + f"{stt_params.task_type!r} yet." + ) + if stt_params.request_prompt or stt_params.hotwords: + logger.warning_once( + "Nemotron-3.5-Transcribe ignores the prompt and hotwords of " + "transcription requests." + ) + tokenizer = cached_tokenizer_from_config(stt_params.model_config) + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": f"{_TRANSCRIPTION_PROMPT} {AUDIO_PLACEHOLDER}", + } + ], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + return TokensPrompt( + prompt_token_ids=tokenizer.encode(prompt, add_special_tokens=False), + multi_modal_data={ + "audio": (stt_params.audio, stt_params.stt_config.sample_rate) + }, + ) + + @classmethod + def get_num_audio_tokens( + cls, + audio_duration_s: float, + stt_config: SpeechToTextConfig, + model_config: ModelConfig, + ) -> int: + audio_config = model_config.hf_config.audio_config + num_samples = max( + int(audio_duration_s * stt_config.sample_rate), + audio_config["preprocessor"]["hop_length"], + ) + return _num_audio_embeddings(audio_config, num_samples) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: Nemotron3_5TranscribeConfig = vllm_config.model_config.hf_config + self.config = config + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config.text_config, + prefix=maybe_prefix(prefix, "language_model"), + ) + with self._mark_tower_model(vllm_config, "audio"): + self.audio_tower = Nemotron3_5TranscribeAudioTower( + config.audio_config, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "audio_tower"), + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + def _parse_audio_input( + self, **kwargs: object + ) -> Nemotron3_5TranscribeAudioInputs | None: + audio_signal = kwargs.pop("audio_signal", None) + if audio_signal is None: + return None + return Nemotron3_5TranscribeAudioInputs( + audio_signal=audio_signal, + audio_signal_length=kwargs.pop("audio_signal_length"), + ) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + audio_input = self._parse_audio_input(**kwargs) + if audio_input is None: + return [] + signals = audio_input.audio_signal + if isinstance(signals, torch.Tensor): + signals = list(signals) + return tuple( + self.audio_tower.encode(signals, audio_input.audio_signal_length.flatten()) + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + return self.language_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = self.language_model.compute_logits(hidden_states) + if logits is not None: + # Rows past the trained vocabulary (the audio placeholder and the + # embedding headroom) are never generated. + logits[..., self.config.output_vocab_size :] = -torch.inf + return logits + + def get_mm_mapping(self) -> MultiModelKeys: + return MultiModelKeys.from_string_field( + language_model="language_model", + connector="audio_tower.proj", + tower_model=["audio_tower.preprocessor", "audio_tower.encoder"], + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + load_audio = isinstance(self.audio_tower, Nemotron3_5TranscribeAudioTower) + audio_weights: list[tuple[str, torch.Tensor]] = [] + + def decoder_source() -> Iterator[tuple[str, torch.Tensor]]: + for name, tensor in self.hf_to_vllm_mapper.apply(weights): + if name.endswith("._extra_state"): + continue + if name.startswith("audio_tower."): + if load_audio: + # Copied so the tensor outlives a reusable read buffer. + audio_weights.append( + (name[len("audio_tower.") :], tensor.detach().clone()) + ) + elif name.startswith("llm.mtp."): + continue # loaded by the multi-token-prediction drafter + elif name.startswith("ctc_timestamp."): + continue # the CTC timestamp head and its tokenizer + elif name.startswith("language_model."): + yield name[len("language_model.") :], tensor + else: + raise ValueError(f"Unexpected checkpoint weight {name!r}.") + + decoder_weights = _decoder_weights( + decoder_source(), self.config.text_config.vocab_size + ) + loaded = { + f"language_model.{name}" + for name in self.language_model.load_weights(decoder_weights) + } + if load_audio: + loaded |= { + f"audio_tower.{name}" + for name in self.audio_tower.load_weights(audio_weights) + } + return loaded + + @classmethod + def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + return NemotronHForCausalLM.get_mamba_state_shape_from_config( + vllm_config.with_hf_config(text_config) + ) + + @classmethod + def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): + text_config = vllm_config.model_config.hf_config.text_config + return NemotronHForCausalLM.get_mamba_state_dtype_from_config( + vllm_config.with_hf_config(text_config) + ) + + @classmethod + def get_mamba_state_copy_func(cls): + return NemotronHForCausalLM.get_mamba_state_copy_func() + + +class Nemotron3_5TranscribeMTP(NemotronHMTP, SupportsMultiModalEmbeddings): + """Multi-token-prediction draft head of Nemotron-3.5-Transcribe. + + One attention-plus-MoE block, trained to be reused for every draft token. + It was trained on the decoder's input stream, so it receives the audio + embeddings at the ``<|audio|>`` positions, not the placeholder embedding. + """ + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self.model.get_input_embeddings(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + if is_multimodal is None: + raise ValueError("is_multimodal is required with multimodal_embeddings.") + return _merge_multimodal_embeddings( + inputs_embeds, multimodal_embeddings, is_multimodal + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = super().compute_logits(hidden_states) + if logits is not None: + logits[..., self.config.speechlm_output_vocab_size :] = -torch.inf + return logits + + def is_unused_checkpoint_weight(self, name: str) -> bool: + # Shard selection passes the raw llm.mtp.* names; load_weights passes + # the mtp.* names that _mtp_weights produces. + return super().is_unused_checkpoint_weight(name.removeprefix("llm.")) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return super().load_weights(_mtp_weights(weights)) diff --git a/vllm/model_executor/models/nemotron3_5_transcribe_audio.py b/vllm/model_executor/models/nemotron3_5_transcribe_audio.py new file mode 100644 index 000000000000..47552bca867f --- /dev/null +++ b/vllm/model_executor/models/nemotron3_5_transcribe_audio.py @@ -0,0 +1,1006 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Audio tower of Nemotron-3.5-Transcribe. + +A native port of the inference path of NeMo's ``AudioPerceptionModule`` with a +``ParallelExpertEncoder``: + + waveform -> log-mel -> Parallel Expert Encoder -> Linear -> decoder width + +The Parallel Expert Encoder runs an ASR ``TransformerEncoder`` and a streaming +Sortformer speaker encoder over the same mel features. Both run in windows of +``window_frames`` encoder frames with left and right context, and the speaker +encoder carries a speaker cache and a FIFO between windows. Thresholded speaker +activity is added to the normalized ASR states through a fixed sinusoidal +kernel. Module names follow NeMo so the checkpoint loads by prefix. +""" + +import math +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any, cast + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.utils.torch_utils import async_tensor_h2d + +from .utils import maybe_prefix + +# Checkpoint tensors that the inference path does not use. The speaker +# encoder's own preprocessor is unused because both encoders read the mel +# features of the perception preprocessor; ``hidden_to_spks`` is a frozen +# legacy head. +_UNUSED_CHECKPOINT_PREFIXES = ( + "encoder.diarization_model.preprocessor.", + "encoder.diarization_model.sortformer_modules.hidden_to_spks.", +) + + +class _TowerLinear(ReplicatedLinear): + """A replicated linear layer over the last dimension of any input shape. + + vLLM's quantized kernels take ``(tokens, features)``; the tower's activations + are ``(batch, frames, features)``. + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Built with return_bias=False, so the layer returns only the output. + out = cast(torch.Tensor, super().forward(x.reshape(-1, x.shape[-1]))) + return out.view(*x.shape[:-1], out.shape[-1]) + + +def _linear( + in_features: int, + out_features: int, + *, + bias: bool, + quant_config: QuantizationConfig | None, + prefix: str, +) -> _TowerLinear: + """A linear layer the checkpoint's quantization config applies to by name.""" + return _TowerLinear( + in_features, + out_features, + bias=bias, + quant_config=quant_config, + prefix=prefix, + return_bias=False, + disable_tp=True, + ) + + +class MelSpectrogram(nn.Module): + """Log-mel features as NeMo ``FilterbankFeatures`` computes them at inference. + + Dither is off and no normalization is applied; frames past each row's + length are set to ``pad_value``. The filterbank and window come from the + checkpoint and are held in the model dtype, as in the reference, and are + upcast to float32 for the computation. + """ + + def __init__(self, cfg: dict[str, Any]) -> None: + super().__init__() + self.n_fft = cfg["n_fft"] + self.hop_length = cfg["hop_length"] + self.win_length = cfg["win_length"] + self.preemph = cfg["preemph"] + self.log_zero_guard_value = cfg["log_zero_guard_value"] + self.mag_power = cfg["mag_power"] + self.pad_value = cfg["pad_value"] + self.register_buffer("fb", torch.zeros(1, cfg["n_mels"], self.n_fft // 2 + 1)) + self.register_buffer("window", torch.zeros(self.win_length)) + + def get_seq_len(self, num_samples: torch.Tensor) -> torch.Tensor: + pad_amount = self.n_fft // 2 * 2 + seq_len = torch.floor_divide( + num_samples + pad_amount - self.n_fft, self.hop_length + ).to(torch.long) + return torch.where(num_samples == 0, torch.zeros_like(seq_len), seq_len) + + @torch.no_grad() + def forward( + self, audio: torch.Tensor, num_samples: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + seq_len = self.get_seq_len(num_samples) + x = audio.float() + if self.preemph: + time_mask = torch.arange(x.shape[1], device=x.device).unsqueeze( + 0 + ) < num_samples.unsqueeze(1) + x = torch.cat((x[:, :1], x[:, 1:] - self.preemph * x[:, :-1]), dim=1) + x = x.masked_fill(~time_mask, 0.0) + x = torch.stft( + x, + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + center=True, + window=self.window.float(), + return_complex=True, + pad_mode="constant", + ) + x = torch.sqrt(torch.view_as_real(x).pow(2).sum(-1)) + if self.mag_power != 1.0: + x = x.pow(self.mag_power) + x = torch.matmul(self.fb.float(), x) + x = torch.log(x + self.log_zero_guard_value) + pad_mask = torch.arange(x.shape[-1], device=x.device).unsqueeze( + 0 + ) >= seq_len.unsqueeze(1) + x = x.masked_fill(pad_mask.unsqueeze(1), self.pad_value) + return x, seq_len + + +class AudioPreprocessor(nn.Module): + """Holds the featurizer under NeMo's ``preprocessor.featurizer`` name.""" + + def __init__(self, cfg: dict[str, Any]) -> None: + super().__init__() + self.featurizer = MelSpectrogram(cfg) + + def forward( + self, audio: torch.Tensor, num_samples: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.featurizer(audio, num_samples) + + +def normalize_per_feature(x: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor: + """NeMo ``normalize_batch(x, seq_len, "per_feature")`` on ``(B, C, T)`` features.""" + batch_size, _, max_time = x.shape + time_steps = ( + torch.arange(max_time, device=x.device) + .unsqueeze(0) + .expand(batch_size, max_time) + ) + valid_mask = time_steps < seq_len.unsqueeze(1) + denominator = valid_mask.sum(dim=1) + reference = x[:, :, 0] if max_time else x.new_zeros((batch_size, x.shape[1])) + reference = reference.masked_fill((denominator == 0).unsqueeze(1), 0.0) + centered = torch.where(valid_mask.unsqueeze(1), x - reference.unsqueeze(2), 0.0) + mean = reference + centered.sum(dim=2) / denominator.clamp_min(1).unsqueeze(1) + variance_sum = torch.sum( + torch.where(valid_mask.unsqueeze(1), x - mean.unsqueeze(2), 0.0) ** 2, dim=2 + ) + std = torch.sqrt(variance_sum / (denominator.unsqueeze(1) - 1.0)) + std = std.masked_fill(std.isnan(), 0.0) + 1e-5 + normalized = (x - mean.unsqueeze(2)) / std.unsqueeze(2) + return normalized.masked_fill(~valid_mask.unsqueeze(1), 0.0) + + +class RotaryEmbedding: + """NeMo ``RotaryPositionalEncoding``: NeoX rotation of the first ``rot_dim`` dims. + + Angles are computed in float32 on the CPU and stored in the activation + dtype before use, as the reference does, so rotations match bit for bit. + """ + + def __init__(self, head_dim: int, rotary_fraction: float, base: float) -> None: + self.rot_dim = int(head_dim * rotary_fraction) + if self.rot_dim < 2 or self.rot_dim % 2: + raise ValueError( + f"Rotary dim must be a positive even number, got {self.rot_dim}." + ) + self.base = base + self._cache: dict[ + tuple[torch.dtype, torch.device], tuple[torch.Tensor, torch.Tensor] + ] = {} + + def cos_sin( + self, length: int, dtype: torch.dtype, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor]: + key = (dtype, device) + cached = self._cache.get(key) + if cached is None or cached[0].shape[0] < length: + inv_freq = 1.0 / ( + self.base + ** ( + torch.arange(0, self.rot_dim, 2, dtype=torch.float32, device="cpu") + / self.rot_dim + ) + ) + positions = torch.arange( + 0, max(length, 1024), dtype=torch.float32, device="cpu" + ) + freqs = torch.outer(positions, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + cached = ( + emb.cos().to(device=device, dtype=dtype), + emb.sin().to(device=device, dtype=dtype), + ) + self._cache[key] = cached + return cached[0][:length], cached[1][:length] + + def apply( + self, t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + def rotate(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return x * cos + torch.cat((-x[..., half:], x[..., :half]), dim=-1) * sin + + if self.rot_dim == t.shape[-1]: + return rotate(t) + return torch.cat( + (rotate(t[..., : self.rot_dim]), t[..., self.rot_dim :]), dim=-1 + ) + + +class FeatureStacking(nn.Module): + """Stack ``factor`` consecutive frames and project them to the model width.""" + + def __init__( + self, + factor: int, + feat_in: int, + feat_out: int, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.factor = factor + self.proj = _linear( + factor * feat_in, + feat_out, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.proj", + ) + + def forward( + self, x: torch.Tensor, lengths: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + x = x.transpose(1, 2) + batch_size, num_frames, channels = x.shape + pad = (self.factor - num_frames % self.factor) % self.factor + if pad: + x = F.pad(x, (0, 0, 0, pad)) + x = x.reshape( + batch_size, (num_frames + pad) // self.factor, channels * self.factor + ) + return self.proj(x), (lengths + self.factor - 1) // self.factor + + +class MultiHeadAttention(nn.Module): + def __init__( + self, + cfg: dict[str, Any], + rope: RotaryEmbedding, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.d_model = cfg["d_model"] + self.n_heads = cfg["n_heads"] + self.head_dim = self.d_model // self.n_heads + self.rope = rope + self.w_qkv = _linear( + self.d_model, + 3 * self.d_model, + bias=cfg["qkv_bias"], + quant_config=quant_config, + prefix=f"{prefix}.w_qkv", + ) + self.out_proj = _linear( + self.d_model, + self.d_model, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.out_proj", + ) + self.qk_norm = cfg["qk_norm"] + if self.qk_norm: + self.q_norm = nn.LayerNorm(self.head_dim) + self.k_norm = nn.LayerNorm(self.head_dim) + + def forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + batch_size, num_frames, _ = x.shape + qkv = self.w_qkv(x).view(batch_size, num_frames, 3, self.n_heads, self.head_dim) + q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) + if self.qk_norm: + q = self.q_norm(q).to(v.dtype) + k = self.k_norm(k).to(v.dtype) + q = self.rope.apply(q, cos.to(q.dtype), sin.to(q.dtype)) + k = self.rope.apply(k, cos.to(k.dtype), sin.to(k.dtype)) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + out = out.transpose(1, 2).reshape(batch_size, num_frames, self.d_model) + return self.out_proj(out) + + +class TransformerBlock(nn.Module): + def __init__( + self, + cfg: dict[str, Any], + rope: RotaryEmbedding, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + d_model = cfg["d_model"] + self.norm1 = nn.LayerNorm(d_model) + self.attn = MultiHeadAttention(cfg, rope, quant_config, f"{prefix}.attn") + self.norm2 = nn.LayerNorm(d_model) + ff_hidden = int(cfg["ff_expansion"] * d_model) + # Indices 0 and 3 match the checkpoint's ``ffn.net.{0,3}``. + self.ffn = nn.Module() + self.ffn.net = nn.Sequential( + _linear( + d_model, + ff_hidden, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.ffn.net.0", + ), + nn.GELU(), + nn.Identity(), + _linear( + ff_hidden, + d_model, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.ffn.net.3", + ), + ) + + def forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + x = x + self.attn(self.norm1(x), attn_mask, cos, sin) + return x + self.ffn.net(self.norm2(x)) + + +class TransformerEncoder(nn.Module): + """NeMo ``TransformerEncoder`` with RoPE and full attention, at inference. + + ``forward`` takes ``(B, C, T)`` mel features, or ``(B, T, d_model)`` + pre-encoded embeddings with ``bypass_pre_encode=True``, and returns + ``(B, d_model, T')`` states with their lengths. Only keys are masked, so + padded frames never influence valid ones. + """ + + def __init__( + self, + cfg: dict[str, Any], + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.d_model = cfg["d_model"] + self.subsampling_factor = cfg["subsampling_factor"] + self.pre_encode = FeatureStacking( + self.subsampling_factor, + cfg["feat_in"], + self.d_model, + quant_config, + f"{prefix}.pre_encode", + ) + self.xscale = math.sqrt(self.d_model) if cfg["xscaling"] else None + self.rope = RotaryEmbedding( + self.d_model // cfg["n_heads"], cfg["rotary_fraction"], cfg["rope_base"] + ) + self.embed_norm = ( + nn.LayerNorm(self.d_model) if cfg["pre_block_norm"] else nn.Identity() + ) + self.layers = nn.ModuleList( + [ + TransformerBlock(cfg, self.rope, quant_config, f"{prefix}.layers.{i}") + for i in range(cfg["n_layers"]) + ] + ) + self.final_norm = nn.LayerNorm(self.d_model) + + def forward( + self, x: torch.Tensor, lengths: torch.Tensor, bypass_pre_encode: bool = False + ) -> tuple[torch.Tensor, torch.Tensor]: + if not bypass_pre_encode: + x, lengths = self.pre_encode(x, lengths) + lengths = lengths.to(torch.int64) + if self.xscale: + x = x * self.xscale + x = self.embed_norm(x) + num_frames = x.shape[1] + positions = torch.arange(num_frames, device=x.device).unsqueeze(0) + # A row with no valid frame (a short clip in a later window) attends to + # frame 0 to stay finite; its outputs are discarded. + key_mask = (positions < lengths.unsqueeze(1)) | ( + (positions == 0) & (lengths.unsqueeze(1) == 0) + ) + attn_mask = key_mask[:, None, None, :] + cos, sin = self.rope.cos_sin(num_frames, x.dtype, x.device) + for layer in self.layers: + x = layer(x, attn_mask, cos, sin) + return self.final_norm(x).transpose(1, 2), lengths + + +def downsample_preds(preds: torch.Tensor, factor: int) -> torch.Tensor: + """Average ``(B, T, S)`` probabilities over windows of ``factor`` frames.""" + if factor == 1 or preds.shape[1] == 0: + return preds + pooled = F.avg_pool1d( + preds.transpose(1, 2), + kernel_size=factor, + stride=factor, + ceil_mode=True, + count_include_pad=False, + ) + return pooled.transpose(1, 2) + + +def align_frames(preds: torch.Tensor, target_len: int) -> torch.Tensor: + """Pad with the last frame, or truncate, ``(B, T, S)`` to ``target_len`` frames.""" + current_len = preds.shape[1] + if current_len < target_len: + if current_len == 0: + raise ValueError( + "Speaker activity cannot be empty when the encoder output is not." + ) + preds = torch.cat( + [preds, preds[:, -1:].repeat(1, target_len - current_len, 1)], dim=1 + ) + return preds[:, :target_len] + + +@dataclass +class SpeakerCacheState: + """Synchronous streaming state of the Sortformer speaker encoder.""" + + spkcache: torch.Tensor + fifo: torch.Tensor + spkcache_preds: torch.Tensor | None = None + spkcache_compressed: bool = False + + +class SortformerModules(nn.Module): + """Speaker head and speaker cache of NeMo ``SortformerModules`` at inference.""" + + def __init__( + self, + cfg: dict[str, Any], + online: dict[str, Any], + upsample_factor: int, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.n_spk = cfg["num_speakers"] + self.upsample_factor = upsample_factor + hidden = cfg["tf_d_model"] + self.encoder_proj = _linear( + cfg["fc_d_model"], + hidden, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.encoder_proj", + ) + self.subpixel_upsample = nn.Conv1d( + hidden, hidden * upsample_factor, kernel_size=3, padding=1 + ) + self.first_hidden_to_hidden = _linear( + hidden, + hidden, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.first_hidden_to_hidden", + ) + self.single_hidden_to_spks = _linear( + hidden, + self.n_spk, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.single_hidden_to_spks", + ) + self.learnable_sil_emb = nn.Parameter(torch.zeros(cfg["fc_d_model"])) + + self.fifo_len = online["fifo_len"] + self.spkcache_len = online["spkcache_len"] + self.spkcache_update_period = online["spkcache_update_period"] + self.sil_frames_per_spk = cfg["spkcache_sil_frames_per_spk"] + self.pred_score_threshold = cfg["pred_score_threshold"] + self.max_index = cfg["max_index"] + self.scores_boost_latest = cfg["scores_boost_latest"] + self.strong_boost_rate = cfg["strong_boost_rate"] + self.weak_boost_rate = cfg["weak_boost_rate"] + self.min_pos_scores_rate = cfg["min_pos_scores_rate"] + if self.spkcache_len < (1 + self.sil_frames_per_spk) * self.n_spk: + raise ValueError("spkcache_len is too small for the number of speakers.") + + def speaker_probs( + self, hidden: torch.Tensor, lengths: torch.Tensor + ) -> torch.Tensor: + """Upsampled speaker probabilities, ``(B, T * upsample_factor, n_spk)``.""" + batch_size, num_frames, width = hidden.shape + mask = torch.arange(num_frames, device=hidden.device).unsqueeze( + 0 + ) < lengths.unsqueeze(1) + # The convolution reads neighbouring frames: zero the padding so a clip's + # last frame does not depend on the clips batched with it. + hidden = hidden.masked_fill(~mask.unsqueeze(-1), 0.0) + upsampled = self.subpixel_upsample(hidden.transpose(1, 2)).transpose(1, 2) + upsampled = upsampled.reshape( + batch_size, num_frames, self.upsample_factor, width + ).reshape(batch_size, num_frames * self.upsample_factor, width) + logits = self.single_hidden_to_spks( + F.relu(self.first_hidden_to_hidden(F.relu(upsampled))) + ) + output_mask = mask.repeat_interleave(self.upsample_factor, dim=1) + return torch.sigmoid(logits) * output_mask.unsqueeze(-1) + + def update( + self, + state: SpeakerCacheState, + chunk: torch.Tensor, + preds: torch.Tensor, + lc: int, + rc: int, + ) -> SpeakerCacheState: + """NeMo ``SortformerModules.streaming_update`` (synchronous).""" + spkcache_len, fifo_len = state.spkcache.shape[1], state.fifo.shape[1] + chunk_len = chunk.shape[1] - lc - rc + fifo_preds = preds[:, spkcache_len : spkcache_len + fifo_len] + chunk = chunk[:, lc : chunk_len + lc] + chunk_preds = preds[ + :, spkcache_len + fifo_len + lc : spkcache_len + fifo_len + chunk_len + lc + ] + fifo = torch.cat([state.fifo, chunk], dim=1) + fifo_preds = torch.cat([fifo_preds, chunk_preds], dim=1) + spkcache, spkcache_preds, compressed = ( + state.spkcache, + state.spkcache_preds, + state.spkcache_compressed, + ) + + if fifo_len + chunk_len > self.fifo_len: + pop_out_len = max( + self.spkcache_update_period, chunk_len - self.fifo_len + fifo_len + ) + pop_out_len = min(pop_out_len, fifo_len + chunk_len) + pop_out_embs, pop_out_preds = ( + fifo[:, :pop_out_len], + fifo_preds[:, :pop_out_len], + ) + fifo = fifo[:, pop_out_len:] + spkcache = torch.cat([spkcache, pop_out_embs], dim=1) + if compressed: + spkcache_preds = torch.cat([spkcache_preds, pop_out_preds], dim=1) + else: + spkcache_preds = torch.cat( + [preds[:, :spkcache_len], pop_out_preds], dim=1 + ) + if spkcache.shape[1] > self.spkcache_len: + spkcache, spkcache_preds = self._compress_spkcache( + spkcache, spkcache_preds + ) + compressed = True + return SpeakerCacheState(spkcache, fifo, spkcache_preds, compressed) + + def _compress_spkcache( + self, emb_seq: torch.Tensor, preds: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Keep the ``spkcache_len`` most informative frames, as NeMo does.""" + batch_size, _, n_spk = preds.shape + per_spk = self.spkcache_len // n_spk - self.sil_frames_per_spk + strong_boost = math.floor(per_spk * self.strong_boost_rate) + weak_boost = math.floor(per_spk * self.weak_boost_rate) + min_pos_scores = math.floor(per_spk * self.min_pos_scores_rate) + + threshold = self.pred_score_threshold + log_probs = torch.log(torch.clamp(preds, min=threshold)) + log_1_probs = torch.log(torch.clamp(1.0 - preds, min=threshold)) + log_1_probs_sum = log_1_probs.sum(dim=2).unsqueeze(-1).expand(-1, -1, n_spk) + scores = log_probs - log_1_probs + log_1_probs_sum - math.log(0.5) + + is_speech = preds > 0.5 + scores = torch.where(is_speech, scores, float("-inf")) + is_pos = scores > 0 + replace = ( + (~is_pos) * is_speech * (is_pos.sum(dim=1).unsqueeze(1) >= min_pos_scores) + ) + scores = torch.where(replace, float("-inf"), scores) + + if self.scores_boost_latest > 0: + scores[:, self.spkcache_len :, :] += self.scores_boost_latest + scores = self._boost_topk_scores(scores, strong_boost, scale_factor=2) + scores = self._boost_topk_scores(scores, weak_boost, scale_factor=1) + if self.sil_frames_per_spk > 0: + pad = torch.full( + (batch_size, self.sil_frames_per_spk, n_spk), + float("inf"), + dtype=scores.dtype, + device=scores.device, + ) + scores = torch.cat([scores, pad], dim=1) + + num_frames = scores.shape[1] + flat = scores.permute(0, 2, 1).reshape(batch_size, -1) + topk_values, topk_indices = torch.topk( + flat, self.spkcache_len, dim=1, sorted=False + ) + topk_indices = torch.where( + topk_values != float("-inf"), topk_indices, self.max_index + ) + topk_indices, _ = torch.sort(topk_indices, dim=1) + is_disabled = topk_indices == self.max_index + topk_indices = torch.remainder(topk_indices, num_frames) + is_disabled += topk_indices >= num_frames - self.sil_frames_per_spk + topk_indices[is_disabled] = 0 + + sil_emb = ( + self.learnable_sil_emb.to(dtype=emb_seq.dtype) + .unsqueeze(0) + .expand(batch_size, -1) + ) + emb_gathered = torch.gather( + emb_seq, 1, topk_indices.unsqueeze(-1).expand(-1, -1, emb_seq.shape[2]) + ) + emb_gathered = torch.where( + is_disabled.unsqueeze(-1), + sil_emb.unsqueeze(1).expand(-1, self.spkcache_len, -1), + emb_gathered, + ) + preds_gathered = torch.gather( + preds, 1, topk_indices.unsqueeze(-1).expand(-1, -1, n_spk) + ) + preds_gathered = torch.where(is_disabled.unsqueeze(-1), 0.0, preds_gathered) + return emb_gathered, preds_gathered + + @staticmethod + def _boost_topk_scores( + scores: torch.Tensor, n_boost_per_spk: int, scale_factor: float + ) -> torch.Tensor: + batch_size, _, n_spk = scores.shape + _, topk_indices = torch.topk( + scores, n_boost_per_spk, dim=1, largest=True, sorted=False + ) + batch_indices = torch.arange(batch_size, device=scores.device).view(-1, 1, 1) + speaker_indices = torch.arange(n_spk, device=scores.device).view(1, 1, -1) + scores[batch_indices, topk_indices, speaker_indices] -= scale_factor * math.log( + 0.5 + ) + return scores + + +class SortformerSpeakerEncoder(nn.Module): + """Streaming inference of the high-resolution Sortformer in NeMo's PEE. + + Mirrors ``SortformerEncLabelModel.forward_streaming_step`` for synchronous + streaming without speaker permutation. Each step returns speaker + probabilities on the encoder frame grid for the window's core frames. + """ + + def __init__( + self, + encoder_cfg: dict[str, Any], + sortformer_cfg: dict[str, Any], + online: dict[str, Any], + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.encoder = TransformerEncoder( + encoder_cfg, quant_config, f"{prefix}.encoder" + ) + self.subsampling_factor = encoder_cfg["subsampling_factor"] + self.sortformer_modules = SortformerModules( + sortformer_cfg, + online, + self.subsampling_factor, + quant_config, + f"{prefix}.sortformer_modules", + ) + self.fc_d_model = sortformer_cfg["fc_d_model"] + self.n_spk = sortformer_cfg["num_speakers"] + + def init_state( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> SpeakerCacheState: + empty = torch.zeros( + (batch_size, 0, self.fc_d_model), dtype=dtype, device=device + ) + return SpeakerCacheState(spkcache=empty, fifo=empty.clone()) + + def step( + self, + features: torch.Tensor, + lengths: torch.Tensor, + state: SpeakerCacheState, + left_offset: int, + right_offset: int, + ) -> tuple[SpeakerCacheState, torch.Tensor]: + """Process one window of ``(B, C, T)`` mel features. + + Returns the new state and the core frames' speaker probabilities on the + encoder grid. + """ + modules = self.sortformer_modules + factor = self.subsampling_factor + chunk_embs, chunk_lengths = self.encoder.pre_encode(features, lengths) + embs = torch.cat([state.spkcache, state.fifo, chunk_embs], dim=1) + context_len = state.spkcache.shape[1] + state.fifo.shape[1] + emb_lengths = context_len + chunk_lengths + states, state_lengths = self.encoder(embs, emb_lengths, bypass_pre_encode=True) + hidden = modules.encoder_proj(states.transpose(1, 2)) + high_res_preds = modules.speaker_probs(hidden, state_lengths) + + lc = round(left_offset / factor) + rc = math.ceil(right_offset / factor) + preds = downsample_preds(high_res_preds, modules.upsample_factor) + valid = torch.arange(preds.shape[1], device=preds.device).view( + 1, -1, 1 + ) < state_lengths.view(-1, 1, 1) + preds = torch.where(valid, preds, 0.0) + state = modules.update(state, chunk_embs, preds, lc, rc) + + chunk_len = chunk_embs.shape[1] - lc - rc + start = (context_len + lc) * modules.upsample_factor + chunk_preds = high_res_preds[ + :, start : start + chunk_len * modules.upsample_factor + ] + return state, downsample_preds(chunk_preds, factor) + + +class ParallelExpertEncoder(nn.Module): + """NeMo ``ParallelExpertEncoder`` online inference with a transformer ASR encoder. + + ``forward`` takes unnormalized ``(B, C, T)`` mel features and returns the + fused ``(B, d_model, T')`` states. + """ + + def __init__( + self, + audio_config: dict[str, Any], + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + online = audio_config["online"] + fusion = audio_config["fusion"] + self.asr_encoder = TransformerEncoder( + audio_config["asr_encoder"], + quant_config, + f"{prefix}.asr_encoder", + ) + self.diarization_model = SortformerSpeakerEncoder( + audio_config["speaker_encoder"], + audio_config["sortformer"], + online, + quant_config, + f"{prefix}.diarization_model", + ) + self.subsampling_factor = self.asr_encoder.subsampling_factor + self.window_feat_len = online["window_frames"] * self.subsampling_factor + self.left_ctx_feat_len = online["left_context_frames"] * self.subsampling_factor + self.right_ctx_feat_len = ( + online["right_context_frames"] * self.subsampling_factor + ) + self.speaker_feature_mode = fusion["speaker_feature_mode"] + self.speaker_activity_threshold = fusion["speaker_activity_threshold"] + self.spk_kernel_scale = fusion["spk_kernel_scale"] + n_spk = self.diarization_model.n_spk + self.asr_norm = nn.LayerNorm(self.asr_encoder.d_model) + self.diar_norm = nn.LayerNorm(n_spk) + self.register_buffer( + "diar_kernel", + self._sinusoid_kernel(n_spk, self.asr_encoder.d_model), + persistent=False, + ) + + @staticmethod + def _sinusoid_kernel(num_speakers: int, width: int) -> torch.Tensor: + position = torch.arange(num_speakers, dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, width, 2, dtype=torch.float32) + * -(math.log(10000.0) / width) + ) + kernel = torch.zeros(num_speakers, width, dtype=torch.float32) + kernel[:, 0::2] = torch.sin(position * div_term) + kernel[:, 1::2] = torch.cos(position * div_term) + return kernel + + def _output_boundary(self, feat_boundary: int) -> int: + return (feat_boundary + self.subsampling_factor - 1) // self.subsampling_factor + + @property + def dtype(self) -> torch.dtype: + """The model dtype, which quantized layers' weights may not have.""" + return self.asr_norm.weight.dtype + + def forward( + self, mel: torch.Tensor, mel_lengths: torch.Tensor, max_feat_len: int + ) -> torch.Tensor: + total_feat_len = min(mel.shape[-1], max_feat_len) + num_windows = max(1, math.ceil(total_feat_len / self.window_feat_len)) + asr_signal = normalize_per_feature(mel, mel_lengths).to(self.dtype) + diar_signal = mel.to(self.dtype) + state = self.diarization_model.init_state(mel.shape[0], self.dtype, mel.device) + + asr_windows: list[torch.Tensor] = [] + speaker_windows: list[torch.Tensor] = [] + for index in range(num_windows): + start = index * self.window_feat_len + end = min(start + self.window_feat_len, total_feat_len) + context_start = max(start - self.left_ctx_feat_len, 0) + context_end = min(end + self.right_ctx_feat_len, total_feat_len) + left_offset, right_offset = start - context_start, context_end - end + window_lengths = (mel_lengths - context_start).clamp( + min=0, max=context_end - context_start + ) + + encoded, _ = self.asr_encoder( + asr_signal[:, :, context_start:context_end], window_lengths + ) + left_drop = left_offset // self.subsampling_factor + core_len = self._output_boundary(end) - self._output_boundary(start) + core_len = max(0, min(core_len, encoded.shape[-1] - left_drop)) + asr_windows.append(encoded[:, :, left_drop : left_drop + core_len]) + + state, window_preds = self.diarization_model.step( + diar_signal[:, :, context_start:context_end], + window_lengths, + state, + left_offset, + right_offset, + ) + speaker_windows.append(align_frames(window_preds, core_len)) + + asr_encoded = torch.cat(asr_windows, dim=2) + speaker_probs = torch.cat(speaker_windows, dim=1) + return self._fuse(asr_encoded, speaker_probs) + + def _fuse( + self, asr_encoded: torch.Tensor, speaker_probs: torch.Tensor + ) -> torch.Tensor: + states = asr_encoded.transpose(1, 2) + speaker_probs = align_frames(speaker_probs, states.shape[1]).to( + device=states.device, dtype=states.dtype + ) + if self.speaker_feature_mode == "thresholded": + speaker_features = (speaker_probs > self.speaker_activity_threshold).to( + states.dtype + ) + else: + speaker_features = speaker_probs + normalized_states = self.asr_norm(states) + normalized_speakers = self.diar_norm(speaker_features) + infusion = torch.matmul( + normalized_speakers, self.diar_kernel.to(normalized_speakers.dtype) + ) + return (normalized_states + self.spk_kernel_scale * infusion).transpose(1, 2) + + +class Nemotron3_5TranscribeAudioTower(nn.Module): + """Waveform to decoder-width embeddings: mel, Parallel Expert Encoder, Linear.""" + + def __init__( + self, + audio_config: dict[str, Any], + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.preprocessor = AudioPreprocessor(audio_config["preprocessor"]) + self.encoder = ParallelExpertEncoder( + audio_config, quant_config, maybe_prefix(prefix, "encoder") + ) + projection = audio_config["projection"] + self.proj = _linear( + projection["input_dim"], + projection["output_dim"], + bias=True, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "proj"), + ) + self.subsampling_factor = self.encoder.subsampling_factor + + def encode( + self, signals: Sequence[torch.Tensor], num_samples_cpu: torch.Tensor + ) -> list[torch.Tensor]: + """Encode each clip into its embeddings, returned in input order. + + Every row of a batch runs every window of the batch's longest clip, so a + clip is batched only with clips that span as many windows. + """ + window = self.encoder.window_feat_len + num_mel_frames = self.preprocessor.featurizer.get_seq_len(num_samples_cpu) + groups: dict[int, list[int]] = {} + for index, num_windows in enumerate( + ((num_mel_frames + window - 1) // window).tolist() + ): + groups.setdefault(num_windows, []).append(index) + + device = self.encoder.asr_norm.weight.device + embeddings: dict[int, torch.Tensor] = {} + for indices in groups.values(): + lengths_cpu = num_samples_cpu[indices] + audio = torch.zeros( + len(indices), int(lengths_cpu.max()), dtype=torch.float32, device=device + ) + for row, index in enumerate(indices): + length = int(lengths_cpu[row]) + audio[row, :length] = ( + signals[index] + .flatten()[:length] + .to(device=device, dtype=torch.float32) + ) + lengths = ( + lengths_cpu + if device.type == "cpu" + else async_tensor_h2d(lengths_cpu, device=device, dtype=torch.long) + ) + for index, emb in zip(indices, self(audio, lengths, lengths_cpu)): + embeddings[index] = emb + return [embeddings[index] for index in range(len(signals))] + + def forward( + self, + audio: torch.Tensor, + num_samples: torch.Tensor, + num_samples_cpu: torch.Tensor, + ) -> list[torch.Tensor]: + """Encode padded ``(B, N)`` float32 audio into one embedding tensor per row. + + ``num_samples_cpu`` holds the lengths on the host, so the window bounds + and the output lengths need no device-to-host copy. + """ + mel, mel_lengths = self.preprocessor(audio, num_samples) + # NeMo's AudioPreprocessor computes in float32 but returns features in + # the module dtype, so the encoders normalize model-dtype features. + mel = mel.to(self.encoder.dtype) + mel_lengths_cpu = self.preprocessor.featurizer.get_seq_len( + num_samples_cpu + ).clamp(max=mel.shape[-1]) + encoded = self.encoder(mel, mel_lengths, int(mel_lengths_cpu.max())) + encoded_lengths = ( + mel_lengths_cpu + self.subsampling_factor - 1 + ) // self.subsampling_factor + embeddings = self.proj(encoded.transpose(1, 2)) + return [ + emb[:length] for emb, length in zip(embeddings, encoded_lengths.tolist()) + ] + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load NeMo ``perception.*`` tensors, named without the ``perception.``.""" + targets = dict(self.named_parameters()) + targets.update(dict(self.named_buffers())) + targets.pop("encoder.diar_kernel", None) + loaded: set[str] = set() + for name, tensor in weights: + if name.endswith("._extra_state") or name.startswith( + _UNUSED_CHECKPOINT_PREFIXES + ): + continue + target = targets.get(name) + if target is None: + raise ValueError(f"Unexpected audio tower weight {name!r}.") + weight_loader = getattr(target, "weight_loader", None) + if weight_loader is not None: + # vLLM's linear layers: quantized parameters and scalar scales. + weight_loader(target, tensor) + elif target.shape != tensor.shape: + raise ValueError( + f"Shape mismatch for {name!r}: {tuple(tensor.shape)} vs " + f"{tuple(target.shape)}." + ) + else: + with torch.no_grad(): + target.copy_(tensor.to(device=target.device, dtype=target.dtype)) + loaded.add(name) + missing = sorted(set(targets) - loaded) + if missing: + raise ValueError( + f"Audio tower weights missing from the checkpoint ({len(missing)}): " + f"{missing[:8]}" + ) + return loaded diff --git a/vllm/model_executor/models/qwen3_dflash2.py b/vllm/model_executor/models/qwen3_dflash2.py index 861a7c5412ea..622f38ce4c40 100644 --- a/vllm/model_executor/models/qwen3_dflash2.py +++ b/vllm/model_executor/models/qwen3_dflash2.py @@ -296,6 +296,10 @@ def _score_edges( anchor_token_ids: torch.Tensor, top_k: int, ) -> torch.Tensor: + # A request still in chunked prefill anchors on its next prompt token, which + # can lie past the draft vocabulary (an audio placeholder, say). Its draft is + # discarded, so only keep the codebook lookup in range. + anchor_token_ids = anchor_token_ids.clamp(0, predecessor_table.shape[0] - 1) successors = successor_table[candidate_ids] predecessor_ids = torch.cat( ( @@ -399,8 +403,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__(vllm_config=vllm_config, prefix=prefix) draft_config = self.config.dflash_config softcap = float(draft_config.get("final_logit_softcapping") or 0.0) + # Candidates index the selector codebooks, which have the draft's + # vocabulary; the target's vocabulary can be larger. self.candidate_logits_processor = LogitsProcessor( - vllm_config.model_config.get_vocab_size(), + min(vllm_config.model_config.get_vocab_size(), self.config.vocab_size), scale=float(draft_config.get("output_multiplier", 1.0)), soft_cap=softcap if softcap > 0 else None, ) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 4050e8c73a58..8e4f7bba7a23 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -546,6 +546,10 @@ "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NemotronH_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), + "Nemotron3_5TranscribeForConditionalGeneration": ( + "nemotron3_5_transcribe", + "Nemotron3_5TranscribeForConditionalGeneration", + ), "NVLM_D": ("nvlm_d", "NVLM_D_Model"), "MuseGlimmerForConditionalGeneration": ("muse_glimmer", "MuseGlimmerForCausalLM"), "MuseSparkForConditionalGeneration": ( @@ -720,6 +724,10 @@ "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"), "Exaone4_5_MTP": ("exaone4_5_mtp", "Exaone4_5_MTP"), "NemotronHMTPModel": ("nemotron_h_mtp", "NemotronHMTP"), + "Nemotron3_5TranscribeMTPModel": ( + "nemotron3_5_transcribe", + "Nemotron3_5TranscribeMTP", + ), "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"), "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"), "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"), diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index e18fc2753cc0..a1354e5f0887 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -156,6 +156,7 @@ def __getitem__(self, key): minimax_m3_mtp="MiniMaxM3MTPConfig", moondream3="Moondream3Config", moss_transcribe_diarize="MossTranscribeDiarizeConfig", + nemo_speechlm="Nemotron3_5TranscribeConfig", eagle="EAGLEConfig", speculators="SpeculatorsConfig", openvla="OpenVLAConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 3042e1f5aa10..1d4797bf3d5c 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -50,6 +50,9 @@ "MossTranscribeDiarizeConfig": ( "vllm.transformers_utils.configs.moss_transcribe_diarize" ), + "Nemotron3_5TranscribeConfig": ( + "vllm.transformers_utils.configs.nemotron3_5_transcribe" + ), "MoonViTConfig": "vllm.transformers_utils.configs.moonvit", "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear", "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl", @@ -120,6 +123,7 @@ "Moondream3TextConfig", "Moondream3VisionConfig", "MossTranscribeDiarizeConfig", + "Nemotron3_5TranscribeConfig", "MoonViTConfig", "KimiLinearConfig", "KimiVLConfig", diff --git a/vllm/transformers_utils/configs/nemotron3_5_transcribe.py b/vllm/transformers_utils/configs/nemotron3_5_transcribe.py new file mode 100644 index 000000000000..f90207dd03e6 --- /dev/null +++ b/vllm/transformers_utils/configs/nemotron3_5_transcribe.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Configuration for NVIDIA Nemotron-3.5-Transcribe. + +The checkpoint's ``config.json`` is a NeMo SpeechLM config: the decoder under +``llm_config``, the speech front end under ``perception`` and the Parallel +Expert Encoder (an ASR encoder plus a Sortformer speaker encoder) under +``pe_encoder_config``. Everything the audio tower needs is resolved into one +flat ``audio_config`` dict. +""" + +import math +from typing import Any + +from transformers import AutoConfig, PreTrainedConfig + +# The trained output vocabulary is followed by this many embedding rows. The +# audio placeholder token lives in that range. +SPEECHLM_EMBED_EXTRA_ROWS = 10 + +_ARCHITECTURE = "Nemotron3_5TranscribeForConditionalGeneration" +_DECODER_ARCHITECTURE = "NemotronHForCausalLM" +_HYBRID_ARCHITECTURES = ("NemotronHForCausalLM", "NemotronHybridForCausalLM") +_SUPPORTED_SPEAKER_FEATURE_MODES = ("thresholded", "continuous") + + +def _encoder_config(cfg: dict[str, Any], name: str) -> dict[str, Any]: + """Resolve a NeMo ``TransformerEncoder`` config, applying its defaults.""" + subsampling = cfg.get("subsampling", "feature_stacking") + if subsampling != "feature_stacking": + raise ValueError( + f"{name}: only feature_stacking subsampling is supported, " + f"got {subsampling!r}." + ) + if cfg.get("self_attention_model", "rel_pos") != "rope": + raise ValueError(f"{name}: only rope self-attention is supported.") + if cfg.get("attn_mode", "full") != "full": + raise ValueError(f"{name}: only full attention is supported.") + if int(cfg.get("causal_tail_len", 0)) != 0: + raise ValueError(f"{name}: causal_tail_len must be 0 at inference.") + feat_out = int(cfg.get("feat_out", -1)) + d_model = int(cfg["d_model"]) + if feat_out > 0 and feat_out != d_model: + raise ValueError( + f"{name}: an output projection (feat_out={feat_out}) is not supported." + ) + return { + "feat_in": int(cfg.get("feat_in", 128)), + "d_model": d_model, + "n_heads": int(cfg["n_heads"]), + "n_layers": int(cfg["n_layers"]), + "ff_expansion": float(cfg.get("ff_expansion", 4.0)), + "qkv_bias": bool(cfg.get("qkv_bias", False)), + "qk_norm": bool(cfg.get("qk_norm", False)), + "pre_block_norm": bool(cfg.get("pre_block_norm", True)), + "subsampling_factor": int(cfg.get("subsampling_factor", 4)), + "rope_base": float(cfg.get("rope_base", 10000.0)), + "rotary_fraction": float(cfg.get("rotary_fraction", 1.0)), + "xscaling": bool(cfg.get("xscaling", False)), + } + + +def _preprocessor_config(cfg: dict[str, Any]) -> dict[str, Any]: + """Resolve a NeMo ``AudioToMelSpectrogramPreprocessor`` config.""" + sample_rate = int(cfg.get("sample_rate", 16000)) + # NeMo normalizes per feature when the key is missing. + if cfg.get("normalize", "per_feature") not in (None, "", "NA"): + raise ValueError( + "The perception preprocessor must not normalize; the encoders normalize " + "their inputs." + ) + if bool(cfg.get("exact_pad", False)): + raise ValueError("exact_pad preprocessing is not supported.") + if int(cfg.get("frame_splicing", 1)) != 1: + raise ValueError("frame_splicing is not supported.") + if cfg.get("window", "hann") != "hann": + raise ValueError("Only the hann window is supported.") + if not bool(cfg.get("log", True)): + raise ValueError("Only log-mel features are supported (log=true).") + if cfg.get("log_zero_guard_type", "add") != "add": + raise ValueError("Only log_zero_guard_type='add' is supported.") + win_length = int(float(cfg.get("window_size", 0.02)) * sample_rate) + return { + "sample_rate": sample_rate, + "win_length": win_length, + "hop_length": int(float(cfg.get("window_stride", 0.01)) * sample_rate), + # NeMo's default: the next power of two at or above the window length. + "n_fft": int(cfg.get("n_fft") or 2 ** math.ceil(math.log2(win_length))), + "n_mels": int(cfg.get("features", 64)), + "preemph": float(cfg.get("preemph", 0.97)), + "log_zero_guard_value": float(cfg.get("log_zero_guard_value", 2**-24)), + "mag_power": float(cfg.get("mag_power", 2.0)), + "pad_value": float(cfg.get("pad_value", 0.0)), + } + + +def _checked_online(online: dict[str, int]) -> dict[str, int]: + """Reject window settings the native encoders do not implement.""" + if online["window_frames"] <= 0: + raise ValueError( + "The encoders run in windows; online_inference_length must be positive, " + f"got {online['window_frames']}." + ) + for key in ("left_context_frames", "right_context_frames", "fifo_len"): + if online[key] < 0: + raise ValueError(f"{key} must not be negative, got {online[key]}.") + for key in ("spkcache_update_period", "spkcache_len"): + if online[key] <= 0: + raise ValueError(f"{key} must be positive, got {online[key]}.") + return online + + +def resolve_audio_config( + perception: dict[str, Any], pe_encoder_config: dict[str, Any] +) -> dict[str, Any]: + """Flatten the NeMo perception and Parallel Expert Encoder configs.""" + if not perception: + raise ValueError("Nemotron-3.5-Transcribe config must declare perception.") + if not pe_encoder_config: + raise ValueError( + "Nemotron-3.5-Transcribe config must declare pe_encoder_config." + ) + asr_type = pe_encoder_config.get("asr_encoder_type", "fastconformer") + if asr_type != "transformer": + raise ValueError( + f"Only a transformer ASR encoder is supported, got {asr_type!r}." + ) + + diar = pe_encoder_config["diarization_model_cfg"] + sortformer = diar.get("sortformer_modules", {}) + asr_encoder = _encoder_config( + pe_encoder_config["asr_encoder_cfg"], "asr_encoder_cfg" + ) + speaker_encoder = _encoder_config(diar["encoder"], "diarization_model_cfg.encoder") + if speaker_encoder["subsampling_factor"] != asr_encoder["subsampling_factor"]: + raise ValueError( + "The speaker and ASR encoders must share one subsampling factor." + ) + if diar.get("transformer_encoder") is not None: + raise ValueError("A Sortformer transformer_encoder stage is not supported.") + if bool(diar.get("async_streaming", False)): + raise ValueError("Asynchronous Sortformer streaming is not supported.") + if not bool(diar.get("high_resolution", False)): + raise ValueError("The speaker encoder must be a high-resolution Sortformer.") + if not bool(sortformer.get("use_learnable_sil_emb", False)): + raise ValueError("The speaker encoder must use a learnable silence embedding.") + if float(sortformer.get("activity_weight", diar.get("activity_weight", 0.0))) > 0.0: + raise ValueError("A Sortformer activity head is not supported.") + diar_normalize = pe_encoder_config.get("diar_normalize_type") + if diar_normalize is None: + diar_normalize = diar.get("preprocessor", {}).get("normalize") + if diar_normalize not in (None, "", "NA"): + raise ValueError( + f"Speaker-encoder input normalization {diar_normalize!r} is not supported." + ) + asr_normalize = pe_encoder_config.get("asr_normalize_type") or "per_feature" + if asr_normalize != "per_feature": + raise ValueError( + f"ASR-encoder input normalization {asr_normalize!r} is not supported." + ) + + mode = str(pe_encoder_config.get("speaker_feature_mode", "thresholded")).lower() + if mode not in _SUPPORTED_SPEAKER_FEATURE_MODES: + raise ValueError(f"Unsupported speaker_feature_mode {mode!r}.") + threshold = pe_encoder_config.get("speaker_activity_threshold") + if mode == "thresholded" and threshold is None: + raise ValueError( + "speaker_feature_mode='thresholded' requires speaker_activity_threshold." + ) + + modality_adapter = perception.get("modality_adapter", {}) + if not str(modality_adapter.get("_target_", "")).endswith("IdentityConnector"): + raise ValueError("Only the identity modality adapter is supported.") + projection_in = int(modality_adapter.get("d_model", asr_encoder["d_model"])) + if projection_in != asr_encoder["d_model"]: + raise ValueError("The projection input must match the ASR encoder width.") + + return { + "preprocessor": _preprocessor_config(perception["preprocessor"]), + "asr_encoder": asr_encoder, + "speaker_encoder": speaker_encoder, + "sortformer": { + "num_speakers": int( + sortformer.get("num_spks", diar.get("max_num_of_spks", 4)) + ), + "fc_d_model": int(sortformer.get("fc_d_model", speaker_encoder["d_model"])), + "tf_d_model": int(sortformer.get("tf_d_model", 192)), + "spkcache_sil_frames_per_spk": int( + sortformer.get("spkcache_sil_frames_per_spk", 3) + ), + "pred_score_threshold": float(sortformer.get("pred_score_threshold", 0.25)), + "max_index": int(sortformer.get("max_index", 99999)), + "scores_boost_latest": float(sortformer.get("scores_boost_latest", 0.05)), + "strong_boost_rate": float(sortformer.get("strong_boost_rate", 0.75)), + "weak_boost_rate": float(sortformer.get("weak_boost_rate", 1.5)), + "min_pos_scores_rate": float(sortformer.get("min_pos_scores_rate", 0.5)), + }, + "online": _checked_online( + { + "window_frames": int( + pe_encoder_config.get("online_inference_length", 500) + ), + "left_context_frames": int( + pe_encoder_config.get("chunk_left_context", 50) + ), + "right_context_frames": int( + pe_encoder_config.get("chunk_right_context", 50) + ), + "fifo_len": int(pe_encoder_config.get("diar_fifo_len", 40)), + "spkcache_update_period": int( + pe_encoder_config.get("diar_spkcache_update_period", 300) + ), + "spkcache_len": int(pe_encoder_config.get("diar_spkcache_len", 188)), + } + ), + "fusion": { + "speaker_feature_mode": mode, + "speaker_activity_threshold": None + if threshold is None + else float(threshold), + "spk_kernel_scale": float(pe_encoder_config.get("spk_kernel_scale", 1.0)), + }, + "projection": { + "input_dim": projection_in, + "output_dim": int(perception["output_dim"]), + }, + } + + +def _build_text_config(llm_config: dict[str, Any] | None) -> PreTrainedConfig: + if llm_config is None: + return AutoConfig.for_model("nemotron_h") + if not isinstance(llm_config, dict): + raise ValueError(f"llm_config must be a dict, got {type(llm_config).__name__}.") + llm_config = dict(llm_config) + model_type = llm_config.pop("model_type", None) + if not model_type: + raise ValueError("llm_config must declare model_type.") + return AutoConfig.for_model(model_type, **llm_config) + + +def _decoder_architectures(text_config: PreTrainedConfig) -> list[str]: + return getattr(text_config, "architectures", None) or [_DECODER_ARCHITECTURE] + + +def _other_architecture_hint( + pe_encoder_config: dict[str, Any] | None, text_config: PreTrainedConfig +) -> str: + # NeMo names checkpoints with this structure after this model; older exports + # carry NeMo's generic SpeechLM name. + if pe_encoder_config and set(_decoder_architectures(text_config)) & set( + _HYBRID_ARCHITECTURES + ): + return ( + "This checkpoint has the Parallel Expert Encoder on a NemotronH decoder: " + f'set "architectures": ["{_ARCHITECTURE}"] in its config.json.' + ) + return "Serve other NeMo SpeechLM models with the NeMo Speech vLLM plugin." + + +class Nemotron3_5TranscribeConfig(PreTrainedConfig): + """Configuration for Nemotron-3.5-Transcribe. + + Reads the checkpoint's NeMo SpeechLM ``config.json`` (see the module + docstring). The decoder config becomes ``text_config``; the audio tower + settings become ``audio_config`` (see :func:`resolve_audio_config`). + + ``nemo_speechlm`` is NeMo's model type for all its SpeechLM checkpoints; the + architecture name selects Nemotron-3.5-Transcribe. + """ + + model_type = "nemo_speechlm" + + def __init__( + self, + llm_config: dict[str, Any] | None = None, + perception: dict[str, Any] | None = None, + pe_encoder_config: dict[str, Any] | None = None, + text_config: dict[str, Any] | PreTrainedConfig | None = None, + audio_config: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + # Transformers may read the text config while the base class + # initializes; give it the real one before calling super().__init__. + if llm_config is not None or text_config is None: + self.text_config = _build_text_config(llm_config) + elif isinstance(text_config, PreTrainedConfig): + self.text_config = text_config + else: + self.text_config = _build_text_config(text_config) + super().__init__(**kwargs) + if self.architectures and self.architectures != [_ARCHITECTURE]: + raise ValueError( + f"Only {_ARCHITECTURE} is supported among NeMo SpeechLM checkpoints, " + f"got architectures={self.architectures!r}. " + + _other_architecture_hint(pe_encoder_config, self.text_config) + ) + + self.llm_config = llm_config + self.perception = perception or {} + self.pe_encoder_config = pe_encoder_config or {} + + decoder_config = self.text_config + architectures = _decoder_architectures(decoder_config) + if len(architectures) != 1 or architectures[0] not in _HYBRID_ARCHITECTURES: + raise ValueError( + f"Expected a NemotronH decoder, got architectures={architectures!r}." + ) + decoder_config.architectures = [_DECODER_ARCHITECTURE] + + if ( + llm_config is not None + and getattr(decoder_config, "speechlm_output_vocab_size", None) is None + ): + decoder_config.speechlm_output_vocab_size = int(decoder_config.vocab_size) + decoder_config.vocab_size = ( + int(decoder_config.vocab_size) + SPEECHLM_EMBED_EXTRA_ROWS + ) + self.output_vocab_size = int( + getattr( + decoder_config, "speechlm_output_vocab_size", decoder_config.vocab_size + ) + ) + + if self.perception: + self.audio_config = resolve_audio_config( + self.perception, self.pe_encoder_config + ) + else: + self.audio_config = audio_config or {} + + def get_text_config(self, decoder: bool = False) -> PreTrainedConfig: + return self.text_config From 65faf4efa0281fc2e04096a28cd35fb8dc7d3218 Mon Sep 17 00:00:00 2001 From: Ronald Date: Thu, 29 Oct 2026 00:10:39 +0800 Subject: [PATCH 067/406] [Refactor][R3] Initialize routed experts capturer before KV cache profiling (#63237) Signed-off-by: Ronald1995 Signed-off-by: Ao Shen Co-authored-by: Ao Shen --- .../test_routed_experts_capture.py | 36 ++++++------ tests/v1/worker/test_gpu_model_runner_v2.py | 58 +++++++++++++++++++ tests/v1/worker/test_gpu_worker.py | 18 ++++++ .../aux_output_connector/worker.py | 16 ++--- vllm/v1/worker/gpu/model_runner.py | 12 +++- vllm/v1/worker/startup_plan.py | 2 +- 6 files changed, 112 insertions(+), 30 deletions(-) diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 7b96d277efdf..0a0fb3e45c52 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -660,15 +660,10 @@ def test_get_aux_output_connector_passes_config(monkeypatch): scheduler_config=SimpleNamespace(max_num_batched_tokens=32) ) model = Mock() - kv_cache_config = Mock() - result = aux_output_worker.get_aux_output_connector(model, config, kv_cache_config) + result = aux_output_worker.get_aux_output_connector(model, config) - constructor.assert_called_once_with( - model=model, - kv_cache_config=kv_cache_config, - vllm_config=config, - ) + constructor.assert_called_once_with(model=model, vllm_config=config) assert result is connector @@ -694,11 +689,13 @@ def test_aux_output_worker_connector_binds_capture_on_non_output_rank(monkeypatc ) model = Mock() connector = aux_output_worker.AuxOutputWorkerConnector( - vllm_config=config, - model=model, - kv_cache_config=SimpleNamespace(kv_cache_groups=[_full_attention_kv_group()]), + vllm_config=config, model=model + ) + connector.initialize_output_storage( + SimpleNamespace(kv_cache_groups=[_full_attention_kv_group()]) ) + assert connector._store is None and connector._buffer is None constructor.assert_called_once_with( max_num_batched_tokens=32, vllm_config=config, @@ -730,7 +727,8 @@ def test_aux_output_worker_connector_default_capacity(monkeypatch): monkeypatch.setattr( aux_output_worker, "BackgroundBlockObjectStore", background_store_constructor ) - monkeypatch.setattr(aux_output_worker, "RoutedExpertsBuffer", Mock()) + buffer_constructor = Mock() + monkeypatch.setattr(aux_output_worker, "RoutedExpertsBuffer", buffer_constructor) config = SimpleNamespace( aux_output_config=SimpleNamespace(max_bytes=None), @@ -739,16 +737,18 @@ def test_aux_output_worker_connector_default_capacity(monkeypatch): scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=32), max_concurrent_batches=2, ) - kwargs = dict( - vllm_config=config, - model=Mock(), - kv_cache_config=SimpleNamespace( + connector = aux_output_worker.AuxOutputWorkerConnector( + vllm_config=config, model=Mock() + ) + store_constructor.assert_not_called() + buffer_constructor.assert_not_called() + + connector.initialize_output_storage( + SimpleNamespace( num_blocks=10, kv_cache_groups=[_full_attention_kv_group(MLAAttentionSpec)], - ), + ) ) - - aux_output_worker.AuxOutputWorkerConnector(**kwargs) assert store_constructor.call_args.kwargs["max_bytes"] == 2560 assert store_constructor.call_args.kwargs["object_nbytes"] == 128 assert background_store_constructor.call_args.kwargs["max_pending_batches"] == 16 diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 9aebddceda3f..6a043f83c78e 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -10,6 +10,7 @@ import pytest import torch +import vllm.distributed.aux_output_connector.worker as aux_output_worker import vllm.v1.worker.gpu.model_runner as model_runner_module from vllm.config import CacheConfig, LoRAConfig, ParallelConfig from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry @@ -71,6 +72,63 @@ def test_non_last_pp_rank_uses_global_batch_for_sample_feedback(monkeypatch): restore_for_sampling.assert_not_called() +@pytest.mark.parametrize( + "enabled,expected_events", + [ + (True, ["allocate", "bind", "dummy_run"]), + (False, ["dummy_run"]), + ], + ids=["aux-output-enabled", "aux-output-disabled"], +) +def test_profile_run_creates_aux_output_connector_before_forward( + monkeypatch, enabled, expected_events +): + events = [] + capturer = object() + model = torch.nn.Module() + vllm_config = SimpleNamespace( + aux_output_config=SimpleNamespace(enabled=enabled), + max_concurrent_batches=1, + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + ) + + def allocate(**kwargs): + assert kwargs == {"max_num_batched_tokens": 8, "vllm_config": vllm_config} + events.append("allocate") + return capturer + + def bind(bind_model, bind_capturer): + assert bind_model is model and bind_capturer is capturer + events.append("bind") + + def dummy_run(*args, **kwargs): + events.append("dummy_run") + return None, None + + monkeypatch.setattr(aux_output_worker, "RoutedExpertsCapturer", allocate) + monkeypatch.setattr(aux_output_worker, "bind_routed_experts_capturer", bind) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.vllm_config = vllm_config + runner.model = model + runner.aux_output_connector = None + runner.supports_mm_inputs = False + runner.is_last_pp_rank = False + runner.max_num_tokens = 8 + runner.encoder_cache = None + runner.pooling_runner = None + monkeypatch.setattr(runner, "_dummy_run", dummy_run) + + runner.profile_run() + + assert events == expected_events + connector = runner.aux_output_connector + if enabled: + assert connector is not None and connector._capturer is capturer + else: + assert connector is None + + def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch): runner = GPUModelRunner.__new__(GPUModelRunner) runner.max_model_len = 262144 diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 3ed2fea99d94..cde1b316dbfd 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import patch @@ -352,6 +353,23 @@ def test_startup_plan_apply_gate(plan_env): assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes +def test_startup_plan_rejects_budget_recorded_before_aux_capture(plan_env): + """Ignore an old-schema budget even when the fingerprint matches.""" + worker = _plan_worker() + maybe_save_startup_plan(worker, 1024) + fingerprint = startup_plan.compute_plan_fingerprint(worker.vllm_config, 0, 1) + path = startup_plan._plan_path(fingerprint) + with open(path) as f: + plan = json.load(f) + plan["schema"] = 1 + with open(path, "w") as f: + json.dump(plan, f) + + maybe_apply_startup_plan(worker) + + assert worker.cache_config.kv_cache_memory_bytes is None + + # Memory accounting of the profiling run (Worker.determine_available_memory). # The fallback reads only the sign of the measured drop and this process's torch diff --git a/vllm/distributed/aux_output_connector/worker.py b/vllm/distributed/aux_output_connector/worker.py index 6044a61cccc3..6c053d17304c 100644 --- a/vllm/distributed/aux_output_connector/worker.py +++ b/vllm/distributed/aux_output_connector/worker.py @@ -88,7 +88,6 @@ def __init__( *, vllm_config: VllmConfig, model: torch.nn.Module, - kv_cache_config: KVCacheConfig, ) -> None: max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens capturer = RoutedExpertsCapturer( @@ -96,6 +95,7 @@ def __init__( ) bind_routed_experts_capturer(model, capturer) self._capturer = capturer + self._vllm_config = vllm_config self._store: BackgroundBlockObjectStore | None = None self._buffer: RoutedExpertsBuffer | None = None self._requests: dict[str, _WorkerRequestState] = {} @@ -107,11 +107,15 @@ def __init__( self._pending_outputs: list[PendingAuxOutput] = [] self._lock = Lock() self._max_concurrent_batches = vllm_config.max_concurrent_batches + + def initialize_output_storage(self, kv_cache_config: KVCacheConfig) -> None: + """Create the output data plane once the KV cache layout is resolved.""" # Every TP rank participates in capture collectives, but only the # executor output rank owns the auxiliary output data plane. if not get_tp_group().is_first_rank: return + vllm_config = self._vllm_config shape_per_token = self._capturer.shape_per_token dtype: np.dtype[Any] = np.dtype(self._capturer.output_dtype_name) scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( @@ -131,7 +135,7 @@ def __init__( shape_per_token, hash_block_size, vllm_config.scheduler_config.max_num_seqs, - max_num_batched_tokens, + vllm_config.scheduler_config.max_num_batched_tokens, vllm_config.max_concurrent_batches, ) @@ -406,10 +410,6 @@ def close(self) -> None: def get_aux_output_connector( - model: torch.nn.Module, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig + model: torch.nn.Module, vllm_config: VllmConfig ) -> AuxOutputWorkerConnector: - return AuxOutputWorkerConnector( - model=model, - kv_cache_config=kv_cache_config, - vllm_config=vllm_config, - ) + return AuxOutputWorkerConnector(model=model, vllm_config=vllm_config) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 6509eaf52c20..568fd2f496af 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -842,9 +842,11 @@ def initialize_kv_cache( # AuxOutput connector requires resolved kv_cache_config. if self.vllm_config.aux_output_config.enabled: - self.aux_output_connector = get_aux_output_connector( - self.model, self.vllm_config, kv_cache_config - ) + if self.aux_output_connector is None: + self.aux_output_connector = get_aux_output_connector( + self.model, self.vllm_config + ) + self.aux_output_connector.initialize_output_storage(kv_cache_config) def _init_kv_zero_meta(self) -> None: """Build KV-block zeroing metadata; invoked from gpu_worker.""" @@ -1042,6 +1044,10 @@ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: @torch.inference_mode() def profile_run(self, randomize_inputs: bool = False) -> None: + if self.vllm_config.aux_output_config.enabled: + self.aux_output_connector = get_aux_output_connector( + self.model, self.vllm_config + ) if self.supports_mm_inputs and self.is_first_pp_rank: mm_config = self.model_config.multimodal_config if mm_config is not None and not mm_config.skip_mm_profiling: diff --git a/vllm/v1/worker/startup_plan.py b/vllm/v1/worker/startup_plan.py index 2c4942070705..a5812d8bf5f5 100644 --- a/vllm/v1/worker/startup_plan.py +++ b/vllm/v1/worker/startup_plan.py @@ -35,7 +35,7 @@ logger = init_logger(__name__) -PLAN_SCHEMA_VERSION = 1 +PLAN_SCHEMA_VERSION = 2 def compute_plan_fingerprint( From 165213f91169009bd78f8262a4bfd3b9a11bd841 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 28 Oct 2026 09:34:00 -0700 Subject: [PATCH 068/406] [Bugfix][Muse] Finalize dummy-loaded MegaMoE weights (#63871) Signed-off-by: Woosuk Kwon Co-authored-by: Codex --- docs/features/quantization/modelopt.md | 3 + tests/quantization/test_blackwell_moe.py | 84 +++++++++++++++++++++++ vllm/models/muse_spark/nvidia/mega_moe.py | 10 +++ vllm/models/muse_spark/nvidia/model.py | 7 +- 4 files changed, 100 insertions(+), 4 deletions(-) diff --git a/docs/features/quantization/modelopt.md b/docs/features/quantization/modelopt.md index 6d59a23565f0..2560f3cd76c6 100644 --- a/docs/features/quantization/modelopt.md +++ b/docs/features/quantization/modelopt.md @@ -55,6 +55,9 @@ For performance experiments with random weights, use `--load-format dummy`. ModelOpt NVFP4 MoE layers initialize positive quantization scales with matching gate/up global scales, including for FlashInfer MegaMoE (`--moe-backend flashinfer_moe_ep_cutedsl`). +Muse Spark's NVIDIA `deep_gemm_mega_moe` backend also supports dummy loading: +it initializes valid NVFP4 scales and finalizes its expert weights and residual +constants before the first forward. Enable expert parallelism for this backend. Dummy weights are for benchmarking, not model quality evaluation. Normal checkpoint loading uses the checkpoint's scales. diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index e249a7ad91ca..15fd10028cb8 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -387,3 +387,87 @@ def check_scales(model): for step in completion.logprobs for value in step.values() ) + + +@pytest.mark.parametrize("nvfp4", [False, True]) +def test_muse_megamoe_dummy_loading(tmp_path, monkeypatch, vllm_runner, nvfp4): + """Dummy-loaded Muse experts generate finite outputs in both weight formats.""" + import math + + from vllm import SamplingParams + from vllm.model_executor.models import ModelRegistry + + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + ModelRegistry.register_model( + "MuseSpark1xForCausalLM", + "vllm.models.muse_spark.nvidia.model:MuseSpark1xForCausalLM", + ) + config = dict( + architectures=["MuseSpark1xForCausalLM"], + model_type="muse_spark_text", + vocab_size=128, + hidden_size=512, + intermediate_size=512, + moe_hidden_size=512, + moe_intermediate_size=512, + num_hidden_layers=1, + num_attention_heads=8, + num_key_value_heads=8, + head_dim=64, + num_local_experts=8, + num_experts_per_tok=2, + num_shared_experts=0, + max_position_embeddings=64, + hidden_act="silu", + attention_bias=False, + layer_types=["full_attention"], + mlp_layer_types=["sparse"], + layer_rope_theta=[10000.0], + rope_parameters={"rope_type": "default", "rope_theta": 10000.0}, + rms_norm_eps=1e-5, + post_norm_eps=1e-8, + residual_gate_temperature=0.3, + final_logit_softcapping=None, + output_multiplier=0.5, + routed_scaling_factor=1.0, + qk_scale_factor=1.0, + tie_word_embeddings=False, + torch_dtype="bfloat16", + ) + if nvfp4: + config["quantization_config"] = { + "quant_method": "modelopt", + "quant_algo": "NVFP4", + "group_size": 16, + "kv_cache_quant_algo": None, + "exclude_modules": ["lm_head", "model.layers.*.self_attn.*"], + } + (tmp_path / "config.json").write_text(json.dumps(config)) + + with vllm_runner( + str(tmp_path), + dtype="bfloat16", + load_format="dummy", + moe_backend="deep_gemm_mega_moe", + enable_expert_parallel=True, + skip_tokenizer_init=True, + enforce_eager=True, + enable_flashinfer_autotune=False, + kernel_config={"enable_jit_warmup": False, "enable_cutedsl_warmup": False}, + max_model_len=32, + max_num_seqs=1, + max_num_batched_tokens=32, + kv_cache_memory_bytes=16 << 20, + ) as llm: + outputs = llm.get_llm().generate( + [{"prompt_token_ids": [1, 2, 3]}], + SamplingParams(temperature=0, max_tokens=2, ignore_eos=True, logprobs=1), + ) + completion = outputs[0].outputs[0] + assert len(completion.token_ids) == 2 + assert all( + math.isfinite(value.logprob) + for step in completion.logprobs + for value in step.values() + ) diff --git a/vllm/models/muse_spark/nvidia/mega_moe.py b/vllm/models/muse_spark/nvidia/mega_moe.py index d5d34c96f7f8..0ad466115451 100644 --- a/vllm/models/muse_spark/nvidia/mega_moe.py +++ b/vllm/models/muse_spark/nvidia/mega_moe.py @@ -138,6 +138,16 @@ def param(*shape: int, dtype: torch.dtype) -> nn.Parameter: self.w2_input_scale = param(num_local_experts, dtype=torch.float32) self.w13_input_scale = param(num_local_experts, 2, dtype=torch.float32) + for scale, value in ( + (self.w13_weight_scale, 1.0), + (self.w2_weight_scale, 1.0), + (self.w13_weight_scale_2, 1e-3), + (self.w2_weight_scale_2, 1e-3), + (self.w13_input_scale, 1e-2), + (self.w2_input_scale, 1e-2), + ): + set_weight_attrs(scale, {"dummy_weight_value": value}) + else: self.w13_weight = param( num_local_experts, diff --git a/vllm/models/muse_spark/nvidia/model.py b/vllm/models/muse_spark/nvidia/model.py index 779f86043225..db1ae0f0062c 100644 --- a/vllm/models/muse_spark/nvidia/model.py +++ b/vllm/models/muse_spark/nvidia/model.py @@ -1930,13 +1930,12 @@ def interleaved_weights() -> Iterable[tuple[str, torch.Tensor]]: yield name, weight loader = AutoWeightsLoader(self) - loaded = loader.load_weights( - interleaved_weights(), mapper=self.hf_to_vllm_mapper - ) + return loader.load_weights(interleaved_weights(), mapper=self.hf_to_vllm_mapper) + + def process_weights_after_loading(self) -> None: self.model.finalize_mega_moe_weights() n_gates = materialize_residual_path_constants(self) logger.info( "Materialized residual-path constants for %d gated residual connections", n_gates, ) - return loaded From af4db2a675f79ed20a8888ee0ec0c71ad581ad25 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Wed, 28 Oct 2026 17:42:08 +0100 Subject: [PATCH 069/406] [Quantization] Enable per-layer fused shared expert selection (#61457) --- tests/kernels/moe/test_fp8_tp_loading.py | 1 + .../layers/test_fused_shared_expert.py | 980 ++++++++++++++++-- .../model_loader/test_reload.py | 5 +- tests/models/test_deepseek_v4_vl_rocm.py | 70 +- tests/quantization/test_auto_gptq.py | 2 + .../test_online_shared_expert_fusion.py | 29 +- .../test_spec_decode_embed_sharing_pp.py | 2 + .../model_executor/layers/fused_moe/config.py | 6 +- .../layers/fused_moe/routed_experts.py | 18 + vllm/model_executor/layers/fused_moe/utils.py | 92 +- vllm/model_executor/models/AXK1.py | 17 +- vllm/model_executor/models/alice_ai.py | 10 +- vllm/model_executor/models/deepseek_mtp.py | 26 +- vllm/model_executor/models/deepseek_v2.py | 20 +- vllm/model_executor/models/glm4_moe.py | 10 +- vllm/model_executor/models/glm4_moe_lite.py | 24 +- .../models/glm4_moe_lite_mtp.py | 30 +- vllm/model_executor/models/glm4_moe_mtp.py | 26 +- vllm/model_executor/models/qwen3_5.py | 12 +- vllm/model_executor/models/qwen3_5_mtp.py | 12 +- vllm/model_executor/models/qwen3_next.py | 10 +- vllm/model_executor/models/qwen3_next_mtp.py | 12 +- vllm/model_executor/models/utils.py | 97 +- vllm/models/deepseek_v32/amd/model.py | 24 +- vllm/models/deepseek_v32/amd/mtp.py | 25 +- vllm/models/deepseek_v4/amd/model.py | 119 +-- vllm/models/deepseek_v4/amd/mtp.py | 12 + vllm/models/deepseek_v4/cpu/model.py | 9 +- vllm/models/deepseek_v41/amd/model.py | 12 + vllm/models/glm5next/common/model.py | 21 +- vllm/models/glm5next/common/mtp.py | 28 +- vllm/models/minimax_m3/amd/model.py | 50 +- vllm/models/qwen4_exp/amd/model.py | 12 +- vllm/models/qwen4_exp/amd/mtp.py | 12 +- vllm/models/qwen4_exp/nvidia/model.py | 12 +- vllm/models/qwen4_exp/nvidia/mtp.py | 12 +- vllm/triton_utils/dispatcher.py | 18 +- 37 files changed, 1418 insertions(+), 459 deletions(-) diff --git a/tests/kernels/moe/test_fp8_tp_loading.py b/tests/kernels/moe/test_fp8_tp_loading.py index 680ccb812cca..0ac056919693 100644 --- a/tests/kernels/moe/test_fp8_tp_loading.py +++ b/tests/kernels/moe/test_fp8_tp_loading.py @@ -68,6 +68,7 @@ def _make_fp8_tp_experts( or Fp8Config(is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128]), expert_map_manager=SimpleNamespace( local_num_experts=num_experts, + num_fused_shared_experts=0, placement_strategy="linear", expert_map=None, expert_mask=None, diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 30b7e2fbd34a..166ccdd197e7 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -17,6 +17,7 @@ import vllm.config as vllm_config_module from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.config.quantization import QuantizationConfigArgs +from vllm.logger import _print_warning_once from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config @@ -31,7 +32,10 @@ from vllm.model_executor.models.qwen3_next import ( _should_replicate_misaligned_shared_expert, ) -from vllm.model_executor.models.utils import PPMissingLayer +from vllm.model_executor.models.utils import ( + get_fused_shared_expert_layers, + maybe_fuse_shared_experts, +) from vllm.models.deepseek_v4 import quant_config as deepseek_v4_quant_config from vllm.models.minimax_m3.amd import model as minimax_m3_model from vllm.platforms import current_platform @@ -148,13 +152,15 @@ class _StubGlm5NextAttention(nn.Module): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__() self.o_proj = SimpleNamespace(reduce_results=True) + # GLM5Next's MTP predictor reads this attribute from MLA layers. + self.mla_attn = SimpleNamespace() def get_fse_test_model_config( model_type: str, quantization_config: dict[str, Any], -) -> tuple[object, type[nn.Module]]: - config: object +) -> tuple[Any, type[nn.Module]]: + config: Any if model_type == "minimax_m3": config = MiniMaxM3TextConfig( hidden_size=128, @@ -166,24 +172,29 @@ def get_fse_test_model_config( head_dim=128, num_local_experts=2, num_experts_per_tok=1, - moe_layer_freq=[0, 1], + moe_layer_freq=[1, 1], sparse_attention_config=None, rotary_dim=64, quantization_config=quantization_config, ) return config, minimax_m3_model.MiniMaxM3Model - if model_type == "deepseek_v4": - from vllm.models.deepseek_v4.amd.model import DeepseekV4Model + if model_type in ("deepseek_v4", "deepseek_v4_cpu"): + if model_type == "deepseek_v4": + from vllm.models.deepseek_v4.amd.model import DeepseekV4Model + else: + from vllm.models.deepseek_v4.cpu.model import DeepseekV4Model config = SimpleNamespace( + model_type="deepseek_v4", vocab_size=256, hidden_size=128, - num_hidden_layers=1, + num_hidden_layers=2, num_attention_heads=1, head_dim=128, max_position_embeddings=128, sliding_window=None, - layer_types=["sliding_attention"], + layer_types=["sliding_attention", "sliding_attention"], + mlp_layer_types=[], compress_rates={}, rope_theta=10000.0, compress_rope_theta=10000.0, @@ -214,13 +225,13 @@ def get_fse_test_model_config( quantization_config=quantization_config, ) return config, DeepseekV4Model - if model_type == "qwen3_5": - from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + if model_type in ("qwen3_next", "qwen3_5", "alice_ai"): + from transformers import Qwen3NextConfig - config = Qwen3_5MoeTextConfig( + common = dict( vocab_size=256, hidden_size=128, - num_hidden_layers=1, + num_hidden_layers=2, num_attention_heads=1, num_key_value_heads=1, head_dim=128, @@ -232,20 +243,113 @@ def get_fse_test_model_config( shared_expert_intermediate_size=32, num_experts_per_tok=1, num_experts=2, - layer_types=["full_attention"], + layer_types=["full_attention", "full_attention"], quantization_config=quantization_config, ) - return config, Qwen3_5Model - if model_type == "deepseek_v2": - from transformers import DeepseekV2Config + if model_type == "qwen3_next": + from vllm.model_executor.models.qwen3_next import Qwen3NextModel + + return Qwen3NextConfig(**common), Qwen3NextModel + if model_type == "qwen3_5": + from vllm.model_executor.models.qwen3_5 import Qwen3_5Model + + return Qwen3_5MoeTextConfig(**common), Qwen3_5Model + from vllm.model_executor.models.alice_ai import AliceAIModel + + config = Qwen3NextConfig(**common, linear_conv_kernel_dim=4) + # AliceAI requires the AttnRes grouping metadata. + config.block_attn_res_block_size = 1 + return config, AliceAIModel + if model_type == "qwen4_exp": + from transformers import Qwen4ExpTextConfig + + config = Qwen4ExpTextConfig( + vocab_size=256, + hidden_size=128, + num_hidden_layers=2, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=1, + linear_num_value_heads=1, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + num_experts_per_tok=1, + num_experts=2, + layer_types=["full_attention", "full_attention"], + eos_token_id=1, + hc_count=2, + hc_lowrank=4, + ple_layer_ids=[], + quantization_config=quantization_config, + ) + from vllm.models.qwen4_exp.amd.model import Qwen4ExpModel + + return config, Qwen4ExpModel + if model_type == "AXK1": + from transformers import AXK1Config - from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model + from vllm.model_executor.models.AXK1 import AXK1Model - config = DeepseekV2Config( + config = AXK1Config( vocab_size=256, hidden_size=128, intermediate_size=32, - num_hidden_layers=1, + moe_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=1, + num_key_value_heads=1, + n_shared_experts=1, + n_routed_experts=2, + num_experts_per_tok=1, + n_group=1, + topk_group=1, + first_k_dense_replace=0, + q_lora_rank=0, + kv_lora_rank=0, + qk_nope_head_dim=0, + qk_rope_head_dim=0, + v_head_dim=0, + quantization_config=quantization_config, + ) + return config, AXK1Model + if model_type == "glm4_moe_lite": + from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig + + from vllm.model_executor.models.glm4_moe_lite import Glm4MoeLiteModel + + config = Glm4MoeLiteConfig( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + moe_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=1, + num_key_value_heads=1, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + q_lora_rank=0, + kv_lora_rank=32, + qk_nope_head_dim=32, + qk_rope_head_dim=16, + v_head_dim=32, + dtype=torch.bfloat16, + quantization_config=quantization_config, + ) + config.first_k_dense_replace = 0 + config.num_nextn_predict_layers = 1 + return config, Glm4MoeLiteModel + if model_type in ("deepseek_v2", "deepseek_v32"): + from transformers import DeepseekV2Config + + common = dict( + vocab_size=256, + hidden_size=128, + intermediate_size=32, + num_hidden_layers=2, num_attention_heads=1, num_key_value_heads=1, head_dim=128, @@ -263,7 +367,18 @@ def get_fse_test_model_config( v_head_dim=0, quantization_config=quantization_config, ) - return config, DeepseekV2Model + config = DeepseekV2Config(**common) + if model_type == "deepseek_v2": + from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model + + # Needed to build the DeepSeek-V2 MTP draft model in this test. + config.num_nextn_predict_layers = 1 + return config, DeepseekV2Model + from vllm.models.deepseek_v32.amd.model import DeepseekV32Model + + config.index_topk = 1 + config.num_nextn_predict_layers = 1 + return config, DeepseekV32Model if model_type == "glm4_moe": from transformers.models.glm4_moe import Glm4MoeConfig @@ -273,7 +388,7 @@ def get_fse_test_model_config( vocab_size=256, hidden_size=128, intermediate_size=32, - num_hidden_layers=1, + num_hidden_layers=2, num_attention_heads=1, num_key_value_heads=1, head_dim=128, @@ -295,14 +410,14 @@ def get_fse_test_model_config( hidden_size=128, intermediate_size=32, moe_intermediate_size=32, - num_hidden_layers=1, + num_hidden_layers=2, num_attention_heads=1, num_key_value_heads=1, n_routed_experts=2, n_shared_experts=1, num_experts_per_tok=1, - layer_types=["linear_attention"], - mlp_layer_types=["sparse"], + layer_types=["linear_attention", "linear_attention"], + mlp_layer_types=["sparse", "sparse"], index_topk=1, index_kpool=1, pad_token_id=None, @@ -311,6 +426,8 @@ def get_fse_test_model_config( mhc=False, quantization_config=quantization_config, ) + config.mla_use_nope = False + config.num_nextn_predict_layers = 1 return config, Glm5NextModel raise ValueError(f"Unsupported FSE test model: {model_type}") @@ -402,9 +519,18 @@ def check_compatibility( ] -def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization( +@pytest.mark.parametrize( + ("reason", "enable_ep"), + [ + ("shared experts are excluded", False), + ("expert parallelism is enabled", True), + ], +) +def test_resolve_layer_fused_shared_expert_logs_skip_reason( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, + reason: str, + enable_ep: bool, ) -> None: monkeypatch.setattr( fused_moe_utils.rocm_aiter_ops, @@ -414,13 +540,20 @@ def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization( monkeypatch.setattr( fused_moe_utils, "is_shared_expert_quant_fse_compatible", - lambda *_, **__: (False, "shared experts are excluded"), + lambda *_, **__: (False, reason), + ) + monkeypatch.setattr( + "vllm.config.get_current_vllm_config_or_none", + lambda: SimpleNamespace( + parallel_config=SimpleNamespace(enable_expert_parallel=enable_ep) + ), ) + _print_warning_once.cache_clear() assert not fused_moe_utils.resolve_layer_fused_shared_expert( - ExpertsInt8Config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp", disable_under_ep=enable_ep ) - assert "shared experts are excluded" in caplog.text + assert reason in caplog.text def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix( @@ -884,45 +1017,409 @@ def test_aiter_fused_moe_validates_shared_expert_arguments() -> None: assert validate(shared_tensor, shared_tensor, shared_tensor, shared_tensor, 0) -def test_is_model_fused_shared_expert_compatible() -> None: - class MoE(nn.Module): - def __init__(self, enabled: bool) -> None: - super().__init__() - self.is_fused_shared_expert_enabled = enabled +def test_deepseek_v4_shared_expert_mapper_fallback() -> None: + """Mixed and all-excluded DSV4 checkpoints map shared experts correctly.""" + from vllm.models.deepseek_v4.amd.model import ( + _make_deepseek_v4_weights_mapper, + _redirect_shared_expert_name, + ) + + tensor = torch.empty(0) + fuse_by_layer = frozenset({6}) + # Mixed: any fused layer puts the mapper in fused mode, which leaves the + # checkpoint .w2 names untouched; the loader redirects per layer. + mapper = _make_deepseek_v4_weights_mapper("fp8", fuse_shared_experts=True) + for raw, expected in [ + ( + "layers.5.ffn.shared_experts.w2.weight", + "model.layers.5.ffn.shared_experts.down_proj.weight", + ), + ( + "layers.5.ffn.shared_experts.w2.scale", + "model.layers.5.ffn.shared_experts.down_proj.weight_scale_inv", + ), + ( + "layers.6.ffn.shared_experts.w2.weight", + "model.layers.6.ffn.experts.4.w2.weight", + ), + ( + "layers.6.ffn.shared_experts.w2.scale", + "model.layers.6.ffn.experts.4.w2.weight_scale_inv", + ), + ]: + [(mapped, _)] = list(mapper.apply([(raw, tensor)])) + assert _redirect_shared_expert_name(mapped, fuse_by_layer, 4) == expected + + # All layers excluded (fused_shared_expert_layers == frozenset()) selects + # the unfused mapper, which renames .w2 -> .down_proj itself. + mapper = _make_deepseek_v4_weights_mapper("fp8", fuse_shared_experts=False) + for raw, expected in [ + ( + "layers.5.ffn.shared_experts.w2.weight", + "model.layers.5.ffn.shared_experts.down_proj.weight", + ), + ( + "layers.5.ffn.shared_experts.w2.scale", + "model.layers.5.ffn.shared_experts.down_proj.weight_scale_inv", + ), + ]: + [(mapped, _)] = list(mapper.apply([(raw, tensor)])) + assert _redirect_shared_expert_name(mapped, frozenset(), 4) == expected + - class Layer(nn.Module): - def __init__(self, enabled: bool) -> None: - super().__init__() - self.mlp = MoE(enabled) +class _FakeMoE(nn.Module): + def __init__(self, enabled: bool) -> None: + super().__init__() + self.is_fused_shared_expert_enabled = enabled - enabled_layers = nn.ModuleList([Layer(True)]) - disabled_layers = nn.ModuleList([Layer(False)]) - mixed_layers = nn.ModuleList([Layer(True), Layer(False)]) - empty_layers = nn.ModuleList() - pipeline_layers = nn.ModuleList([Layer(True), PPMissingLayer()]) - assert fused_moe_utils.is_model_fused_shared_expert_compatible( - enabled_layers, MoE, "mlp" +def _make_fake_moe_model( + layer_names: list[str], enabled: list[bool], use_dict: bool +) -> nn.Module: + layers = [] + for flag in enabled: + layer = nn.Module() + layer.mlp = _FakeMoE(flag) + layers.append(layer) + model = nn.Module() + model.layers = ( + nn.ModuleDict(dict(zip(layer_names, layers, strict=True))) + if use_dict + else nn.ModuleList(layers) ) - assert not fused_moe_utils.is_model_fused_shared_expert_compatible( - disabled_layers, MoE, "mlp" + return model + + +def test_fused_shared_expert_layers_matches_container_index_space() -> None: + """The fused-layer set follows the module tree, not a hard-coded offset. + + ``nn.ModuleList`` MTP layers (Qwen3 Next/3.5, Qwen4 Exp) are 0-based, + while ``nn.ModuleDict`` MTP layers (GLM4, DeepSeek) are keyed by absolute + layer index. Both must redirect without a per-loader offset. + """ + list_model = _make_fake_moe_model(["0", "1"], [True, False], use_dict=False) + dict_model = _make_fake_moe_model(["28", "29"], [False, True], use_dict=True) + list_layers = get_fused_shared_expert_layers(list_model, _FakeMoE) + dict_layers = get_fused_shared_expert_layers(dict_model, _FakeMoE) + assert list_layers == frozenset({0}) + assert dict_layers == frozenset({29}) + + weight = torch.arange(8.0).reshape(4, 2) + + def redirected(name: str, fused_shared_expert_layers: frozenset[int]) -> bool: + [(out_name, _)] = list( + maybe_fuse_shared_experts( + [(name, weight)], + n_routed_experts=4, + n_shared_experts=1, + fused_shared_expert_layers=fused_shared_expert_layers, + ckpt_prefix="mlp.shared_expert", + ) + ) + return ".experts." in out_name + + # ModuleList: 0-based names, including the Qwen4 Exp MTP case where the + # quant-config prefix uses an absolute index but the weight names do not. + assert redirected("layers.0.mlp.shared_expert.down_proj.weight", list_layers) + assert not redirected("layers.1.mlp.shared_expert.down_proj.weight", list_layers) + # ModuleDict: absolute names, including the MTP layer offset. + assert redirected("layers.29.mlp.shared_expert.down_proj.weight", dict_layers) + assert not redirected("layers.28.mlp.shared_expert.down_proj.weight", dict_layers) + + +@pytest.mark.parametrize( + ("enable_expert_parallel", "match"), + [ + (True, "expert parallelism is not tested"), + (False, "nixl_ep"), + ], +) +def test_mixed_fused_shared_experts_rejected_with_ep_or_nixl_ep( + enable_expert_parallel: bool, match: str +) -> None: + """Mixed FSE must be rejected with EP or a shared-workspace all2all backend. + + With EP, per-layer geometry differences are not validated. ``nixl_ep`` + additionally keeps a single process-wide buffer initialized from the first + MoE layer's ``num_experts_per_rank`` and ignores later geometries. A mixed + model has fused layers with ``n_routed + n_shared`` experts and non-fused + layers with ``n_routed``, so reusing that buffer would silently map experts + to the wrong slots. Construction must fail cleanly instead of running the + untested mixed path. + """ + model = _make_fake_moe_model(["0", "1"], [True, False], use_dict=False) + + vllm_config = VllmConfig() + vllm_config.parallel_config.enable_expert_parallel = enable_expert_parallel + vllm_config.parallel_config.data_parallel_size = 2 + vllm_config.parallel_config.all2all_backend = "nixl_ep" + + with ( + set_current_vllm_config(vllm_config), + pytest.raises(NotImplementedError, match=match), + ): + get_fused_shared_expert_layers(model, _FakeMoE) + + +def _assert_mtp_fused_shared_expert_redirect( + mtp: nn.Module, ckpt_prefix: str, moe: nn.Module, n_routed_experts: int +) -> None: + """Load a fused MTP shared expert and verify it lands in the appended slot. + + The MTP loaders see different name/index spaces (0-based ``ModuleList`` or + absolute ``ModuleDict``); using the live module map must redirect either. + """ + packed_down = torch.full((128, 16), 3, dtype=torch.uint8) + down_scale = torch.full((128, 1), 127, dtype=torch.uint8) + mtp.load_weights( + [ + (f"{ckpt_prefix}.down_proj.weight", packed_down), + (f"{ckpt_prefix}.down_proj.weight_scale", down_scale), + ] ) - assert not fused_moe_utils.is_model_fused_shared_expert_compatible( - empty_layers, MoE, "mlp" + torch.testing.assert_close( + moe.experts.routed_experts.w2_weight[n_routed_experts][:128, :16].float(), + packed_down.float(), ) - assert fused_moe_utils.is_model_fused_shared_expert_compatible( - pipeline_layers, MoE, "mlp" + + +def _assert_mtp_shared_expert_not_fused( + mtp: nn.Module, ckpt_prefix: str, moe: nn.Module, n_routed_experts: int +) -> None: + """Load an unfused MTP shared expert and verify no fused slot is appended. + + Without AITER the shared expert stays a separate MLP, so the routed tensor must + keep exactly ``n_routed_experts`` rows. + """ + packed_down = torch.full((128, 16), 3, dtype=torch.uint8) + down_scale = torch.full((128, 1), 127, dtype=torch.uint8) + mtp.load_weights( + [ + (f"{ckpt_prefix}.down_proj.weight", packed_down), + (f"{ckpt_prefix}.down_proj.weight_scale", down_scale), + ] ) - with pytest.raises(NotImplementedError, match="1 enabled and 1 disabled layers"): - fused_moe_utils.is_model_fused_shared_expert_compatible( - mixed_layers, MoE, "mlp" + assert moe.experts.routed_experts.w2_weight.shape[0] == n_routed_experts + + +def _make_deepseek_v4_shared_expert_checkpoint( + layer: nn.Module, idx: int, fused: bool, exclude: bool +) -> tuple[list[tuple[str, torch.Tensor]], list[tuple[torch.Tensor, torch.Tensor]]]: + """Build DSV4's FP8 shared-expert / MXFP4 routed-slot checkpoint pairs.""" + ffn = layer.ffn + if not fused: + dtype = torch.bfloat16 if exclude else torch.float8_e4m3fn + hidden = ffn.shared_experts.gate_up_proj.weight.shape[1] + intermediate = ffn.shared_experts.down_proj.weight.shape[1] + gate_up = torch.full((2 * intermediate, hidden), idx + 1, dtype=dtype) + down = torch.full((hidden, intermediate), idx + 1, dtype=dtype) + return ( + [ + (f"layers.{idx}.ffn.shared_experts.gate_up_proj.weight", gate_up), + (f"layers.{idx}.ffn.shared_experts.down_proj.weight", down), + ], + [ + (ffn.shared_experts.gate_up_proj.weight, gate_up), + (ffn.shared_experts.down_proj.weight, down), + ], + ) + + routed = ffn.experts.routed_experts + slot = routed.w13_weight.shape[0] - 1 + w13_rows, hidden_half = routed.w13_weight[slot].shape + intermediate = w13_rows // 2 + hidden = hidden_half * 2 + w1 = torch.full((intermediate, hidden_half), idx + 1, dtype=torch.uint8) + w3 = torch.full((intermediate, hidden_half), idx + 1, dtype=torch.uint8) + w2 = torch.full((hidden, intermediate // 2), idx + 1, dtype=torch.uint8) + return ( + [ + (f"layers.{idx}.ffn.shared_experts.w1.weight", w1), + (f"layers.{idx}.ffn.shared_experts.w3.weight", w3), + (f"layers.{idx}.ffn.shared_experts.w2.weight", w2), + ], + [ + (routed.w13_weight[slot].chunk(2, dim=0)[0], w1), + (routed.w13_weight[slot].chunk(2, dim=0)[1], w3), + (routed.w2_weight[slot], w2), + ], + ) + + +def _build_mtp_model( + mtp_cls: Any, + vllm_config: VllmConfig, + config: object, + *, + with_text_config: bool = False, + with_target_model: bool = False, +) -> nn.Module: + draft_model_config: dict[str, object] = {"hf_config": config} + if with_text_config: + draft_model_config["hf_text_config"] = config + speculative_config: dict[str, object] = { + "draft_model_config": SimpleNamespace(**draft_model_config), + "method": "mtp", + "parallel_drafting": False, + "enable_adaptive_verification": False, + "use_dspark": lambda: False, + } + if with_target_model: + speculative_config["target_model_config"] = vllm_config.model_config + vllm_config.speculative_config = SimpleNamespace(**speculative_config) + return mtp_cls(vllm_config=vllm_config) + + +_MTP_FACTORIES: dict[str, tuple[str, str, dict[str, bool]]] = { + "deepseek_v4": ("vllm.models.deepseek_v4.amd.mtp", "DeepSeekV4MTP", {}), + "glm4_moe": ( + "vllm.model_executor.models.glm4_moe_mtp", + "Glm4MoeMTP", + {}, + ), + "glm4_moe_lite": ( + "vllm.model_executor.models.glm4_moe_lite_mtp", + "Glm4MoeLiteMTP", + {}, + ), + "deepseek_v32": ( + "vllm.models.deepseek_v32.amd.mtp", + "DeepseekV32MTP", + {}, + ), + "glm5_next": ( + "vllm.models.glm5next.common.mtp", + "Glm5NextMTP", + {"with_text_config": True}, + ), + "deepseek_v2": ( + "vllm.model_executor.models.deepseek_mtp", + "DeepSeekMTP", + {}, + ), + "qwen3_next": ( + "vllm.model_executor.models.qwen3_next_mtp", + "Qwen3NextMTP", + {"with_text_config": True, "with_target_model": True}, + ), + "qwen3_5": ( + "vllm.model_executor.models.qwen3_5_mtp", + "Qwen3_5MoeMTP", + {"with_text_config": True, "with_target_model": True}, + ), +} + + +def _build_mtp_for_model( + model_type: str, + vllm_config: VllmConfig, + config: object, + enabled: bool, +) -> nn.Module | None: + factory = _MTP_FACTORIES.get(model_type) + if not enabled or factory is None: + return None + module_name, class_name, kwargs = factory + module = importlib.import_module(module_name) + return _build_mtp_model(getattr(module, class_name), vllm_config, config, **kwargs) + + +# TODO: move to conftest.py. +class _SingleRankGroup: + """No-op single-rank stand-in for a compiled ``GroupCoordinator``.""" + + world_size = 1 + rank_in_group = 0 + local_rank = 0 + is_first_rank = True + is_last_rank = True + + @property + def device_group(self): + return self + + @property + def cpu_group(self): + return self + + @property + def device_communicator(self): + return self + + def rank(self) -> int: + return 0 + + def size(self) -> int: + return 1 + + def all_reduce(self, input_: torch.Tensor, *args: Any, **kwargs: Any): + return input_ + + def all_gather(self, input_: torch.Tensor, dim: int = 0): + return input_ + + def all_gatherv(self, inputs: Any, dim: int = 0, sizes: Any = None): + return inputs + + def reduce_scatter(self, input_: torch.Tensor, dim: int = 0): + return input_ + + def reduce_scatterv(self, input_: Any, dim: int = 0, sizes: Any = None): + return input_ + + def broadcast(self, input_: torch.Tensor, src: int = 0): + return input_ + + def send(self, *args: Any, **kwargs: Any) -> None: + return None + + def recv(self, *args: Any, **kwargs: Any) -> None: + return None + + def barrier(self) -> None: + return None + + def destroy(self) -> None: + return None + + +@pytest.fixture +def single_rank_parallel_state(monkeypatch: pytest.MonkeyPatch): + """Single rank parallel groups without initializing a process group. + + This is faster than the `dist_init` fixture. + """ + for name in ("_WORLD", "_TP", "_ETP", "_DCP", "_PP", "_DP", "_EP", "_PCP"): + monkeypatch.setattr( + f"vllm.distributed.parallel_state.{name}", _SingleRankGroup() ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0 + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1 + ) @pytest.mark.parametrize( "model_type", - ["minimax_m3", "deepseek_v4", "qwen3_5", "glm4_moe", "deepseek_v2", "glm5_next"], + [ + "minimax_m3", + "deepseek_v4", + "deepseek_v4_cpu", + "qwen3_next", + "qwen3_5", + "alice_ai", + "glm4_moe", + "deepseek_v2", + "AXK1", + "glm4_moe_lite", + "deepseek_v32", + "glm5_next", + "qwen4_exp", + ], ) +@pytest.mark.parametrize("shared_expert_precision", ["all_mxfp4", "mixed"]) @pytest.mark.parametrize( ("use_fse", "exclude"), [ @@ -933,17 +1430,82 @@ def __init__(self, enabled: bool) -> None: ) def test_models_fse_init( model_type: str, + shared_expert_precision: str, use_fse: bool, exclude: list[str], - dist_init: None, + single_rank_parallel_state: None, monkeypatch: pytest.MonkeyPatch, ) -> None: - """Model construction resolves FSE consistently with Quark quantization.""" + """Construction and loading preserve each shared expert's precision.""" + _print_warning_once.cache_clear() + per_layer_fse = model_type != "deepseek_v4_cpu" + moe_attr = { + "minimax_m3": "block_sparse_moe", + "deepseek_v4": "ffn", + }.get(model_type, "mlp") + shared_name = ( + "shared_experts" + if model_type + in { + "minimax_m3", + "deepseek_v4", + "glm4_moe", + "deepseek_v2", + "AXK1", + "glm4_moe_lite", + "deepseek_v32", + "glm5_next", + } + else "shared_expert" + ) + if shared_expert_precision == "mixed" and not per_layer_fse: + pytest.skip( + "The mixed Quark two-layer setup does not apply to this architecture" + ) quantization_config: dict[str, Any] = ( - get_deepseek_v4_quark_config(["layers.0.ffn.shared_experts"] if exclude else []) + get_deepseek_v4_quark_config( + [f"layers.{idx}.ffn.shared_experts" for idx in range(2)] if exclude else [] + ) if model_type == "deepseek_v4" - else {**_QUARK_FSE_CONFIG, "exclude": exclude} - ) + else {**deepcopy(_QUARK_FSE_CONFIG), "exclude": exclude} + ) + if shared_expert_precision == "mixed": + mxfp8_config = deepcopy(quantization_config["global_quant_config"]) + mxfp8_config["input_tensors"].update( + dtype="fp8_e4m3", + symmetric=True, + scale_type=None, + scale_format=None, + scale_calculation_mode=None, + observer_cls="PerGroupMinMaxObserver", + zero_point_type="int32", + ) + mxfp8_config["weight"].update( + dtype="fp8_e4m3", + qscheme="per_block", + ch_axis=None, + group_size=None, + block_size=[1, 32], + symmetric=True, + scale_type="float8_e8m0fnu", + scale_format=None, + scale_calculation_mode=None, + observer_cls="PerBlock2DMinMaxObserver", + zero_point_type="int32", + ) + if model_type == "deepseek_v4": + quantization_config.setdefault("layer_quant_config", {})[ + "layers.1.ffn.shared_experts.w1" + ] = deepcopy(mxfp8_config) + else: + quantization_config.setdefault("layer_quant_config", {}).update( + { + f"model.layers.1.{moe_attr}.{shared_name}.{projection}": deepcopy( + mxfp8_config + ) + for projection in ("gate_proj", "up_proj", "down_proj") + } + ) config, model_constructor = get_fse_test_model_config( model_type, quantization_config @@ -961,11 +1523,12 @@ def test_models_fse_init( quantization_config=None, runner_type="generate", is_moe=True, + use_mla=False, logits_processors=None, rswa_window=None, ) vllm_config.parallel_config.enable_expert_parallel = False - if model_type == "deepseek_v4": + if model_type in ("deepseek_v4", "deepseek_v4_cpu"): from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config vllm_config.cache_config.cache_dtype = "fp8_ds_mla" @@ -975,6 +1538,10 @@ def test_models_fse_init( ) else: vllm_config.quant_config = QuarkConfig(quantization_config) + if per_layer_fse: + vllm_config.quant_config.packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"] + } if model_type == "glm5_next": from vllm.models.glm5next.common import model as glm5_next_model @@ -985,18 +1552,40 @@ def test_models_fse_init( monkeypatch.setattr( glm5_next_model, "_fused_shared_experts_tuned", lambda _: True ) + from vllm.models.glm5next.common import mtp as glm5_next_mtp_model + + # The MTP predictor asserts on the concrete MLA attention class and + # reads ``mla_attn``, so patch its imported name to the same stub. + monkeypatch.setattr( + glm5_next_mtp_model, "Glm5NextMLAAttention", _StubGlm5NextAttention + ) + if model_type == "glm4_moe_lite": + from vllm.model_executor.models import glm4_moe_lite as glm4_moe_lite_model + + # Only the MoE/shared-expert wiring is under test here; the MLA backend + # does not accept this test's default dtype. + for attention in ("Glm4MoeLiteAttention", "Glm4MoeLiteMLAAttention"): + monkeypatch.setattr(glm4_moe_lite_model, attention, _StubGlm5NextAttention) + if model_type == "deepseek_v32": + from vllm.models.deepseek_v32.amd import model as deepseek_v32_model + + # The sparse-MLA attention needs a real MLA backend; only the MoE + # wiring is under test here. + monkeypatch.setattr( + deepseek_v32_model, "DeepseekV32MLAAttention", _StubGlm5NextAttention + ) + if model_type == "qwen4_exp": + from vllm.models.qwen4_exp.amd import model as qwen4_exp_model + + # Only the MoE/shared-expert wiring is under test here. + monkeypatch.setattr( + qwen4_exp_model, "Qwen3NextAttention", _StubGlm5NextAttention + ) import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops - if model_type == "deepseek_v4": - from vllm.models.deepseek_v4.amd import model as deepseek_v4_model - - warning_logger = deepseek_v4_model.logger - elif model_type == "minimax_m3": - warning_logger = minimax_m3_model.logger - else: - warning_logger = fused_moe_utils.logger + warning_logger = fused_moe_utils.logger with monkeypatch.context() as mp: mp.setenv("VLLM_ROCM_USE_AITER", str(use_fse)) @@ -1010,33 +1599,196 @@ def test_models_fse_init( importlib.reload(envs) mp.setattr(envs, "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", use_fse) rocm_aiter_ops.refresh_env_variables() + from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe + + mp.setattr(rocm_aiter_moe, "aiter_topK_meta_data", None) + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() aiter_fse_enabled = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) - # These AMD-specific models currently use the raw environment flag. - # and do not rely on `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`. - fse_enabled = use_fse and ( - aiter_fse_enabled or model_type in {"deepseek_v4", "minimax_m3"} - ) + fse_enabled = use_fse and aiter_fse_enabled with ( patch.object(warning_logger, "warning") as warning, set_current_vllm_config(vllm_config), ): - model = model_constructor(vllm_config=vllm_config) - mtp = None - if model_type == "deepseek_v4" and use_fse and not exclude: - from vllm.models.deepseek_v4.amd.mtp import ( - DeepSeekV4MTP, + model = model_constructor( + vllm_config=vllm_config, + **({"prefix": "model"} if per_layer_fse else {}), + ) + mtps = { + model_type: _build_mtp_for_model( + model_type, vllm_config, config, use_fse and not exclude ) - - vllm_config.speculative_config = SimpleNamespace( - draft_model_config=SimpleNamespace(hf_config=config), - method="mtp", - parallel_drafting=False, - enable_adaptive_verification=False, - use_dspark=lambda: False, + } + mtp = mtps.get("deepseek_v4") + if per_layer_fse: + expected_policy = { + f"layers.0.{moe_attr}": bool(fse_enabled and not exclude), + f"layers.1.{moe_attr}": bool( + fse_enabled + and not exclude + and shared_expert_precision == "all_mxfp4" + ), + } + assert model.fused_shared_expert_layers == frozenset( + idx for idx in range(2) if expected_policy[f"layers.{idx}.{moe_attr}"] + ) + if any(expected_policy.values()) and not all(expected_policy.values()): + with pytest.raises(NotImplementedError, match="EPLB"): + get_fused_shared_expert_layers( + model, + type(getattr(model.layers[0], moe_attr)), + enable_eplb=True, + ) + for idx, layer in enumerate(model.layers): + fused = expected_policy[f"layers.{idx}.{moe_attr}"] + moe = getattr(layer, moe_attr) + assert (getattr(moe, shared_name) is None) is fused + assert moe.experts.expert_map_manager.num_fused_shared_experts == int( + fused + ) + assert ( + moe.experts.moe_config.is_fused_shared_expert_requested + is aiter_fse_enabled + ) + if not fused and not exclude and model_type != "deepseek_v4": + scheme = getattr(moe, shared_name).gate_up_proj.scheme + expected_dtype = ( + "mxfp8_e4m3" + if idx == 1 and shared_expert_precision == "mixed" + else "mxfp4" + ) + assert scheme.weight_dtype == expected_dtype + + checkpoint = [] + expected_weights = [] + for idx, layer in enumerate(model.layers): + fused = expected_policy[f"layers.{idx}.{moe_attr}"] + if model_type == "deepseek_v4": + dsv4_checkpoint, dsv4_expected = ( + _make_deepseek_v4_shared_expert_checkpoint( + layer, idx, fused, exclude + ) + ) + checkpoint.extend(dsv4_checkpoint) + expected_weights.extend(dsv4_expected) + continue + high_precision = idx == 1 and shared_expert_precision == "mixed" + dtype = ( + torch.bfloat16 + if exclude + else torch.float8_e4m3fn + if high_precision + else torch.uint8 ) - mtp = DeepSeekV4MTP(vllm_config=vllm_config) - assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude) + pack_factor = 1 if exclude or high_precision else 2 + for projection in ("gate_proj", "up_proj", "down_proj"): + shape = ( + (128, 32 // pack_factor) + if projection == "down_proj" + else (32, 128 // pack_factor) + ) + weight = torch.full(shape, idx + 1).to(dtype) + name = f"layers.{idx}.{moe_attr}.{shared_name}.{projection}" + checkpoint.append((f"{name}.weight", weight)) + if not exclude: + scale_shape = (128, 1) if projection == "down_proj" else (32, 4) + checkpoint.append( + ( + f"{name}.weight_scale", + torch.full(scale_shape, 127, dtype=torch.uint8), + ) + ) + if fused: + routed = getattr(layer, moe_attr).experts.routed_experts + target = ( + routed.w2_weight[2] + if projection == "down_proj" + else routed.w13_weight[2].chunk(2, dim=0)[ + int(projection == "up_proj") + ] + ) + else: + shared = getattr(getattr(layer, moe_attr), shared_name) + target = ( + shared.down_proj.weight + if projection == "down_proj" + else shared.gate_up_proj.weight.chunk(2, dim=0)[ + int(projection == "up_proj") + ] + ) + expected_weights.append((target, weight)) + if shared_name == "shared_expert": + checkpoint.append( + ( + f"layers.{idx}.{moe_attr}.shared_expert_gate.weight", + torch.full((1, 128), idx + 1, dtype=torch.bfloat16), + ) + ) + with set_current_vllm_config(vllm_config): + loaded = model.load_weights(checkpoint) + base = config.num_hidden_layers + mtp_cases = [ + ( + mtp_model, + f"model.layers.{base}.mlp.{shared}", + mtp_model.model.layers[str(base)].mtp_block.mlp, + getattr(config, "n_routed_experts", 0), + ) + for mtp_model, shared in ( + (mtps.get("glm4_moe"), "shared_experts"), + (mtps.get("glm4_moe_lite"), "shared_experts"), + (mtps.get("deepseek_v2"), "shared_experts"), + (mtps.get("deepseek_v32"), "shared_experts"), + (mtps.get("glm5_next"), "shared_experts"), + ) + if mtp_model is not None + ] + mtp_cases += [ + ( + mtp_model, + "mtp.layers.0.mlp.shared_expert", + mtp_model.model.layers[0].mlp, + getattr(config, "num_experts", 0), + ) + for mtp_model in ( + mtps.get("qwen3_next"), + mtps.get("qwen3_5"), + ) + if mtp_model is not None + ] + for mtp_model, ckpt_prefix, moe, n_routed in mtp_cases: + if fse_enabled and not exclude: + _assert_mtp_fused_shared_expert_redirect( + mtp_model, ckpt_prefix, moe, n_routed + ) + else: + _assert_mtp_shared_expert_not_fused( + mtp_model, ckpt_prefix, moe, n_routed + ) + for idx in range(2): + prefix = f"layers.{idx}.{moe_attr}" + expected_name = ( + f"{prefix}.experts.routed_experts.w2_weight" + if expected_policy[prefix] + else f"{prefix}.{shared_name}.down_proj.weight" + ) + assert expected_name in loaded + if shared_name == "shared_expert": + assert f"{prefix}.shared_expert_gate.weight" in loaded + for target, weight in expected_weights: + torch.testing.assert_close( + target[: weight.shape[0], : weight.shape[1]].float(), weight.float() + ) + elif model_type == "deepseek_v4_cpu": + # CPU never supports fused shared experts; the separate MLP is + # always built and no routed slot is appended. + for layer in model.layers: + assert layer.ffn.shared_experts is not None + assert ( + layer.ffn.experts.expert_map_manager.num_fused_shared_experts == 0 + ) + else: + assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude) if model_type == "glm5_next": assert (model.layers[0].mlp.shared_experts is None) is ( fse_enabled and not exclude @@ -1046,21 +1798,53 @@ def test_models_fse_init( # so should just raise a warning. if mtp is not None: assert not mtp.model.layers[ - "1" + str(config.num_hidden_layers) ].mtp_block.ffn.is_fused_shared_expert_enabled + if aiter_fse_enabled: + warning.assert_called_once() + assert ( + "DeepSeek-V4/V4.1 shared experts at " + in (warning.call_args.args[1]) + ) + else: + # Without AITER nothing fuses, so no skip reason is reported. + warning.assert_not_called() + if aiter_fse_enabled and shared_expert_precision == "mixed" and not exclude: warning.assert_called_once() - assert ( - "DeepSeek-V4/V4.1 shared experts at mtp.0.ffn.shared_experts" - in (warning.call_args.args[1]) - ) - if aiter_fse_enabled and exclude: + if aiter_fse_enabled and shared_expert_precision == "mixed" and not exclude: warning.assert_called_once() - assert ( - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled" - in (warning.call_args.args[0]) - ) - assert "excludes shared experts" in warning.call_args.args[1] + if model_type == "deepseek_v4": + assert "are not MXFP4" in warning.call_args.args[1] + else: + assert ( + "different quantization configurations" in warning.call_args.args[1] + ) + for mtp_model in ( + mtps.get("glm4_moe"), + mtps.get("deepseek_v2"), + mtps.get("glm4_moe_lite"), + mtps.get("deepseek_v32"), + mtps.get("glm5_next"), + ): + if mtp_model is not None: + assert mtp_model.fused_shared_expert_layers == ( + frozenset({config.num_hidden_layers}) + if fse_enabled + else frozenset() + ) + if aiter_fse_enabled and exclude: + if model_type == "deepseek_v4_cpu": + # CPU has no FSE path, so the resolver never warns. + warning.assert_not_called() + else: + assert warning.call_count == 1 + assert ( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled" + in (warning.call_args.args[0]) + ) + assert "excludes shared experts" in warning.call_args.args[1] + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() importlib.reload(envs) rocm_aiter_ops.refresh_env_variables() diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index eac5245bd988..01763dfb831c 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -720,6 +720,7 @@ def test_padded_moe_reload_releases_each_layer( has_bias=has_bias, tp_rank=tp_rank, tp_shard_with_padding=False, + num_experts=experts, moe_parallel_config=SimpleNamespace(tp_size=2), ) model = torch.nn.ModuleList() @@ -735,7 +736,9 @@ def test_padded_moe_reload_releases_each_layer( layer.moe_config = config layer.quant_config = None layer.quant_method = method - layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i) + layer.expert_map_manager = SimpleNamespace( + map_global_to_local=lambda i: i, num_fused_shared_experts=0 + ) layer._loaded_expert_biases = set() layer._fused_shared_expert_quantizer = None method.create_weights( diff --git a/tests/models/test_deepseek_v4_vl_rocm.py b/tests/models/test_deepseek_v4_vl_rocm.py index 762f1cd6c0e0..418a2b47c3a5 100644 --- a/tests/models/test_deepseek_v4_vl_rocm.py +++ b/tests/models/test_deepseek_v4_vl_rocm.py @@ -249,19 +249,71 @@ def fake_factory(**kwargs): def test_rocm_fse_rejects_data_parallel_deployment( monkeypatch: pytest.MonkeyPatch, ) -> None: - from vllm.models.deepseek_v4.amd.model import _fuse_shared_experts_enabled + from vllm.config import set_current_vllm_config + from vllm.model_executor.layers.fused_moe.utils import ( + resolve_layer_fused_shared_expert, + ) + from vllm.models.deepseek_v4.amd import model as rocm_model - config = SimpleNamespace(n_shared_experts=1) + class FakeGate(nn.Module): + def __init__(self, **kwargs) -> None: + super().__init__() - monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1") - with pytest.raises(ValueError, match="data parallelism"): - _fuse_shared_experts_enabled(config, ParallelConfig(data_parallel_size=2)) + monkeypatch.setattr(rocm_model, "GateLinear", FakeGate) + monkeypatch.setattr(rocm_model, "FusedMoEFactory", lambda **kwargs: nn.Identity()) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(rocm_model, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + rocm_model.rocm_aiter_ops, + "is_fusion_moe_shared_experts_enabled", + lambda: True, + ) - # Expert parallelism already disables FSE: no raise, FSE stays off. - assert not _fuse_shared_experts_enabled( - config, ParallelConfig(data_parallel_size=2, enable_expert_parallel=True) + config = SimpleNamespace( + hidden_size=16, + n_routed_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + swiglu_limit=None, + norm_topk_prob=True, + scoring_func="sqrtsoftplus", + num_hash_layers=0, + vocab_size=32, + topk_method="noaux_tc", + vision_n_layers=0, + n_shared_experts=1, ) - assert _fuse_shared_experts_enabled(config, ParallelConfig()) + + def build_vllm_config(parallel_config: ParallelConfig) -> SimpleNamespace: + return SimpleNamespace( + model_config=SimpleNamespace(hf_config=config), + quant_config=None, + parallel_config=parallel_config, + kernel_config=KernelConfig(), + ) + + # Data parallelism without expert parallelism fails fast at MoE construction. + with pytest.raises(ValueError, match="not supported with data parallelism"): + rocm_model.DeepseekV4MoE( + build_vllm_config(ParallelConfig(data_parallel_size=2)), + prefix="model.layers.0.ffn", + ) + + # Under expert parallelism the resolver disables FSE instead of raising. + with set_current_vllm_config( + build_vllm_config( + ParallelConfig(data_parallel_size=2, enable_expert_parallel=True) + ) + ): + assert not resolve_layer_fused_shared_expert( + None, "model.layers.0.ffn", disable_under_ep=True + ) + + # Nothing blocks FSE once expert parallelism is off. + with set_current_vllm_config(build_vllm_config(ParallelConfig())): + assert resolve_layer_fused_shared_expert( + None, "model.layers.0.ffn", disable_under_ep=True + ) def test_rocm_mtp_forwards_input_ids_for_vision_routing( diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index 4f60a974caed..f0df2aefe10b 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -166,8 +166,10 @@ class Loader: moe_config = SimpleNamespace( is_act_and_mul=True, tp_rank=0, + num_experts=1, moe_parallel_config=SimpleNamespace(tp_size=1), ) + expert_map_manager = SimpleNamespace(num_fused_shared_experts=0) _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim) _narrow_expert_data_for_padding = staticmethod( RoutedExperts._narrow_expert_data_for_padding diff --git a/tests/quantization/test_online_shared_expert_fusion.py b/tests/quantization/test_online_shared_expert_fusion.py index 394ba04a1b30..c10ba5e38046 100644 --- a/tests/quantization/test_online_shared_expert_fusion.py +++ b/tests/quantization/test_online_shared_expert_fusion.py @@ -573,18 +573,19 @@ def record_warning(message: str, *args: object, **_kwargs: object) -> None: "cannot be enabled" in warning for warning in logged_warnings ) - if architecture == "MiniMaxM3SparseForCausalLM": - # TODO: Use resolve_layer_fused_shared_expert in MiniMax M3 and remove - # this branch once its virtual fused projections are registered. - assert logged_messages == ["Quantizing 0 layers of types: "] - else: - assert logged_messages == [ - "Quantizing 2 layers of types: " - f"mlp.{shared_expert_name}.down_proj: 1 " - f"(from targets: {target_pattern}, mxfp4); " - f"mlp.{shared_expert_name}.gate_up_proj: 1 " - f"(from targets: {target_pattern}, mxfp4)" - ] + # MiniMax M3 now resolves fused shared experts per layer through + # `resolve_layer_fused_shared_expert`, so its virtual fused projections are + # registered and quantized like any other architecture. + moe_prefix = ( + "block_sparse_moe" if architecture == "MiniMaxM3SparseForCausalLM" else "mlp" + ) + assert logged_messages == [ + "Quantizing 2 layers of types: " + f"{moe_prefix}.{shared_expert_name}.down_proj: 1 " + f"(from targets: {target_pattern}, mxfp4); " + f"{moe_prefix}.{shared_expert_name}.gate_up_proj: 1 " + f"(from targets: {target_pattern}, mxfp4)" + ] def test_online_gate_up_proj_target_disables_fse( @@ -645,6 +646,6 @@ def record_warning(message: str, *args: object, **_kwargs: object) -> None: ] assert logged_warnings == [ "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be " - "enabled - skipping for this layer: online quantization targets only " - "part of the shared expert at model.layers.0.mlp.shared_experts." + "enabled - skipping shared-expert fusion: online quantization targets only " + "part of the shared expert at ." ] diff --git a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py index 0fef7bb47bdf..901a4d184c73 100644 --- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py +++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from vllm.config import ParallelConfig from vllm.model_executor.model_loader.mtp_validation import ( disable_mtp_completeness_check, ) @@ -158,6 +159,7 @@ def m3_config(): speculative_config=SimpleNamespace( method="mtp", draft_model_config=SimpleNamespace(hf_config=hf_config) ), + parallel_config=ParallelConfig(), quant_config=None, cache_config=None, kernel_config=SimpleNamespace(monokernel=False), diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index e8bd822becab..83df5d16c16e 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1338,7 +1338,7 @@ class FusedMoEConfig: # Use the allocated width as the checkpoint TP stride, including for scales. tp_shard_with_padding: bool = False rocm_aiter_fmoe_enabled: bool = False - aiter_fmoe_shared_expert_enabled: bool = False + is_fused_shared_expert_requested: bool = False # Whether to force MXFP4 (a4w4) MoE activations for DeepSeek V4.1 on # ROCm/AITER. Opt-in via VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1; rejected # for any other model type. Resolved here, not in the forward path, @@ -1371,7 +1371,7 @@ def __post_init__(self): if self.is_act_and_mul: self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled() - self.aiter_fmoe_shared_expert_enabled = ( + self.is_fused_shared_expert_requested = bool( rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() ) @@ -1402,7 +1402,7 @@ def __post_init__(self): assert self.rocm_aiter_fmoe_enabled, ( "Mori needs to be used with aiter fused_moe for now." ) - assert not self.aiter_fmoe_shared_expert_enabled, ( + assert not self.is_fused_shared_expert_requested, ( "Mori does not support fusion shared expert now. " "Turn it off by setting VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0" ) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 59aea0599b4e..3c4a1ffc7707 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -725,6 +725,24 @@ def weight_loader( quant_method_name = base_quant_method.__class__.__name__ global_expert_id = expert_id + num_global_experts = ( + self.moe_config.num_experts + + self.expert_map_manager.num_fused_shared_experts + ) + # `num_experts` counts logical and redundant experts; fused shared + # slots are appended after them. + # Expert substitution checkpoints name retained experts by their + # logical ID, which is not bounded by the number of allocated slots. + expert_substitution = getattr(self, "expert_substitution", None) + if expert_substitution is not None: + num_global_experts = ( + expert_substitution.num_logical_experts + + self.expert_map_manager.num_fused_shared_experts + ) + assert global_expert_id < num_global_experts, ( + f"Expert id {global_expert_id} is out of range for " + f"{self.layer_name}; expected < {num_global_experts}" + ) expert_id = self._map_global_expert_id_to_local_expert_id(global_expert_id) use_global_sf = ( diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 27df12668584..6a4911e8387f 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools -from collections.abc import Iterable from math import prod from typing import TYPE_CHECKING +import regex as re import torch import torch.nn.functional as F -from torch import nn import vllm.envs as envs from vllm import _custom_ops as ops @@ -40,7 +39,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( per_tensor_dequantize, ) -from vllm.model_executor.models.utils import PPMissingLayer from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv @@ -52,11 +50,23 @@ logger = init_logger(__name__) +# Matches any layer-qualified module path (e.g. +# ``model.layers.5.mlp.shared_expert`` or +# ``model.language_model.model.layers.5.mlp.shared_expert``) so skip reasons +# are reported once per cause instead of once per layer. +_LAYER_MODULE_PATH_RE = re.compile(r"(?:[\w]+\.)*(?:layers|mtp)\.\d+\.[\w.]*") + + +def _normalize_skip_reason(reason: str) -> str: + return _LAYER_MODULE_PATH_RE.sub("", reason) + + def resolve_layer_fused_shared_expert( quant_config: "QuantizationConfig | None", prefix: str, shared_expert_name: str = "shared_experts", moe_backend: str | None = None, + disable_under_ep: bool = False, ) -> bool: """Resolve whether AITER fused shared-expert execution is enabled. @@ -65,6 +75,9 @@ def resolve_layer_fused_shared_expert( prefix: MoE module prefix. shared_expert_name: Shared-expert module name under ``prefix``. moe_backend: MoE backend that runs the fused shared expert. + disable_under_ep: Whether this model family does not support FSE under + expert parallelism. Kept as an explicit call-site policy instead of + branching on model type here. Returns: Whether AITER fused shared experts are enabled. @@ -76,16 +89,31 @@ def resolve_layer_fused_shared_expert( # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported # that returns None if AITER is not available. fse_requested = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) - fse_compatible, fse_reason = ( - is_shared_expert_quant_fse_compatible( + + # TODO: should it be disabled for all models? + skip_reason: str | None = None + if fse_requested and disable_under_ep: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + parallel_config = ( + vllm_config.parallel_config if vllm_config is not None else None + ) + if parallel_config is not None and parallel_config.enable_expert_parallel: + fse_requested = False + skip_reason = "expert parallelism is enabled" + + if fse_requested: + fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible( quant_config, f"{prefix}.experts", f"{prefix}.{shared_expert_name}", moe_backend=moe_backend, ) - if fse_requested - else (True, None) - ) + if not fse_compatible: + skip_reason = fse_reason or "unknown reason" + else: + fse_compatible = True is_fused_shared_expert_enabled = fse_requested and fse_compatible # In online shared_expert quantization is used, register it to @@ -121,52 +149,18 @@ def resolve_layer_fused_shared_expert( quant_key_str, target_pattern, ) - if fse_requested and not is_fused_shared_expert_enabled: - logger.warning( + if skip_reason is not None: + # Keep the full per-layer reason at debug level; the warning below is + # deduplicated by cause so a long model logs it once. + logger.debug("Skipping shared-expert fusion at %s: %s", prefix, skip_reason) + logger.warning_once( "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " - "cannot be enabled - skipping for this layer: %s.", - fse_reason, + "cannot be enabled - skipping shared-expert fusion: %s.", + _normalize_skip_reason(skip_reason), ) return is_fused_shared_expert_enabled -def is_model_fused_shared_expert_compatible( - layers: nn.ModuleList | Iterable[nn.Module], - moe_cls: type[nn.Module], - moe_name: str, -) -> bool: - """Resolve one fused-shared-expert state for a model's MoE layers.""" - - def get_moe_layer(layer: nn.Module) -> nn.Module | None: - for name in moe_name.split("."): - layer = getattr(layer, name, None) - if layer is None: - return None - return layer - - moe_layers = ( - moe_layer - for layer in layers - if not isinstance(layer, PPMissingLayer) - and (moe_layer := get_moe_layer(layer)) is not None - and isinstance(moe_layer, moe_cls) - ) - - enabled = [ - getattr(layer, "is_fused_shared_expert_enabled", False) for layer in moe_layers - ] - enabled_count = sum(enabled) - disabled_count = len(enabled) - enabled_count - if enabled_count > 0 and disabled_count > 0: - raise NotImplementedError( - "Fused shared experts must be enabled for all MoE layers; found " - f"{enabled_count} enabled and {disabled_count} disabled layers. " - "Per-layer fused shared experts is not yet supported. Please open " - "an issue." - ) - return enabled_count > 0 and disabled_count == 0 - - @triton.jit def _count_expert_num_tokens( topk_ids_ptr, diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 6a12fc88b255..cf0828e9951d 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -49,7 +49,6 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -84,6 +83,8 @@ from .utils import ( AutoWeightsLoader, PPMissingLayer, + extract_layer_index, + get_fused_shared_expert_layers, get_spec_layer_idx_from_weight_name, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -742,10 +743,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, AXK1MoE, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -839,10 +840,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts + ( (self.config.n_shared_experts or 0) - if self.is_fused_shared_expert_enabled + if self.fused_shared_expert_layers else 0 ), num_redundant_experts=self.num_redundant_experts, @@ -858,8 +860,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is not None: continue # skip spec decode layers for main model + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/alice_ai.py b/vllm/model_executor/models/alice_ai.py index a0c5f53ecf92..9dfdab545916 100644 --- a/vllm/model_executor/models/alice_ai.py +++ b/vllm/model_executor/models/alice_ai.py @@ -20,7 +20,6 @@ ) from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import GemmaRMSNorm @@ -57,6 +56,7 @@ AutoWeightsLoader, WeightsMapper, extract_layer_index, + get_fused_shared_expert_layers, make_layers, maybe_fuse_shared_experts, maybe_prefix, @@ -500,8 +500,10 @@ def get_layer(prefix: str) -> AliceAIDecoderBlock: layer_fn=get_layer, prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - layers=self.layers, moe_cls=AliceAISparseMoeBlock, moe_name="mlp" + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, + AliceAISparseMoeBlock, + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.norm = GemmaRMSNorm( hidden_size=config.hidden_size, eps=config.rms_norm_eps @@ -567,7 +569,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights=weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=self.config.num_experts, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 4f6125a2279b..c870660935bb 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -7,15 +7,13 @@ import torch.nn as nn from transformers import PreTrainedConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import tensor_model_parallel_all_gather from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -41,6 +39,8 @@ ) from .interfaces import SupportsPP from .utils import ( + extract_layer_index, + get_fused_shared_expert_layers, get_pp_missing_layer_names, get_spec_layer_idx_from_weight_name, make_empty_intermediate_tensors_factory, @@ -242,6 +242,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config self.quant_config = vllm_config.quant_config + self.vllm_config = vllm_config self.model = DeepSeekMultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) @@ -267,10 +268,10 @@ def set_moe_parameters(self): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV2MoE, - "mtp_block.mlp", + enable_eplb=self.vllm_config.parallel_config.enable_eplb, ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -316,17 +317,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ] stacked_params_mapping.extend(indexer_fused_mapping) + fused_shared_expert_layers = self.fused_shared_expert_layers expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if fused_shared_expert_layers else 0), ) pp_missing_layer_names = get_pp_missing_layer_names(self) @@ -353,8 +352,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight) loaded_params.add(name) continue + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers ) name = self._rewrite_spec_layer_name(spec_layer, name) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index a1a919fae644..c3befbc1a3a5 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -59,7 +59,6 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm @@ -97,6 +96,7 @@ from vllm.model_executor.models.utils import ( AutoWeightsLoader, extract_layer_index, + get_fused_shared_expert_layers, sequence_parallel_chunk, ) from vllm.platforms import current_platform @@ -1476,10 +1476,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV2MoE, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -1629,12 +1629,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0), num_redundant_experts=self.num_redundant_experts, ) @@ -1659,8 +1656,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ): continue # this layer has no indexer; drop its checkpoint weights + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers ) if _try_load_fp8_indexer_wk( diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 4e541a49413e..c43f12e41fad 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -46,7 +46,6 @@ GateLinear, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -69,6 +68,7 @@ AutoWeightsLoader, PPMissingLayer, WeightsMapper, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -444,10 +444,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Glm4MoE, - "mlp", + enable_eplb=enable_eplb, ) if get_pp_group().is_last_rank: @@ -495,7 +495,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: skip_spec_layers(weights, self.config), n_routed_experts=self.config.n_routed_experts, n_shared_experts=self.config.n_shared_experts or 1, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, ) loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/glm4_moe_lite.py b/vllm/model_executor/models/glm4_moe_lite.py index 019380540d08..367fbbbcd96a 100644 --- a/vllm/model_executor/models/glm4_moe_lite.py +++ b/vllm/model_executor/models/glm4_moe_lite.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: from transformers.models.glm4_moe_lite import Glm4MoeLiteConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import ( @@ -43,9 +44,6 @@ from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -76,6 +74,8 @@ from .utils import ( AutoWeightsLoader, PPMissingLayer, + extract_layer_index, + get_fused_shared_expert_layers, get_spec_layer_idx_from_weight_name, is_pp_missing_parameter, make_layers, @@ -276,10 +276,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Glm4MoeLite, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -364,12 +364,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0), ) params_dict = dict(self.named_parameters()) @@ -382,8 +379,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is not None: continue # skip spec decode layers for main model + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/glm4_moe_lite_mtp.py b/vllm/model_executor/models/glm4_moe_lite_mtp.py index f5a82b72314b..a59eeb8f60fb 100644 --- a/vllm/model_executor/models/glm4_moe_lite_mtp.py +++ b/vllm/model_executor/models/glm4_moe_lite_mtp.py @@ -30,14 +30,12 @@ import torch.nn as nn from transformers import PreTrainedConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.model_executor.layers.fused_moe import ( MoERunner, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -62,7 +60,12 @@ Glm4MoeLiteDecoderLayer, ) from .interfaces import SupportsPP -from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix +from .utils import ( + extract_layer_index, + get_fused_shared_expert_layers, + get_spec_layer_idx_from_weight_name, + maybe_prefix, +) class SharedHead(nn.Module): @@ -241,10 +244,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Glm4MoeLite, - "mtp_block.mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -279,17 +282,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), ] + fused_shared_expert_layers = self.fused_shared_expert_layers expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if fused_shared_expert_layers else 0), ) params_dict = dict(self.named_parameters()) @@ -300,8 +301,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is None: continue + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers ) name = self._rewrite_spec_layer_name(spec_layer, name) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/glm4_moe_mtp.py b/vllm/model_executor/models/glm4_moe_mtp.py index 330f891d1cf3..0464322cae02 100644 --- a/vllm/model_executor/models/glm4_moe_mtp.py +++ b/vllm/model_executor/models/glm4_moe_mtp.py @@ -31,14 +31,12 @@ import torch.nn as nn from transformers import PreTrainedConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import CacheConfig, ParallelConfig, VllmConfig from vllm.model_executor.layers.fused_moe import ( MoERunner, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -54,7 +52,12 @@ Glm4MoE, Glm4MoeDecoderLayer, ) -from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix +from .utils import ( + extract_layer_index, + get_fused_shared_expert_layers, + get_spec_layer_idx_from_weight_name, + maybe_prefix, +) class SharedHead(nn.Module): @@ -215,10 +218,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.moe_mlp_layers.append(layer.mlp) self.moe_layers.append(layer.mlp.experts) self.extract_moe_parameters(example_moe) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Glm4MoE, - "mtp_block.mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -257,8 +260,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) + fused_shared_expert_layers = self.fused_shared_expert_layers + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts = self.config.n_routed_experts - if self.is_fused_shared_expert_enabled and self.config.n_shared_experts: + if fused_shared_expert_layers and self.config.n_shared_experts: num_experts += self.config.n_shared_experts expert_params_mapping = fused_moe_make_expert_params_mapping( self, @@ -282,8 +287,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue name = self._rewrite_spec_layer_name(spec_layer, name) + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers ) for param_name, weight_name, shard_id in stacked_params_mapping: diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index f105948f9283..6e370f180da5 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -42,9 +42,6 @@ get_pp_group, ) from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3_5RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -99,6 +96,7 @@ WeightsMapper, _merge_multimodal_embeddings, extract_layer_index, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -263,10 +261,10 @@ def get_layer(prefix: str): self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen3NextSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size @@ -286,7 +284,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if "moe" in self.config.model_type: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=self.config.num_experts, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 2acac6b5b7c5..aa7e6c06e338 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -12,9 +12,6 @@ from vllm.config import VllmConfig from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -44,6 +41,7 @@ AutoWeightsLoader, PPMissingLayer, _merge_multimodal_embeddings, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, maybe_fuse_shared_experts, maybe_prefix, @@ -124,10 +122,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): for idx in range(self.num_mtp_layers) ) vllm_config.quant_config = original_quant - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen3NextSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size @@ -193,7 +191,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0), n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index ee5c05d0a4ae..f6df90594a69 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -24,7 +24,6 @@ GateLinear, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate @@ -75,6 +74,7 @@ PPMissingLayer, WeightsMapper, extract_layer_index, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -672,10 +672,10 @@ def get_layer(prefix: str): self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen3NextSparseMoeBlock, - "mlp", + enable_eplb=parallel_config.enable_eplb, ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size @@ -756,7 +756,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0), n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index 4ae8f21840a1..47d038b8e287 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -12,9 +12,6 @@ from vllm.config import VllmConfig from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -33,6 +30,7 @@ from .utils import ( AutoWeightsLoader, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, maybe_fuse_shared_experts, maybe_prefix, @@ -95,10 +93,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): for idx in range(self.num_mtp_layers) ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen3NextSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size @@ -164,7 +162,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=self.config.num_experts, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index a86496b48a9b..3386a76a0ea6 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -12,7 +12,7 @@ import torch.nn as nn from torch.nn.modules.module import register_module_module_registration_hook -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -502,13 +502,88 @@ def _check_skipped_aliases(self, autoloaded_weights: set[str]) -> None: ) +def get_fused_shared_expert_layers( + module: nn.Module, + moe_cls: type[nn.Module], + enable_eplb: bool = False, +) -> frozenset[int]: + """Resolve fused-shared-expert layer indices from the live module tree. + + The index is taken from each MoE module's qualified name, so it follows the + same index space as the caller's weight names: 0-based for modules stored in + an ``nn.ModuleList``, and absolute for ``nn.ModuleDict`` MTP layers keyed by + layer index. No loader needs to hard-code an offset. + + Also rejects mixed fusion when EPLB is enabled or the ``nixl_ep`` all2all + backend is selected, and logs the resolved policy. + """ + fused_layers: set[int] = set() + moe_layers: set[int] = set() + for name, mod in module.named_modules(): + if not isinstance(mod, moe_cls): + continue + idx = extract_layer_index(name) + moe_layers.add(idx) + if getattr(mod, "is_fused_shared_expert_enabled", False): + fused_layers.add(idx) + + # EPLB bookkeeping is per rank/model: `EplbState` builds each model's + # `expert_weights`/`num_moe_layers` from the local `MixtureOfExperts` (one + # pipeline stage here), while main and draft models have separate + # `EplbModelState`s keyed by model config. The uniformity requirement is + # therefore stage/model-local by design: mixed within a stage is rejected, + # mixed across PP stages (or main vs draft) is allowed. + if 0 < len(fused_layers) < len(moe_layers): + if enable_eplb: + raise NotImplementedError( + "Mixed per-layer shared-expert fusion within a pipeline " + "stage/model is not supported with EPLB" + ) + vllm_config = get_current_vllm_config_or_none() + parallel_config = ( + vllm_config.parallel_config if vllm_config is not None else None + ) + if parallel_config is not None and parallel_config.enable_expert_parallel: + raise NotImplementedError( + "Mixed per-layer shared-expert fusion with expert parallelism is " + "not tested. Please open an issue if you need this combination." + ) + # These backends size one shared buffer from the first MoE layer's + # geometry (all2all.py NixlEPAll2AllManager._init_buffer/get_handle). + all2all_backend = ( + parallel_config.all2all_backend if parallel_config is not None else None + ) + if all2all_backend == "nixl_ep": + raise NotImplementedError( + "Mixed per-layer shared-expert fusion is not supported with " + f"all2all_backend={all2all_backend!r}" + ) + + if fused_layers: + logger.info( + "AITER fused shared experts: enabled on MoE layer(s) %s; " + "disabled on MoE layer(s) %s.", + sorted(fused_layers), + sorted(moe_layers - fused_layers) or "none", + ) + elif moe_layers: + from vllm._aiter_ops import rocm_aiter_ops + + if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled(): + logger.info( + "AITER fused shared experts: disabled on all %d MoE layer(s).", + len(moe_layers), + ) + return frozenset(fused_layers) + + def maybe_fuse_shared_experts( weights: Iterable[tuple[str, torch.Tensor]], *, n_routed_experts: int, n_shared_experts: int, + fused_shared_expert_layers: frozenset[int], ckpt_prefix: str = "mlp.shared_experts", - enabled: bool | None = None, ) -> Iterable[tuple[str, torch.Tensor]]: """Route AITER fused-shared-expert checkpoint weights into fused slots. @@ -524,21 +599,16 @@ def maybe_fuse_shared_experts( n_routed_experts: Number of routed experts; offsets the fused slots. n_shared_experts: Number of shared experts packed into the tensor. ckpt_prefix: Checkpoint module name of the shared experts. - enabled: Whether AITER fused-shared-experts is active. Defaults to - `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`; pass an - explicit value only when the model gates on something more (e.g. - quant-spec compatibility) and it must match its construction-time - decision. + fused_shared_expert_layers: Layer indices whose shared expert was fused + at construction time, as returned by + `get_fused_shared_expert_layers`. Only weights belonging to these + layers are routed to fused slots. Yields: `(name, tensor)` pairs with shared experts routed to fused slots. """ - if enabled is None: - from vllm._aiter_ops import rocm_aiter_ops - - enabled = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() - if not enabled: + if not fused_shared_expert_layers: yield from weights return @@ -549,6 +619,9 @@ def maybe_fuse_shared_experts( if prefix not in name: yield name, loaded_weight continue + if extract_layer_index(name) not in fused_shared_expert_layers: + yield name, loaded_weight + continue # gate/up split on the output dim; down_proj on the input dim. split_dim = 1 if ("down_proj.weight" in name and loaded_weight.ndim > 1) else 0 total = loaded_weight.shape[split_dim] diff --git a/vllm/models/deepseek_v32/amd/model.py b/vllm/models/deepseek_v32/amd/model.py index 225e9b0ffe10..2338b0d030ef 100644 --- a/vllm/models/deepseek_v32/amd/model.py +++ b/vllm/models/deepseek_v32/amd/model.py @@ -7,14 +7,12 @@ import torch +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_rank from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, @@ -32,6 +30,8 @@ ) from vllm.model_executor.models.utils import ( PPMissingLayer, + extract_layer_index, + get_fused_shared_expert_layers, get_pp_missing_layer_names, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -172,10 +172,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV2MoE, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -246,12 +246,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if self.fused_shared_expert_layers else 0), num_redundant_experts=self.num_redundant_experts, ) @@ -269,8 +266,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue if get_spec_layer_idx_from_weight_name(self.config, name) is not None: continue + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers ) if _try_load_fp8_indexer_wk( name, diff --git a/vllm/models/deepseek_v32/amd/mtp.py b/vllm/models/deepseek_v32/amd/mtp.py index d5066b39de8d..b773035e9199 100644 --- a/vllm/models/deepseek_v32/amd/mtp.py +++ b/vllm/models/deepseek_v32/amd/mtp.py @@ -8,14 +8,12 @@ import torch import torch.nn as nn +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.distributed import tensor_model_parallel_all_reduce from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -34,6 +32,8 @@ ) from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( + extract_layer_index, + get_fused_shared_expert_layers, get_pp_missing_layer_names, make_empty_intermediate_tensors_factory, maybe_prefix, @@ -183,10 +183,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], self.config.hidden_size ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV2MoE, - "mtp_block.mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) def set_moe_parameters(self): @@ -266,17 +266,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), ] + fused_shared_expert_layers = self.fused_shared_expert_layers expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts - + ( - self.config.n_shared_experts - if self.is_fused_shared_expert_enabled - else 0 - ), + + (self.config.n_shared_experts if fused_shared_expert_layers else 0), ) pp_missing_layer_names = get_pp_missing_layer_names(self) @@ -303,8 +301,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, loaded_weight) loaded_params.add(name) continue + # Only parse the layer index when FSE is enabled. is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + "mlp.shared_experts" in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers ) name = self._rewrite_spec_layer_name(spec_layer, name) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index b0d798033f11..b60189b0b57d 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -13,7 +13,7 @@ import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops -from vllm.config import ParallelConfig, VllmConfig +from vllm.config import VllmConfig from vllm.distributed import ( get_pp_group, get_tensor_model_parallel_rank, @@ -35,7 +35,6 @@ fused_topk_bias, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.layernorm import RMSNorm @@ -54,9 +53,6 @@ ) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod -from vllm.model_executor.layers.quantization.utils.config_utils import ( - is_shared_expert_quant_fse_compatible, -) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -72,6 +68,7 @@ PPMissingLayer, WeightsMapper, extract_layer_index, + get_fused_shared_expert_layers, is_pp_missing_parameter, make_layers, maybe_prefix, @@ -525,31 +522,6 @@ def forward_modular( return routed + shared_out -def _fuse_shared_experts_enabled(config, parallel_config: ParallelConfig) -> bool: - if ( - getattr(config, "n_shared_experts", None) - and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - and not parallel_config.enable_expert_parallel - and parallel_config.data_parallel_size > 1 - ): - # Fused shared experts are not supported under data parallelism: - # the fused path cannot load the FP8 shared expert into the MXFP4 - # routed slot, which fails obscurely during weight loading. Fail - # fast with an actionable message instead. - raise ValueError( - "DeepSeek-V4 fused shared experts " - "(VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1) are not supported " - "with data parallelism (data_parallel_size > 1). Set " - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 to run the shared " - "expert as an unfused FP8 MLP." - ) - return bool( - getattr(config, "n_shared_experts", None) - and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - and not parallel_config.enable_expert_parallel - ) - - def _validate_aiter_mega_moe_config(vllm_config: VllmConfig) -> None: parallel_config = vllm_config.parallel_config if not rocm_aiter_ops.is_fused_moe_enabled(): @@ -658,28 +630,31 @@ def __init__( self.n_shared_experts = config.n_shared_experts - # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` - # is checked to enable FSE for DeepSeek-v4, despite AITER not being used. - # This should be cleaned up and use `resolve_layer_fused_shared_expert`. self.fuse_heterogeneous_shared_expert = fuse_heterogeneous_shared_expert - fse_requested = ( - _fuse_shared_experts_enabled(config, vllm_config.parallel_config) + self.is_fused_shared_expert_enabled = False + if ( + getattr(config, "n_shared_experts", None) and not self.fuse_heterogeneous_shared_expert - ) - fse_compatible = False - if fse_requested: - fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible( - quant_config, - f"{prefix}.experts", - f"{prefix}.shared_experts", - ) - if not fse_compatible: - logger.warning( - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " - "cannot be enabled: %s.", - fse_reason, + ): + # Fused shared experts are not supported under data parallelism: + # the fused path cannot load the FP8 shared expert into the MXFP4 + # routed slot. Fail fast with an actionable message. + if ( + not vllm_config.parallel_config.enable_expert_parallel + and vllm_config.parallel_config.data_parallel_size > 1 + and rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + ): + raise ValueError( + "DeepSeek-V4 fused shared experts " + "(VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1) are not " + "supported with data parallelism " + "(data_parallel_size > 1). Set " + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 to run the " + "shared expert as an unfused FP8 MLP." ) - self.is_fused_shared_expert_enabled = fse_requested and fse_compatible + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix, disable_under_ep=True + ) self.use_mega_moe = vllm_config.kernel_config.moe_backend == "aiter_mega_moe" fuse_mega_moe_shared_expert = ( @@ -1160,10 +1135,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, DeepseekV4MoE, - "ffn", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -1344,16 +1319,12 @@ def _resolve_param_name(name: str) -> str: # at init) as the single source of truth, so the redirect below cannot # diverge from how the module was built if per-layer quantization ever # mixes fused and non-fused layers. - fuse_by_layer = { - extract_layer_index(mod_name): mod.is_fused_shared_expert_enabled - for mod_name, mod in self.named_modules() - if isinstance(mod, DeepseekV4MoE) - } + fuse_by_layer = self.fused_shared_expert_layers n_routed = self.config.n_routed_experts # The redirect below maps the single shared-expert tensor group to one # appended slot; multiple shared experts would need per-expert slicing # (see deepseek_v2.py). DeepSeek-V4 has n_shared_experts == 1. - if any(fuse_by_layer.values()) and self.config.n_shared_experts != 1: + if fuse_by_layer and self.config.n_shared_experts != 1: raise NotImplementedError( "deepseek-v4 fused shared-expert loading supports only " f"n_shared_experts == 1, got {self.config.n_shared_experts}" @@ -1364,13 +1335,7 @@ def _resolve_param_name(name: str) -> str: # into appended routed-expert slot ``.ffn.experts.{n_routed}`` # so the MXFP4-quantized shared expert loads through the routed # expert loader (grouped GEMM). Single shared expert only. - if ".ffn.shared_experts.w" in name and fuse_by_layer.get( - extract_layer_index(name), False - ): - name = name.replace( - ".ffn.shared_experts.w", - f".ffn.experts.{n_routed}.w", - ) + name = _redirect_shared_expert_name(name, fuse_by_layer, n_routed) for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). @@ -1448,7 +1413,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # redirected shared-expert weights route through the expert loader. n_shared = getattr(self.config, "n_shared_experts", 0) or 0 num_experts = self.config.n_routed_experts + ( - n_shared if self.is_fused_shared_expert_enabled else 0 + n_shared if self.fused_shared_expert_layers else 0 ) return fused_moe_make_expert_params_mapping( self, @@ -1460,6 +1425,24 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: ) +def _redirect_shared_expert_name( + name: str, fuse_by_layer: frozenset[int], n_routed_experts: int +) -> str: + """Route a DSV4 shared-expert checkpoint name to its final parameter name. + + Fused layers go to the appended routed-expert slot. When any layer is fused + the weight mapper leaves the checkpoint ``.w2`` names untouched, so + non-fused layers still need the native ``down_proj`` parameter name. + """ + if ".ffn.shared_experts.w" not in name: + return name + if extract_layer_index(name) in fuse_by_layer: + return name.replace( + ".ffn.shared_experts.w", f".ffn.experts.{n_routed_experts}.w" + ) + return name.replace(".ffn.shared_experts.w2", ".ffn.shared_experts.down_proj") + + def _make_deepseek_v4_weights_mapper( expert_dtype: str, fuse_shared_experts: bool = False ) -> WeightsMapper: @@ -1533,10 +1516,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - if expert_dtype != "fp4" or self.model.is_fused_shared_expert_enabled: + if expert_dtype != "fp4" or self.model.fused_shared_expert_layers: self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper( expert_dtype, - fuse_shared_experts=self.model.is_fused_shared_expert_enabled, + fuse_shared_experts=bool(self.model.fused_shared_expert_layers), ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index d49769c843d9..1144bbde9604 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -302,6 +302,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = DeepSeekV4MultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + if any( + getattr( + layer.mtp_block.ffn, + "is_fused_shared_expert_enabled", + False, + ) + for layer in self.model.layers.values() + ): + raise NotImplementedError( + "DeepSeek-V4 MTP fused shared experts are not supported yet; " + "please open an issue." + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py index 63ec337fa48b..2e6f76c09e07 100644 --- a/vllm/models/deepseek_v4/cpu/model.py +++ b/vllm/models/deepseek_v4/cpu/model.py @@ -21,9 +21,6 @@ GateLinear, fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -493,11 +490,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, - DeepseekV4MoE, - "ffn", - ) + self.is_fused_shared_expert_enabled = False if get_pp_group().is_last_rank: self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index b9094c697f50..68c2f47ec3ef 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -1344,6 +1344,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + if any( + getattr( + getattr(layer, "ffn", None), + "is_fused_shared_expert_enabled", + False, + ) + for layer in self.model.layers + ): + raise NotImplementedError( + "DeepSeek-V4.1 AMD fused shared experts are not supported yet; " + "please open an issue." + ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( config.vocab_size, diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index 6b0cbe4ad61b..e944ad38f151 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -8,6 +8,7 @@ from torch import nn from transformers import Glm5NextTextConfig +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import ParallelConfig, VllmConfig from vllm.distributed import ( get_ep_group, @@ -24,7 +25,6 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm @@ -76,6 +76,8 @@ from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + extract_layer_index, + get_fused_shared_expert_layers, init_vllm_registered_model, is_pp_missing_parameter, make_layers, @@ -749,8 +751,10 @@ def get_layer(prefix: str): # The active slice is fixed after construction; cache it so forward # doesn't rebuild the slice (a fresh list) every step. self._active_layers = self.layers[self.start_layer : self.end_layer] - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, Glm5NextMoE, "mlp" + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, + Glm5NextMoE, + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -894,13 +898,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: 0, ) num_fused_shared = _num_fused_shared_experts( - self.config.n_shared_experts, self.is_fused_shared_expert_enabled + self.config.n_shared_experts, + bool(self.fused_shared_expert_layers), ) expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=num_redundant_experts, ) @@ -930,7 +936,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue - if self.is_fused_shared_expert_enabled: + # Only parse the layer index when FSE is enabled. + if ( + "mlp.shared_experts." in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers + ): name = _fused_shared_expert_name(name, self.config.n_routed_experts) # Handle FP8 indexer WK: dequantize to BF16 for fusion with diff --git a/vllm/models/glm5next/common/mtp.py b/vllm/models/glm5next/common/mtp.py index 6e6950ca700f..deb3c61a1dfb 100644 --- a/vllm/models/glm5next/common/mtp.py +++ b/vllm/models/glm5next/common/mtp.py @@ -6,13 +6,11 @@ import torch import torch.nn as nn +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -24,7 +22,11 @@ ) 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.model_executor.models.utils import ( + extract_layer_index, + get_fused_shared_expert_layers, + 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 @@ -212,8 +214,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) self.set_moe_parameters() - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.model.layers.values(), Glm5NextMoE, "mtp_block.mlp" + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, + Glm5NextMoE, + enable_eplb=vllm_config.parallel_config.enable_eplb, ) def set_moe_parameters(self): @@ -296,14 +300,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("wk_weights_proj", "wk", 0), ("wk_weights_proj", "weights_proj", 1), ] + fused_shared_expert_layers = self.fused_shared_expert_layers num_fused_shared = _num_fused_shared_experts( - self.config.n_shared_experts, self.is_fused_shared_expert_enabled + self.config.n_shared_experts, + bool(fused_shared_expert_layers), ) expert_params_mapping = fused_moe_make_expert_params_mapping( self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts=self.config.n_routed_experts + num_fused_shared, num_redundant_experts=self.num_redundant_experts, ) @@ -329,7 +336,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if spec_layer is None: continue name = self._rewrite_spec_layer_name(spec_layer, name) - if self.is_fused_shared_expert_enabled: + # Only parse the layer index when FSE is enabled. + if ( + "mlp.shared_experts." in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in fused_shared_expert_layers + ): name = _fused_shared_expert_name(name, self.config.n_routed_experts) if _try_load_fp8_indexer_wk( diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 7b7459770d87..424b24e51a87 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -51,7 +51,7 @@ fused_moe_make_expert_params_mapping, ) from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, + resolve_layer_fused_shared_expert, ) from vllm.model_executor.layers.fusion.quant_activation import ( QuantizedActivation, @@ -65,9 +65,6 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.quantization.utils.config_utils import ( - is_shared_expert_quant_fse_compatible, -) from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -89,6 +86,8 @@ AutoWeightsLoader, PPMissingLayer, WeightsMapper, + extract_layer_index, + get_fused_shared_expert_layers, init_vllm_registered_model, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -405,32 +404,11 @@ def __init__( # MoE call as the last expert slot, so we don't build a separate module. # On gfx950 with aiter MoE the append is fused inside aiter's grouped # top-k kernel; otherwise it goes through the vLLM top-k bias router. - # It is disabled under expert parallelism (the shared slot is appended to - # the routed top-k, which the EP expert-mapping path does not handle). - # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1` - # is checked to enable FSE for MiniMax-M3, despite AITER not being used. - # This should be cleaned up and use `resolve_layer_fused_shared_expert`. - fse_requested = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - self.is_fused_shared_expert_enabled = False - if ( - fse_requested - and bool(getattr(config, "n_shared_experts", None)) - and not get_current_vllm_config().parallel_config.enable_expert_parallel - ): - fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible( - quant_config, - f"{prefix}.experts", - f"{prefix}.shared_experts", + if bool(getattr(config, "n_shared_experts", None)): + self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( + quant_config, prefix, disable_under_ep=True ) - if not fse_compatible: - logger.warning( - "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " - "cannot be enabled: %s.", - fse_reason, - ) - else: - self.is_fused_shared_expert_enabled = True # When additionally on gfx950 with an active aiter MoE backend, the shared # expert is appended inside aiter's an active aiter MoE backend, the shared @@ -1538,10 +1516,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, MiniMaxM3MoE, - "block_sparse_moe", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) if get_pp_group().is_last_rank: @@ -1622,8 +1600,9 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Checkpoint experts use w1=gate, w2=down, w3=up. When fusing the shared # expert, include the appended slot (id == num_local_experts). n_shared = getattr(self.config, "n_shared_experts", 0) or 0 + # NOTE: global despite per-layer FSE; non-fused layers skip the extras. num_experts = self.config.num_local_experts + ( - n_shared if self.is_fused_shared_expert_enabled else 0 + n_shared if self.fused_shared_expert_layers else 0 ) return fused_moe_make_expert_params_mapping( self, @@ -1672,7 +1651,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # down->w2) so it loads via the routed expert loader. Runs before the # stacked/dense mappings so shared_experts.gate_proj/up_proj are not # captured by the dense gate_up_proj mapping. - if self.is_fused_shared_expert_enabled and ".shared_experts." in name: + # Only parse the layer index when FSE is enabled. + if ( + ".shared_experts." in name + and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()) + and extract_layer_index(name) in self.fused_shared_expert_layers + ): sid = self.config.num_local_experts name = name.replace(".shared_experts.gate_proj.", f".experts.{sid}.w1.") name = name.replace(".shared_experts.up_proj.", f".experts.{sid}.w3.") diff --git a/vllm/models/qwen4_exp/amd/model.py b/vllm/models/qwen4_exp/amd/model.py index 5715d038dd1a..675b9a9a6cff 100644 --- a/vllm/models/qwen4_exp/amd/model.py +++ b/vllm/models/qwen4_exp/amd/model.py @@ -12,9 +12,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import get_pp_group -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( @@ -64,6 +61,7 @@ WeightsMapper, _merge_multimodal_embeddings, extract_layer_index, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -446,10 +444,10 @@ def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen4ExpSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) intermediate_size = config.hidden_size * config.hc_count self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( @@ -620,7 +618,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0) or 0, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py index 0805d102b0c9..7a73ac8f33ed 100644 --- a/vllm/models/qwen4_exp/amd/mtp.py +++ b/vllm/models/qwen4_exp/amd/mtp.py @@ -23,9 +23,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig, replace, set_current_vllm_config from vllm.distributed import get_pp_group -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import GemmaRMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -41,6 +38,7 @@ PPMissingLayer, WeightsMapper, get_draft_quant_config, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, maybe_fuse_shared_experts, maybe_prefix, @@ -229,10 +227,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) for idx in range(self.num_mtp_layers) ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen4ExpSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.pre_fc_norm_embedding = GemmaRMSNorm( @@ -373,7 +371,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0) or 0, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index ed9aaa927733..08ad38361e14 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -13,9 +13,6 @@ from vllm.config import VllmConfig from vllm.distributed import get_pp_group from vllm.forward_context import get_forward_context, is_forward_context_available -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( @@ -65,6 +62,7 @@ WeightsMapper, _merge_multimodal_embeddings, extract_layer_index, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, make_layers, maybe_fuse_shared_experts, @@ -446,10 +444,10 @@ def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen4ExpSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) intermediate_size = config.hidden_size * config.hc_count self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( @@ -656,7 +654,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0) or 0, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index fb15d076467a..30401ca56dcb 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -24,9 +24,6 @@ from vllm.config import VllmConfig, replace, set_current_vllm_config from vllm.distributed import get_pp_group from vllm.forward_context import get_forward_context, is_forward_context_available -from vllm.model_executor.layers.fused_moe.utils import ( - is_model_fused_shared_expert_compatible, -) from vllm.model_executor.layers.layernorm import GemmaRMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -42,6 +39,7 @@ PPMissingLayer, WeightsMapper, get_draft_quant_config, + get_fused_shared_expert_layers, make_empty_intermediate_tensors_factory, maybe_fuse_shared_experts, maybe_prefix, @@ -228,10 +226,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) for idx in range(self.num_mtp_layers) ) - self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( - self.layers, + self.fused_shared_expert_layers = get_fused_shared_expert_layers( + self, Qwen4ExpSparseMoeBlock, - "mlp", + enable_eplb=vllm_config.parallel_config.enable_eplb, ) self.pre_fc_norm_embedding = GemmaRMSNorm( @@ -384,7 +382,7 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = maybe_fuse_shared_experts( weights, - enabled=self.is_fused_shared_expert_enabled, + fused_shared_expert_layers=self.fused_shared_expert_layers, n_routed_experts=getattr(self.config, "num_experts", 0) or 0, n_shared_experts=1, ckpt_prefix="mlp.shared_expert", diff --git a/vllm/triton_utils/dispatcher.py b/vllm/triton_utils/dispatcher.py index c7e2f9e46986..7d1c2e79aa0e 100644 --- a/vllm/triton_utils/dispatcher.py +++ b/vllm/triton_utils/dispatcher.py @@ -198,8 +198,22 @@ def lookup(value: Any) -> Callable[..., Any] | None: # kernel as an instance attribute; read the instance dict directly # so module-level __getattr__ hooks are never triggered, and so # this scan needs no dependency on the warmup machinery. - if hasattr(value, "__dict__"): - replacement = lookup(vars(value).get("kernel")) + # + # `hasattr(value, "__dict__")` is not a sufficient guard: proxy + # objects such as llvmlite's `binding.ffi.lib` synthesise every + # attribute through `__getattr__`, so `__dict__` resolves to a + # `_lib_fn_wrapper` and `.get()` would raise AttributeError. Only + # treat a real dict as an instance namespace. + try: + instance_dict = vars(value) + except TypeError: + # Slots-only objects have no instance dict. + instance_dict = None + replacement = ( + lookup(instance_dict.get("kernel")) + if isinstance(instance_dict, dict) + else None + ) if replacement is None: continue value.kernel = replacement From 4da2a82a24b3e5917901ca77254272d86c185f3d Mon Sep 17 00:00:00 2001 From: Ao Tang Date: Wed, 28 Oct 2026 13:14:43 -0400 Subject: [PATCH 070/406] [Bugfix][Attention] FlashInfer sparse MLA: launch at most 65535 tokens per trtllm-gen call (#63791) Signed-off-by: Ao Tang Co-authored-by: Claude Opus 5.5 --- .../test_flashinfer_sparse_mla_workspace.py | 17 ++++ .../v1/attention/test_sparse_mla_backends.py | 97 ++++++++++++++++++- .../backends/mla/flashinfer_mla_sparse.py | 97 ++++++++----------- 3 files changed, 152 insertions(+), 59 deletions(-) diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py index 5b63c2c609f9..ea1b1a232787 100644 --- a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py +++ b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py @@ -62,6 +62,23 @@ def test_non_dcp_size_is_unchanged(): assert computed == _DEFAULT_WORKSPACE_BUFFER_SIZE +def test_dcp_size_is_capped_at_one_launch(): + # Launches cover at most 65535 tokens, so larger token budgets reuse the slab. + at_limit = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=0, + dcp_world_size=2, + num_heads_per_rank=8, + max_num_batched_tokens=65535, + ) + above_limit = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=0, + dcp_world_size=2, + num_heads_per_rank=8, + max_num_batched_tokens=131072, + ) + assert above_limit == at_limit + + def test_combined_pcp_dcp_rejects_flashinfer_sparse(): config = SimpleNamespace( parallel_config=SimpleNamespace( diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index e15dc7ab5402..8a55a6771b32 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -287,6 +287,94 @@ def fake_flashinfer(**kwargs): assert recorded_scale != kv_lora_rank**-0.5 +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="trtllm-gen sparse MLA requires SM10x", +) +@pytest.mark.parametrize( + ("qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"), + [(128, 64, 128), (256, 0, 256)], +) +def test_flashinfer_sparse_mla_attends_tokens_past_grid_limit( + default_vllm_config, dist_init, qk_nope_head_dim, qk_rope_head_dim, v_head_dim +): + """trtllm-gen maps each token to gridDim.z (limit 65535); a larger step, + e.g. a 64K-token prefill, must still attend every token.""" + device = torch.device("cuda") + num_tokens, num_heads, kv_lora_rank, topk, block_size = ( + 65536 + 4099, + 16, + 512, + 128, + 64, + ) + head_size = kv_lora_rank + qk_rope_head_dim + impl = FlashInferMLASparseImpl( + num_heads=num_heads, + head_size=head_size, + scale=head_size**-0.5, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim, + v_head_dim=v_head_dim, + kv_b_proj=ColumnParallelLinear( + kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim), bias=False + ), + ) + gen = torch.Generator(device=device).manual_seed(0) + kv_cache = ( + torch.rand(64, block_size, head_size, generator=gen, device=device) - 0.5 + ).to(torch.bfloat16) + q = ( + torch.rand(num_tokens, num_heads, head_size, generator=gen, device=device) - 0.5 + ).to(torch.bfloat16) + # Mixed active lengths, valid indices first and -1 padding after, as the + # index converters produce. + valid_counts = torch.randint( + 1, topk + 1, (num_tokens,), generator=gen, device=device, dtype=torch.int32 + ) + valid_counts[::3] = topk + topk_indices = torch.randint( + 0, + kv_cache.shape[0] * block_size, + (num_tokens, topk), + generator=gen, + device=device, + dtype=torch.int32, + ) + slots = torch.arange(topk, device=device) + topk_indices.masked_fill_(slots[None, :] >= valid_counts[:, None], -1) + + out, lse = impl._forward_mqa_kernel( + q, + kv_cache, + topk_indices, + valid_counts, + layer=SimpleNamespace(), + block_size=block_size, + is_decode=False, + ) + + assert lse is None + kv_rows = kv_cache.view(-1, head_size).float() + # Both sides of the 65535 boundary, the ends, and a sweep in between. + rows = sorted({0, 65534, 65535, 65536, num_tokens - 1, *range(0, num_tokens, 997)}) + for t in rows: + kv = kv_rows[topk_indices[t, : valid_counts[t]].long()] + scores = q[t].float() @ kv.T * impl.scale + ref = torch.softmax(scores, dim=-1) @ kv[:, :kv_lora_rank] + torch.testing.assert_close(out[t].float(), ref, atol=2e-2, rtol=2e-2) + + def test_hisparse_routes_prefill_to_sparse_mqa(): layer = SimpleNamespace(hisparse_cache=object()) @@ -4322,10 +4410,11 @@ def fake_fa4(**kwargs): def fake_trtllm(**kwargs): trtllm_calls.append(kwargs) rows = kwargs["query"].shape[0] - out = torch.full( - (rows, 1, num_heads, 512), 2.0, dtype=torch.bfloat16, device=device - ) - return (out, lse(rows)) if kwargs.get("return_lse") else out + # Like FlashInfer, write into the caller-owned ``out`` / ``lse``. + out = kwargs["out"].fill_(2.0) + if not kwargs.get("return_lse"): + return out + return out, kwargs["lse"].copy_(lse(rows)) monkeypatch.setattr(flashattn_sparse_module, "flash_attn_varlen_func", fake_fa4) monkeypatch.setattr( diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 1f8338dd1de3..3231a9249c24 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -42,6 +42,9 @@ logger = init_logger(__name__) +# CUDA's gridDim.z limit; trtllm-gen MLA decode launches one z-slice per request. +_TRTLLM_GEN_MAX_REQUESTS = 65535 + class _FlashInferMLASparseBackendBase(AttentionBackend): """Common metadata for concrete FlashInfer sparse MLA backends.""" @@ -366,14 +369,16 @@ def compute_trtllm_sparse_mla_workspace_bytes( 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. + returned unchanged. A launch covers at most ``_TRTLLM_GEN_MAX_REQUESTS`` + tokens, and launches run in order on one stream, so larger batches reuse + the same slab. """ 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 + * min(max_num_batched_tokens, _TRTLLM_GEN_MAX_REQUESTS) * _TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN + _TRTLLM_GEN_SOFTMAX_GUARD_BYTES ) @@ -612,7 +617,6 @@ def _forward_mqa_kernel( # deferred until MTP is validated end-to-end for this backend. query = q.unsqueeze(1) block_tables = topk_indices.unsqueeze(1) - seq_lens_arg = seq_lens # page_table width = topk buffer width, which kpool widens past # index_topk (topk_tokens) and rounds up to a multiple of 128. The @@ -623,45 +627,52 @@ def _forward_mqa_kernel( # mismatches the page_table when index_kpool > 1. sparse_topk_capacity = topk_indices.shape[1] - extra_kwargs: dict[str, torch.Tensor] = {} needs_empty_query_guard = self.need_to_return_lse_for_decode or isinstance( self.index_group, HiSparseMLAIndexGroup ) + topk_lens = None if self.is_nope_mla: # Resident TP queries have nonempty selections. Preserve empty-query # handling for DCP's local selections and host-backed HiSparse. topk_lens = seq_lens if needs_empty_query_guard: topk_lens = prepare_sparse_mla_safe_lengths(topk_indices, seq_lens) - extra_kwargs["sparse_mla_top_k_lens"] = topk_lens - - kernel_out = trtllm_batch_decode_with_kv_cache_mla( - query=query, - kv_cache=kv_cache.unsqueeze(1), - workspace_buffer=self._workspace_buffer, - qk_nope_head_dim=self.qk_nope_head_dim, - kv_lora_rank=self.kv_lora_rank, - qk_rope_head_dim=self.qk_rope_head_dim, - block_tables=block_tables, - seq_lens=seq_lens_arg, - max_seq_len=sparse_topk_capacity, - bmm1_scale=self.bmm1_scale, - bmm2_scale=self.bmm2_scale, - sparse_mla_top_k=sparse_topk_capacity, - return_lse=self.need_to_return_lse_for_decode, - **extra_kwargs, + + num_tokens, num_heads = q.shape[0], q.shape[1] + out = q.new_empty( + (num_tokens, num_heads, self.kv_lora_rank), dtype=torch.bfloat16 ) + lse = None if self.need_to_return_lse_for_decode: - assert isinstance(kernel_out, tuple) - o, lse = kernel_out - else: - assert isinstance(kernel_out, torch.Tensor) - o = kernel_out - lse = None - - out = o.view(-1, o.shape[-2], o.shape[-1]) - if lse is not None: - lse = self._normalize_lse(lse, out.shape[0], out.shape[1]) + lse = q.new_empty((num_tokens, num_heads), dtype=torch.float32) + + # One trtllm-gen request per token, and requests map to gridDim.z: a + # launch over more than 65535 tokens fails without raising and leaves + # ``out`` unwritten (flashinfer-ai/flashinfer#3125). Launch row chunks. + for start in range(0, num_tokens, _TRTLLM_GEN_MAX_REQUESTS): + end = min(start + _TRTLLM_GEN_MAX_REQUESTS, num_tokens) + extra_kwargs: dict[str, torch.Tensor] = {} + if topk_lens is not None: + extra_kwargs["sparse_mla_top_k_lens"] = topk_lens[start:end] + if lse is not None: + extra_kwargs["lse"] = lse[start:end] + trtllm_batch_decode_with_kv_cache_mla( + query=query[start:end], + kv_cache=kv_cache.unsqueeze(1), + workspace_buffer=self._workspace_buffer, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=block_tables[start:end], + seq_lens=seq_lens[start:end], + max_seq_len=sparse_topk_capacity, + out=out[start:end].unsqueeze(1), + bmm1_scale=self.bmm1_scale, + bmm2_scale=self.bmm2_scale, + sparse_mla_top_k=sparse_topk_capacity, + return_lse=lse is not None, + **extra_kwargs, + ) if self.is_nope_mla and needs_empty_query_guard: empty_queries = seq_lens == 0 if lse is not None: @@ -674,27 +685,3 @@ def _forward_mqa_kernel( out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) return out, lse - - @staticmethod - def _normalize_lse( - lse: torch.Tensor, - num_tokens: int, - num_heads: int, - ) -> torch.Tensor: - # FlashInfer returns the decode LSE either as 2D (num_tokens, num_heads) - # or 3D ((num_tokens, num_heads, 1) / (num_tokens, 1, num_heads)). - # Collapse all of these to the (num_tokens, num_heads) the shared DCP - # reducer expects. - if lse.dim() == 3: - if lse.shape[-1] == 1: - lse = lse.squeeze(-1) - elif lse.shape[1] == 1: - lse = lse.squeeze(1) - elif lse.shape[0] * lse.shape[1] == num_tokens: - lse = lse.reshape(num_tokens, lse.shape[-1]) - if lse.shape != (num_tokens, num_heads): - raise RuntimeError( - "Unexpected FlashInfer sparse MLA LSE shape: " - f"{tuple(lse.shape)}, expected ({num_tokens}, {num_heads})." - ) - return lse From 8321d1119bc2a2048f5575421db3507de8b5f3d5 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:26:28 +0000 Subject: [PATCH 071/406] [CI] Exclude nvidia-cuda-cccl from the ROCm test lockfile (#64302) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .pre-commit-config.yaml | 1 + requirements/test/rocm.txt | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c940db4ae844..34d5fcd6a1c3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -74,6 +74,7 @@ repos: --no-emit-package, cupy-cuda12x, # nvidia packages (unsuffixed / unified naming) --no-emit-package, nvidia-cublas, + --no-emit-package, nvidia-cuda-cccl, --no-emit-package, nvidia-cuda-cupti, --no-emit-package, nvidia-cuda-nvrtc, --no-emit-package, nvidia-cuda-runtime, diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index ba87f1c3a336..e3af4cbc3946 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13 +# uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cccl --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13 absl-py==2.4.0 # via rouge-score accelerate==1.13.0 @@ -1446,12 +1446,12 @@ zipp==3.23.0 # cuda-pathfinder # cuda-toolkit # nvidia-cublas +# nvidia-cuda-cccl # nvidia-cuda-cupti # nvidia-cuda-nvrtc # nvidia-cuda-runtime # nvidia-cufft # nvidia-cufile -# nvidia-curand # nvidia-cusolver # nvidia-cusparse # nvidia-nvjitlink From f034f9ed1ff69b3ec95988b66ceb28b4639aa1f0 Mon Sep 17 00:00:00 2001 From: bofengl Date: Wed, 28 Oct 2026 10:52:06 -0700 Subject: [PATCH 072/406] [Perf][Kernel][DSv4.1] Set the Q pad-copy threshold from a 64-16384 token sweep A 64-16384-token sweep of the fused pad vs the KV-only launch + copy on B300 (the description's bench script with a finer grid, plus an eager column that fully exposes the copy's extra launch) puts the GPU-time crossover at ~256 tokens and the eager, gap-exposed breakeven at ~2048. Set _Q_PAD_COPY_MIN_TOKENS to 2048 and decouple it from the op's NUM_TOKEN_CUTOFF: the copy path now wins even when nothing hides its second launch. The q_head_padded=0 call never reaches the op's cutoff branch, so no other behavior changes. Co-Authored-By: Claude Fable 5 Signed-off-by: bofengl --- vllm/models/deepseek_v41/attention.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 43351fab460f..615c45213074 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -97,9 +97,12 @@ _SHARD_FUSED_WQA_WKV_MIN_TOKENS = 2048 # From this many tokens, eager steps build the padded interleaved Q with -# _pad_interleaved_q instead of in the fused KV-insert op. Kept equal to the -# op's NUM_TOKEN_CUTOFF, where it switches to its reduced-grid kernel. -_Q_PAD_COPY_MIN_TOKENS = 1024 +# _pad_interleaved_q instead of in the fused KV-insert op. Set from a +# 64-16384-token sweep on B300: by GPU time alone the copy wins from about +# 256 tokens, but it costs one extra eager launch, and with that launch gap +# fully exposed the crossover moves to about 2048. This value keeps the +# large-step win without betting on launch overlap. +_Q_PAD_COPY_MIN_TOKENS = 2048 @triton.jit From c105af429241fe07853114e2f6420197c21554b3 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Wed, 28 Oct 2026 11:31:42 -0700 Subject: [PATCH 073/406] [Perf][DSA] Rank the prefill indexer logits with DeepSelect in BF16 (#62240) Signed-off-by: Juntian Liu Co-authored-by: Claude Opus 5.5 --- .../layers/sparse_attn_indexer.py | 57 ++++++++++++------- 1 file changed, 37 insertions(+), 20 deletions(-) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 1c89724c2f60..eeedc9edb51f 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -33,7 +33,9 @@ ) from vllm.model_executor.layers.indexer_topk import ( RADIX_TOPK_WORKSPACE_SIZE, + deep_select_topk, get_indexer_topk, + is_deep_select_supported, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, @@ -720,6 +722,7 @@ def sparse_attn_indexer( # DeepGEMM returns compressed logits: row i starts at K position # cu_seqlen_ks[i]. None means every row starts at column zero. row_starts: torch.Tensor | None = None + use_deep_select = False if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) @@ -754,17 +757,24 @@ def sparse_attn_indexer( q_slice_cast, q_scale_slice, weights[row_start:row_end] ) row_ends = cu_seqlen_ke - cu_seqlen_ks - logits = as_fp32_logits( - fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - chunk_weights, - cu_seqlen_ks, - cu_seqlen_ke, - chunk.max_seqlen_k, - ), - row_ends, + logits = fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks, + cu_seqlen_ke, + chunk.max_seqlen_k, ) + # DeepSelect ranks the bf16 logits in place of their exact fp32 + # copy; fp32 logits keep the per-row top-k. + use_deep_select = ( + topk_backend == "deep_select" + and logits.dtype == torch.bfloat16 + and dcp_world_size == 1 + and is_deep_select_supported(logits, topk_tokens) + ) + if not use_deep_select: + logits = as_fp32_logits(logits, row_ends) num_rows = logits.shape[0] if candidate_blocks is not None: # Two-level selection (v4.1): the candidate source @@ -790,16 +800,23 @@ def sparse_attn_indexer( chunk_candidates, candidate_block_size, ) - ops.top_k_per_row_prefill( - logits, - torch.zeros_like(row_ends) if row_starts is None else row_starts, - row_ends, - topk_indices, - num_rows, - logits.stride(0), - logits.stride(1), - topk_tokens, - ) + if use_deep_select: + deep_select_topk( + logits, topk_tokens, end=row_ends, output_idx=topk_indices + ) + else: + ops.top_k_per_row_prefill( + logits, + torch.zeros_like(row_ends) + if row_starts is None + else row_starts, + row_ends, + topk_indices, + num_rows, + logits.stride(0), + logits.stride(1), + topk_tokens, + ) if deinterleave_idx is None: # Under the PCP path the top-k already ran over the whole From 0fddee1f5869465a6235d2076f82f1ea70318064 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:34:16 +0000 Subject: [PATCH 074/406] [CI] Bump GitHub Actions versions and replace node20 pre-commit/action (#64328) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .github/workflows/buf.yml | 4 ++-- .github/workflows/issue_autolabel.yml | 6 +++--- .github/workflows/macos-smoke-test.yml | 4 ++-- .github/workflows/notify-ci-authorized.yml | 4 ++-- .github/workflows/pre-commit.yml | 12 ++++++++---- .github/workflows/proto-crate.yml | 10 +++++----- .github/workflows/recipes-preview.yml | 2 +- .github/workflows/run-ci-command.yml | 4 ++-- .github/workflows/stale.yml | 2 +- .pre-commit-config.yaml | 2 +- 10 files changed, 27 insertions(+), 23 deletions(-) diff --git a/.github/workflows/buf.yml b/.github/workflows/buf.yml index dd695d28d529..15c7ce0e1179 100644 --- a/.github/workflows/buf.yml +++ b/.github/workflows/buf.yml @@ -17,8 +17,8 @@ jobs: proto: runs-on: ubuntu-latest steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - - uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: bufbuild/buf-action@85aebf73123b5c15fd5528aaecbf9129cddf7fa7 # v1.6.0 with: version: 1.72.0 input: rust/proto diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 43a458df7678..03eeb71a5bda 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -15,7 +15,7 @@ jobs: steps: - name: Label issues based on keywords id: label-step - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | // Configuration: Add new labels and keywords here @@ -633,7 +633,7 @@ jobs: - name: CC users for labeled issues if: steps.label-step.outputs.labels_added != '[]' - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | // Configuration: Map labels to GitHub users to CC @@ -714,7 +714,7 @@ jobs: - name: Request missing ROCm info from issue author if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug') - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | const body = (context.payload.issue.body || '').toLowerCase(); diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml index 2e63c5676dd8..8de27d848e8b 100644 --- a/.github/workflows/macos-smoke-test.yml +++ b/.github/workflows/macos-smoke-test.yml @@ -28,9 +28,9 @@ jobs: timeout-minutes: 30 steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3 with: enable-cache: true cache-dependency-glob: | diff --git a/.github/workflows/notify-ci-authorized.yml b/.github/workflows/notify-ci-authorized.yml index c260cc67dba0..1396032d51cb 100644 --- a/.github/workflows/notify-ci-authorized.yml +++ b/.github/workflows/notify-ci-authorized.yml @@ -34,11 +34,11 @@ jobs: github.event.workflow_run.conclusion == 'success')) runs-on: [self-hosted, linux, x64, vllm-runners] steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ github.event.repository.default_branch }} persist-credentials: false - - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3 with: python-version: "3.12" - name: Notify PR author that CI is available diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 861d4a5e498b..1410a8944828 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -49,14 +49,18 @@ jobs: if: github.repository_owner == 'vllm-project' && always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped') runs-on: [self-hosted, linux, x64, vllm-runners] steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3 with: python-version: "3.12" - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json" - run: echo "::add-matcher::.github/workflows/matchers/markdownlint.json" - run: echo "::add-matcher::.github/workflows/matchers/mypy.json" - run: echo "::add-matcher::.github/workflows/matchers/shellcheck.json" - - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 + # Replaces pre-commit/action, whose latest release (v3.0.1) uses the node20 actions/cache@v4 + - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: - extra_args: --all-files --hook-stage manual + path: ~/.cache/pre-commit + key: pre-commit-uv|${{ runner.os }}|${{ env.UV_PYTHON }}|${{ hashFiles('.pre-commit-config.yaml') }} + - run: >- + uvx pre-commit run --show-diff-on-failure --color=always --all-files --hook-stage manual diff --git a/.github/workflows/proto-crate.yml b/.github/workflows/proto-crate.yml index 43be5c8c9214..ad20db39fa66 100644 --- a/.github/workflows/proto-crate.yml +++ b/.github/workflows/proto-crate.yml @@ -24,8 +24,8 @@ jobs: semver: runs-on: ubuntu-latest steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - - uses: obi1kenobi/cargo-semver-checks-action@6b69fcf40e9b5fb17adeb57e4b6ecd020649a239 # v2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: obi1kenobi/cargo-semver-checks-action@6b69fcf40e9b5fb17adeb57e4b6ecd020649a239 # v2.9 with: manifest-path: rust/proto/Cargo.toml package: vllm-proto @@ -33,7 +33,7 @@ jobs: check: runs-on: ubuntu-latest steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - run: rustup show active-toolchain - name: Check release version if: github.event_name == 'push' @@ -56,13 +56,13 @@ jobs: contents: read id-token: write steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ github.sha }} - run: rustup show active-toolchain - name: Verify package before authentication run: cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run - - uses: rust-lang/crates-io-auth-action@c6f97d42243bad5fab37ca0427f495c86d5b1a18 # v1.0.5 + - uses: rust-lang/crates-io-auth-action@3268eb413133e2960fa4ae0e963ca76380c54cd9 # v1.0.6 id: auth - name: Publish env: diff --git a/.github/workflows/recipes-preview.yml b/.github/workflows/recipes-preview.yml index 42127e99b79e..efaca4681a67 100644 --- a/.github/workflows/recipes-preview.yml +++ b/.github/workflows/recipes-preview.yml @@ -46,7 +46,7 @@ jobs: steps: - name: Create a token for vllm-project/recipes id: token - uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0 + uses: actions/create-github-app-token@64ab729a6294000f5aa4a4bc0b73d1f32bc45510 # v3.3.0 with: app-id: ${{ vars.RECIPES_APP_ID }} private-key: ${{ secrets.RECIPES_APP_PRIVATE_KEY }} diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml index d3336772717f..877c2d91b94b 100644 --- a/.github/workflows/run-ci-command.yml +++ b/.github/workflows/run-ci-command.yml @@ -22,11 +22,11 @@ jobs: startsWith(github.event.comment.body, '/amd-ci')) runs-on: [self-hosted, linux, x64, vllm-runners] steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ github.event.repository.default_branch }} persist-credentials: false - - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3 with: python-version: "3.12" - name: Authorize and run CI command diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index ba807fab7c3a..d2c5533d31f6 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -15,7 +15,7 @@ jobs: actions: write runs-on: ubuntu-latest steps: - - uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10.3.0 + - uses: actions/stale@4391f3da665fdf50b6810c1a66712fb9ba21aa93 # v11.0.0 with: # Increasing this value ensures that changes to this workflow # propagate to all issues and PRs in days rather than months diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 34d5fcd6a1c3..a2709f09142f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -28,7 +28,7 @@ repos: rev: v0.21.0 hooks: - id: markdownlint-cli2 - language_version: lts + language_version: 24.22.1 # Node 26 (current lts) needs libatomic, which the CI runners lack args: [--config, tools/pre_commit/.markdownlint.yaml, --fix] - repo: https://github.com/rhysd/actionlint rev: v1.7.7 From 366beb221557d9b43cb2975a380c7b9962ea20f9 Mon Sep 17 00:00:00 2001 From: Xin Li <119016172+xinli-sw@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:16:24 -0400 Subject: [PATCH 075/406] [Doc] MiniMax-M3.1-Flash: add the GB300 AgentX presets (#64284) Signed-off-by: Xin Li Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> --- .../models/MiniMaxAI/MiniMax-M3.1-Flash.yaml | 268 ++++++++++++++++++ 1 file changed, 268 insertions(+) diff --git a/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml b/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml index 25ea0663eff1..466d22b808e7 100644 --- a/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml +++ b/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml @@ -378,6 +378,274 @@ strategy_overrides: VLLM_KV_CACHE_LAYOUT: "BLHNC" presets: + # GB300 AgentX under Dynamo: aggregated TP4 workers, and DEP2 prefill + # workers feeding one DEP decode worker. Each preset restates its engine + # flags, so it does not depend on the pd_cluster GB300 block above. + gb300_agentx_agg: + label: "AgentX TP4 (Dynamo)" + hardware: gb300 + strategy: single_node_tp + variants: [default] + workload: "Agentic (InferenceX AgentX)" + description: "Aggregated TP4 vLLM workers, one per 4-GPU GB300 tray, behind the Dynamo KV router with session affinity. Seven-token DSpark draft on FLASH_ATTN and prefill CUDA graphs up to 8192 tokens. KV stays in HBM, or spills to host DRAM through SimpleCPUOffloadConnector." + source: { label: "InferenceX #4468", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4468", key: "minimaxm31flash-fp4-gb300-dynamo-vllm-agentic-dspark-agg", revision: 41f7dfe5e5f4a6ad1acb56ca5f2f47412fbdcbb7 } + image: + vllm: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134" + dynamo: "1.6.0.dev20261023" + pip: ["ai-dynamo-runtime==1.6.0.dev20261023", "ai-dynamo==1.6.0.dev20261023"] + tp: 4 + select: + features: { text_only: true } + args: + - "--enable-prefix-caching" + - "--max-model-len" + - "1048576" + - "--max-num-batched-tokens" + - "32768" + env: + # Pins the prefix-block hashes, so they match across processes and restarts. + PYTHONHASHSEED: "0" + # A cold start (JIT builds, graph capture, autotune) takes more than 10 minutes. + VLLM_ENGINE_READY_TIMEOUT_S: "7200" + VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600" + VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT: "600" + json_args: + "--compilation-config": { cudagraph_capture_sizes: [8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 320, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960, 1024, 1536, 2048, 3072, 4096, 6144, 8192] } + "--speculative-config": { attention_backend: "FLASH_ATTN" } + # --max-num-seqs is about twice the sessions a worker serves. + points: + gpu_kv: + label: "GPU-resident KV" + description: "KV cache in HBM only: --max-num-seqs 32 for about 16 sessions per worker, 64 for up to 32." + default: true + args: ["--max-num-seqs", "32", "--gpu-memory-utilization", "0.9"] + select: { kv_offload: "off" } + dram_kv: + label: "Host-DRAM KV offload" + description: "Adds a host-DRAM KV tier (860 GB per worker) for about 40-56 sessions per worker." + args: ["--max-num-seqs", "96", "--gpu-memory-utilization", "0.86"] + env: { VLLM_USE_SIMPLE_KV_OFFLOAD: "1" } + select: { kv_offload: simple } + json_args: + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 860000000000, lazy_offload: true } } + runbook: | + The GB300 aggregated AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb300-dynamo-vllm-agentic-dspark-agg`, [SemiAnalysisAI/InferenceX#4468](https://github.com/SemiAnalysisAI/InferenceX/pull/4468)) serves the default checkpoint text-only, with one TP4 `dynamo.vllm` worker per 4-GPU tray behind the Dynamo frontend. Add trays to scale out: the KV router spreads conversations across them. The lane measured 29.2k tok/s/GPU at 416 tok/s/user with 16 sessions per tray (GPU-resident KV, `--max-num-seqs 32`), and 95.0k tok/s/GPU at 180 tok/s/user with 56 sessions per tray (host-DRAM KV offload, `--max-num-seqs 96`). + + **Image and versions.** The lane used an arm64 CUDA 13 build of the `vllm/vllm-openai:minimax-m31-flash` image (tag `minimax-m31-flash-arm64-cu134`); any build of that image family for GB300 will do. Dynamo `1.6.0.dev20261023` was installed on top: `pip install --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime==1.6.0.dev20261023 ai-dynamo==1.6.0.dev20261023`. + + **Tuning rules.** + + - `--max-num-seqs` is about twice the sessions one worker serves: 32 up to 16 sessions, 64 at 32-48, 96 at 56. + - **GPU-resident KV** (`--gpu-memory-utilization 0.9`) is enough up to about 32 sessions per tray. Beyond that, multi-turn contexts outgrow HBM and evicted prefixes are recomputed. + - **Host-DRAM KV offload** adds `SimpleCPUOffloadConnector` with lazy offload (`VLLM_USE_SIMPLE_KV_OFFLOAD=1`). `cpu_bytes_to_use` is the budget for the whole worker, split across its four ranks: 860 GB pins 215 GB per rank. Size it from the host DRAM that is free after the engine processes are up. Lower `--gpu-memory-utilization` to 0.86 to leave room for the transfer buffers. + - The CUDA graph sizes go up to 8192 tokens, so prefill chunks of up to 8192 tokens also replay graphs. `--max-num-batched-tokens 32768` still allows larger mixed steps, which run piecewise. + + **Running under Dynamo.** Start etcd and NATS (32 MB max payload) once: + + ```bash + etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://:2379 + printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf + ``` + + On each tray, export `ETCD_ENDPOINTS=http://:2379` and `NATS_SERVER=nats://:4222` together with the generated environment and `DYN_TCP_CONNECT_TIMEOUT=120`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT=1024`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT=1024` and `DYN_BACKEND_ADMISSION_REQUEST_QUEUE_LIMIT=2048`. Then run the generated flags with `python3 -m dynamo.vllm --model MiniMaxAI/MiniMax-M3.1-Flash --served-model-name MiniMaxAI/MiniMax-M3.1-Flash --trust-remote-code` in place of `vllm serve MiniMaxAI/MiniMax-M3.1-Flash`. `dynamo.vllm` rejects vLLM's API-server flags, so leave out `--tool-call-parser minimax_m3 --enable-auto-tool-choice` and add: + + ```bash + --dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3 \ + --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' + ``` + + Start the frontend: + + ```bash + export ETCD_ENDPOINTS=http://:2379 NATS_SERVER=nats://:4222 + export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TCP_CONNECT_TIMEOUT=120 + python3 -m dynamo.frontend --http-port 8000 --router-mode kv --router-session-affinity-ttl-secs 14400 + ``` + + The KV router uses the workers' KV events to send each request to the worker that holds its prefix. Session affinity keeps a conversation on that worker for up to four hours; send a stable conversation ID in the `X-Dynamo-Session-ID` request header. + + The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04), so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler (`rejection_sample_method: block`), as in the command above. + + gb300_agentx_pd: + label: "AgentX PD (Dynamo)" + hardware: gb300 + strategy: pd_cluster + variants: [default] + workload: "Agentic (InferenceX AgentX)" + description: "DEP2 prefill workers, two per GB300 tray, feed one DEP8 decode worker (throughput) or DEP16 (lower latency) over NIXL. Both roles share a Mooncake KV store in host DRAM. The Dynamo KV router routes with session affinity. Seven-token DSpark draft on FLASHINFER_TS, DeepGEMM MegaMoE on both roles." + source: { label: "InferenceX #4468", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4468", key: "minimaxm31flash-fp4-gb300-dynamo-vllm-agentic-dspark-disagg", revision: 41f7dfe5e5f4a6ad1acb56ca5f2f47412fbdcbb7 } + image: + vllm: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134" + dynamo: "1.6.0.dev20261023" + pip: ["ai-dynamo-runtime==1.6.0.dev20261023", "ai-dynamo==1.6.0.dev20261023", "mooncake-transfer-engine-cuda13==0.3.13.post1"] + # Dynamo workers take engine flags only; the tool-call parser moves to + # --dyn-tool-call-parser (see the runbook). + select: + pd_router: dynamo + features: { text_only: true, tool_calling: false } + args: + - "--enable-expert-parallel" + - "--moe-backend" + - "deep_gemm_mega_moe" + - "--numa-bind" + - "--enable-prefix-caching" + - "--max-model-len" + - "1048576" + # The RDMA names are a GB300 tray's four NICs; use the host's own. + env: + PYTHONHASHSEED: "0" + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_ENGINE_READY_TIMEOUT_S: "7200" + VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600" + NCCL_NET_GDR_C2C: "1" + NCCL_IB_HCA: "mlx5_0,mlx5_1,mlx5_2,mlx5_3" + NVSHMEM_REMOTE_TRANSPORT: "none" + NVSHMEM_ENABLE_NIC_PE_MAPPING: "1" + NVSHMEM_HCA_PE_MAPPING: "mlx5_0:1:1,mlx5_1:1:1,mlx5_2:1:1,mlx5_3:1:1" + UCX_TLS: "rc,cuda_copy" + UCX_NET_DEVICES: "mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1" + UCX_CUDA_IPC_ENABLE_MNNVL: "n" + UCX_MEMTYPE_CACHE: "n" + UCX_MEMTYPE_REG_WHOLE: "n" + UCX_RNDV_PIPELINE_ERROR_HANDLING: "y" + NIXL_UCX_SGL_ENABLE: "1" + WITH_NVIDIA_PEERMEM: "0" + MC_SLICE_SIZE: "1048576" + MC_MAX_MR_SIZE: "4294967296" + MC_ENABLE_DEST_DEVICE_AFFINITY: "1" + DYN_TCP_CONNECT_TIMEOUT: "120" + DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT: "1024" + DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT: "1024" + DYN_BACKEND_ADMISSION_REQUEST_QUEUE_LIMIT: "2048" + roles: + prefill: + # One DEP2 worker on two GPUs; two per tray, as many trays as the load needs. + nodes: 1 + parallelism: dep + args: + - "--data-parallel-size" + - "2" + - "--data-parallel-size-local" + - "2" + - "--max-num-seqs" + - "64" + - "--max-num-batched-tokens" + - "32768" + - "--long-prefill-token-threshold" + - "4096" + - "--gpu-memory-utilization" + - "0.85" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE"}' + - "--max-cudagraph-capture-size" + - "2048" + - "--kv-events-config" + - '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' + - "--kv-transfer-config" + - '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"backends":["UCX"],"enforce_handshake_compat":false,"kv_lease_duration":120}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false}}]}}' + env: + VLLM_MOONCAKE_LOAD_RECV_THREADS: "8" + decode: + nodes: 2 + parallelism: dep + # 48 sequences per DP rank; decode graphs cover 48 x (1 + 7) tokens. + # Decode also writes the KV it generates into the Mooncake store. + remove_args: ["--no-enable-prefix-caching"] + args: + - "--max-num-seqs" + - "48" + - "--max-num-batched-tokens" + - "8192" + - "--watermark" + - "0.05" + - "--gpu-memory-utilization" + - "0.95" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY"}' + - "--max-cudagraph-capture-size" + - "384" + - "--numa-bind-nodes" + - "0" + - "0" + - "1" + - "1" + - "--kv-transfer-config" + - '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"backends":["UCX"],"enforce_handshake_compat":false,"kv_lease_duration":120}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false,"save_decode_cache":true}}]}}' + env: + VLLM_MOONCAKE_LOAD_RECV_THREADS: "4" + VLLM_MAX_ASYNC_KV_ADMITS_PER_STEP: "4" + VLLM_RANDOMIZE_DP_DUMMY_INPUTS: "1" + VLLM_ALLREDUCE_USE_SYMM_MEM: "0" + NCCL_CUMEM_ENABLE: "1" + NCCL_MNNVL_ENABLE: "1" + NCCL_NVLS_ENABLE: "1" + points: + dep8: + label: "DEP8 decode (throughput)" + description: "One DEP8 decode worker across two trays." + default: true + dep16: + label: "DEP16 decode (lower latency)" + description: "One DEP16 decode worker across four trays." + roles: { decode: { nodes: 4 } } + runbook: | + The GB300 disaggregated AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb300-dynamo-vllm-agentic-dspark-disagg`, [SemiAnalysisAI/InferenceX#4468](https://github.com/SemiAnalysisAI/InferenceX/pull/4468)) serves the default checkpoint text-only with Dynamo and vLLM. Six prefill workers, each DEP2 on two GPUs (two per 4-GPU tray, three trays), feed one decode worker that is DEP8 across two trays or DEP16 across four. KV moves from prefill to decode over NIXL (UCX over RDMA), and every rank contributes host DRAM to a shared Mooncake KV store, so a returning conversation's prefix is loaded from the store instead of recomputed. Both roles use expert parallelism with DeepGEMM MegaMoE, attention data parallelism, the V2 model runner and the seven-token DSpark draft on `FLASHINFER_TS`. + + | Prefill workers (DEP2) | Decode | GPUs | Sessions | tok/s/GPU | tok/s/user | TTFT p90 | + |---:|---|---:|---:|---:|---:|---:| + | 6 | DEP8 | 20 | 280 | 90.0k | 241 | 3.5 s | + | 6 | DEP8 | 20 | 320 | 97.1k | 236 | 4.3 s | + | 6 | DEP8 | 20 | 360 | 103.9k | 230 | 5.8 s | + | 6 | DEP16 | 28 | 300 | 74.0k | 299 | 2.2 s | + + The preset renders one prefill worker and the decode worker. Run more prefill workers with the same command. Two workers share a tray on GPUs 0-1 and 2-3 (`--device-ids 0,1` and `--device-ids 2,3` on `dynamo.vllm`). Give each its own `--data-parallel-rpc-port`, `VLLM_NIXL_SIDE_CHANNEL_PORT` and `--kv-events-config` endpoint port. + + **Image and versions.** The lane used an arm64 CUDA 13 build of the `vllm/vllm-openai:minimax-m31-flash` image (tag `minimax-m31-flash-arm64-cu134`), with Dynamo `1.6.0.dev20261023` and Mooncake `mooncake-transfer-engine-cuda13==0.3.13.post1` installed on every worker and on the Mooncake master host: `pip install --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime==1.6.0.dev20261023 ai-dynamo==1.6.0.dev20261023 mooncake-transfer-engine-cuda13==0.3.13.post1`. + + ### Tuning rules + + - **Prefill pool.** The prefill pool sets TTFT, so scale prefill workers with load. The lane's points kept TTFT p90 between 2.2 and 5.8 s with one DEP2 prefill worker per 45-60 concurrent agentic sessions, so 90-120 sessions per prefill tray (six workers at 280-360 sessions). Add prefill workers when TTFT p90 climbs while decode still has room. Two DEP2 workers per tray give the router twice as many prefill targets as one DEP4 worker. + - **DEP8 or DEP16 decode.** DEP8 gives the most tokens per GPU: up to 45 sessions per decode GPU (360 sessions) at 230 tok/s/user. DEP16 halves the batch per rank and gives about 25% more tok/s/user (299 at 300 sessions, against 241 at 280 on DEP8) at about 18% fewer tokens per GPU. `--max-num-seqs 48` per DP rank is a cap, not the operating point. The decode graph cap is `--max-num-seqs` x (1 + 7 draft tokens) = 384. `--watermark 0.05` keeps 5% of the decode KV cache free, so extra requests wait in the queue instead of preempting running ones. Keep decode `--max-num-batched-tokens` at 8192: a preempted request recomputes its context in chunks of that size. + - **Prefill batching.** `--long-prefill-token-threshold 4096` chunks long prompts into 4096-token pieces. Short prompts then share each step with them instead of waiting behind a 100K-token context, which lowers the TTFT tail. The recipe's 1P1D section above uses 8192. `--max-num-batched-tokens 32768` bounds a prefill step. Prefill captures CUDA graphs up to 2048 tokens. + - **Mooncake pool size.** `global_segment_size` is registered by every DP rank, prefill and decode alike, so a tray pins four times that value. The lane used 160 GB per rank (640 GB per tray). Size it as (host DRAM free per tray after the engines are up - headroom for NIXL and the 4 GB local buffers) / ranks per tray. The decode ranks also write the KV they generate into the store (`save_decode_cache: true`), so the next turn's prefill can load the previous answer instead of recomputing it. + - **Routing.** The KV router (`--router-mode kv`) reads the prefill workers' KV events and sends each request to the worker that already holds its prefix. Session affinity (`--router-session-affinity-ttl-secs 14400`) keeps a conversation on that worker for up to four hours. Send a stable conversation ID in the `X-Dynamo-Session-ID` header. The Mooncake store catches the prefixes that miss after a reroute. + - **Transport.** Cross-tray NIXL needs `UCX_TLS=rc,cuda_copy` with `UCX_NET_DEVICES` set to the tray's RDMA ports. Without them, UCX can fall back to TCP and transfers stall. Point `UCX_MODULE_DIR` at the UCX bundled with the NIXL wheel (`nixl_cu13.libs/ucx` in the image's `dist-packages`). `kv_lease_duration: 120` holds a prefill's blocks for up to 120 s until decode reads them (the default is 30 s). `VLLM_MAX_ASYNC_KV_ADMITS_PER_STEP=4` caps how many requests with finished KV loads decode admits per step. + - **NUMA.** `--numa-bind` pins each rank to its GPU's NUMA node. Decode passes the tray's map, `--numa-bind-nodes 0 0 1 1`, explicitly. + - `--moe-backend deep_gemm_mega_moe` is the same as the lane's `--kernel-config '{"moe_backend":"deep_gemm_mega_moe"}'`. `--prefill-schedule-interval 1`, also set by the lane, is vLLM's default. + + ### Services + + Start etcd, NATS and the Mooncake master once: + + ```bash + etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://:2379 + printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf + mooncake_master --rpc_port=50051 \ + --default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10 + ``` + + Every worker reads the store configuration from the JSON file that `MOONCAKE_CONFIG_PATH` points to (export it with the rest of the environment): + + ```json + {"master_server_address": ":50051", "metadata_server": "P2PHANDSHAKE", + "global_segment_size": "160GB", "local_buffer_size": "4GB", "protocol": "rdma", + "device_name": "mlx5_0,mlx5_1,mlx5_2,mlx5_3", "enable_offload": false} + ``` + + ### Workers + + Export `ETCD_ENDPOINTS=http://:2379`, `NATS_SERVER=nats://:4222` and the generated environment on every tray, then run the generated `python3 -m dynamo.vllm` commands. Add `--served-model-name MiniMaxAI/MiniMax-M3.1-Flash --trust-remote-code --dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3` to both roles. `dynamo.vllm` rejects vLLM's API-server flags `--tool-call-parser` and `--enable-auto-tool-choice`. Multi-node DEP decode runs one process per tray with `--data-parallel-start-rank` set to 0, 4, 8 and so on, and a shared `--data-parallel-address` and `--data-parallel-rpc-port`. `--data-parallel-hybrid-lb` lets every tray's Dynamo runtime register its own DP ranks with the frontend. + + ### Frontend + + ```bash + export ETCD_ENDPOINTS=http://:2379 NATS_SERVER=nats://:4222 + export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TCP_CONNECT_TIMEOUT=120 + python3 -m dynamo.frontend --http-port 8000 --router-mode kv --router-session-affinity-ttl-secs 14400 + ``` + + The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04), so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler (`rejection_sample_method: block`), as in the command above. + b300_agentx_tp4: label: "AgentX TP4" hardware: b300 From 3a7552056c22e9e8008c56204034df22e6544743 Mon Sep 17 00:00:00 2001 From: Xin Li <119016172+xinli-sw@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:17:12 -0400 Subject: [PATCH 076/406] [Doc] MiniMax-M3.1-Flash: add the GB200 AgentX presets (#64283) Signed-off-by: Xin Li --- .../models/MiniMaxAI/MiniMax-M3.1-Flash.yaml | 266 ++++++++++++++++++ 1 file changed, 266 insertions(+) diff --git a/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml b/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml index 466d22b808e7..d351b692e9f1 100644 --- a/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml +++ b/recipes/models/MiniMaxAI/MiniMax-M3.1-Flash.yaml @@ -756,6 +756,272 @@ presets: The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04) so every point is measured at the same acceptance rate. The command above keeps block verification. + # GB200 AgentX under Dynamo: aggregated TP4 workers, and DEP4 prefill trays + # feeding one DEP decode worker. Each preset restates its engine flags, so it + # does not depend on the pd_cluster GB200 block above. + gb200_agentx_agg: + label: "AgentX TP4 (Dynamo)" + hardware: gb200 + strategy: single_node_tp + variants: [default] + workload: "Agentic (InferenceX AgentX)" + description: "Aggregated TP4 vLLM workers, one per 4-GPU GB200 tray, behind the Dynamo KV router with session affinity. Seven-token DSpark draft on FLASH_ATTN and prefill CUDA graphs up to 8192 tokens. KV stays in HBM, or spills to host DRAM through SimpleCPUOffloadConnector." + source: { label: "InferenceX #4464", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4464", key: "minimaxm31flash-fp4-gb200-dynamo-vllm-agentic-dspark-agg", revision: 25845c731734bde9298eece5f9ed296e3b0745ca } + image: + vllm: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134" + dynamo: "1.6.0.dev20261023" + pip: ["ai-dynamo-runtime==1.6.0.dev20261023", "ai-dynamo==1.6.0.dev20261023"] + tp: 4 + select: + features: { text_only: true } + args: + - "--enable-prefix-caching" + - "--max-model-len" + - "1048576" + - "--max-num-batched-tokens" + - "32768" + env: + # Pins the prefix-block hashes, so they match across processes and restarts. + PYTHONHASHSEED: "0" + # A cold start (JIT builds, graph capture, autotune) takes more than 10 minutes. + VLLM_ENGINE_READY_TIMEOUT_S: "7200" + json_args: + "--compilation-config": { cudagraph_capture_sizes: [8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 320, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960, 1024, 1536, 2048, 3072, 4096, 6144, 8192] } + "--speculative-config": { attention_backend: "FLASH_ATTN" } + # --max-num-seqs is about twice the sessions a worker serves. + points: + gpu_kv: + label: "GPU-resident KV" + description: "KV cache in HBM only, for up to about 16 sessions per worker." + default: true + args: ["--max-num-seqs", "32", "--gpu-memory-utilization", "0.88"] + select: { kv_offload: "off" } + dram_kv: + label: "Host-DRAM KV offload" + description: "Adds a host-DRAM KV tier (360 GB per worker) for about 32-40 sessions per worker." + args: ["--max-num-seqs", "64", "--gpu-memory-utilization", "0.82", "--kv-cache-metrics"] + select: { kv_offload: simple } + json_args: + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: null, cpu_bytes_to_use: 360000000000, lazy_offload: true } } + runbook: | + The GB200 aggregated AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb200-dynamo-vllm-agentic-dspark-agg`, [SemiAnalysisAI/InferenceX#4464](https://github.com/SemiAnalysisAI/InferenceX/pull/4464)) serves the default checkpoint text-only, with one TP4 `dynamo.vllm` worker per 4-GPU tray behind the Dynamo frontend. Add trays to scale out: the KV router spreads conversations across them. At 16 sessions per tray (GPU-resident KV) the lane measured 29.0k tok/s/GPU at 402 tok/s/user. + + **Image and versions.** The lane used an arm64 CUDA 13 build of the `vllm/vllm-openai:minimax-m31-flash` image (tag `minimax-m31-flash-arm64-cu134`); any build of that image family for GB200 will do. Dynamo `1.6.0.dev20261023` was installed on top: `pip install --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime==1.6.0.dev20261023 ai-dynamo==1.6.0.dev20261023`. + + **Tuning rules.** + + - `--max-num-seqs` is about twice the sessions one worker serves: 32 up to 16 sessions, 64 at 32-40. + - **GPU-resident KV** (`--gpu-memory-utilization 0.88`) is enough up to about 16 sessions per tray. Beyond that, multi-turn contexts outgrow HBM and evicted prefixes are recomputed. + - **Host-DRAM KV offload** adds `SimpleCPUOffloadConnector` with lazy offload. `cpu_bytes_to_use` is the budget for the whole worker, split across its four ranks: 360 GB pins 90 GB per rank. Size it from the host DRAM that is free after the engine processes are up. Lower `--gpu-memory-utilization` to 0.82 to leave room for the transfer buffers. `--kv-cache-metrics` exports KV residency metrics, which help with that sizing; it can be left out. + - The CUDA graph sizes go up to 8192 tokens, so prefill chunks of up to 8192 tokens also replay graphs. `--max-num-batched-tokens 32768` still allows larger mixed steps, which run piecewise. + + **Running under Dynamo.** Start etcd and NATS (32 MB max payload) once: + + ```bash + etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://:2379 + printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf + ``` + + On each tray, export `ETCD_ENDPOINTS=http://:2379` and `NATS_SERVER=nats://:4222` together with the generated environment and `DYN_TCP_CONNECT_TIMEOUT=120`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT=1024`, `DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT=1024` and `DYN_BACKEND_ADMISSION_REQUEST_QUEUE_LIMIT=2048`. Then run the generated flags with `python3 -m dynamo.vllm --model MiniMaxAI/MiniMax-M3.1-Flash --served-model-name MiniMaxAI/MiniMax-M3.1-Flash --trust-remote-code` in place of `vllm serve MiniMaxAI/MiniMax-M3.1-Flash`. `dynamo.vllm` rejects vLLM's API-server flags, so leave out `--tool-call-parser minimax_m3 --enable-auto-tool-choice` and add: + + ```bash + --dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3 \ + --kv-events-config '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' + ``` + + Start the frontend: + + ```bash + export ETCD_ENDPOINTS=http://:2379 NATS_SERVER=nats://:4222 + export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TCP_CONNECT_TIMEOUT=120 + python3 -m dynamo.frontend --http-port 8000 --router-mode kv --router-session-affinity-ttl-secs 14400 + ``` + + The KV router uses the workers' KV events to send each request to the worker that holds its prefix. Session affinity keeps a conversation on that worker for up to four hours; send a stable conversation ID in the `X-Dynamo-Session-ID` request header. + + The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04), so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler (`rejection_sample_method: block`), as in the command above. + + gb200_agentx_pd: + label: "AgentX PD (Dynamo)" + hardware: gb200 + strategy: pd_cluster + variants: [default] + workload: "Agentic (InferenceX AgentX)" + description: "DEP4 prefill workers, one per GB200 tray, feed one DEP8 decode worker (throughput) or DEP16 (lower latency) over NIXL. Both roles share a Mooncake KV store in host DRAM. The Dynamo KV router routes with session affinity. Seven-token DSpark draft on FLASHINFER_TS, DeepGEMM MegaMoE on both roles." + source: { label: "InferenceX #4464", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4464", key: "minimaxm31flash-fp4-gb200-dynamo-vllm-agentic-dspark-disagg", revision: 25845c731734bde9298eece5f9ed296e3b0745ca } + image: + vllm: "vllm/vllm-openai:minimax-m31-flash-arm64-cu134" + dynamo: "1.6.0.dev20261023" + pip: ["ai-dynamo-runtime==1.6.0.dev20261023", "ai-dynamo==1.6.0.dev20261023", "mooncake-transfer-engine-cuda13==0.3.13.post1"] + # Dynamo workers take engine flags only; the tool-call parser moves to + # --dyn-tool-call-parser (see the runbook). + select: + pd_router: dynamo + features: { text_only: true, tool_calling: false } + args: + - "--enable-expert-parallel" + - "--moe-backend" + - "deep_gemm_mega_moe" + - "--numa-bind" + - "--enable-prefix-caching" + - "--max-model-len" + - "1048576" + - "--load-format" + - "fastsafetensors" + # The RDMA names are a GB200 tray's four NICs; use the host's own. + env: + PYTHONHASHSEED: "0" + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_ENGINE_READY_TIMEOUT_S: "7200" + VLLM_HTTP_TIMEOUT_KEEP_ALIVE: "600" + VLLM_RANDOMIZE_DP_DUMMY_INPUTS: "1" + NCCL_NET_GDR_C2C: "1" + NCCL_IB_HCA: "mlx5_0,mlx5_1,mlx5_2,mlx5_3" + NVSHMEM_REMOTE_TRANSPORT: "none" + NVSHMEM_ENABLE_NIC_PE_MAPPING: "1" + NVSHMEM_HCA_PE_MAPPING: "mlx5_0:1:1,mlx5_1:1:1,mlx5_2:1:1,mlx5_3:1:1" + UCX_TLS: "rc,cuda_copy" + UCX_NET_DEVICES: "mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1" + UCX_CUDA_IPC_ENABLE_MNNVL: "n" + UCX_MEMTYPE_CACHE: "n" + UCX_MEMTYPE_REG_WHOLE: "n" + UCX_RNDV_PIPELINE_ERROR_HANDLING: "y" + NIXL_UCX_SGL_ENABLE: "1" + WITH_NVIDIA_PEERMEM: "0" + MC_SLICE_SIZE: "1048576" + MC_MAX_MR_SIZE: "4294967296" + MC_ENABLE_DEST_DEVICE_AFFINITY: "1" + DYN_TCP_CONNECT_TIMEOUT: "120" + DYN_BACKEND_ADMISSION_ENGINE_REQUEST_LIMIT: "1024" + DYN_BACKEND_ADMISSION_ENGINE_REQUEST_WAITING_LIMIT: "1024" + DYN_BACKEND_ADMISSION_REQUEST_QUEUE_LIMIT: "2048" + roles: + prefill: + # One DEP4 worker per tray; run as many trays as the load needs. + nodes: 1 + parallelism: dep + args: + - "--data-parallel-size" + - "4" + - "--data-parallel-size-local" + - "4" + - "--max-num-seqs" + - "64" + - "--max-num-batched-tokens" + - "16384" + - "--long-prefill-token-threshold" + - "4096" + - "--gpu-memory-utilization" + - "0.88" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE"}' + - "--max-cudagraph-capture-size" + - "2048" + - "--kv-events-config" + - '{"publisher":"zmq","topic":"kv-events","enable_kv_cache_events":true,"endpoint":"tcp://*:5200"}' + - "--kv-transfer-config" + - '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"backends":["UCX"],"enforce_handshake_compat":false,"kv_lease_duration":120}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false}}]}}' + env: + VLLM_MOONCAKE_LOAD_RECV_THREADS: "8" + decode: + nodes: 2 + parallelism: dep + # 64 sequences per DP rank; decode graphs cover 64 x (1 + 7) tokens. + # The decode ranks only add capacity to the Mooncake store. + remove_args: ["--no-enable-prefix-caching"] + args: + - "--max-num-seqs" + - "64" + - "--max-num-batched-tokens" + - "8192" + - "--watermark" + - "0.05" + - "--gpu-memory-utilization" + - "0.88" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY"}' + - "--max-cudagraph-capture-size" + - "512" + - "--numa-bind-nodes" + - "0" + - "0" + - "1" + - "1" + - "--kv-transfer-config" + - '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"backends":["UCX"],"enforce_handshake_compat":false,"kv_lease_duration":120}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_consumer","kv_load_failure_policy":"recompute","kv_connector_extra_config":{"load_async":true,"lookup_async":true,"enable_offload":false,"save_decode_cache":false,"enable_lookup":false}}]}}' + env: + VLLM_MOONCAKE_LOAD_RECV_THREADS: "4" + VLLM_MAX_ASYNC_KV_ADMITS_PER_STEP: "4" + VLLM_ALLREDUCE_USE_SYMM_MEM: "0" + NCCL_CUMEM_ENABLE: "1" + NCCL_MNNVL_ENABLE: "1" + NCCL_NVLS_ENABLE: "1" + points: + dep8: + label: "DEP8 decode (throughput)" + description: "One DEP8 decode worker across two trays." + default: true + dep16: + label: "DEP16 decode (lower latency)" + description: "One DEP16 decode worker across four trays." + roles: { decode: { nodes: 4 } } + runbook: | + The GB200 disaggregated AgentX lane in [SemiAnalysis InferenceX](https://inferencex.semianalysis.com/) (`minimaxm31flash-fp4-gb200-dynamo-vllm-agentic-dspark-disagg`, [SemiAnalysisAI/InferenceX#4464](https://github.com/SemiAnalysisAI/InferenceX/pull/4464)) serves the default checkpoint text-only with Dynamo and vLLM. Several prefill workers, each DEP4 on one 4-GPU tray, feed one decode worker that is DEP8 across two trays or DEP16 across four. KV moves from prefill to decode over NIXL (UCX over RDMA), and every rank contributes host DRAM to a shared Mooncake KV store, so a returning conversation's prefix is loaded from the store instead of recomputed. Both roles use expert parallelism with DeepGEMM MegaMoE, attention data parallelism, the V2 model runner and the seven-token DSpark draft on `FLASHINFER_TS`. + + | Prefill workers (DEP4) | Decode | GPUs | Sessions | tok/s/GPU | tok/s/user | TTFT p90 | + |---:|---|---:|---:|---:|---:|---:| + | 3 | DEP8 | 20 | 256 | 80.8k | 244 | 4.4 s | + | 3 | DEP16 | 28 | 280 | 68.6k | 303 | 2.6 s | + | 4 | DEP16 | 32 | 352 | 70.1k | 288 | 4.1 s | + + The preset renders one prefill worker and the decode worker. Run more prefill workers with the same command, one per tray. + + **Image and versions.** The lane used an arm64 CUDA 13 build of the `vllm/vllm-openai:minimax-m31-flash` image (tag `minimax-m31-flash-arm64-cu134`), with Dynamo `1.6.0.dev20261023` and Mooncake `mooncake-transfer-engine-cuda13==0.3.13.post1` installed on every worker and on the Mooncake master host: `pip install --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime==1.6.0.dev20261023 ai-dynamo==1.6.0.dev20261023 mooncake-transfer-engine-cuda13==0.3.13.post1`. + + ### Tuning rules + + - **Prefill pool.** The prefill pool sets TTFT, so scale prefill workers with load. The lane's points kept TTFT p90 between 2.6 and 4.4 s with one DEP4 prefill tray per 85-95 concurrent agentic sessions (3 trays at 256-280 sessions, 4 at 352). Add a prefill tray when TTFT p90 climbs while decode still has room. + - **DEP8 or DEP16 decode.** DEP8 gives the most tokens per GPU: 32 sessions per decode GPU at 244 tok/s/user. DEP16 halves the batch per rank and gives about 25% more tok/s/user (303 at 280 sessions) at about 15% fewer tokens per GPU. `--max-num-seqs 64` per DP rank is a cap, not the operating point. The decode graph cap is `--max-num-seqs` x (1 + 7 draft tokens) = 512. `--watermark 0.05` keeps 5% of the decode KV cache free, so extra requests wait in the queue instead of preempting running ones. Keep decode `--max-num-batched-tokens` at 8192: a preempted request recomputes its context in chunks of that size. + - **Prefill batching.** `--long-prefill-token-threshold 4096` chunks long prompts into 4096-token pieces. Short prompts then share each step with them instead of waiting behind a 100K-token context, which lowers the TTFT tail. The recipe's 1P1D section above uses 8192. `--max-num-batched-tokens 16384` bounds a prefill step. Prefill captures CUDA graphs up to 2048 tokens. + - **Mooncake pool size.** `global_segment_size` is registered by every DP rank, prefill and decode alike, so a tray pins four times that value. The lane used 150 GB per rank (600 GB per tray). Size it as (host DRAM free per tray after the engines are up - headroom for NIXL and the 4 GB local buffers) / ranks per tray. The decode ranks only add capacity to the pool (`kv_consumer` with `enable_lookup: false` and `save_decode_cache: false`): the prefill workers write and read every block. + - **Routing.** The KV router (`--router-mode kv`) reads the prefill workers' KV events and sends each request to the worker that already holds its prefix. Session affinity (`--router-session-affinity-ttl-secs 14400`) keeps a conversation on that worker for up to four hours. Send a stable conversation ID in the `X-Dynamo-Session-ID` header. The Mooncake store catches the prefixes that miss after a reroute. + - **Transport.** Cross-tray NIXL needs `UCX_TLS=rc,cuda_copy` with `UCX_NET_DEVICES` set to the tray's RDMA ports. Without them, UCX can fall back to TCP and transfers stall. Point `UCX_MODULE_DIR` at the UCX bundled with the NIXL wheel (`nixl_cu13.libs/ucx` in the image's `dist-packages`). `kv_lease_duration: 120` holds a prefill's blocks for up to 120 s until decode reads them (the default is 30 s). `VLLM_MAX_ASYNC_KV_ADMITS_PER_STEP=4` caps how many requests with finished KV loads decode admits per step. + - **NUMA.** `--numa-bind` pins each rank to its GPU's NUMA node. Decode passes the tray's map, `--numa-bind-nodes 0 0 1 1`, explicitly. + - `--moe-backend deep_gemm_mega_moe` is the same as the lane's `--kernel-config '{"moe_backend":"deep_gemm_mega_moe"}'`. `--prefill-schedule-interval 1`, also set by the lane, is vLLM's default. + + ### Services + + Start etcd, NATS and the Mooncake master once: + + ```bash + etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://:2379 + printf 'max_payload: 33554432\n' > nats.conf && nats-server -c nats.conf + mooncake_master --rpc_port=50051 \ + --default_kv_lease_ttl=60000 --eviction_high_watermark_ratio=0.95 --eviction_ratio=0.10 + ``` + + Every worker reads the store configuration from the JSON file that `MOONCAKE_CONFIG_PATH` points to (export it with the rest of the environment): + + ```json + {"master_server_address": ":50051", "metadata_server": "P2PHANDSHAKE", + "global_segment_size": "150GB", "local_buffer_size": "4GB", "protocol": "rdma", + "device_name": "mlx5_0,mlx5_1,mlx5_2,mlx5_3", "enable_offload": false} + ``` + + ### Workers + + Export `ETCD_ENDPOINTS=http://:2379`, `NATS_SERVER=nats://:4222` and the generated environment on every tray, then run the generated `python3 -m dynamo.vllm` commands. Add `--served-model-name MiniMaxAI/MiniMax-M3.1-Flash --trust-remote-code --dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3` to both roles. `dynamo.vllm` rejects vLLM's API-server flags `--tool-call-parser` and `--enable-auto-tool-choice`. Multi-node DEP decode runs one process per tray with `--data-parallel-start-rank` set to 0, 4, 8 and so on, and a shared `--data-parallel-address` and `--data-parallel-rpc-port`. `--data-parallel-hybrid-lb` lets every tray's Dynamo runtime register its own DP ranks with the frontend. + + ### Frontend + + ```bash + export ETCD_ENDPOINTS=http://:2379 NATS_SERVER=nats://:4222 + export DYN_ROUTER_ACTIVE_REQUEST_EXPIRY_SECS=3600 DYN_TCP_CONNECT_TIMEOUT=120 + python3 -m dynamo.frontend --http-port 8000 --router-mode kv --router-session-affinity-ttl-secs 14400 + ``` + + The InferenceX throughput points replaced block verification with synthetic rejection at the committed golden DSpark acceptance length (4.04), so every point is measured at the same acceptance rate. Production serving keeps the default rejection sampler (`rejection_sample_method: block`), as in the command above. + guide: | ## Overview From d929755b7a475d75bae0c8c095b75cfb398389fa Mon Sep 17 00:00:00 2001 From: hshrivastava-droid Date: Wed, 28 Oct 2026 12:17:21 -0700 Subject: [PATCH 077/406] [Doc] DeepSeek-V4-Pro: add B200 InferenceX AgentX presets (#64316) Signed-off-by: hshrivastava-droid Co-authored-by: Claude Sonnet 5.5 --- .../models/deepseek-ai/DeepSeek-V4-Pro.yaml | 337 ++++++++++++++++++ 1 file changed, 337 insertions(+) diff --git a/recipes/models/deepseek-ai/DeepSeek-V4-Pro.yaml b/recipes/models/deepseek-ai/DeepSeek-V4-Pro.yaml index d57a3fc7c916..e5bde595f970 100644 --- a/recipes/models/deepseek-ai/DeepSeek-V4-Pro.yaml +++ b/recipes/models/deepseek-ai/DeepSeek-V4-Pro.yaml @@ -275,6 +275,343 @@ strategy_overrides: - "--enable-sleep-mode" env: {} +# Tuned presets (InferenceX AgentX lanes); see SCHEMA.md "presets". +presets: + b200_agentx_tp: + label: "AgentX TP8" + hardware: b200 + strategy: single_node_tp + tp: 8 + variants: [default] + workload: "Agentic (InferenceX AgentX)" + description: "TP8 on one 8-GPU B200 node with DSpark (six draft tokens), FlashInfer sparse MLA and a fixed 50 GiB KV cache per GPU; no host KV offload." + source: { label: "InferenceX #4554", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4554", key: "dsv4-fp4-b200-vllm-agentic-mtp", revision: 2c044c3 } + image: "vllm/vllm-openai:nightly-8e2beaa3bff595577a069d49d96bde9074294ec2@sha256:3207cb99c0f59b92f7d47a74da8c6eb5467565d811a4f8d9df51df5188f59c70" + select: + features: { spec_decoding: dspark } + # The lane's FlashInfer sparse MLA attention config replaces the builder's + # standalone MXFP4 indexer flag; DSpark drafts six tokens, not seven; the pinned + # image predates --disable-flashinfer-autotune, so keep the --no-enable form. + remove_args: ["--attention_config.indexer_kv_dtype=mxfp4", "--disable-flashinfer-autotune"] + args: + - "--no-enable-flashinfer-autotune" + - "--max-model-len" + - "1048576" + - "--gpu-memory-utilization" + - "0.90" + - "--max-num-batched-tokens" + - "16384" + - "--numa-bind" + - "--enable-cumem-allocator" + - "--disable-uvicorn-access-log" + - "--no-disable-hybrid-kv-cache-manager" + - "--attention-config" + - '{"backend":"FLASHINFER_MLA_SPARSE_DSV4","use_prefill_query_quantization":true,"indexer_kv_dtype":"mxfp4"}' + json_args: + "--speculative-config": { num_speculative_tokens: 6 } + env: + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_FLOAT32_MATMUL_PRECISION: "high" + TORCH_CUDA_ARCH_LIST: "10.0" + VLLM_ENGINE_READY_TIMEOUT_S: "7200" + VLLM_RPC_TIMEOUT: "600000" + # --max-num-seqs equals the concurrency. Graphs capture every decode batch of up to + # that many sequences (target plus six drafts each) and mixed batches of 100-500 tokens. + points: + c1: + label: "1" + concurrency: 1 + args: + - "--max-num-seqs" + - "1" + - "--kv-cache-memory-bytes" + - "53687091200" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,100,200,300,400,500]}' + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + c4: + label: "4 (low latency)" + concurrency: 4 + image: "vllm/vllm-openai@sha256:d428631bb917a33d22308265d1b8be91522f862bf9fd133bee62dde3313664ae" + args: + - "--max-num-seqs" + - "4" + - "--kv-cache-memory-bytes" + - "42949672960" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,14,21,28,100,200,300,400,500,640,800,1024,1280,1600,2048,2560,3200,4096,5120,6400,8192,10240,12800,16384],"max_cudagraph_capture_size":16384}' + - "--sparse-indexer-topk-backend" + - "deep_select" + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + c6: + label: "6" + concurrency: 6 + args: + - "--max-num-seqs" + - "6" + - "--kv-cache-memory-bytes" + - "53687091200" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,14,21,28,35,42,100,200,300,400,500]}' + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + c10: + label: "10" + concurrency: 10 + image: "vllm/vllm-openai@sha256:d428631bb917a33d22308265d1b8be91522f862bf9fd133bee62dde3313664ae" + args: + - "--max-num-seqs" + - "10" + - "--kv-cache-memory-bytes" + - "53687091200" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,100,200,300,400,500]}' + - "--prefix-cache-retention-interval" + - "32768" + - "--prefill-schedule-interval" + - "64" + json_args: + "--attention-config": { indexer_tp_shard_min_rows: 4096 } + c14: + label: "14" + concurrency: 14 + args: + - "--max-num-seqs" + - "14" + - "--kv-cache-memory-bytes" + - "53687091200" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,100,200,300,400,500]}' + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + runbook: | + The TP8 points of the B200 AgentX lane (`dsv4-fp4-b200-vllm-agentic-mtp` and its low-concurrency + companion, [SemiAnalysisAI/InferenceX#4554](https://github.com/SemiAnalysisAI/InferenceX/pull/4554)) serve the + 0813 checkpoint with DSpark (six draft tokens) as one aggregated vLLM worker on one 8-GPU B200 node: TP8, no + expert parallelism, no host KV offload. Each GPU holds a fixed 50 GiB KV cache and `--max-num-seqs` equals the + concurrency. + + Concurrency 4 is the low-latency variant: a 40 GiB KV cache, the `deep_select` sparse indexer top-k backend and + graph capture up to 16384 tokens. Concurrency 10 sets `--prefix-cache-retention-interval 32768` and + `--prefill-schedule-interval 64`, shards the indexer at 4096 rows (`indexer_tp_shard_min_rows`) and leaves + `VLLM_DSV4_MEGA_FP8_COMBINE` unset. Concurrency 4 and 10 were measured on a second pinned image, shown on + those points. + + Left out of the command: the lane also exported `VLLM_PREFIX_CACHE_RETENTION_INTERVAL`, which vLLM removed in + v0.30.0 ([vllm#55353](https://github.com/vllm-project/vllm/pull/55353)) in favour of the flag above, so it had + no effect on the image it ran on; and `--served-model-name`, which the builder replaces with the checkpoint ID. + The throughput points in InferenceX swap block verification for synthetic rejection at a committed + acceptance length so every point is measured at the same acceptance rate; the command does not set it. + + b200_agentx_dep: + label: "AgentX DEP" + hardware: b200 + strategy: single_node_dep + variants: [default] + workload: "Agentic (InferenceX AgentX)" + description: "Data-parallel ranks with expert parallelism and a Simple CPU KV tier behind a consistent-hash vLLM Router: DEP8 for concurrency 64 to 192, TP4 x DP2 at 32 and TP2 x DP4 at 48 and 64." + source: { label: "InferenceX #4554", href: "https://github.com/SemiAnalysisAI/InferenceX/pull/4554", key: "dsv4-fp4-b200-vllm-agentic-mtp", revision: 2c044c3 } + image: "vllm/vllm-openai:nightly-8e2beaa3bff595577a069d49d96bde9074294ec2@sha256:3207cb99c0f59b92f7d47a74da8c6eb5467565d811a4f8d9df51df5188f59c70" + select: + kv_offload: simple + features: { spec_decoding: dspark } + remove_args: ["--attention_config.indexer_kv_dtype=mxfp4"] + args: + - "--max-model-len" + - "1048576" + - "--gpu-memory-utilization" + - "0.90" + - "--numa-bind" + - "--enable-cumem-allocator" + - "--disable-uvicorn-access-log" + - "--enable-ep-weight-filter" + - "--no-disable-hybrid-kv-cache-manager" + - "--no-enable-flashinfer-autotune" + - "--prefill-schedule-interval" + - "16" + - "--long-prefill-token-threshold" + - "512" + - "--moe-backend" + - "deep_gemm_mega_moe" + - "--attention-config" + - '{"backend":"FLASHINFER_MLA_SPARSE_DSV4","use_prefill_query_quantization":true,"indexer_kv_dtype":"mxfp4"}' + json_args: + "--speculative-config": { num_speculative_tokens: 6 } + "--kv-transfer-config": { kv_connector_extra_config: { cpu_bytes_to_use_per_rank: 241500000000, enable_cross_layers_blocks: "true", lazy_offload: false } } + env: + VLLM_USE_V2_MODEL_RUNNER: "1" + VLLM_FLOAT32_MATMUL_PRECISION: "high" + TORCH_CUDA_ARCH_LIST: "10.0" + VLLM_ENGINE_READY_TIMEOUT_S: "7200" + VLLM_RPC_TIMEOUT: "600000" + PYTORCH_ALLOC_CONF: "expandable_segments:True" + PYTHONHASHSEED: "42" + # --max-num-seqs is per DP rank; --max-num-batched-tokens is an 8192-token prefill budget + # plus seven verification slots per sequence; graphs capture decode batches of 7 to 7 x seqs. + points: + dep8_c64: + label: "64 (DEP8)" + concurrency: 64 + args: + - "--tensor-parallel-size" + - "1" + - "--data-parallel-size" + - "8" + - "--max-num-seqs" + - "16" + - "--max-num-batched-tokens" + - "8304" + - "--kv-cache-memory-bytes" + - "37580963840" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112],"mode":0}' + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + dep8_c96: + label: "96 (DEP8)" + concurrency: 96 + args: + - "--tensor-parallel-size" + - "1" + - "--data-parallel-size" + - "8" + - "--max-num-seqs" + - "24" + - "--max-num-batched-tokens" + - "8360" + - "--kv-cache-memory-bytes" + - "33285996544" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168],"mode":0}' + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + dep8_c128: + label: "128 (DEP8)" + concurrency: 128 + args: + - "--tensor-parallel-size" + - "1" + - "--data-parallel-size" + - "8" + - "--max-num-seqs" + - "32" + - "--max-num-batched-tokens" + - "8416" + - "--kv-cache-memory-bytes" + - "33285996544" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224],"mode":0}' + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + dep8_c160: + label: "160 (DEP8)" + concurrency: 160 + args: + - "--tensor-parallel-size" + - "1" + - "--data-parallel-size" + - "8" + - "--max-num-seqs" + - "40" + - "--max-num-batched-tokens" + - "8472" + - "--kv-cache-memory-bytes" + - "32212254720" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224,231,238,245,252,259,266,273,280],"mode":0}' + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + dep8_c192: + label: "192 (DEP8)" + concurrency: 192 + args: + - "--tensor-parallel-size" + - "1" + - "--data-parallel-size" + - "8" + - "--max-num-seqs" + - "48" + - "--max-num-batched-tokens" + - "8528" + - "--kv-cache-memory-bytes" + - "32212254720" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224,231,238,245,252,259,266,273,280,287,294,301,308,315,322,329,336],"mode":0}' + env: + VLLM_DSV4_MEGA_FP8_COMBINE: "1" + tp4dp2_c32: + label: "32 (TP4 x DP2)" + concurrency: 32 + args: + - "--tensor-parallel-size" + - "4" + - "--data-parallel-size" + - "2" + - "--max-num-seqs" + - "32" + - "--max-num-batched-tokens" + - "8416" + - "--kv-cache-memory-bytes" + - "51539607552" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224],"mode":0}' + tp2dp4_c48: + label: "48 (TP2 x DP4)" + concurrency: 48 + image: "vllm/vllm-openai@sha256:d428631bb917a33d22308265d1b8be91522f862bf9fd133bee62dde3313664ae" + args: + - "--tensor-parallel-size" + - "2" + - "--data-parallel-size" + - "4" + - "--max-num-seqs" + - "24" + - "--max-num-batched-tokens" + - "8360" + - "--kv-cache-memory-bytes" + - "37580963840" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168],"mode":0}' + tp2dp4_c64: + label: "64 (TP2 x DP4)" + concurrency: 64 + image: "vllm/vllm-openai@sha256:d428631bb917a33d22308265d1b8be91522f862bf9fd133bee62dde3313664ae" + args: + - "--tensor-parallel-size" + - "2" + - "--data-parallel-size" + - "4" + - "--max-num-seqs" + - "32" + - "--max-num-batched-tokens" + - "8416" + - "--kv-cache-memory-bytes" + - "37580963840" + - "--compilation-config" + - '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[7,14,21,28,35,42,49,56,63,70,77,84,91,98,105,112,119,126,133,140,147,154,161,168,175,182,189,196,203,210,217,224],"mode":0}' + runbook: | + The data-parallel points of the B200 AgentX lane (`dsv4-fp4-b200-vllm-agentic-mtp` and its low-concurrency + companion, [SemiAnalysisAI/InferenceX#4554](https://github.com/SemiAnalysisAI/InferenceX/pull/4554)) serve the + 0813 checkpoint with DSpark (six draft tokens) as one engine on one 8-GPU B200 node: DEP8 (TP1 x DP8) for + concurrency 64 to 192, TP4 x DP2 at 32 and TP2 x DP4 at 48 and 64, all with expert parallelism and the Simple + CPU KV tier. Each rank gets 241.5 GB of host memory (1932 GB across the node), so size it to the machine. + `--max-num-seqs` is per DP rank, and `--max-num-batched-tokens` is an 8192-token prefill budget plus seven + verification slots per sequence. The TP4 x DP2 and TP2 x DP4 points leave `VLLM_DSV4_MEGA_FP8_COMBINE` unset. + TP2 x DP4 was measured on a second pinned image, shown on those points. + + The lane fronts the server with vLLM Router 0.1.14 (`--policy consistent_hash`, `--intra-node-data-parallel-size` + equal to the server's DP size, as on the router tab), plus `--request-timeout-secs 14400` so long agentic + requests outlive the default and `--disable-retries` so a failed request is not replayed. Clients send a stable + per-conversation `X-Session-ID` so the router keeps a conversation on the rank that holds its prefix. + + Left out of the command: `VLLM_PREFIX_CACHE_RETENTION_INTERVAL` (removed in vLLM v0.30.0, + [vllm#55353](https://github.com/vllm-project/vllm/pull/55353), so it had no effect), `--served-model-name` + (the builder replaces it with the checkpoint ID) and the router's Prometheus port. The throughput points in + InferenceX swap block verification for synthetic rejection at a committed acceptance length so every point is + measured at the same acceptance rate; the command does not set it. + guide: | ## Overview From 892cb1981653e503ecc1a87c55809627ae24f492 Mon Sep 17 00:00:00 2001 From: Luyi Xiao <81631911+luyixiao95@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:04:07 -0700 Subject: [PATCH 078/406] [MyPy][8/N] Enable strict checking for model test subdirectories (#63042) Signed-off-by: Luyi Xiao <81631911+luyixiao95@users.noreply.github.com> Co-authored-by: OpenAI Codex --- .../inkling/test_contract_validation.py | 6 ++-- tests/models/inkling/test_mtp_input_fusion.py | 1 + .../amd_specialized_decode/test_runtime.py | 6 +++- tests/models/kimi_k3/test_amd_attn_res.py | 2 ++ .../kimi_k3/test_amd_latent_moe_runner.py | 3 ++ tests/models/kimi_k3/test_eagle3.py | 15 +++++--- tests/models/kimi_k3/test_kda.py | 3 +- tests/models/kimi_k3/test_kda_metadata.py | 14 +++++--- tests/models/kimi_k3/test_latent_moe_tail.py | 3 +- .../kimi_k3/test_mla_prefill_context.py | 17 ++++----- .../models/kimi_k3/test_sequence_parallel.py | 10 +++--- .../amd_specialized_decode/test_ops.py | 8 ++++- .../amd_specialized_decode/test_runtime.py | 8 +++-- tests/models/qwen4_exp/test_ple.py | 3 ++ tests/models/qwen4_exp/test_qsa_amd.py | 12 ++++--- tests/models/qwen4_exp/test_qsa_reference.py | 1 + tests/models/qwen4_exp/test_tpsp.py | 5 ++- .../models/transformers/fusers/test_linear.py | 36 +++++++++++++------ tests/models/transformers/fusers/test_moe.py | 5 +++ .../transformers/fusers/test_rms_norm.py | 13 +++++-- tests/models/transformers/test_backend.py | 7 ++-- tools/pre_commit/mypy.py | 8 +++++ 22 files changed, 137 insertions(+), 49 deletions(-) diff --git a/tests/models/inkling/test_contract_validation.py b/tests/models/inkling/test_contract_validation.py index 4bdb75dcc499..e5c13903f4e8 100644 --- a/tests/models/inkling/test_contract_validation.py +++ b/tests/models/inkling/test_contract_validation.py @@ -37,8 +37,10 @@ def test_vision_scale_plan_breaks_assignment_ties_in_order(): @pytest.mark.parametrize( ("config_cls", "kwargs", "missing"), [ - (InklingAudioConfig, {"decoder_dmodel": 16}, "n_mel_bins"), - (InklingVisionConfig, {"decoder_dmodel": 16}, "vision_encoder_type"), + pytest.param(InklingAudioConfig, {"decoder_dmodel": 16}, "n_mel_bins"), + pytest.param( + InklingVisionConfig, {"decoder_dmodel": 16}, "vision_encoder_type" + ), ], ) def test_enabled_tower_requires_architecture_fields(config_cls, kwargs, missing): diff --git a/tests/models/inkling/test_mtp_input_fusion.py b/tests/models/inkling/test_mtp_input_fusion.py index 7ab33dd45ae8..0e17eb523413 100644 --- a/tests/models/inkling/test_mtp_input_fusion.py +++ b/tests/models/inkling/test_mtp_input_fusion.py @@ -85,6 +85,7 @@ def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None: # Gather + embed_norm (base model / MTP prefill embed path). out = embed_rmsnorm(ids, table, w, EPS) + assert isinstance(out, torch.Tensor) assert out.shape == (t, n) assert torch.equal(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb) diff --git a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py index 266b947dbe01..2ea326b752e4 100644 --- a/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py +++ b/tests/models/kimi_k3/amd_specialized_decode/test_runtime.py @@ -331,7 +331,11 @@ def _decode(monkeypatch, mode=CUDAGraphMode.FULL, md=None): decode = object.__new__(SpecializedDecode) decode.config = DecodeConfig(q=4) decode.runner = SimpleNamespace(kda_prefix="kda") - decode.kernels = {w: {1: _Kernel(), 2: _Kernel()} for w in (4, 8)} + # Launch-only stand-ins for the GPU kernels. + decode.kernels = { + w: {1: _Kernel(), 2: _Kernel()} # type: ignore[dict-item] + for w in (4, 8) + } decode._kv = (1, 2) monkeypatch.setattr(SpecializedDecode, "_kv_key", lambda self: (1, 2)) context = SimpleNamespace( diff --git a/tests/models/kimi_k3/test_amd_attn_res.py b/tests/models/kimi_k3/test_amd_attn_res.py index 871d53faae12..622e8f989782 100644 --- a/tests/models/kimi_k3/test_amd_attn_res.py +++ b/tests/models/kimi_k3/test_amd_attn_res.py @@ -99,6 +99,7 @@ def test_amd_attn_res_matches_reference( 0.0, ) + assert isinstance(actual, torch.Tensor) torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) torch.testing.assert_close(prefix, original_prefix, atol=0, rtol=0) torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0) @@ -185,6 +186,7 @@ def test_amd_attn_res_fused_contract( output_eps, ) + assert isinstance(actual, torch.Tensor) torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0) if write_block: diff --git a/tests/models/kimi_k3/test_amd_latent_moe_runner.py b/tests/models/kimi_k3/test_amd_latent_moe_runner.py index e47b5fc08e87..237c3e20bc52 100644 --- a/tests/models/kimi_k3/test_amd_latent_moe_runner.py +++ b/tests/models/kimi_k3/test_amd_latent_moe_runner.py @@ -131,6 +131,7 @@ def _check_matches_replicated( row_sharded: bool = False, ) -> None: transform, up_weight = _build_transform(device, row_sharded=row_sharded) + assert transform.norm is not None runner = _tail_runner(transform, tp_world, use_ep) group = get_tp_group().device_group @@ -174,6 +175,7 @@ def _check_unfused_tail_gathers_full_hidden( ) -> None: transform, up_weight = _build_transform(device, row_sharded=True) assert transform.up_proj.weight.shape[0] == HIDDEN_SIZE // tp_world + assert transform.norm is not None # all-gather reconstructs the full projection when all ranks feed the same # input, so use one seed for all ranks. @@ -194,6 +196,7 @@ def _check_writes_only_its_own_shard( """Two ranks that swapped both weight rows and write offsets still sum to the right total, so inspect each slice before the final collective.""" transform, up_weight = _build_transform(device) + assert transform.norm is not None runner = _tail_runner(transform, tp_world) group = get_tp_group().device_group diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index 858622fa07da..562ade2b8b6a 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -36,7 +36,8 @@ def test_kimi_linear_advertises_eagle3_support(): def test_kimi_k3_uses_shared_eagle3_layer_configuration(): - target = object.__new__(KimiK3ForConditionalGeneration) + # Skip model initialization; the test supplies the required attributes below. + target = object.__new__(KimiK3ForConditionalGeneration) # type: ignore[type-abstract] torch.nn.Module.__init__(target) model = _make_kimi_linear_model() object.__setattr__(model, "layers", [None] * 93) @@ -79,12 +80,14 @@ def test_kimi_linear_forward_extracts_standard_aux_hidden_states(monkeypatch): lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), ) - output, aux_hidden_states = model.forward( + result = model.forward( input_ids=None, positions=torch.tensor([0]), intermediate_tensors=None, inputs_embeds=initial_hidden_states, ) + assert isinstance(result, tuple) + output, aux_hidden_states = result expected_layer_output = layer_hidden_states + layer_residual torch.testing.assert_close(output, expected_layer_output) @@ -128,12 +131,14 @@ def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): final_attn_res = Mock(return_value=final_hidden_states) monkeypatch.setattr(kimi_model, "attn_res", final_attn_res) - output, aux_hidden_states = model.forward( + result = model.forward( input_ids=None, positions=torch.tensor([0]), intermediate_tensors=None, inputs_embeds=initial_hidden_states, ) + assert isinstance(result, tuple) + output, aux_hidden_states = result torch.testing.assert_close(output, final_hidden_states) torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) @@ -187,12 +192,14 @@ def test_attn_res_stream_capture_receives_the_layer_outputs_in_order(monkeypatch capture = Mock(return_value=captured) monkeypatch.setattr(KimiLinearModel, "_capture_aux_hidden_stream", capture) - _, aux_hidden_states = model.forward( + result = model.forward( input_ids=None, positions=torch.tensor([0]), intermediate_tensors=None, inputs_embeds=initial_hidden_states, ) + assert isinstance(result, tuple) + _, aux_hidden_states = result layer_idx, got_prefix, got_pending, got_residual = capture.call_args.args assert layer_idx == 0 diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 2dfc08eb0497..becf78b59445 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -6,6 +6,7 @@ Uses torch.rand for q/k/v to match FLA's test pattern. """ +from collections.abc import Callable from types import SimpleNamespace import pytest @@ -77,7 +78,7 @@ "nvidia": fused_recurrent_kda_packed_decode, "amd": fused_recurrent_kda_packed_decode_amd, } -SPEC_DECODE_IMPLS = { +SPEC_DECODE_IMPLS: dict[str, Callable[..., tuple[torch.Tensor, torch.Tensor]]] = { "nvidia": fused_recurrent_kda_nvidia, "amd": fused_recurrent_kda_amd, } diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 61c19e8961de..20fae682c40a 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -3,6 +3,7 @@ from dataclasses import fields from types import SimpleNamespace +from typing import TypeVar from unittest.mock import Mock, patch import pytest @@ -26,7 +27,6 @@ from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM from vllm.platforms import current_platform from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig -from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionBackend, GDNAttentionMetadata, @@ -58,6 +58,8 @@ "flashinfer_prefill_seq_order", } +_BuilderT = TypeVar("_BuilderT", bound=GDNAttentionMetadataBuilder) + def _assert_matches_shared_gdn( reference: GDNAttentionMetadata, actual: KimiK3KDAMetadata @@ -92,7 +94,7 @@ def _assert_matches_shared_gdn( def _make_builder( - builder_cls: type[AttentionMetadataBuilder], + builder_cls: type[_BuilderT], num_speculative_tokens: int, full_cuda_graph: bool, device: torch.device = DEVICE, @@ -105,7 +107,7 @@ def _make_builder( disable_eagle_block_drop: bool = False, cache_hit_alignment_tokens: int | None = None, hash_block_size: int | None = None, -) -> AttentionMetadataBuilder: +) -> _BuilderT: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE, @@ -659,6 +661,8 @@ def test_recoverssm_spec_uses_one_state_slot_and_current_window( ) context = builder.recoverssm_context assert context is not None + commit = context.commit + assert isinstance(commit, Mock) actual = builder.build( 0, common_attn_metadata, @@ -696,7 +700,7 @@ def test_recoverssm_spec_uses_one_state_slot_and_current_window( is common_attn_metadata.compute_num_computed_tokens() ) assert postprocess.block_size == BLOCK_SIZE - args = context.commit.call_args.args + args = commit.call_args.args assert args[0] is num_accepted_tokens torch.testing.assert_close(args[1], commit_metadata.state_indices[:, 0]) torch.testing.assert_close(args[2], commit_metadata.query_start_loc) @@ -807,7 +811,7 @@ def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn(): ], ) def test_cudagraph_capture_metadata_avoids_device_to_host_copy( - builder_cls: type[AttentionMetadataBuilder], monkeypatch: pytest.MonkeyPatch + builder_cls: type[GDNAttentionMetadataBuilder], monkeypatch: pytest.MonkeyPatch ): """`build_for_cudagraph_capture` is not capture-only: a DP rank with nothing scheduled re-stages its FULL-graph metadata through it on every dummy step, diff --git a/tests/models/kimi_k3/test_latent_moe_tail.py b/tests/models/kimi_k3/test_latent_moe_tail.py index dd3136641121..45235a67a923 100644 --- a/tests/models/kimi_k3/test_latent_moe_tail.py +++ b/tests/models/kimi_k3/test_latent_moe_tail.py @@ -101,7 +101,8 @@ def fake_tail_initialize(**kwargs): latent_moe_runner.LatentMoERunner() - assert moe_config.use_deferred_moe_finalize + # The stand-in deliberately borrows the real config's property. + assert moe_config.use_deferred_moe_finalize # type: ignore[misc] assert moe_config.defer_moe_finalize_max_num_tokens == 128 assert initialized_with["experts_per_token"] == 16 diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py index e8cfb014f63c..6c5362e35520 100644 --- a/tests/models/kimi_k3/test_mla_prefill_context.py +++ b/tests/models/kimi_k3/test_mla_prefill_context.py @@ -126,7 +126,6 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]: class _FusedLayer: """Only the attributes K3's context loop reads.""" - _compute_prefill_context = MultiHeadLatentAttention._compute_prefill_context _gather_context_latent = MultiHeadLatentAttention._gather_context_latent _attn_read_kv_cache = MultiHeadLatentAttention._attn_read_kv_cache @@ -144,7 +143,6 @@ def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None: class _ReferenceImpl: """Only the attributes the generic context loop reads.""" - _compute_prefill_context = MLACommonBaseImpl._compute_prefill_context _concat_k_nope_k_pe = MLACommonBaseImpl._concat_k_nope_k_pe _use_flashinfer_concat_mla_k = False @@ -163,7 +161,7 @@ def _build_prefill_metadata( workspace_dtype: torch.dtype, q_data_type: torch.dtype, backend: _RecordingPrefillBackend, -) -> MLACommonPrefillMetadata: +) -> tuple[MLACommonPrefillMetadata, int]: query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32) query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0) workspace = torch.empty( @@ -253,13 +251,13 @@ def test_fused_context_matches_generic_impl( ).to(q_data_type) layer = _FusedLayer(kv_b_proj, kv_cache, kv_cache_dtype, k_scale) - fused_out, fused_lse = layer._compute_prefill_context( - q, SimpleNamespace(prefill=prefill_fused) + fused_out, fused_lse = MultiHeadLatentAttention._compute_prefill_context( + layer, q, SimpleNamespace(prefill=prefill_fused) ) impl = _ReferenceImpl(kv_b_proj, kv_cache_dtype) - ref_out, ref_lse = impl._compute_prefill_context( - q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale + ref_out, ref_lse = MLACommonBaseImpl._compute_prefill_context( + impl, q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale ) assert len(backend_fused.calls) == len(backend_ref.calls) @@ -281,6 +279,7 @@ def test_fused_context_matches_generic_impl( # Every chunk but the continuation should have been written in place, i.e. # straight into the returned accumulator, with no intermediate copy. + assert prefill_fused.chunked_context is not None wrote_in_place = [ out is not None and out.data_ptr() == fused_out[chunk.token_slice].data_ptr() for out, chunk in zip( @@ -318,7 +317,9 @@ def test_fused_context_rejects_an_unquantized_query() -> None: dtype=torch.bfloat16, ) with pytest.raises(AssertionError, match="new-token epilogue"): - layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) + MultiHeadLatentAttention._compute_prefill_context( + layer, q, SimpleNamespace(prefill=prefill) + ) # -------------------------------------------------------------------------- diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py index 3057f85c7e44..09564b4221d3 100644 --- a/tests/models/kimi_k3/test_sequence_parallel.py +++ b/tests/models/kimi_k3/test_sequence_parallel.py @@ -142,9 +142,10 @@ def test_kimi_decoder_layer_keeps_moe_states_sequence_sharded(monkeypatch): layer.input_layernorm = _IdentityNorm() layer.post_attention_layernorm = _IdentityNorm() layer.mlp = _RecordingMoE() - layer._run_self_attn = MethodType( - lambda self, positions, hidden_states: hidden_states, + monkeypatch.setattr( layer, + "_run_self_attn", + MethodType(lambda self, positions, hidden_states: hidden_states, layer), ) _mock_sequence_parallel_collectives(monkeypatch) @@ -188,9 +189,10 @@ def test_kimi_attn_residual_states_stay_sequence_sharded(monkeypatch): layer.self_attention_res_proj = _Projection() layer.mlp_res_proj = _Projection() layer.mlp = _RecordingMoE() - layer._run_self_attn = MethodType( - lambda self, positions, hidden_states: hidden_states, + monkeypatch.setattr( layer, + "_run_self_attn", + MethodType(lambda self, positions, hidden_states: hidden_states, layer), ) _mock_sequence_parallel_collectives(monkeypatch) diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_ops.py b/tests/models/minimax_m3/amd_specialized_decode/test_ops.py index c72a1fde334b..27d2aa2b70b4 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_ops.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_ops.py @@ -144,12 +144,18 @@ def _pairs(region, n): return region[: n * 8].view(torch.int32).view(n, 2)[:, 0].contiguous() +class _TestFfnRunner(FfnRunner): + peer_bufs: list[torch.Tensor] + + def _runner(weights, s, tp=1, widths=None, **build): dev = torch.device("cuda") bufs = [torch.zeros(ffn.peer_bytes(tp), dtype=torch.uint8, device=dev)] bufs += [torch.zeros_like(bufs[0]) for _ in range(tp - 1)] peers = torch.tensor([b.data_ptr() for b in bufs], dtype=torch.int64, device=dev) - runner = FfnRunner(tp, 0, bufs[0].data_ptr(), peers, dev, widths or (s,), **build) + runner = _TestFfnRunner( + tp, 0, bufs[0].data_ptr(), peers, dev, widths or (s,), **build + ) runner.warm(weights) runner.peer_bufs = bufs return runner diff --git a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py index 326ebde2555f..57863815b88e 100644 --- a/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py +++ b/tests/models/minimax_m3/amd_specialized_decode/test_runtime.py @@ -139,9 +139,13 @@ def _vllm_config(layers=60, **parallel): def test_config_refusals_report_every_reason(): - assert specialized_decode_config(_vllm_config()).last_layer == 59 + full_config = specialized_decode_config(_vllm_config()) + assert full_config is not None + assert full_config.last_layer == 59 # a truncated checkpoint: the sparse layers end at its last layer - assert specialized_decode_config(_vllm_config(layers=5)).last_layer == 4 + truncated_config = specialized_decode_config(_vllm_config(layers=5)) + assert truncated_config is not None + assert truncated_config.last_layer == 4 config = _vllm_config(tensor_parallel_size=2, enable_expert_parallel=True) config.cache_config.block_size = 64 diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index b4382ebd15b2..a50387614305 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -139,6 +139,9 @@ def check(actual, expected): assert metadata.num_decodes == sum(query_lens[i] == 1 for i in non_spec_rows) assert metadata.num_prefills == sum(query_lens[i] > 1 for i in non_spec_rows) if spec_rows: + assert metadata.spec_state_indices_tensor is not None + assert metadata.num_accepted_tokens is not None + assert metadata.spec_query_start_loc is not None torch.testing.assert_close( metadata.spec_state_indices_tensor[: len(spec_rows)], common.block_table_tensor[spec_rows, 0], diff --git a/tests/models/qwen4_exp/test_qsa_amd.py b/tests/models/qwen4_exp/test_qsa_amd.py index c8ebf0190432..58f33b6cd53c 100644 --- a/tests/models/qwen4_exp/test_qsa_amd.py +++ b/tests/models/qwen4_exp/test_qsa_amd.py @@ -139,7 +139,9 @@ def _make_attn(*, use_fused: bool) -> SimpleNamespace: ) -def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer: +def _make_indexer( + monkeypatch: pytest.MonkeyPatch, *, use_fused: bool, num_tokens: int = 2 +) -> QSAIndexer: """A QSAIndexer whose unfused path runs for real, with no projection weights. Both branches have to be reachable: a fused-path test is only meaningful if @@ -190,7 +192,9 @@ def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer: slot_mapping=torch.arange(num_tokens), k_work_metadata=torch.zeros(1, 2, dtype=torch.int32), ) - indexer._metadata = lambda: (raw_metadata, compressed_metadata) + monkeypatch.setattr( + indexer, "_metadata", lambda: (raw_metadata, compressed_metadata) + ) indexer.test_metadata = (raw_metadata, compressed_metadata) return indexer @@ -199,7 +203,7 @@ def _make_indexer(*, use_fused: bool, num_tokens: int = 2) -> QSAIndexer: def test_qsa_fused_prepare_replaces_the_unfused_chain( monkeypatch: pytest.MonkeyPatch, ) -> None: - indexer = _make_indexer(use_fused=True) + indexer = _make_indexer(monkeypatch, use_fused=True) attn = _make_attn(use_fused=True) calls = [] main_outputs = (torch.empty(0), torch.empty(0)) @@ -260,7 +264,7 @@ def fake_prepare(*args, **kwargs): def test_qsa_unsupported_config_keeps_the_unfused_chain( monkeypatch: pytest.MonkeyPatch, ) -> None: - indexer = _make_indexer(use_fused=False) + indexer = _make_indexer(monkeypatch, use_fused=False) monkeypatch.setattr( indexer_qsa, "qsa_prepare", diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index a55cc028f9db..8cea1f83d420 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -431,6 +431,7 @@ def test_clearing_qsa_key_cache_releases_its_storage() -> None: references_before_bind = torch._C._storage_Use_Count(kv.untyped_storage()._cdata) raw_cache.bind_kv_cache(kv) assert raw_cache.key_cache.shape[-1] == 64 + assert raw_cache.rope_position_cache is not None assert raw_cache.rope_position_cache.dtype == torch.int64 clear_layer_kv_caches([raw_cache]) diff --git a/tests/models/qwen4_exp/test_tpsp.py b/tests/models/qwen4_exp/test_tpsp.py index 33e01ab24f76..1d879470ffbc 100644 --- a/tests/models/qwen4_exp/test_tpsp.py +++ b/tests/models/qwen4_exp/test_tpsp.py @@ -48,7 +48,9 @@ from vllm.models.qwen4_exp.amd import model as qwen4_model from vllm.models.qwen4_exp.amd.model import Qwen4ExpDecoderLayer else: - from vllm.models.qwen4_exp.nvidia import model as qwen4_model + from vllm.models.qwen4_exp.nvidia import ( # type: ignore[no-redef] + model as qwen4_model, + ) from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer @@ -220,6 +222,7 @@ def _check_moe_sp(rank: int, config: VllmConfig) -> None: full_hidden, moe_input, injection = layer.mlp_hyper_connection.combine_and_mix( full_hidden, attn_output, injection ) + assert injection is not None expected = tuple( x.clone() for x in (full_hidden, layer.mlp(moe_input), injection) ) diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index bcad62765a76..ecb7ff0b08c8 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/tests/models/transformers/fusers/test_linear.py @@ -4,6 +4,7 @@ import ast import inspect +from collections.abc import MutableMapping from types import MethodType, SimpleNamespace import pytest @@ -838,9 +839,14 @@ def test_glu_identifies_down_projection(): matches the column-parallel merged gate/up; `None` when there is no such projection to force (fusion of gate/up still applies).""" with torch.device("meta"): - assert get_fuser(GLUMLP(), GLUFuser).down_name == "down_proj" - assert get_fuser(ReversedGLUMLP(), GLUFuser).down_name == "down_proj" - assert get_fuser(NoDownGLU(), GLUFuser).down_name is None + for cls, down_name in [ + (GLUMLP, "down_proj"), + (ReversedGLUMLP, "down_proj"), + (NoDownGLU, None), + ]: + fuser = get_fuser(cls(), GLUFuser) + assert fuser is not None + assert fuser.down_name == down_name @pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention]) @@ -891,14 +897,19 @@ def test_detects_and_rewrites_qkv(attn_cls, kv_heads): def test_qkv_identifies_output_projection(): with torch.device("meta"): - assert get_fuser(FakeAttention(), QKVFuser).o_name == "o_proj" - assert get_fuser(ReversedFakeAttention(), QKVFuser).o_name == "o_proj" - assert get_fuser(ExtraProjAttention(), QKVFuser).o_name == "o_proj" - # Norm children (q_norm/k_norm) must not disturb o_proj identification. - assert get_fuser(QKNormAttention(), QKVFuser).o_name == "o_proj" - assert get_fuser(PerHeadQKNormAttention(), QKVFuser).o_name == "o_proj" - # A module between o_proj and the return is transparent. - assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj" + for cls in [ + FakeAttention, + ReversedFakeAttention, + ExtraProjAttention, + # Norm children must not disturb o_proj identification. + QKNormAttention, + PerHeadQKNormAttention, + # A module between o_proj and the return is transparent. + ResidDropoutAttention, + ]: + fuser = get_fuser(cls(), QKVFuser) + assert fuser is not None + assert fuser.o_name == "o_proj" @pytest.mark.parametrize( @@ -1265,6 +1276,7 @@ def test_fuser_is_cached_per_class_and_structure(): fuser_a = get_fuser(GLUMLP(), GLUFuser) fuser_b = get_fuser(GLUMLP(), GLUFuser) assert fuser_a is fuser_b + assert get_fusers.cache is not None assert any(key[0] is GLUMLP for key in get_fusers.cache) @@ -1297,8 +1309,10 @@ def test_weight_mappings_are_scoped_to_fused_prefixes(): with torch.device("meta"): glu_fuser = get_fuser(GLUMLP(), GLUFuser) qkv_fuser = get_fuser(FakeAttention(), QKVFuser) + assert glu_fuser is not None and qkv_fuser is not None mapper = WeightsMapper() + assert isinstance(mapper.orig_to_new_stacked, MutableMapping) for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"): mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix)) mapper.orig_to_new_stacked.update( diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index 284b1d6a094e..7a8d801765c4 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -493,6 +493,7 @@ def test_moe_fuser_reads_router_dtype_from_the_gate(): with torch.device("meta"): fp32 = MoEBlockFuser.match(MoEBlock(Fp32Router), "experts") default = MoEBlockFuser.match(MoEBlock(TopKRouter), "experts") + assert fp32 is not None and default is not None assert fp32.router_dtype == torch.float32 assert default.router_dtype is None @@ -566,6 +567,7 @@ def test_moe_fuser_unrecognized_router_moves_shared_expert(): with torch.device("meta"): block = MoEBlockSharedBiasedRouter() fuser = MoEBlockFuser.match(block, "experts") + assert fuser is not None assert fuser.scoring_func is None assert fuser.shared_name == "shared_expert" @@ -580,6 +582,7 @@ def test_moe_fuser_unrecognized_router_moves_shared_expert(): def test_moe_fuser_divided_shared_output(block_cls, scoring_func): with torch.device("meta"): fuser = MoEBlockFuser.match(block_cls(), "experts", divides_shared_output=True) + assert fuser is not None assert fuser.scoring_func == scoring_func @@ -600,6 +603,7 @@ def test_moe_fuser_scaled_output_with_vllm_routing(): """A scaled block fuses because vLLM routing applies `routed_scaling_factor`.""" with torch.device("meta"): fuser = MoEBlockFuser.match(MoEBlockScaled(), "experts") + assert fuser is not None assert fuser.scoring_func == "softmax" assert fuser.shared_name == "shared_expert" @@ -608,4 +612,5 @@ def test_moe_fuser_scaled_output_with_vllm_routing(): def test_moe_fuser_looks_through_casts_and_identity(block_cls): with torch.device("meta"): fuser = MoEBlockFuser.match(block_cls(), "experts") + assert fuser is not None assert fuser.shared_name == "shared_expert" diff --git a/tests/models/transformers/fusers/test_rms_norm.py b/tests/models/transformers/fusers/test_rms_norm.py index 256f83649a48..142c8075f454 100644 --- a/tests/models/transformers/fusers/test_rms_norm.py +++ b/tests/models/transformers/fusers/test_rms_norm.py @@ -255,6 +255,7 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c with torch.device("meta"): module = cls() fuser = get_fuser(module, RMSNormFuser) + assert fuser is not None built = fuser.fuse(module, "norm", default_vllm_config) from vllm.model_executor.models.transformers.fusers.rms_norm import ( TPAwareNormMixin, @@ -283,7 +284,9 @@ def test_pow_spelled_norm_fuses_to_equivalent_math(cls, default_vllm_config): module = cls(16) with torch.no_grad(): module.weight.copy_(torch.randn(16)) - built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config) + fuser = get_fuser(module, RMSNormFuser) + assert fuser is not None + built = fuser.fuse(module, "norm", default_vllm_config) with torch.no_grad(): built.weight.copy_(module.weight) x = torch.randn(4, 16) @@ -298,7 +301,9 @@ def test_weightless_norm_has_no_hidden_size(default_vllm_config): matching the unfused norm and vLLM's native `RMSNorm(hidden_size=head_dim)`. """ module = WeightlessRMSNorm(72) - built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config) + fuser = get_fuser(module, RMSNormFuser) + assert fuser is not None + built = fuser.fuse(module, "norm", default_vllm_config) assert built.hidden_size == 0 built.tp_size = 2 # emulate TP=2 without a real process group @@ -329,6 +334,7 @@ def test_eps_is_derived_per_instance(default_vllm_config): for eps in (1e-5, 1e-6): module = RMSNorm(16, eps=eps) fuser = get_fuser(module, RMSNormFuser) + assert fuser is not None built = fuser.fuse(module, "norm", default_vllm_config) assert built.variance_epsilon == eps @@ -341,6 +347,7 @@ def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config): for eps in (1e-5, 1e-6): module = cls(16, eps=eps) fuser = get_fuser(module, RMSNormFuser) + assert fuser is not None assert fuser.eps_attr == cls.attr built = fuser.fuse(module, "norm", default_vllm_config) assert built.variance_epsilon == eps @@ -353,6 +360,7 @@ def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplo with caplog.at_level("DEBUG", logger=logger), torch.device("meta"): module = LiteralEpsRMSNorm() fuser = get_fuser(module, RMSNormFuser) + assert fuser is not None built = fuser.fuse(module, "norm", default_vllm_config) assert fuser.eps_attr is None assert built.variance_epsilon == 1e-4 @@ -366,6 +374,7 @@ def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config): module = AmbiguousEpsRMSNorm(16, eps=1e-6) before = dict(vars(module)) fuser = get_fuser(module, RMSNormFuser) + assert fuser is not None assert fuser.eps_attr == "variance_epsilon" assert vars(module) == before diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index dd7cae111fb1..ef76ecd754d4 100644 --- a/tests/models/transformers/test_backend.py +++ b/tests/models/transformers/test_backend.py @@ -694,8 +694,11 @@ def test_attention_layer_index_is_the_modules_own(): attentions = [ model.get_submodule(f"layers.{i}.self_attn.{j}") for i in (0, 1) for j in (0, 1) ] - fusers = [AttentionFuser.match(None, attn) for attn in attentions] - assert all(fuser is not None for fuser in fusers) + fusers: list[AttentionFuser] = [] + for attn in attentions: + fuser = AttentionFuser.match(None, attn) + assert fuser is not None + fusers.append(fuser) assert [f.layer_index(a) for f, a in zip(fusers, attentions)] == [0, 1, 2, 3] # Not dispatching to vLLM yet, as a vision tower never would be. The vLLM diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 5c90eee54963..61dd74669a9d 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -49,9 +49,17 @@ "tests/kernels", "tests/lora", "tests/model_executor", + "tests/models/deepseek_v32", + "tests/models/glm5next", + "tests/models/inkling", + "tests/models/kimi_k3", "tests/models/language", + "tests/models/minimax_m3", "tests/models/quantization", + "tests/models/qwen4_exp", + "tests/models/step5", "tests/models/test_initialization.py", + "tests/models/transformers", "tests/multimodal", "tests/parser", "tests/plugins/bge_m3_sparse_plugin", From cdd2750cd5e7d90849282a16bc46c21bbc58bcbd Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 28 Oct 2026 13:20:41 -0700 Subject: [PATCH 079/406] [MRV2][DBO] Support DBO with speculative decoding and DCP (#61464) Signed-off-by: Nick Hill Signed-off-by: Lucas Wilkinson Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Lucas Wilkinson Co-authored-by: vllm-agent --- tests/test_config.py | 17 +++++++ tests/v1/distributed/test_dbo.py | 46 +++++++++++++++++- ...test_gpu_target_dependent_ar_speculator.py | 36 ++++++++++++++ tests/v1/worker/test_gpu_ubatch_slicing.py | 47 +++++++++++++++++++ vllm/config/vllm.py | 40 ++++++++++++---- vllm/v1/worker/gpu/cudagraph_utils.py | 2 +- vllm/v1/worker/gpu/model_runner.py | 3 +- .../gpu/spec_decode/extract_hidden_states.py | 4 ++ vllm/v1/worker/gpu/spec_decode/speculator.py | 43 +++++++++++++++++ .../target_dependent_ar/speculator.py | 14 +++++- vllm/v1/worker/gpu/ubatch_utils.py | 47 +++++++++++++++---- 11 files changed, 276 insertions(+), 23 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index a53843e2a93a..97c20a19fbf7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1926,6 +1926,23 @@ def test_deepseek_engram_rejects_microbatching(explicit, enable_dbo, ubatch_size VllmConfig._resolve_and_verify_engram_config(config) +@pytest.mark.parametrize("threshold,expected", [(7, 8), (8, 8), (16, 16)]) +def test_dbo_thresholds_cover_a_request_per_microbatch(threshold, expected): + # Verifying 3 drafts: 4 tokens per request, so 2 microbatches need 8. + parallel_config = ParallelConfig( + enable_dbo=True, + dbo_decode_token_threshold=threshold, + dbo_prefill_token_threshold=threshold, + ) + config = cast( + VllmConfig, + SimpleNamespace(uniform_decode_query_len=4, parallel_config=parallel_config), + ) + VllmConfig._raise_dbo_thresholds_for_spec_decode(config) + assert parallel_config.dbo_decode_token_threshold == expected + assert parallel_config.dbo_prefill_token_threshold == expected + + def test_engram_explicit_config_requires_supported_model(): """Explicit all-false settings still opt into model validation.""" with pytest.raises(ValueError, match="requires a model with supported Engram"): diff --git a/tests/v1/distributed/test_dbo.py b/tests/v1/distributed/test_dbo.py index 632bdfff6755..6f2dfc6a2658 100644 --- a/tests/v1/distributed/test_dbo.py +++ b/tests/v1/distributed/test_dbo.py @@ -58,11 +58,51 @@ ) def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available): """Test DBO with DP+EP using GSM8K evaluation.""" + _run_dbo_gsm8k(all2all_backend, num_gpus_available) + + +# Combinations the V2 runner supports only (DBO picks V1 unless told otherwise). +# ngram verifies multi-token queries without a draft model; DCP shards the KV +# cache across each DP rank's TP group. +V2_FEATURES = { + "spec_decode": ( + [ + "--speculative-config", + '{"method": "ngram", "num_speculative_tokens": 2, ' + '"prompt_lookup_max": 4, "prompt_lookup_min": 2}', + ], + 1, + ), + "dcp": (["--tensor-parallel-size", "2", "--decode-context-parallel-size", "2"], 2), +} + + +@pytest.mark.skipif(not has_deep_ep(), reason="These tests require deep_ep to run") +@pytest.mark.parametrize("feature", list(V2_FEATURES)) +def test_dbo_v2_features_gsm8k(feature: str, num_gpus_available): + """Test V2-runner DBO combined with spec decode or DCP.""" + extra_args, tp_size = V2_FEATURES[feature] + _run_dbo_gsm8k( + "deepep_low_latency", + num_gpus_available, + extra_args=extra_args, + tp_size=tp_size, + env_dict={"VLLM_USE_V2_MODEL_RUNNER": "1"}, + ) + + +def _run_dbo_gsm8k( + all2all_backend: str, + num_gpus_available: int, + extra_args: list[str] | None = None, + tp_size: int = 1, + env_dict: dict[str, str] | None = None, +): lm_eval = pytest.importorskip("lm_eval") - required_gpus = DP_SIZE + required_gpus = DP_SIZE * tp_size if num_gpus_available < required_gpus: - pytest.skip(f"Need at least {required_gpus} GPUs (DP={DP_SIZE})") + pytest.skip(f"Need at least {required_gpus} GPUs (DP={DP_SIZE}, TP={tp_size})") # Server arguments for DBO + DP + EP server_args = [ @@ -83,11 +123,13 @@ def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available): "32", "--all2all-backend", all2all_backend, + *(extra_args or []), ] with RemoteOpenAIServer( MODEL_NAME, server_args, + env_dict=env_dict, max_wait_seconds=600, # Allow time for model loading with DP+EP ) as remote_server: # Use host and port directly from RemoteOpenAIServer diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 5e2a3f060dbf..27c841a75754 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py @@ -200,6 +200,7 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected) data_parallel_size=1, data_parallel_rank=0, decode_context_parallel_size=1, + use_ubatching=False, ), ) @@ -679,3 +680,38 @@ def test_decode_metadata_update_needs_draft_decode_steps( speculator._configure_decode_metadata_update() assert speculator.supports_decode_metadata_update is expected + + +def test_rebuilt_prefill_metadata_uses_target_batch_slots(): + # Under DBO the target builds only per-microbatch metadata, so the drafter + # rebuilds the full batch's on its own builders from the target's slots. + speculator = object.__new__(_TestSpeculator) + target = ({"draft": "target metadata"}, {"draft": torch.zeros(6)}) + assert speculator._build_drafter_attn_metadata(None, *target) == target + + speculator.can_reuse_target_attn_metadata = False + speculator.attn_groups = [[Mock()]] + speculator.model_state = Mock() + speculator.model_state.prepare_attn.return_value = {"draft": "metadata"} + speculator.kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(layer_names=["target", "draft"])] + ) + speculator.block_tables = SimpleNamespace( + input_block_tables=[torch.arange(8, dtype=torch.int32).view(4, 2)], + slot_mappings=torch.arange(8).view(1, 8), + ) + input_batch = SimpleNamespace(num_reqs_after_padding=2, num_tokens_after_padding=6) + + attn_metadata, slot_mappings = speculator._build_drafter_attn_metadata( + input_batch, *target, CUDAGraphMode.FULL + ) + + assert attn_metadata == {"draft": "metadata"} + args = speculator.model_state.prepare_attn.call_args.args + batch, cg_mode, (block_table,), target_slots, attn_groups, _ = args + assert batch is input_batch and cg_mode == CUDAGraphMode.FULL + assert attn_groups is speculator.attn_groups + assert block_table.shape == (2, 2) + assert target_slots.data_ptr() == speculator.block_tables.slot_mappings.data_ptr() + assert target_slots.shape == (1, 6) + assert slot_mappings["draft"].tolist() == list(range(6)) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index 406b21f7f2a4..1bf8d3d6c057 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/tests/v1/worker/test_gpu_ubatch_slicing.py @@ -48,6 +48,7 @@ UBatchState, _slice_input_batch, create_ubatch_slices, + merge_ubatch_outputs, slice_model_inputs, ) from vllm.v1.worker.ubatch_utils import ( @@ -662,6 +663,19 @@ def test_ubatch_runner_overlaps_and_matches_single_batch(): ) +def test_merge_keeps_aux_hidden_states_per_layer(): + # EAGLE3 and DFlash targets return (hidden_states, [aux per layer]). + hidden = torch.arange(6.0).view(6, 1) + aux = [torch.arange(6.0).view(6, 1) + 10 * i for i in range(3)] + merged_hidden, merged_aux = merge_ubatch_outputs( + [(hidden[:2], [x[:2] for x in aux]), (hidden[2:], [x[2:] for x in aux])] + ) + torch.testing.assert_close(merged_hidden, hidden) + assert len(merged_aux) == len(aux) + for merged, expected in zip(merged_aux, aux): + torch.testing.assert_close(merged, expected) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU") @pytest.mark.parametrize("failed_ubatch", [0, 1]) def test_ubatch_runner_names_the_microbatch_that_failed(failed_ubatch): @@ -1002,3 +1016,36 @@ def test_microbatched_graph_needs_every_rank_to_reach_the_split(): assert desc.cg_mode == CUDAGraphMode.NONE assert desc.num_ubatches == 2 assert dp_sync_state is not None and dp_sync_state.eager + + +@pytest.mark.parametrize("num_reqs", [3, 4, 5]) +def test_uniform_spec_decode_batches_split_on_request_boundaries(num_reqs: int): + # Verifying 2 drafts per request: 3 tokens each. A midpoint split would cut + # a request in two whenever num_reqs is odd, leaving neither microbatch a + # uniform decode. + query_len = 3 + input_batch = _make_input_batch( + [query_len] * num_reqs, [64] * num_reqs, _make_buffers() + ) + slices = create_ubatch_slices(input_batch, 2, uniform_token_count=query_len) + first, second = slices + assert first.token_slice.stop % query_len == 0 + assert first.request_slice.stop == second.request_slice.start + assert first.num_tokens == query_len * (num_reqs // 2) + assert second.token_slice.stop == query_len * num_reqs + + +def test_single_token_or_unknown_uniform_count_keeps_midpoint_split(): + input_batch = _make_input_batch([3] * 5, [64] * 5, _make_buffers()) + for uniform_token_count in (None, 1): + first, _ = create_ubatch_slices( + input_batch, 2, uniform_token_count=uniform_token_count + ) + assert first.token_slice.stop == 15 // 2 + + +def test_too_few_requests_to_align_keeps_midpoint_split(): + # One request cannot be split on a request boundary. + input_batch = _make_input_batch([4], [64], _make_buffers()) + first, _ = create_ubatch_slices(input_batch, 2, uniform_token_count=4) + assert first.token_slice.stop == 2 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 69b4d64469a2..b2f33c272da0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2204,6 +2204,8 @@ def has_blocked_weights(): " and install the matching kernels." ) + self._raise_dbo_thresholds_for_spec_decode() + if not self.instance_id: self.instance_id = random_uuid()[:5] @@ -3220,12 +3222,35 @@ def _validate_batch_sharded_sampling(self) -> None: f"{'; '.join(blockers)}." ) + def _raise_dbo_thresholds_for_spec_decode(self) -> None: + """The speculative counterpart of ParallelConfig's DBO threshold check. + + A uniform decode batch below one request per microbatch cannot be split + on a request boundary, and an idle DP rank's one-request dummy batch + must stay below the thresholds too. Such batches gain nothing from + microbatching, so raise the thresholds rather than reject the config. + """ + parallel_config = self.parallel_config + min_ubatch_tokens = parallel_config.num_ubatches * self.uniform_decode_query_len + for name in ("dbo_decode_token_threshold", "dbo_prefill_token_threshold"): + threshold = getattr(parallel_config, name) + if threshold < min_ubatch_tokens: + logger.warning_once( + "Raising %s from %d to %d, the number of microbatches times " + "the decode query length.", + name, + threshold, + min_ubatch_tokens, + ) + setattr(parallel_config, name, min_ubatch_tokens) + def _get_dbo_unsupported_features(self) -> list[str]: """Collect what the V2 model runner cannot combine with DBO. - The V2 runner microbatches a plain decoder forward pass. Anything that - slices or replays the batch differently (drafting, adapters, pipeline - stages, context parallelism, encoders) is not handled yet. + The V2 runner microbatches a decoder forward pass (drafting runs whole, + after it). Anything else that slices or replays the batch differently + (adapters, pipeline stages, prefill context parallelism, encoders) is + not handled yet. """ # TODO: DBO with model runner V2 is under development. # It should be enabled with explicit VLLM_USE_V2_MODEL_RUNNER environ. @@ -3239,15 +3264,10 @@ def _get_dbo_unsupported_features(self) -> list[str]: if self.lora_config is not None: unsupported.append("dual batch overlap with LoRA") - if self.speculative_config is not None: - unsupported.append("dual batch overlap with speculative decoding") if parallel_config.pipeline_parallel_size > 1: unsupported.append("dual batch overlap with pipeline parallelism") - if ( - parallel_config.decode_context_parallel_size > 1 - or parallel_config.prefill_context_parallel_size > 1 - ): - unsupported.append("dual batch overlap with context parallelism") + if parallel_config.prefill_context_parallel_size > 1: + unsupported.append("dual batch overlap with prefill context parallelism") if model_config is not None and ( model_config.is_multimodal_model or model_config.is_encoder_decoder ): diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 33555fcc7aee..995657462bf2 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -721,7 +721,7 @@ def create_forward_fn( input_batch, input_block_tables, block_tables.get_dummy_slot_mappings(num_tokens), - cg_mode=CUDAGraphMode.FULL, + desc, for_capture=True, ) # Capture with dummy rows marked as padding. diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 568fd2f496af..74ae2c3c8624 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1976,7 +1976,8 @@ def execute_model( input_batch, block_tables, slot_mappings, - cg_mode=batch_desc.cg_mode, + batch_desc, + dp_sync_state, for_capture=dummy_run and batch_desc.cg_mode == CUDAGraphMode.FULL, ) elif not (dummy_run and skip_attn_for_dummy_run): diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py index e02d7ee11d70..8b4d8d7c3dae 100644 --- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py +++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py @@ -134,6 +134,10 @@ def propose( num_tokens = stacked_hidden_states.shape[0] self.hidden_states[:num_tokens].copy_(stacked_hidden_states) + attn_metadata, slot_mappings = self._build_drafter_attn_metadata( + input_batch, attn_metadata, slot_mappings + ) + draft_attn_metadata = { name: attn_metadata[name] for name in self.draft_attn_layer_names } diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 87fbc3704f10..b3d520b9dc19 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -25,6 +25,7 @@ ) from vllm.v1.worker.gpu.attn_utils import ( build_attn_metadata, + build_slot_mappings_by_layer, init_attn_backend, ) from vllm.v1.worker.gpu.block_table import BlockTables @@ -106,6 +107,10 @@ def propose( class DraftModelSpeculator(BaseSpeculator): + # Draft prefill runs on the target's batch, so it can reuse the target + # forward's attention metadata. Otherwise the drafter rebuilds it. + can_reuse_target_attn_metadata: bool = True + def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.device = device @@ -145,6 +150,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.dp_size = vllm_config.parallel_config.data_parallel_size self.dp_rank = vllm_config.parallel_config.data_parallel_rank self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size + if vllm_config.parallel_config.use_ubatching: + # A microbatched target forward builds only per-microbatch metadata. + self.can_reuse_target_attn_metadata = False self.eplb_state: EplbState | None = None @@ -344,6 +352,41 @@ def set_attn( self.target_attn_groups = target_attn_groups self._configure_decode_metadata_update() + def _build_drafter_attn_metadata( + self, + input_batch: InputBatch, + attn_metadata: dict[str, Any], + slot_mappings: dict[str, torch.Tensor], + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + for_capture: bool = False, + ) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: + """Attention metadata and slot mappings for draft prefill on the target + batch: the target forward's when reusable, else rebuilt on the draft + builders from the block tables and slot mappings the target prepared. + Must run before draft decode steps overwrite the slot mappings. + """ + if self.can_reuse_target_attn_metadata: + return attn_metadata, slot_mappings + block_tables = tuple( + block_table[: input_batch.num_reqs_after_padding] + for block_table in self.block_tables.input_block_tables + ) + target_slot_mappings = self.block_tables.slot_mappings[ + :, : input_batch.num_tokens_after_padding + ] + attn_metadata = self.model_state.prepare_attn( + input_batch, + cg_mode, + block_tables, + target_slot_mappings, + self.attn_groups, + self.kv_cache_config, + for_capture=for_capture, + ) + return attn_metadata, build_slot_mappings_by_layer( + target_slot_mappings, self.kv_cache_config + ) + def _configure_decode_metadata_update(self) -> None: unsupported_backends = sorted( { diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 95aa5959972e..1c7db7c6970e 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -133,7 +133,9 @@ def capture(self) -> None: self.model_state, self.target_input_buffers, self.block_tables, - self.target_attn_groups, + self.target_attn_groups + if self.can_reuse_target_attn_metadata + else self.attn_groups, self.kv_cache_config, progress_bar_desc="Capturing prefill CUDA graphs", ) @@ -280,6 +282,16 @@ def propose( self._prepare_eplb_forward(num_tokens) + if not skip_attn_for_dummy_run: + attn_metadata, slot_mappings = self._build_drafter_attn_metadata( + input_batch, + attn_metadata, + slot_mappings, + prefill_batch_desc.cg_mode, + for_capture=dummy_run + and prefill_batch_desc.cg_mode == CUDAGraphMode.FULL, + ) + self.on_prefill_begin(num_reqs) if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: # Replay the full graph for draft prefill. diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 608231a3ed42..d8216c22ed10 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/vllm/v1/worker/gpu/ubatch_utils.py @@ -6,7 +6,7 @@ from collections.abc import Callable from contextlib import ExitStack from dataclasses import replace -from typing import Any, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple import numpy as np import torch @@ -36,6 +36,10 @@ from vllm.v1.worker.ubatching import make_ubatch_contexts from vllm.v1.worker.utils import AttentionGroup +if TYPE_CHECKING: + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + from vllm.v1.worker.gpu.dp_utils import DPSyncState + logger = init_logger(__name__) @@ -46,7 +50,9 @@ class UBatchState(NamedTuple): forward_contexts: list[ForwardContext] -def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: +def create_ubatch_slices( + input_batch: InputBatch, num_ubatches: int, uniform_token_count: int | None = None +) -> UBatchSlices: """Split a DP-padded batch into the slices its microbatches run on. Splitting at the midpoint of the DP-padded token count leaves the trailing @@ -54,13 +60,24 @@ def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSl request slice onto the last request leaves it holding that request with zero query tokens -- the shape a straddling request already takes -- so it stays well-formed with no work to do, like a dummy run. + + A uniform batch of multi-token queries (spec decode verification) is split + on a request boundary instead, so each microbatch stays a uniform decode. + `uniform_token_count` must be agreed across DP ranks. """ + split_point = None + if uniform_token_count is not None and uniform_token_count > 1: + num_tokens = input_batch.num_tokens_after_padding + reqs_per_ubatch = num_tokens // (uniform_token_count * num_ubatches) + if reqs_per_ubatch > 0: + split_point = reqs_per_ubatch * uniform_token_count _, ubatch_slices = maybe_create_ubatch_slices( True, input_batch.num_scheduled_tokens, input_batch.num_tokens_after_padding, input_batch.num_reqs_after_padding, num_ubatches, + split_point=split_point, ) assert ubatch_slices is not None return [ @@ -263,8 +280,9 @@ def slice_model_inputs( def merge_ubatch_outputs(outputs: list[Any]) -> Any: """Reassemble the full-batch model output from the per-microbatch ones. - Models return hidden states, a tuple of those plus auxiliary hidden states - (EAGLE3), or IntermediateTensors on non-final pipeline ranks. + Models return hidden states, a tuple of those plus a list of auxiliary + hidden states (EAGLE3, DFlash), or IntermediateTensors on non-final + pipeline ranks. """ first = outputs[0] if isinstance(first, IntermediateTensors): @@ -275,7 +293,9 @@ def merge_ubatch_outputs(outputs: list[Any]) -> Any: } ) if isinstance(first, tuple): - return tuple(torch.cat(parts, dim=0) for parts in zip(*outputs)) + return tuple(merge_ubatch_outputs(list(parts)) for parts in zip(*outputs)) + if isinstance(first, list): + return [merge_ubatch_outputs(list(parts)) for parts in zip(*outputs)] return torch.cat(outputs, dim=0) @@ -337,7 +357,8 @@ def prepare( input_batch: InputBatch, block_tables: tuple[torch.Tensor, ...], slot_mappings: torch.Tensor, - cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + batch_desc: "BatchExecutionDescriptor", + dp_sync_state: "DPSyncState | None" = None, for_capture: bool = False, ) -> UBatchState: """Split the batch into the microbatches the step will run on. @@ -346,9 +367,19 @@ def prepare( contexts the threads install, not in the caller's. FULL uses padded sizes; `for_capture` refreshes metadata from dummy - block tables. + block tables. Eager steps split by the DP-agreed uniform token count in + `dp_sync_state` (see `create_ubatch_slices`); FULL replays the split + its graph captured with. """ - ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches) + cg_mode = batch_desc.cg_mode + uniform_token_count = ( + batch_desc.uniform_token_count + if cg_mode == CUDAGraphMode.FULL or dp_sync_state is None + else dp_sync_state.uniform_token_count + ) + ubatch_slices = create_ubatch_slices( + input_batch, self.num_ubatches, uniform_token_count + ) attn_metadata = [] slot_mappings_by_layer = [] From e10ef091baf01b3ba17b93e49e0a94228c36cbfb Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 29 Oct 2026 05:04:16 +0800 Subject: [PATCH 080/406] [Quantization] Keep deleted and renamed params registered after process_weights_after_loading (#62407) Signed-off-by: aoshen02 Signed-off-by: Aaron Hao Signed-off-by: Aaron Hao Co-authored-by: Claude Opus 5.5 Co-authored-by: Aaron Hao Co-authored-by: Aaron Hao --- .../kernels/quantization/test_triton_w4a16.py | 8 +- tests/kernels/quantization/test_zen_da8w4.py | 6 +- tests/model_executor/test_utils.py | 114 +++++++++++++++++- tests/quantization/test_auto_round.py | 6 +- tests/quantization/test_compressed_tensors.py | 2 +- tests/quantization/test_keep_param_objects.py | 102 ++++++++++++++-- tests/quantization/test_moe_wna16.py | 5 +- tests/quantization/test_online.py | 4 +- tests/quantization/test_per_token_kv_cache.py | 6 +- .../linear/mixed_precision/MPLinearKernel.py | 8 ++ .../kernels/linear/mixed_precision/conch.py | 4 +- .../kernels/linear/mixed_precision/cpu.py | 12 +- .../kernels/linear/mixed_precision/cutlass.py | 3 +- .../linear/mixed_precision/dynamic_4bit.py | 4 +- .../kernels/linear/mixed_precision/marlin.py | 6 +- .../kernels/linear/mixed_precision/xpu.py | 4 - .../linear/mixed_precision/zentorch.py | 39 +++--- .../kernels/linear/nvfp4/humming.py | 18 ++- .../linear/scaled_mm/ScaledMMLinearKernel.py | 3 +- .../kernels/linear/scaled_mm/cpu.py | 4 +- .../kernels/linear/scaled_mm/cutlass.py | 1 - .../kernels/linear/scaled_mm/marlin.py | 3 +- .../kernels/linear/scaled_mm/triton.py | 1 - .../layers/fused_moe/oracle/fp8.py | 14 ++- .../layers/fused_moe/oracle/mxfp4.py | 46 ++++--- .../layers/fused_moe/oracle/nvfp4.py | 9 +- .../layers/fused_moe/routed_experts.py | 10 +- .../compressed_tensors/compressed_tensors.py | 8 -- .../compressed_tensors_moe_w4a4_mxfp4.py | 17 ++- .../compressed_tensors_moe_w4a4_nvfp4.py | 22 ++-- .../compressed_tensors_moe_w4a8_fp8.py | 9 +- .../compressed_tensors_moe_w4a8_int8.py | 57 +++------ .../compressed_tensors_moe_w4a8_mxfp4.py | 5 +- .../schemes/compressed_tensors_w4a4_mxfp4.py | 5 +- .../schemes/compressed_tensors_w4a4_nvfp4.py | 5 +- .../schemes/compressed_tensors_w8a16_fp8.py | 10 +- .../schemes/compressed_tensors_wNa8o8.py | 18 +-- .../layers/quantization/fbgemm_fp8.py | 3 +- .../layers/quantization/humming.py | 46 +++---- .../inc/schemes/inc_mxfp4_linear.py | 5 +- .../quantization/inc/schemes/inc_mxfp4_moe.py | 16 ++- .../inc/schemes/inc_wna16_linear.py | 9 +- .../layers/quantization/kv_cache.py | 22 +--- .../layers/quantization/modelopt.py | 27 +++-- .../layers/quantization/mxfp4.py | 28 ++--- .../layers/quantization/online/moe_base.py | 9 +- .../layers/quantization/online/mxfp4.py | 8 +- .../layers/quantization/quark/quark_moe.py | 47 ++++---- .../quantization/quark/schemes/quark_nvfp4.py | 7 +- .../quark/schemes/quark_w8a8_int8.py | 7 +- .../quantization/utils/humming/linear.py | 108 +++++++++++++---- .../layers/quantization/utils/humming/moe.py | 60 ++++----- vllm/model_executor/offloader/prefetch.py | 17 ++- vllm/model_executor/utils.py | 103 +++++++++++++++- vllm/models/deepseek_v4/cpu/cpu_sparse.py | 13 +- 55 files changed, 736 insertions(+), 397 deletions(-) diff --git a/tests/kernels/quantization/test_triton_w4a16.py b/tests/kernels/quantization/test_triton_w4a16.py index bd8bd3cae23a..0b37f67ae457 100644 --- a/tests/kernels/quantization/test_triton_w4a16.py +++ b/tests/kernels/quantization/test_triton_w4a16.py @@ -319,13 +319,15 @@ class DummyLayer(torch.nn.Module): ), ) - kernel.process_weights_after_loading(layer) - - # Expected transformed layouts. + # Expected transformed layouts, taken before processing: `replace_parameter` + # writes a same-size result into the loaded parameter's own storage, which + # the checkpoint tensors above still alias. expected_w_kn8 = _pack_int4_along_n(w_int4_kn) # [K, N//8] expected_scales_gn = scales_ckpt_nkg.t().contiguous() # [K//G, N] expected_zeros_gn8 = zeros_ckpt_n8kg.t().contiguous() # [K//G, N//8] + kernel.process_weights_after_loading(layer) + assert tuple(layer.weight_packed.shape) == (K, N // 8) assert tuple(layer.weight_scale.shape) == (K // G, N) assert tuple(layer.weight_zero_point.shape) == (K // G, N // 8) diff --git a/tests/kernels/quantization/test_zen_da8w4.py b/tests/kernels/quantization/test_zen_da8w4.py index 924b7a865f4e..c82f1f9dd9b3 100644 --- a/tests/kernels/quantization/test_zen_da8w4.py +++ b/tests/kernels/quantization/test_zen_da8w4.py @@ -272,9 +272,9 @@ def test_da8w4_process_weights_layout(mock_zentorch_ops, monkeypatch): ) assert layer._zentorch_da8w4_packed.is_contiguous() assert layer._zentorch_da8w4_scale.is_contiguous() - # The checkpoint parameters are released once repacked. - assert layer.weight_packed.numel() == 0 - assert layer.weight_scale.numel() == 0 + # The repacked tensors live in the checkpoint parameters' storage. + assert layer._zentorch_da8w4_packed.data_ptr() == layer.weight_packed.data_ptr() + assert layer._zentorch_da8w4_scale.data_ptr() == layer.weight_scale.data_ptr() def test_da8w4_apply_weights_matches_dequantized_reference( diff --git a/tests/model_executor/test_utils.py b/tests/model_executor/test_utils.py index 6c098a95f663..b25b042d3e4a 100644 --- a/tests/model_executor/test_utils.py +++ b/tests/model_executor/test_utils.py @@ -13,7 +13,14 @@ PackedvLLMParameter, permute_param_layout_, ) -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import ( + get_live_param, + kernel_parameters, + move_parameter, + replace_parameter, + retire_parameter, + set_kernel_parameter, +) @pytest.fixture @@ -44,6 +51,7 @@ def test_replace_parameter_preserves_custom_attribute(wrap_in_parameter: bool) - "weight", torch.nn.Parameter(torch.zeros(4, 4), requires_grad=False) ) original = layer.weight + original_ptr = original.data_ptr() new_data: torch.Tensor = torch.ones(4, 4).t() if wrap_in_parameter: @@ -63,18 +71,30 @@ def test_replace_parameter_preserves_custom_attribute(wrap_in_parameter: bool) - assert layer.weight is original assert layer.weight.stride() == new_data.stride() - assert layer.weight.data_ptr() == new_data.data_ptr() + # same byte size: copied into the loaded storage, address kept + assert layer.weight.data_ptr() == original_ptr @pytest.mark.parametrize( "new_data,copies", [ (torch.ones(4, 4, dtype=torch.float32, device="cpu"), True), - (torch.ones(4, 4, dtype=torch.float32, device="cpu").t(), False), + (torch.ones(4, 4, dtype=torch.float32, device="cpu").t(), True), (torch.ones(8, 4, dtype=torch.float32, device="cpu"), False), - (torch.ones(4, 4, dtype=torch.bfloat16, device="cpu"), False), + # A result that fits (equal or smaller) goes into the loaded storage, + # which never shrinks. + (torch.ones(4, 4, dtype=torch.bfloat16, device="cpu"), True), + (torch.ones(4, 2, dtype=torch.float32, device="cpu"), True), + (torch.tensor(1.0, device="cpu"), True), + ], + ids=[ + "same-layout-copies", + "new-strides-copies", + "larger-shape", + "smaller-dtype-copies", + "smaller-shape-copies", + "scalar-copies", ], - ids=["same-layout-copies", "new-strides", "new-shape", "new-dtype"], ) def test_replace_parameter_keeps_object_by_default( single_rank_tp: None, new_data: torch.Tensor, copies: bool @@ -102,6 +122,42 @@ def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: assert torch.equal(layer.weight.data, new_data) # Copying in place keeps addresses captured in CUDA graphs valid. assert (layer.weight.data_ptr() == old_data_ptr) == copies + assert layer.weight.untyped_storage().nbytes() >= 4 * 4 * 4 + + +def test_replace_parameter_online_quantization_frees_source() -> None: + layer = torch.nn.Module() + layer.quant_method = type("Online", (), {"uses_meta_device": True})() + layer.register_parameter("weight", torch.nn.Parameter(torch.zeros(4, 4))) + + replace_parameter(layer, "weight", torch.ones(4, 4, dtype=torch.bfloat16)) + + assert layer.weight.untyped_storage().nbytes() == 4 * 4 * 2 + + +def test_humming_replace_keeps_loaded_params(single_rank_tp: None) -> None: + """A new tensor in a leftover's storage, or taking a large leftover's place, + is written to that leftover param, registered under both names.""" + from vllm.model_executor.layers.quantization.utils.humming.linear import ( + replace_layer_parameters, + ) + + layer = torch.nn.Module() + packed = torch.nn.Parameter(torch.ones(64, 8), requires_grad=False) + scale_inv = torch.nn.Parameter(torch.ones(4), requires_grad=False) + layer.register_parameter("weight_packed", packed) + layer.register_parameter("weight_scale_inv", scale_inv) + tensors = { + "weight": torch.full((32, 8), 2.0), + "weight_scale": scale_inv.data.view(2, 2), + } + + replace_layer_parameters(layer, tensors, list(layer._parameters)) + + assert layer.weight is packed and layer.weight_scale is scale_inv + assert torch.equal(layer.weight, torch.full((32, 8), 2.0)) + assert packed.untyped_storage().nbytes() == 64 * 8 * packed.element_size() + assert set(kernel_parameters(layer)) == {"weight", "weight_scale"} def test_replace_parameter_preserves_weight_loader() -> None: @@ -240,6 +296,54 @@ def test_replace_parameter_registers_new_param_for_requantization_source() -> No assert torch.equal(layer.weight, torch.ones(4, 4)) +def test_move_parameter_aliases_one_param(single_rank_tp: None) -> None: + """Both names refer to the loaded param; a later replace updates both.""" + + def weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + pass + + layer = torch.nn.Module() + src = ModelWeightParameter( + data=torch.ones(4, 4), input_dim=1, output_dim=0, weight_loader=weight_loader + ) + layer.register_parameter("weight_packed", src) + + move_parameter(layer, "weight_packed", "weight") + replace_parameter(layer, "weight", torch.full((2, 8), 2.0)) + + assert layer.weight_packed is src and layer.weight is src + assert torch.equal(src.data, torch.full((2, 8), 2.0)) + assert src.weight_loader is weight_loader + assert src._input_dim == 1 and src._output_dim == 0 + assert list(kernel_parameters(layer)) == ["weight"] + + # Kernels writing through the checkpoint name make it the runtime name again. + set_kernel_parameter(layer, "weight_packed", torch.ones(2, 8)) + assert list(kernel_parameters(layer)) == ["weight_packed"] + + +def test_move_parameter_rejects_another_param(single_rank_tp: None) -> None: + layer = torch.nn.Module() + layer.register_parameter("a", torch.nn.Parameter(torch.ones(2))) + layer.register_parameter("b", torch.nn.Parameter(torch.ones(2))) + with pytest.raises(ValueError): + move_parameter(layer, "a", "b") + + +def test_retired_param_keeps_data_but_is_not_live(single_rank_tp: None) -> None: + layer = torch.nn.Module() + layer.register_parameter("scale", torch.nn.Parameter(torch.ones(2))) + layer.register_parameter("weight", torch.nn.Parameter(torch.ones(4))) + + retire_parameter(layer, "scale") + + assert torch.equal(layer.scale, torch.ones(2)) + assert get_live_param(layer, "scale") is None + assert list(kernel_parameters(layer)) == ["weight"] + set_kernel_parameter(layer, "scale", torch.zeros(2)) + assert get_live_param(layer, "scale") is layer.scale + + def test_replace_parameter_with_own_parameter_keeps_layout(single_rank_tp) -> None: """Handing back the registered parameter keeps it unchanged.""" diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 8869b8d177ed..51b47eb0f19e 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -1053,7 +1053,7 @@ def process_weights_after_loading(self, layer) -> None: method.process_weights_after_loading(layer) assert layer.weight.data.data_ptr() == packed_data.data_ptr() - assert not hasattr(layer, "weight_packed") + assert layer.weight_packed is layer.weight assert captured["processed_layer"] is layer @@ -1141,8 +1141,8 @@ def test_inc_mxfp4_moe_method_preserves_checkpoint_packing( assert layer.w13_weight.data.data_ptr() == w13_packed_data.data_ptr() assert layer.w2_weight.data.data_ptr() == w2_packed_data.data_ptr() - assert not hasattr(layer, "w13_weight_packed") - assert not hasattr(layer, "w2_weight_packed") + assert layer.w13_weight_packed is layer.w13_weight + assert layer.w2_weight_packed is layer.w2_weight assert captured["quant_config_kwargs"]["w1_scale"] is layer.w13_weight_scale assert captured["quant_config_kwargs"]["w2_scale"] is layer.w2_weight_scale assert captured["kernel_kwargs"]["moe_quant_config"] is expected_quant_config diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 8120cc35f128..a2b7ec315307 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -1109,7 +1109,7 @@ def test_mxfp4_moe_with_fp8_inputs_shares_native_mxfp4_method(monkeypatch): layer.quant_method.process_weights_after_loading(layer) assert layer.w13_weight is w13 - assert not hasattr(layer, "w13_weight_packed") + assert layer.w13_weight_packed is w13 @pytest.mark.skipif( diff --git a/tests/quantization/test_keep_param_objects.py b/tests/quantization/test_keep_param_objects.py index 94b8110c3b7a..5d98511a8667 100644 --- a/tests/quantization/test_keep_param_objects.py +++ b/tests/quantization/test_keep_param_objects.py @@ -11,6 +11,9 @@ from vllm.model_executor.layers.fused_moe import FusedMoEFactory from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsConfig, +) from vllm.model_executor.layers.quantization.fp8 import Fp8Config from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, @@ -25,6 +28,7 @@ from vllm.model_executor.model_loader.reload.meta import materialize_meta_tensor from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.platforms import current_platform +from vllm.utils.import_utils import has_humming pytestmark = pytest.mark.skipif( not current_platform.is_device_capability_family(100), @@ -46,14 +50,28 @@ def _load(layer: torch.nn.Module) -> None: param.data.copy_(torch.rand(param.shape) + 0.5) +def _nbytes(param: torch.Tensor) -> int: + return param.numel() * param.element_size() + + def _dims(param: torch.Tensor) -> tuple: return tuple( getattr(param, k, None) for k in ("input_dim", "output_dim", "packed_dim") ) -def _linear(quant_config: QuantizationConfig, output_size: int = SIZE): - layer = ReplicatedLinear(SIZE, output_size, bias=False, quant_config=quant_config) +def _linear( + quant_config: QuantizationConfig, + output_size: int = SIZE, + params_dtype: torch.dtype | None = None, +): + layer = ReplicatedLinear( + SIZE, + output_size, + bias=False, + quant_config=quant_config, + params_dtype=params_dtype, + ) return layer, layer.quant_method @@ -102,7 +120,19 @@ def _online_moe(): return layer, method +def _ct(fmt: str, group: str, group_size: int, input_activations: bool): + args = dict(num_bits=4, type="float", strategy=group, group_size=group_size) + args |= dict(symmetric=True, dynamic=False) + acts = args | dict(dynamic="local") if input_activations else None + group_0 = dict(targets=["Linear"], weights=args, input_activations=acts) + return CompressedTensorsConfig.from_config( + dict(format=fmt, config_groups=dict(group_0=group_0), ignore=[]) + ) + + NVFP4 = dict(is_checkpoint_nvfp4_serialized=True, kv_cache_quant_algo=None) +CT_NVFP4 = ("nvfp4-pack-quantized", "tensor_group", 16, True) +CT_MXFP4 = ("mxfp4-pack-quantized", "group", 32, False) CASES: list[Callable] = [ lambda: _linear(Fp8Config(weight_block_size=[128, 128])), @@ -117,11 +147,18 @@ def _online_moe(): output_size=SIZE + 64, ), lambda: _linear(ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})), + # The consumed per-tensor scales stay registered with their data. + lambda: _linear(ModelOptNvFp4Config(**NVFP4, exclude_modules=[])), + # The packed weights are also registered under their runtime names. + lambda: _linear(_ct(*CT_NVFP4)), + lambda: _linear(_ct(*CT_MXFP4)), _online_linear, lambda: _moe(None), lambda: _moe(Fp8Config(weight_block_size=[128, 128])), lambda: _moe(ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"})), lambda: _moe(ModelOptNvFp4Config(**NVFP4, exclude_modules=[])), + lambda: _moe(_ct(*CT_NVFP4)), + lambda: _moe(_ct(*CT_MXFP4)), lambda: _moe(Mxfp4Config()), _online_moe, ] @@ -129,34 +166,81 @@ def _online_moe(): "fp8-block-linear", "modelopt-fp8-pb-wo-linear", "modelopt-mxfp8-linear", + "modelopt-nvfp4-linear", + "ct-nvfp4-linear", + "ct-mxfp4-linear", "online-mxfp4-linear", "unquantized-moe", "fp8-block-moe", "modelopt-mxfp8-moe", "modelopt-nvfp4-moe", + "ct-nvfp4-moe", + "ct-mxfp4-moe", "mxfp4-moe", "online-mxfp4-moe", ] -@pytest.mark.parametrize("build", CASES, ids=IDS) -def test_process_weights_after_loading_keeps_param_objects( - build: Callable, default_vllm_config, dist_init -) -> None: - default_vllm_config.model_config = ModelConfig() +def _check_keeps_param_objects(build: Callable, vllm_config) -> None: + vllm_config.model_config = ModelConfig() with torch.device("cuda"): layer, method = build() _load(layer) before = { - name: (param, set(param.__dict__), _dims(param)) + name: (param, set(param.__dict__), _dims(param), _nbytes(param)) for name, param in layer.named_parameters(recurse=False) } method.process_weights_after_loading(layer) after = dict(layer.named_parameters(recurse=False, remove_duplicate=False)) - for name, (param, attrs, dims) in before.items(): + for name, (param, attrs, dims, nbytes) in before.items(): assert after.get(name) is param, name assert attrs <= param.__dict__.keys(), name # The loader shards by these on reload, so they keep the checkpoint layout. assert _dims(param) == dims, name assert callable(getattr(param, "weight_loader", None)), name + # Online quantization frees the high-precision source by design. + if not method.uses_meta_device: + kept = after[name] + offset = kept.storage_offset() * kept.element_size() + assert kept.untyped_storage().nbytes() - offset >= nbytes, name + # Two names share storage only as one param object. + objs = {id(p): p for p in after.values() if p.numel()} + ptrs = [p.untyped_storage().data_ptr() for p in objs.values()] + assert len(ptrs) == len(set(ptrs)) + + +@pytest.mark.parametrize("build", CASES, ids=IDS) +def test_process_weights_after_loading_keeps_param_objects( + build: Callable, default_vllm_config, dist_init +) -> None: + _check_keeps_param_objects(build, default_vllm_config) + + +FORCED_CASES: list[Callable] = [ + lambda: _linear( + ModelOptNvFp4Config(**NVFP4, exclude_modules=[]), params_dtype=torch.bfloat16 + ), + lambda: _linear(_ct(*CT_NVFP4), params_dtype=torch.bfloat16), + lambda: _moe(ModelOptNvFp4Config(**NVFP4, exclude_modules=[])), + lambda: _moe(_ct(*CT_NVFP4)), +] +FORCED_IDS = [ + "modelopt-nvfp4-linear", + "ct-nvfp4-linear", + "modelopt-nvfp4-moe", + "ct-nvfp4-moe", +] + + +@pytest.mark.parametrize("backend", ["marlin", "humming"]) +@pytest.mark.parametrize("build", FORCED_CASES, ids=FORCED_IDS) +def test_forced_backend_keeps_param_objects( + build: Callable, backend: str, default_vllm_config, dist_init +) -> None: + """Weight-only backends leave the input scales unused but registered.""" + if backend == "humming" and not has_humming(): + pytest.skip("humming is not installed") + default_vllm_config.kernel_config.linear_backend = backend + default_vllm_config.kernel_config.moe_backend = backend + _check_keeps_param_objects(build, default_vllm_config) diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index c295a195a5f3..c320a7e7e9be 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -384,8 +384,9 @@ def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel(): assert method.moe_kernel is not None assert set(layer.weight_schemas) == {"w13", "w2"} assert set(layer.humming_configs) == {"w13", "w2"} - assert not hasattr(layer, "w13_weight_packed") - assert not hasattr(layer, "w2_weight_packed") + # The converted weights are written to the loaded params (one object each). + assert layer.w13_weight_packed is layer.w13_weight + assert layer.w2_weight_packed is layer.w2_weight assert layer.w13_weight.dtype is torch.int32 assert layer.w2_weight.dtype is torch.int32 diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index b572bb425a6d..9f72cbef5c6c 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -2562,8 +2562,8 @@ def test_online_int8_moe_w2_scale_matches_unsharded(monkeypatch) -> None: def quantize(w2_in, moe_tp_size): layer = torch.nn.Module() - layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False) - layer.w2_weight = torch.nn.Parameter(w2_in, requires_grad=False) + layer.w13_weight = torch.nn.Parameter(w13.clone(), requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2_in.clone(), requires_grad=False) layer.num_experts = layer.local_num_experts = w13.shape[0] method = SimpleNamespace(moe=SimpleNamespace(tp_size=moe_tp_size)) Int8OnlineMoEMethod._quantize_weights(method, layer) diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index e66a577737bf..2647e8b5a8e2 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -525,10 +525,8 @@ def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str): assert layer._k_scale_float == 1.0 assert layer._v_scale_float == 1.0 - assert not hasattr(layer, "k_scale") - assert not hasattr(layer, "v_scale") - assert not hasattr(layer, "q_scale") - assert not hasattr(layer, "prob_scale") + for name in ("k_scale", "v_scale", "q_scale", "prob_scale"): + assert hasattr(layer, name) # =========================================================================== diff --git a/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py b/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py index 2657c2a39cf8..056a89c43348 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py @@ -8,6 +8,7 @@ import torch from vllm.model_executor.layers.quantization.utils import replace_parameter +from vllm.model_executor.utils import retire_parameter from vllm.scalar_type import ScalarType @@ -69,6 +70,13 @@ def apply_weights( ) -> torch.Tensor: raise NotImplementedError + def _retire_unused_zp(self, layer: torch.nn.Module) -> None: + name = self.w_zp_name + if name is not None and getattr(layer, name, None) is not None: + # Loaded but unused: keep the param and its data. + retire_parameter(layer, name) + self.w_zp_name = None + def _transform_param( self, layer: torch.nn.Module, name: str | None, fn: Callable ) -> None: diff --git a/vllm/model_executor/kernels/linear/mixed_precision/conch.py b/vllm/model_executor/kernels/linear/mixed_precision/conch.py index 7f9910a76a23..efb1473a9945 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/conch.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/conch.py @@ -113,8 +113,8 @@ def transform_w_zp(x): self._transform_param(layer, self.w_s_name, transform_w_s) if self.config.zero_points: self._transform_param(layer, self.w_zp_name, transform_w_zp) - elif self.w_zp_name is not None: - layer.register_parameter(self.w_zp_name, None) + else: + self._retire_unused_zp(layer) def apply_weights( self, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py index e2f8dbac7fd9..b869e5cc85ef 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py @@ -101,7 +101,7 @@ def _process_gptq_weights_w4a8(self, layer: torch.nn.Module): packed_zp = getattr(layer, self.w_zp_name) else: # w4a8 kernel always requires zp, allocate a fake zp - assert self.w_zp_name + self.w_zp_name = "w_zp" replace_parameter( layer, self.w_zp_name, @@ -138,13 +138,13 @@ def _process_gptq_weights_w4a8(self, layer: torch.nn.Module): if layer.bias is not None: layer.bias.data = layer.bias.float() - packed_weight.data = blocked_w - scales.data = blocked_s - packed_zp.data = blocked_zp + replace_parameter(layer, self.w_q_name, blocked_w) + replace_parameter(layer, self.w_s_name, blocked_s) + replace_parameter(layer, self.w_zp_name, blocked_zp) def process_weights_after_loading(self, layer: torch.nn.Module): - if (not self.config.zero_points) and (self.w_zp_name is not None): - setattr(layer, self.w_zp_name, None) + if not self.config.zero_points: + self._retire_unused_zp(layer) weights = getattr(layer, self.w_q_name) # Require GPTQ pack format diff --git a/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py b/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py index ea867cafd063..64c920ef53c4 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/cutlass.py @@ -12,6 +12,7 @@ convert_packed_uint4b8_to_signed_int4_inplace, ) from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.scalar_type import scalar_types @@ -88,7 +89,7 @@ def transform_w_s(x): w_s = getattr(layer, self.w_s_name) fp8_scales, chan_scales = convert_bf16_scales_to_fp8(self.quant_fp8, w_s.data) - w_s.data = fp8_scales + replace_parameter(layer, self.w_s_name, fp8_scales) # register per-channel scales layer.register_parameter( diff --git a/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py b/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py index d0515027628e..0867968360a4 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py @@ -5,6 +5,7 @@ import torch from vllm.model_executor.layers.quantization.utils import replace_parameter +from vllm.model_executor.utils import retire_parameter from vllm.platforms import CpuArchEnum, current_platform from vllm.scalar_type import scalar_types @@ -111,7 +112,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module): replace_parameter( layer, self.w_q_name, torch.nn.Parameter(w, requires_grad=False) ) - setattr(layer, self.w_s_name, None) + # Scales are packed into the weight; the loaded param keeps its data. + retire_parameter(layer, self.w_s_name) def apply_weights( self, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py index 60687fc8602d..44362f7dfddd 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/marlin.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/marlin.py @@ -21,6 +21,7 @@ unpack_cols, ) from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.scalar_type import scalar_types @@ -169,7 +170,10 @@ def transform_w_s(x): ), ) else: - setattr(layer, self.w_zp_name, marlin_make_empty(device)) + # Marlin takes an empty zp; a loaded unused one keeps its data. + self._retire_unused_zp(layer) + self.w_zp_name = "w_zp" + replace_parameter(layer, self.w_zp_name, marlin_make_empty(device)) self._transform_param(layer, self.w_q_name, transform_w_q) self._transform_param(layer, self.w_s_name, transform_w_s) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py index 3cfbba243ab3..fa8f4afbb90d 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/xpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/xpu.py @@ -181,10 +181,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: torch.nn.Parameter(packed, requires_grad=False), ) - # Free the original unpacked int8 weight (still registered as "weight") - # to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory. - delattr(layer, "weight") - def apply_weights( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py index c62a70507682..5975ddb479db 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/zentorch.py @@ -8,10 +8,12 @@ """ import torch +from torch.nn import Parameter from vllm import envs from vllm.logger import init_logger from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.scalar_type import scalar_types @@ -21,6 +23,17 @@ logger = init_logger(__name__) +def _store_in_param(layer: torch.nn.Module, attr: str, name: str | None) -> None: + """Back kernel tensor ``layer.`` by the storage of loaded param ``name``.""" + tensor = getattr(layer, attr) + if name is None or tensor is None: + return + if not isinstance(getattr(layer, name, None), Parameter): + return + replace_parameter(layer, name, tensor) + setattr(layer, attr, getattr(layer, name).data) + + def _import_unpack_from_int32(): """Import compressed-tensors' ``unpack_from_int32`` across versions.""" try: @@ -144,8 +157,7 @@ def _maybe_process_da8w4_weights(self, layer: torch.nn.Module) -> bool: if (in_features // num_groups) % 4 != 0 or in_features % 2 != 0: return False - if self.w_zp_name is not None: - setattr(layer, self.w_zp_name, None) + self._retire_unused_zp(layer) unpack_from_int32 = _import_unpack_from_int32() @@ -165,10 +177,8 @@ def _maybe_process_da8w4_weights(self, layer: torch.nn.Module) -> bool: # CT stores scales as [N, G]; DA8W4 wants per-group {G, N}. layer._zentorch_da8w4_scale = weight_scale.t().to(torch.bfloat16).contiguous() - for param_name in (self.w_q_name, self.w_s_name): - param = getattr(layer, param_name, None) - if param is not None and hasattr(param, "data"): - param.data = torch.empty(0) + _store_in_param(layer, "_zentorch_da8w4_packed", self.w_q_name) + _store_in_param(layer, "_zentorch_da8w4_scale", self.w_s_name) layer._zentorch_kind = "compressed_tensors_w4a8_da8w4" layer._zentorch_da8w4 = True @@ -205,8 +215,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: super().process_weights_after_loading(layer) return - if (not self.config.zero_points) and (self.w_zp_name is not None): - setattr(layer, self.w_zp_name, None) + if not self.config.zero_points: + self._retire_unused_zp(layer) weight_q = getattr(layer, self.w_q_name) weight_s = getattr(layer, self.w_s_name) @@ -255,16 +265,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer._zentorch_woq_scale = weight_scale.t().contiguous() layer._zentorch_woq_zero_point = zp_tc - for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name): - if param_name is None: - continue - param = getattr(layer, param_name, None) - if param is None: - continue - if hasattr(param, "data"): - param.data = torch.empty(0) - else: - setattr(layer, param_name, torch.empty(0)) + _store_in_param(layer, "_zentorch_woq_packed", self.w_q_name) + _store_in_param(layer, "_zentorch_woq_scale", self.w_s_name) + _store_in_param(layer, "_zentorch_woq_zero_point", self.w_zp_name) layer._zentorch_kind = "compressed_tensors_w4a16_gptq" layer._zentorch_processed_weights = True diff --git a/vllm/model_executor/kernels/linear/nvfp4/humming.py b/vllm/model_executor/kernels/linear/nvfp4/humming.py index d6d9cb441b9e..676c3a6c86e7 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/humming.py +++ b/vllm/model_executor/kernels/linear/nvfp4/humming.py @@ -11,7 +11,11 @@ quant_key_to_input_schema, ) from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import ( + move_parameter, + replace_parameter, + set_kernel_parameter, +) from vllm.platforms import current_platform from vllm.utils.humming import create_humming_locks from vllm.utils.import_utils import has_humming @@ -53,16 +57,18 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: "strategy": "group", "group_size": 16, } - # CT pack-quantized reads `weight_packed`; the scheme renamed it to `weight`. - if not hasattr(layer, "weight_packed"): - layer.weight_packed = layer.weight - del layer.weight + # CT pack-quantized reads `weight_packed`: make it the weight's runtime name + # (a CT checkpoint already has it as an alias of `weight`). + move_parameter(layer, "weight", "weight_packed") # The CT linear scheme inverts the global scale (1/scale) for # marlin/cutlass; humming wants the original. replace_parameter(layer, "weight_global_scale", 1.0 / layer.weight_global_scale) input_quant_key = None if hasattr(layer, "input_global_scale"): - replace_parameter( + # Quark keeps its loaded (retired) `input_scale_2`: revive that one. + if "input_scale_2" not in layer._parameters: + move_parameter(layer, "input_global_scale", "input_scale_2") + set_kernel_parameter( layer, "input_scale_2", layer.input_global_scale.reshape(-1) ) input_quant_key = kNvfp4Dynamic diff --git a/vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py b/vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py index e632f6f13f28..20726608b76c 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py @@ -200,6 +200,7 @@ def _get_layer_params(self, layer) -> _Int8ParamsT: getattr(layer, w_q), getattr(layer, w_s), getattr(layer, i_s, None), - getattr(layer, i_zp, None), + # A loaded zero point of a symmetric scheme stays registered, unused. + None if self.config.input_symmetric else getattr(layer, i_zp, None), getattr(layer, azp_adj, None), ) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index 00a1bf06d617..32ce9296e20d 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -142,8 +142,8 @@ def process_weights_for_onednn(self, layer: torch.nn.Module) -> None: 32, ) # weight is prepacked and maintained by the dnnl_handler, - # release the original weight - setattr(layer, w_q_name, None) + # release the original weight storage + replace_parameter(layer, w_q_name, weight.new_empty(0)) del weight def process_weights_for_sgl(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py index 2d4c5fecbdbc..366e62c2d5e2 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cutlass.py @@ -79,7 +79,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: i_s_name, torch.nn.Parameter(input_scale.max(), requires_grad=False), ) - setattr(layer, i_zp_name, None) else: input_zero_point = getattr(layer, i_zp_name) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 1df6ad55412d..c1f5ad949a1b 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -73,11 +73,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter(layer, scale_name, weight_scale.data) # Non-block: callers must pass weight in (K, N) layout. - layer.input_scale = None + # A loaded input_scale stays registered; weight-only Marlin ignores it. prepare_fp8_layer_for_marlin( layer, self.size_k_first, input_dtype=self.marlin_input_dtype ) - del layer.input_scale def apply_weights( self, diff --git a/vllm/model_executor/kernels/linear/scaled_mm/triton.py b/vllm/model_executor/kernels/linear/scaled_mm/triton.py index b9f7e55d5f07..d641c92f0612 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/triton.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/triton.py @@ -73,7 +73,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: i_s_name, torch.nn.Parameter(i_s.max(), requires_grad=False), ) - setattr(layer, i_zp_name, None) else: input_zero_point = getattr(layer, i_zp_name) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 0aed7b16f2ef..8de9f1a21eb3 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -39,7 +39,11 @@ kFp8Dynamic128Sym, kFp8Static128BlockSym, ) -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import ( + move_parameter, + replace_parameter, + set_kernel_parameter, +) from vllm.platforms import current_platform from vllm.utils.deep_gemm import should_auto_disable_deep_gemm from vllm.utils.math_utils import round_up @@ -639,6 +643,14 @@ def convert_to_fp8_moe_kernel_format( convert_to_humming_moe_kernel_format( layer, quant_config=_humming_fp8_weight_schema(layer, w13, w13_scale) ) + if scale_name == "weight_scale_inv": + # One param under both names: the loaded `_inv` holds the runtime scale. + for prefix in ("w13", "w2"): + scale = layer._parameters.pop(f"{prefix}_weight_scale").data + move_parameter( + layer, f"{prefix}_{scale_name}", f"{prefix}_weight_scale" + ) + set_kernel_parameter(layer, f"{prefix}_weight_scale", scale) w13 = layer.w13_weight w2 = layer.w2_weight w13_scale = layer.w13_weight_scale diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 9de3732f00f1..3e838b4e1d53 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -82,6 +82,19 @@ ) +def store_params_in_triton(layer: torch.nn.Module, **tensors) -> None: + """Back each loaded param and its triton_kernels tensor by one storage, so the + param keeps its loaded bytes without a second copy of the data.""" + for name, tensor in tensors.items(): + replace_parameter(layer, name, tensor.storage.data) + tensor.storage.data = getattr(layer, name).data + + +def triton_weight(layer: torch.nn.Module, name: str): + """The triton_kernels tensor stored for weight `name`, else the param.""" + return getattr(layer, f"{name}_triton", getattr(layer, name)) + + def _pack_deepgemm_mxfp4_scales( w13_weight: torch.Tensor, w2_weight: torch.Tensor, @@ -1433,8 +1446,7 @@ def swap_every_two_rows(x, axis=-1): flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2), ) - del layer.w13_weight - del layer.w2_weight + store_params_in_triton(layer, w13_weight=w13_weight, w2_weight=w2_weight) return ( w13_weight, @@ -1467,13 +1479,14 @@ def swap_every_two_rows(x, axis=-1): **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex) ) - # The original mxfp4 block scales have been swizzled into the - # precision configs above and are no longer read by the kernel, so - # drop the now-dead weight/scale Parameters to free their memory. - del layer.w13_weight - del layer.w2_weight - del layer.w13_weight_scale - del layer.w2_weight_scale + # The swizzled weights and scales live in the params' storage. + store_params_in_triton( + layer, + w13_weight=w13_weight, + w2_weight=w2_weight, + w13_weight_scale=w13_scale, + w2_weight_scale=w2_scale, + ) return ( w13_weight, @@ -1929,13 +1942,14 @@ def shuffle_weight(w: torch.Tensor) -> torch.Tensor: **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex) ) - # The original mxfp4 block scales have been swizzled into the - # precision configs above and are no longer read by the kernel, so - # drop the now-dead weight/scale Parameters to free their memory. - del layer.w13_weight - del layer.w2_weight - del layer.w13_weight_scale - del layer.w2_weight_scale + # The swizzled weights and scales live in the params' storage. + store_params_in_triton( + layer, + w13_weight=w13_weight, + w2_weight=w2_weight, + w13_weight_scale=w13_scale, + w2_weight_scale=w2_scale, + ) return ( w13_weight, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 80aa9122f447..62dceb99b324 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -36,6 +36,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, ) +from vllm.model_executor.utils import move_parameter logger = init_logger(__name__) @@ -441,10 +442,10 @@ def convert_to_nvfp4_moe_kernel_format( "strategy": "group", "group_size": 16, } - # CT pack-quantized reads `weight_packed`; the method renamed it to - # `weight`. Re-alias (convert replaces all params anyway). - layer.w13_weight_packed = layer.w13_weight - layer.w2_weight_packed = layer.w2_weight + # CT pack-quantized reads `weight_packed`; the method aliased it as + # `weight`. Make the loaded `*_weight_packed` name the runtime one. + move_parameter(layer, "w13_weight", "w13_weight_packed") + move_parameter(layer, "w2_weight", "w2_weight_packed") else: quant_config = {"quant_method": "modelopt", "quant_algo": "nvfp4"} diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 3c4a1ffc7707..3813220908e5 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -31,6 +31,7 @@ QuantizationConfig, resolve_quant_method, ) +from vllm.model_executor.utils import kernel_parameters from vllm.utils.math_utils import cdiv if TYPE_CHECKING: @@ -1347,7 +1348,14 @@ def _maybe_make_contiguous( torch.transpose(p.data, 1, 2), requires_grad=False ) - weights = list(self.named_parameters()) + # Kernel views of the params: live ones only, under their runtime names + # (an aliased checkpoint name such as `w13_weight_packed` is reported as + # `w13_weight`, which the name filters below expect). + weights = [ + (f"{prefix}.{name}" if prefix else name, p) + for prefix, module in self.named_modules() + for name, p in kernel_parameters(module).items() + ] # EPLB must use shards when original weight storage has been released. # Use explicit metadata: some tensor subclasses lack untyped_storage(). diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index f5b73ade5338..3d17901029bd 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -1271,11 +1271,3 @@ def _to_scalar(tensor: torch.Tensor) -> float: # Sync host (cpu) scale copies read by AITER fused kernels. layer._k_scale_cpu.fill_(layer._k_scale_float) layer._v_scale_cpu.fill_(layer._v_scale_float) - - # Discard all placeholders. - del layer.k_scale - del layer.v_scale - del layer.q_scale - del layer.k_zero_point - del layer.v_zero_point - del layer.q_zero_point diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index 227382ff1d03..7bdcda4bf7b1 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -43,7 +43,11 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( prepare_moe_fp4_layer_for_marlin, ) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import ( + move_parameter, + replace_parameter, + set_weight_attrs, +) from vllm.platforms import current_platform logger = init_logger(__name__) @@ -168,15 +172,8 @@ def get_fused_moe_quant_config( ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: - layer.w13_weight = torch.nn.Parameter( - layer.w13_weight_packed.data, requires_grad=False - ) - delattr(layer, "w13_weight_packed") - - layer.w2_weight = torch.nn.Parameter( - layer.w2_weight_packed.data, requires_grad=False - ) - delattr(layer, "w2_weight_packed") + move_parameter(layer, "w13_weight_packed", "w13_weight") + move_parameter(layer, "w2_weight_packed", "w2_weight") if self.use_cutlass_mxfp4: # Swizzle weight scales from flat checkpoint layout [E, N, K//32] diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index e7fc3f1df365..86c7f8c7ba81 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -33,7 +33,11 @@ kNvfp4Dynamic, kNvfp4Static, ) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import ( + move_parameter, + replace_parameter, + set_weight_attrs, +) # NVFP4 backends that have been verified to support EPLB for this quantization recipe. _EPLB_SUPPORTED_NVFP4_BACKENDS = frozenset( @@ -206,19 +210,9 @@ def dequantize_weight( def process_weights_after_loading(self, layer: RoutedExperts) -> None: """Convert NVFP4 MoE weights into kernel format and setup the kernel.""" - # NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod - # requires this naming convention. However, the name change breaks - # reloading because the state dict no longer matches disk. Once we - # remove MKM, we should revert this change to ensure compatibility. - layer.w13_weight = torch.nn.Parameter( - layer.w13_weight_packed.data, requires_grad=False - ) - delattr(layer, "w13_weight_packed") - - layer.w2_weight = torch.nn.Parameter( - layer.w2_weight_packed.data, requires_grad=False - ) - delattr(layer, "w2_weight_packed") + # MKM requires wN_weight; the loaded wN_weight_packed params stay registered. + move_parameter(layer, "w13_weight_packed", "w13_weight") + move_parameter(layer, "w2_weight_packed", "w2_weight") # Use a single gscale for w13. if self.moe.is_act_and_mul and not torch.allclose( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py index 4a2223aab8bb..67feb5b11c34 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py @@ -25,7 +25,11 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 CompressedTensorsMoEMethod, ) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import ( + move_parameter, + replace_parameter, + set_weight_attrs, +) class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod): @@ -149,8 +153,7 @@ def create_weights( def process_weights_after_loading(self, layer: RoutedExperts) -> None: for name in ("w13", "w2"): - setattr(layer, f"{name}_weight", getattr(layer, f"{name}_weight_packed")) - delattr(layer, f"{name}_weight_packed") + move_parameter(layer, f"{name}_weight_packed", f"{name}_weight") ( w13_weight, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py index f3952afe0dc1..c071fbc93df0 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py @@ -30,7 +30,12 @@ QuantKey, ScaleDesc, ) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import ( + move_parameter, + replace_parameter, + retire_parameter, + set_weight_attrs, +) logger = init_logger(__name__) @@ -203,47 +208,17 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) ) - # Register packed weights as parameters - replace_parameter( - layer, - "w13_weight_packed", - torch.nn.Parameter(w13_packed, requires_grad=False), - ) - replace_parameter( - layer, - "w2_weight_packed", - torch.nn.Parameter(w2_packed, requires_grad=False), - ) - - replace_parameter( - layer, "w13_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False) - ) - replace_parameter( - layer, "w2_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False) - ) - replace_parameter( - layer, - "w13_weight_scale", - torch.nn.Parameter(w13_weight_scale, requires_grad=False), - ) - replace_parameter( - layer, - "w2_weight_scale", - torch.nn.Parameter(w2_weight_scale, requires_grad=False), - ) - + # The packed weights (with scales and biases folded in) go into the loaded + # weight storage, one param under both names; the rest keep their data. + for name, packed in (("w13", w13_packed), ("w2", w2_packed)): + replace_parameter(layer, f"{name}_weight", packed) + del layer._parameters[f"{name}_weight_packed"] # never-loaded placeholder + move_parameter(layer, f"{name}_weight", f"{name}_weight_packed") + for name in ("w13_weight_scale", "w2_weight_scale"): + retire_parameter(layer, name) if self.has_bias: - replace_parameter( - layer, - "w13_bias", - torch.nn.Parameter(w13_bias, requires_grad=False), - ) - if self.has_bias: - replace_parameter( - layer, - "w2_bias", - torch.nn.Parameter(w2_bias, requires_grad=False), - ) + retire_parameter(layer, "w13_bias") + retire_parameter(layer, "w2_bias") quant_config = self.get_fused_moe_quant_config(layer) assert quant_config is not None diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py index ae2c88044f1f..d22b7ac1dd2b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py @@ -3,6 +3,7 @@ from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod +from vllm.model_executor.utils import move_parameter _PACKED_WEIGHTS = ("w13_weight", "w2_weight") @@ -28,7 +29,5 @@ def create_weights(self, layer: RoutedExperts, *args, **kwargs): def process_weights_after_loading(self, layer: RoutedExperts) -> None: for name in _PACKED_WEIGHTS: - param = getattr(layer, f"{name}_packed") - delattr(layer, f"{name}_packed") - layer.register_parameter(name, param) + move_parameter(layer, f"{name}_packed", name) super().process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py index 4a410bce6ac8..f0eead077962 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py @@ -3,7 +3,6 @@ from collections.abc import Callable import torch -from torch.nn.parameter import Parameter from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( @@ -17,6 +16,7 @@ GroupQuantScaleParameter, ModelWeightParameter, ) +from vllm.model_executor.utils import move_parameter __all__ = ["CompressedTensorsW4A4Mxfp4"] @@ -88,8 +88,7 @@ def create_weights( layer.register_parameter("weight_scale", weight_scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.weight = Parameter(layer.weight_packed.data, requires_grad=False) - del layer.weight_packed + move_parameter(layer, "weight_packed", "weight") self.kernel.process_weights_after_loading(layer) def apply_weights( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py index 78924681e9ae..7a06072bd901 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py @@ -17,7 +17,7 @@ ModelWeightParameter, PerTensorScaleParameter, ) -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import move_parameter, replace_parameter logger = init_logger(__name__) @@ -94,8 +94,7 @@ def create_weights( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Rename CT checkpoint names to standardized names - layer.weight = layer.weight_packed - del layer.weight_packed + move_parameter(layer, "weight_packed", "weight") # Check for mismatched weight global scales if torch.unique(layer.weight_global_scale).numel() != 1: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py index 16a0d9242a18..cdc027fdb367 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py @@ -121,13 +121,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if self.strategy == QuantizationStrategy.BLOCK: assert self.is_static_input_scheme is False # MarlinFP8ScaledMMLinearKernel uses "weight_scale_inv" for block - # quant, while CT registers the scale as "weight_scale". - # Rename by deleting the old parameter and adding the new one so - # that prepare_fp8_layer_for_marlin (which prefers "weight_scale" - # over "weight_scale_inv") picks up "weight_scale_inv" correctly. - weight_scale_data = layer.weight_scale.data - del layer._parameters["weight_scale"] - replace_parameter(layer, "weight_scale_inv", weight_scale_data) + # quant, while CT registers the scale as "weight_scale". Alias the + # loaded parameter (same object) under both names. + layer.weight_scale_inv = layer.weight_scale else: if self.strategy == QuantizationStrategy.TENSOR: # For fused modules with per-tensor scales, expand each scale diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py index dd1b986dc3cd..28e191bf276e 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py @@ -31,6 +31,7 @@ ModelWeightParameter, PackedvLLMParameter, ) +from vllm.model_executor.utils import move_parameter, replace_parameter from vllm.scalar_type import scalar_types __all__ = ["CompressedTensorsWNA8O8Int", "fake_quant_static_int8"] @@ -211,22 +212,14 @@ def _pack_int_quantized_weight(self, layer: torch.nn.Module) -> None: weight = layer.weight out_features, in_features = weight.shape packed = pack_to_int32(weight.data.contiguous(), self.num_bits) - delattr(layer, "weight") + # One param under both names; the packed data goes into the loaded storage. + move_parameter(layer, "weight", "weight_packed") + replace_parameter(layer, "weight_packed", packed.contiguous()) + weight.packed_dim, weight.packed_factor = 1, self.pack_factor def _noop_loader(*_, **__): return None - layer.register_parameter( - "weight_packed", - PackedvLLMParameter( - data=packed.contiguous(), - input_dim=1, - output_dim=0, - packed_dim=1, - packed_factor=self.pack_factor, - weight_loader=_noop_loader, - ), - ) layer.register_parameter( "weight_shape", BasevLLMParameter( @@ -241,7 +234,6 @@ def _take_act_scale(layer, name: str) -> torch.Tensor | None: if param is None: return None scale = param.data.clone() - delattr(layer, name) return None if float(scale.reshape(-1)[0]) == 0.0 else scale def apply_weights( diff --git a/vllm/model_executor/layers/quantization/fbgemm_fp8.py b/vllm/model_executor/layers/quantization/fbgemm_fp8.py index 06589628215d..5afd493fc3f7 100644 --- a/vllm/model_executor/layers/quantization/fbgemm_fp8.py +++ b/vllm/model_executor/layers/quantization/fbgemm_fp8.py @@ -169,9 +169,8 @@ def process_weights_after_loading(self, layer: Module) -> None: # The param keeps its checkpoint dims; Humming reads the (K, N) layout. layer.processed_weight_dims = (0, 1) if self.quant_config.use_marlin: + # Activations are not quantized for marlin (input_scale_ub unused). prepare_fp8_layer_for_marlin(layer) - # Activations not quantized for marlin. - del layer.input_scale_ub self.fp8_linear.process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index d0c225cebceb..ea4c9f6c475b 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -41,6 +41,11 @@ select_humming_moe_experts, weight_schema_to_quant_key, ) +from vllm.model_executor.layers.quantization.utils.humming.linear import ( + get_live_param, + kernel_parameters, + replace_layer_parameters, +) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.parameter import ( BasevLLMParameter, @@ -542,21 +547,14 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: param_dtype=layer.param_dtype, ) tensors.update(input_tensors) - - for name, _ in list(layer.named_parameters()): - delattr(layer, name) - - for name, tensor in tensors.items(): - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) - + replace_layer_parameters(layer, tensors, list(layer._parameters)) del tensors # force requant (origin quant setting -> fp16/bf16 -> new_quant setting) assert isinstance(self.weight_schema, _hm.HummingWeightSchema) force_requant = self.force_weight_schema is not None if force_requant and self.weight_schema != self.force_weight_schema: - source_tensors = layer.state_dict() + source_tensors = kernel_parameters(layer) tensors = self.weight_schema.requant_tensors( tensors=source_tensors, target_weight_schema=self.force_weight_schema, @@ -568,14 +566,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: self.weight_schema = self.force_weight_schema - for name, _ in list(layer.named_parameters()): - if name != "bias": - delattr(layer, name) - - for name, tensor in tensors.items(): - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) - + names = [name for name in layer._parameters if name != "bias"] + replace_layer_parameters(layer, tensors, names) del tensors self.input_schema = check_and_fallback_input_schema( @@ -584,7 +576,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: param_dtype=layer.param_dtype, allow_fallback=self.quant_config.allow_input_schema_fallback, ) - source_tensors = dict(layer.named_parameters()) + source_tensors = kernel_parameters(layer) self.layer_config, tensors = _hm.prepare_and_transform_humming_tensors( tensors=source_tensors, shape_n=layer.output_partition_sizes_sum, @@ -603,11 +595,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if "bias" in tensors: zero_bias = torch.zeros_like(tensors["bias"]) layer.register_buffer("zero_bias", zero_bias) - for name, _ in list(layer.named_parameters()): - delattr(layer, name) - for name, tensor in tensors.items(): - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) + replace_layer_parameters(layer, tensors, list(layer._parameters)) self.compute_config = get_humming_linear_compute_config() self.locks = _hm.create_humming_locks(layer.weight.device) @@ -623,12 +611,12 @@ def apply( self.layer_config, inputs=flatten_inputs, weight=layer.weight, - weight_scale=getattr(layer, "weight_scale", None), - zero_point=getattr(layer, "zero_point", None), - bias=getattr(layer, "bias" if bias is not None else "zero_bias", None), - weight_scale_2=getattr(layer, "weight_scale_2", None), - input_scale=getattr(layer, "input_scale", None), - input_scale_2=getattr(layer, "input_scale_2", None), + weight_scale=get_live_param(layer, "weight_scale"), + zero_point=get_live_param(layer, "zero_point"), + bias=get_live_param(layer, "bias" if bias is not None else "zero_bias"), + weight_scale_2=get_live_param(layer, "weight_scale_2"), + input_scale=get_live_param(layer, "input_scale"), + input_scale_2=get_live_param(layer, "input_scale_2"), hadamard_block_size=self.weight_schema.hadamard_block_size, locks=self.locks, compute_config=self.compute_config, diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py index a4902fae60c3..06143e73756f 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py @@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Any import torch -from torch.nn.parameter import Parameter from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic @@ -12,6 +11,7 @@ GroupQuantScaleParameter, ModelWeightParameter, ) +from vllm.model_executor.utils import move_parameter from .inc_scheme import INCLinearScheme @@ -84,8 +84,7 @@ def create_weights( layer.register_parameter("weight_scale", weight_scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.weight = Parameter(layer.weight_packed.data, requires_grad=False) - del layer.weight_packed + move_parameter(layer, "weight_packed", "weight") self.kernel.process_weights_after_loading(layer) def apply_weights( diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py index e51a51aeaa76..ae5fd4fe6b55 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py @@ -35,7 +35,11 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( prepare_moe_fp4_layer_for_marlin, ) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import ( + move_parameter, + replace_parameter, + set_weight_attrs, +) from vllm.platforms import current_platform logger = init_logger(__name__) @@ -165,14 +169,8 @@ def get_fused_moe_quant_config( ) def process_weights_after_loading(self, layer: RoutedExperts) -> None: - layer.w13_weight = torch.nn.Parameter( - layer.w13_weight_packed.data, requires_grad=False - ) - delattr(layer, "w13_weight_packed") - layer.w2_weight = torch.nn.Parameter( - layer.w2_weight_packed.data, requires_grad=False - ) - delattr(layer, "w2_weight_packed") + move_parameter(layer, "w13_weight_packed", "w13_weight") + move_parameter(layer, "w2_weight_packed", "w2_weight") if self.use_cutlass_mxfp4: # Swizzle weight scales from flat checkpoint layout [E, N, K//32] diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py index 02a14f51e791..cfd5fdad4581 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py @@ -14,7 +14,7 @@ GroupQuantScaleParameter, PackedvLLMParameter, ) -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import replace_parameter, retire_parameter from vllm.scalar_type import scalar_types from .inc_scheme import INCLinearScheme @@ -434,9 +434,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.ark_weight_type = ark_linear.wdt layer.ark_scale_type = ark_linear.sdt - if hasattr(layer, "qzeros"): - del layer.qzeros - del layer.scales + # ARK holds its own qzeros/scales; the loaded params keep their data. + if getattr(layer, "qzeros", None) is not None: + retire_parameter(layer, "qzeros") + retire_parameter(layer, "scales") def apply_weights( self, diff --git a/vllm/model_executor/layers/quantization/kv_cache.py b/vllm/model_executor/layers/quantization/kv_cache.py index ba503a93f1a5..5428ce17ca86 100644 --- a/vllm/model_executor/layers/quantization/kv_cache.py +++ b/vllm/model_executor/layers/quantization/kv_cache.py @@ -77,11 +77,8 @@ def apply(self, layer: torch.nn.Module) -> torch.Tensor: def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if is_weights_pre_processed(): - # Scale buffers were exported with post-processed values; drop the - # consumed scale placeholders and rebuild the host-side copies. - for name in ("k_scale", "v_scale", "q_scale", "prob_scale"): - if hasattr(layer, name): - delattr(layer, name) + # Scale buffers were exported with post-processed values; rebuild the + # host-side copies. if hasattr(layer, "_k_scale"): layer._k_scale_float = layer._k_scale.item() layer._v_scale_float = layer._v_scale.item() @@ -90,7 +87,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer._v_scale_cpu = layer._v_scale.detach().to("cpu") return - # skip if there are no weights to process (for example, weight reloading) + # skip if there are no weights to process if not hasattr(layer, "q_scale"): assert not hasattr(layer, "k_scale") assert not hasattr(layer, "v_scale") @@ -104,10 +101,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer._v_scale.copy_(1.0) layer._k_scale_float = 1.0 layer._v_scale_float = 1.0 - del layer.k_scale - del layer.v_scale - del layer.q_scale - del layer.prob_scale return # If the kv-cache is not quantized, we enforce the k/v_scale to be 1.0 @@ -171,13 +164,13 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if layer.q_scale > 0.0: q_scale = layer.q_scale if current_platform.is_fp8_fnuz(): - q_scale *= 2 + q_scale = q_scale * 2 else: q_scale = 1.0 if layer.prob_scale > 0.0: prob_scale = layer.prob_scale if current_platform.is_fp8_fnuz(): - prob_scale *= 2 + prob_scale = prob_scale * 2 else: prob_scale = 1.0 @@ -206,8 +199,3 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: "issues. Please make sure q/prob scaling factors are " "available in the fp8 checkpoint." ) - - del layer.k_scale - del layer.v_scale - del layer.q_scale - del layer.prob_scale diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 29db26053521..72b190e2b663 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -115,6 +115,7 @@ from vllm.model_executor.utils import ( is_weights_pre_processed, replace_parameter, + retire_parameter, set_weight_attrs, ) from vllm.platforms import current_platform @@ -855,6 +856,8 @@ def __init__( super().__init__(moe_config) self.quant_config = quant_config self.use_a16 = quant_config.quant_method == "W4A16_NVFP4" + # Whether the selected backend reads the activation input scales. + self.uses_input_scales = True activation_key = None if self.use_a16 else kNvfp4Dynamic if quant_config.moe_activation_override is not None: @@ -1093,11 +1096,18 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w13_weight_scale", w13_scale) replace_parameter(layer, "w13_weight_scale_2", w13_scale_2) - replace_parameter(layer, "w13_input_scale", a13_scale) replace_parameter(layer, "w2_weight", w2) replace_parameter(layer, "w2_weight_scale", w2_scale) replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) - replace_parameter(layer, "w2_input_scale", a2_scale) + # Weight-only backends (Marlin, Humming) return no input scales: the + # loaded ones keep their data, retired, and the kernel config gets None. + self.uses_input_scales = a13_scale is not None + if self.uses_input_scales: + replace_parameter(layer, "w13_input_scale", a13_scale) + replace_parameter(layer, "w2_input_scale", a2_scale) + else: + retire_parameter(layer, "w13_input_scale") + retire_parameter(layer, "w2_input_scale") self._build_moe_kernel(layer) @@ -1132,8 +1142,8 @@ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfi w2_scale=layer.w2_weight_scale, w13_scale_2=layer.w13_weight_scale_2, w2_scale_2=layer.w2_weight_scale_2, - a13_scale=layer.w13_input_scale, - a2_scale=layer.w2_input_scale, + a13_scale=layer.w13_input_scale if self.uses_input_scales else None, + a2_scale=layer.w2_input_scale if self.uses_input_scales else None, swiglu_limit=getattr(layer, "swiglu_limit", None), swiglu_alpha=getattr(layer, "swiglu_alpha", None), swiglu_beta=getattr(layer, "swiglu_beta", None), @@ -2092,7 +2102,8 @@ def process(self, layer, role) -> None: # Raw max, no reciprocation — Marlin/cutlass want ModelOpt's amax/2688. weight_global_scale = layer.weight_scale_2.max().to(torch.float32) replace_parameter(layer, "weight_global_scale", weight_global_scale) - del layer.weight_scale_2 + # Kernels read weight_global_scale; the loaded param keeps its data. + retire_parameter(layer, "weight_scale_2") class KNvfp4Dynamic(QuantKeyScheme): @@ -2130,7 +2141,8 @@ def process(self, layer, role) -> None: "input_global_scale_inv", (1.0 / layer.input_global_scale).to(torch.float32), ) - del layer.input_scale + # Kernels read input_global_scale; the loaded param keeps its data. + retire_parameter(layer, "input_scale") class KFp8StaticTensor(QuantKeyScheme): @@ -2550,7 +2562,8 @@ def post_process(self, layer) -> None: "Dropping input_scale from a weight-only W4A16 NVFP4 " "checkpoint (deprecated)." ) - del layer.input_scale + # Weight-only: no kernel reads it; the loaded param keeps its data. + retire_parameter(layer, "input_scale") _DROP_INPUT_SCALE = _DropInputScale() diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index c93977d86c65..4b9c188503c5 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -31,6 +31,7 @@ mxfp4_round_up_hidden_size_and_intermediate_size, select_deepseek_v4_mxfp4_moe_backend, select_mxfp4_moe_backend, + triton_weight, ) from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization import QuantizationMethods @@ -366,8 +367,8 @@ def _setup_kernel( replace_parameter(layer, "w13_weight_scale", w13_scale) replace_parameter(layer, "w2_weight_scale", w2_scale) else: - layer.w13_weight = w13 - layer.w2_weight = w2 + layer.w13_weight_triton = w13 + layer.w2_weight_triton = w2 self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale @@ -444,8 +445,8 @@ def apply( assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, + w1=triton_weight(layer, "w13_weight"), + w2=triton_weight(layer, "w2_weight"), topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -467,8 +468,8 @@ def apply_monolithic( assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, + w1=triton_weight(layer, "w13_weight"), + w2=triton_weight(layer, "w2_weight"), router_logits=router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, @@ -851,8 +852,8 @@ def _setup_kernel( replace_parameter(layer, "w13_weight_scale", w13_scale) replace_parameter(layer, "w2_weight_scale", w2_scale) else: - layer.w13_weight = w13 - layer.w2_weight = w2 + layer.w13_weight_triton = w13 + layer.w2_weight_triton = w2 self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale @@ -982,9 +983,6 @@ def convert( del w_all, s_all assert out_packed is not None assert out_scale is not None - weight.data = torch.empty(0, dtype=torch.uint8, device=out_packed.device) - scale.data = torch.empty(0, dtype=torch.uint8, device=out_packed.device) - torch.accelerator.empty_cache() return out_packed, out_scale w13, w13_scale = convert(layer.w13_weight, layer.w13_weight_scale) @@ -1099,8 +1097,8 @@ def apply( assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, + w1=triton_weight(layer, "w13_weight"), + w2=triton_weight(layer, "w2_weight"), topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -1122,8 +1120,8 @@ def apply_monolithic( assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, + w1=triton_weight(layer, "w13_weight"), + w2=triton_weight(layer, "w2_weight"), router_logits=router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index eb175ebd33d8..efee149146dd 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -12,6 +12,7 @@ ) from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import triton_weight from vllm.model_executor.layers.quantization.online.requantization import ( OnlineRequantizationMixin, ) @@ -237,8 +238,8 @@ def apply_monolithic( assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( x, - layer.w13_weight, - layer.w2_weight, + triton_weight(layer, "w13_weight"), + triton_weight(layer, "w2_weight"), router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, @@ -264,8 +265,8 @@ def apply( assert self.moe_kernel is not None return self.moe_kernel.apply( x, - layer.w13_weight, - layer.w2_weight, + triton_weight(layer, "w13_weight"), + triton_weight(layer, "w2_weight"), topk_weights, topk_ids, activation=layer.activation, diff --git a/vllm/model_executor/layers/quantization/online/mxfp4.py b/vllm/model_executor/layers/quantization/online/mxfp4.py index 70a64340f694..9de7da09fa91 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp4.py +++ b/vllm/model_executor/layers/quantization/online/mxfp4.py @@ -235,8 +235,8 @@ def _setup_kernel( ): # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor # Store on layer for apply(), scales are PrecisionConfig - layer.w13_weight = w13 - layer.w2_weight = w2 + layer.w13_weight_triton = w13 + layer.w2_weight_triton = w2 self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale else: @@ -269,8 +269,8 @@ def get_fused_moe_quant_config( # AITER_MXFP4_FP8` to read `self.w13_precision_config`/ # `self.w2_precision_config` instead of `layer.w13_weight_scale`/ # `layer.w2_weight_scale`. For those backends, - # `convert_weight_to_mxfp4_moe_kernel_format` deletes the latter, so - # `getattr` below would raise `AttributeError`. This is currently + # `convert_weight_to_mxfp4_moe_kernel_format` empties the latter, so + # the scales read below would be empty. This is currently # unreachable because `activation_quant_key` is hardcoded to # `kMxfp4Dynamic` # TODO: When supporting online activation quant key override, diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index a17f7b28da87..25e19e2bb8f8 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -55,6 +55,7 @@ make_mxfp4_moe_quant_config, mxfp4_round_up_hidden_size_and_intermediate_size, select_mxfp4_moe_backend, + triton_weight, ) from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( select_mxfp8_moe_backend, @@ -124,7 +125,11 @@ normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize, ) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.model_executor.utils import ( + replace_parameter, + retire_parameter, + set_weight_attrs, +) from vllm.platforms import current_platform if TYPE_CHECKING: @@ -1239,16 +1244,7 @@ def create_weights( layer.w13_bias, layer.w2_bias = None, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Discard zero points (INT8 fused MoE kernel uses symmetric quant) - for attr in ( - "w13_input_zero_point", - "w2_input_zero_point", - "w13_weight_zero_point", - "w2_weight_zero_point", - ): - if hasattr(layer, attr): - delattr(layer, attr) - + # Loaded zero points stay registered; the INT8 fused MoE kernel ignores them. # For static input scales, collapse the per-expert scales to a single # value (the legacy fused_experts path expects one scale per layer). if self.static_input_scales: @@ -1849,8 +1845,8 @@ def _setup_kernel(self, layer: RoutedExperts): ): # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor # Store on layer for apply(), scales are PrecisionConfig - layer.w13_weight = w13 - layer.w2_weight = w2 + layer.w13_weight_triton = w13 + layer.w2_weight_triton = w2 self.w13_precision_config = w13_scale self.w2_precision_config = w2_scale else: @@ -1970,8 +1966,8 @@ def apply( assert self.moe_kernel is not None return self.moe_kernel.apply( hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, + w1=triton_weight(layer, "w13_weight"), + w2=triton_weight(layer, "w2_weight"), topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -1992,8 +1988,8 @@ def apply_monolithic( assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, + w1=triton_weight(layer, "w13_weight"), + w2=triton_weight(layer, "w2_weight"), router_logits=router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, @@ -2204,6 +2200,8 @@ def __init__( super().__init__(moe, kNvfp4Static, activation_quant_key) self.quant_config = quant_config self.group_size = 16 + # Whether the selected backend reads the activation input scales. + self.uses_input_scales = True # Select experts implementation. self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( @@ -2374,12 +2372,19 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w13_weight_scale", w13_scale) replace_parameter(layer, "w13_weight_scale_2", w13_scale_2) - replace_parameter(layer, "w13_input_scale_2", a13_scale) replace_parameter(layer, "w2_weight", w2) replace_parameter(layer, "w2_weight_scale", w2_scale) replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) - replace_parameter(layer, "w2_input_scale_2", a2_scale) + # Weight-only backends (Marlin, Humming) return no input scales: the + # loaded ones keep their data, retired, and the kernel config gets None. + self.uses_input_scales = a13_scale is not None + if self.uses_input_scales: + replace_parameter(layer, "w13_input_scale_2", a13_scale) + replace_parameter(layer, "w2_input_scale_2", a2_scale) + else: + retire_parameter(layer, "w13_input_scale_2") + retire_parameter(layer, "w2_input_scale_2") # Setup modular kernel. self.moe_quant_config = self.get_fused_moe_quant_config(layer) @@ -2402,8 +2407,8 @@ def get_fused_moe_quant_config( w2_scale=layer.w2_weight_scale, w13_scale_2=layer.w13_weight_scale_2, w2_scale_2=layer.w2_weight_scale_2, - a13_scale=layer.w13_input_scale_2, - a2_scale=layer.w2_input_scale_2, + a13_scale=layer.w13_input_scale_2 if self.uses_input_scales else None, + a2_scale=layer.w2_input_scale_2 if self.uses_input_scales else None, swiglu_limit=getattr(layer, "swiglu_limit", None), swiglu_alpha=getattr(layer, "swiglu_alpha", None), swiglu_beta=getattr(layer, "swiglu_beta", None), diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py index 4634777c185d..38b373ad7d29 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py @@ -23,7 +23,7 @@ ModelWeightParameter, PerTensorScaleParameter, ) -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import replace_parameter, retire_parameter __all__ = ["QuarkNVFP4"] @@ -126,7 +126,8 @@ def create_weights( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: input_global_scale = layer.input_scale_2.max().to(torch.float32) replace_parameter(layer, "input_global_scale", input_global_scale) - del layer.input_scale_2 + # Kernels read input_global_scale; the loaded param keeps its data. + retire_parameter(layer, "input_scale_2") weight_global_scale = layer.weight_scale_2.to(torch.float32) @@ -142,7 +143,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: weight_global_scale = weight_global_scale.max() replace_parameter(layer, "weight_global_scale", weight_global_scale) - del layer.weight_scale_2 + retire_parameter(layer, "weight_scale_2") replace_parameter( layer, "alpha", layer.input_global_scale * layer.weight_global_scale diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py index 364f5eaf614f..4eac3b91c27c 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py @@ -152,12 +152,7 @@ def _scale_weight_loader( # Checkpoints are serialized in quark format, which is # different from the format the kernel may want. Handle repacking here. def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.register_parameter("weight_zero_point", None) - delattr(layer, "weight_zero_point") - if self.input_symmetric: - layer.register_parameter("input_zero_point", None) - delattr(layer, "input_zero_point") - + # The loaded (symmetric) zero points stay registered; kernels ignore them. self.kernel.process_weights_after_loading(layer) def apply_weights( diff --git a/vllm/model_executor/layers/quantization/utils/humming/linear.py b/vllm/model_executor/layers/quantization/utils/humming/linear.py index 821f23aeda85..f8557ccbe4a3 100644 --- a/vllm/model_executor/layers/quantization/utils/humming/linear.py +++ b/vllm/model_executor/layers/quantization/utils/humming/linear.py @@ -12,10 +12,80 @@ from vllm.model_executor.layers.quantization.utils.humming.schema import ( check_and_fallback_input_schema, ) +from vllm.model_executor.utils import ( + get_live_param, + kernel_parameters, + move_parameter, + replace_parameter, + retire_parameter, + set_kernel_parameter, +) if TYPE_CHECKING: from vllm.utils.humming import BaseInputSchema, LayerConfig +# A param the conversion leaves unused keeps its loaded data, retired, when it is +# at most this fraction of the largest converted tensor (scalars, per-tensor and +# per-channel scales); a larger one backs a new tensor (its data is in them). +_KEEP_UNUSED_FRACTION = 1 / 256 + + +def _nbytes(t: torch.Tensor) -> int: + return t.numel() * t.element_size() + + +def replace_layer_parameters( + layer: torch.nn.Module, tensors: dict[str, torch.Tensor], names: list[str] +) -> None: + """Write ``tensors`` into ``layer``, keeping existing parameter objects. + + A param in ``names`` that gets no new tensor stays registered: an alias of a + replaced param follows it; one whose storage a new tensor is in, or a large + one, is also registered under a new name (see `move_parameter`); any other + keeps its data, retired.""" + params = layer._parameters + replaced = {id(p) for name in tensors if (p := params.get(name)) is not None} + leftovers = { + name: p + for name in names + if name not in tensors + and (p := params.get(name)) is not None + and id(p) not in replaced + } + + def adopt(src: str, dst: str) -> None: + move_parameter(layer, src, dst) + for name in [n for n, p in leftovers.items() if p is params[src]]: + del leftovers[name] + + new = [name for name in tensors if name not in params] + for name in new: + ptr = tensors[name].untyped_storage().data_ptr() + src = next( + (n for n, p in leftovers.items() if p.untyped_storage().data_ptr() == ptr), + None, + ) + if ptr and src is not None: + adopt(src, name) + limit = _KEEP_UNUSED_FRACTION * max(map(_nbytes, tensors.values()), default=0) + large = sorted( + (n for n, p in leftovers.items() if _nbytes(p) > limit), + key=lambda n: _nbytes(leftovers[n]), + reverse=True, + ) + free = sorted( + (n for n in new if n not in params), + key=lambda n: _nbytes(tensors[n]), + reverse=True, + ) + for src, dst in zip(large, free): + if src in leftovers: + adopt(src, dst) + for name in leftovers: + retire_parameter(layer, name) + for name, tensor in tensors.items(): + set_kernel_parameter(layer, name, tensor) + def convert_linear_layer_to_humming_standard( layer: LinearBase, name_map: dict[str, str] @@ -24,15 +94,15 @@ def convert_linear_layer_to_humming_standard( layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names and layout humming's weight schema expects (``weight`` / ``weight_scale``).""" for name, checkpoint_name in name_map.items(): - tensor = getattr(layer, checkpoint_name) - delattr(layer, checkpoint_name) + param = getattr(layer, checkpoint_name) + tensor = param.data if name == "weight": # (input_dim, output_dim) after PWAL, if it transposed the weight. input_dim, output_dim = getattr( layer, "processed_weight_dims", - (getattr(tensor, "input_dim", 1), getattr(tensor, "output_dim", 0)), + (getattr(param, "input_dim", 1), getattr(param, "output_dim", 0)), ) if input_dim == 0 and output_dim == 1: @@ -42,19 +112,17 @@ def convert_linear_layer_to_humming_standard( tensor = tensor.view(tensor.size(0), -1).view(torch.int32) elif name in ["weight_scale", "zero_point"]: - if getattr(tensor, "output_dim", 0) == 1: + if getattr(param, "output_dim", 0) == 1: tensor = tensor.transpose(0, 1).contiguous() if tensor.ndim == 1: tensor = tensor.unsqueeze(1) tensor = tensor.view(torch.int32) if name == "zero_point" else tensor - if isinstance(tensor, torch.nn.Parameter): - param = tensor - else: - param = torch.nn.Parameter(tensor, requires_grad=False) - - setattr(layer, name, param) + if name != checkpoint_name: + # One param under both names: the loader keeps its name. + move_parameter(layer, checkpoint_name, name) + replace_parameter(layer, name, tensor) def prepare_humming_linear_layer_config( @@ -91,7 +159,7 @@ def prepare_humming_linear_layer_config( shape_n_stacks = layer.output_partition_sizes # Step 1: convert weight and input schemas to humming standard format - source_tensors = dict(layer.named_parameters(recurse=False)) + source_tensors = kernel_parameters(layer) weight_schema, tensors = weight_schema.convert_humming( tensors=source_tensors, shape_n_stacks=shape_n_stacks, @@ -148,11 +216,7 @@ def prepare_humming_linear_layer_config( if "bias" in tensors: zero_bias = torch.zeros_like(tensors["bias"]) layer.register_buffer("zero_bias", zero_bias) - for name, _ in list(layer.named_parameters(recurse=False)): - delattr(layer, name) - for name, tensor in tensors.items(): - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) + replace_layer_parameters(layer, tensors, list(layer._parameters)) return config @@ -179,19 +243,19 @@ def apply_humming_linear( from vllm.utils.humming import humming_forward mode = layer_config.input_quant_mode - input_scale = getattr(layer, "input_scale", None) if mode.should_quantize else None + input_scale = get_live_param(layer, "input_scale") if mode.should_quantize else None input_scale_2 = ( - getattr(layer, "input_scale_2", None) if mode.has_secondary_scale else None + get_live_param(layer, "input_scale_2") if mode.has_secondary_scale else None ) flatten_inputs = x.reshape(-1, x.size(-1)) output = humming_forward( layer_config, inputs=flatten_inputs, weight=layer.weight, - weight_scale=getattr(layer, "weight_scale", None), - zero_point=getattr(layer, "zero_point", None), - bias=getattr(layer, "zero_bias" if skip_bias_add else "bias", None), - weight_scale_2=getattr(layer, "weight_scale_2", None), + weight_scale=get_live_param(layer, "weight_scale"), + zero_point=get_live_param(layer, "zero_point"), + bias=get_live_param(layer, "zero_bias" if skip_bias_add else "bias"), + weight_scale_2=get_live_param(layer, "weight_scale_2"), input_scale=input_scale, input_scale_2=input_scale_2, hadamard_block_size=layer.weight_schema.hadamard_block_size, diff --git a/vllm/model_executor/layers/quantization/utils/humming/moe.py b/vllm/model_executor/layers/quantization/utils/humming/moe.py index 418e9acbe70a..44b5a9020443 100644 --- a/vllm/model_executor/layers/quantization/utils/humming/moe.py +++ b/vllm/model_executor/layers/quantization/utils/humming/moe.py @@ -24,6 +24,11 @@ HummingGroupedExperts, HummingIndexedExperts, ) +from vllm.model_executor.layers.quantization.utils.humming.linear import ( + get_live_param, + kernel_parameters, + replace_layer_parameters, +) from vllm.model_executor.layers.quantization.utils.humming.schema import ( humming_is_layer_skipped, ) @@ -168,14 +173,14 @@ def get_humming_moe_quant_config( weight_dtype=str(weight_schema.b_dtype), weight_group_shape=weight_group_shape, activation_group_shape=activation_group_shape, - w1_scale=getattr(layer, "w13_weight_scale", None), - w1_gscale=getattr(layer, "w13_weight_scale_2", None), - w1_zp=getattr(layer, "w13_zero_point", None), - w1_bias=getattr(layer, "w13_bias", None), - w2_scale=getattr(layer, "w2_weight_scale", None), - w2_gscale=getattr(layer, "w2_weight_scale_2", None), - w2_zp=getattr(layer, "w2_zero_point", None), - w2_bias=getattr(layer, "w2_bias", None), + w1_scale=get_live_param(layer, "w13_weight_scale"), + w1_gscale=get_live_param(layer, "w13_weight_scale_2"), + w1_zp=get_live_param(layer, "w13_zero_point"), + w1_bias=get_live_param(layer, "w13_bias"), + w2_scale=get_live_param(layer, "w2_weight_scale"), + w2_gscale=get_live_param(layer, "w2_weight_scale_2"), + w2_zp=get_live_param(layer, "w2_zero_point"), + w2_bias=get_live_param(layer, "w2_bias"), gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, @@ -189,7 +194,7 @@ def get_humming_moe_quant_config( ) for name in ("input_scale", "input_scale_2"): key = f"{prefix}_{name}" - value = getattr(layer, f"{sublayer}_{name}", None) + value = get_live_param(layer, f"{sublayer}_{name}") setattr(config, key, value) return config @@ -301,12 +306,17 @@ def _extract_sublayer_tensors( layer: "RoutedExperts", sublayer_name: str, ) -> dict[str, torch.Tensor]: - """Extract tensors for a specific sublayer from the layer's state dict.""" - return dict( - (key.removeprefix(sublayer_name + "_"), value) - for key, value in layer.state_dict().items() + """Extract a sublayer's live params and persistent buffers.""" + buffers = { + name: buffer + for name, buffer in layer._buffers.items() + if buffer is not None and name not in layer._non_persistent_buffers_set + } + return { + key.removeprefix(sublayer_name + "_"): value + for key, value in (kernel_parameters(layer) | buffers).items() if key.startswith(sublayer_name + "_") - ) + } def _replace_layer_parameters( @@ -324,19 +334,15 @@ def _replace_layer_parameters( preserve_bias: If True, don't delete bias parameters """ - # Delete old parameters - for name, _ in list(layer.named_parameters()): - if not name.startswith(sublayer_name + "_"): - continue - if preserve_bias and name == sublayer_name + "_bias": - continue - delattr(layer, name) - - # Set new parameters - for name, tensor in tensors.items(): - param_name = f"{sublayer_name}_{name}" - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, param_name, param) + names = [ + name + for name in layer._parameters + if name.startswith(sublayer_name + "_") + and not (preserve_bias and name == sublayer_name + "_bias") + ] + replace_layer_parameters( + layer, {f"{sublayer_name}_{n}": t for n, t in tensors.items()}, names + ) def _convert_sublayer_to_humming( diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index b87535ac376e..98cc0f06f49f 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -467,7 +467,7 @@ def sync_cpu_storage(self): param_offloader.sync_cpu_storage() # Remove offloaders whose parameter was deleted during - # process_weights_after_loading (e.g. k_scale / v_scale). + # process_weights_after_loading (e.g. online requantization sources). deleted = [ name for name, offloader in self._param_offloaders.items() @@ -639,9 +639,8 @@ def __init__(self, module: nn.Module, param_name: str): self._cpu_storage: torch.Tensor | None = None self._gpu_buffer: torch.Tensor | None = None # Store reference to GPU buffer # Set to True if the underlying nn.Parameter was deleted by - # process_weights_after_loading (e.g. transient KV-cache scale params - # such as k_scale/v_scale created by BaseKVCacheMethod.create_weights - # and deleted after copying into permanent _k_scale buffers). + # process_weights_after_loading (e.g. online requantization source + # params released after conversion). self._param_deleted: bool = False # Offload to CPU immediately to free GPU memory during model loading @@ -750,12 +749,10 @@ def sync_cpu_storage(self) -> None: 2. device_loading_context creates NEW CPU tensors when moving back 3. Our old _cpu_storage would have pre-processed or stale data - If the parameter no longer exists on the module (e.g. transient - KV-cache scale parameters such as k_scale/v_scale that are created - by BaseKVCacheMethod.create_weights() and then deleted by - process_weights_after_loading() after copying their values into - permanent _k_scale buffers), the offloader marks itself as deleted - and skips the sync. The caller (_ModuleOffloader.sync_cpu_storage) + If the parameter no longer exists on the module (e.g. online + requantization source parameters released by + process_weights_after_loading()), the offloader marks itself as + deleted and skips the sync. The caller (_ModuleOffloader.sync_cpu_storage) is responsible for removing these stale entries. """ try: diff --git a/vllm/model_executor/utils.py b/vllm/model_executor/utils.py index e73e4a7ec569..f5bf8d059973 100644 --- a/vllm/model_executor/utils.py +++ b/vllm/model_executor/utils.py @@ -68,6 +68,34 @@ def set_weight_attrs( setattr(weight, key, value) +def _extent(t: torch.Tensor) -> int: + """Number of elements spanned by ``t`` in its storage.""" + if t.numel() == 0: + return 0 + return 1 + sum((n - 1) * s for n, s in zip(t.shape, t.stride())) + + +def _in_own_storage(old: torch.Tensor, new: torch.Tensor) -> torch.Tensor: + """``new`` copied into the bytes ``old`` spans if it fits, so the loaded + storage is kept (with its ``data_ptr``); else ``new`` itself. An equal-size + result fits too: a repack to the same byte count (e.g. FP8 → Marlin int32) + then keeps the address captured CUDA graphs read.""" + storage = old.untyped_storage() + offset, size = old.storage_offset() * old.element_size(), new.element_size() + if ( + new.numel() == 0 + or old.device != new.device + or new.untyped_storage().data_ptr() == storage.data_ptr() + or _extent(new) != new.numel() # overlapping or with gaps + or _extent(old) != old.numel() # its gaps may hold another tensor's data + or offset % size + or new.numel() * size > _extent(old) * old.element_size() + ): + return new + view = new.new_empty(0).set_(storage, offset // size, new.shape, new.stride()) + return view.copy_(new) + + def replace_parameter( layer: torch.nn.Module, param_name: str, @@ -79,7 +107,8 @@ def replace_parameter( The existing parameter object is kept (and so its class and loader attributes for weight reloading): ``new_tensor`` is copied into it when shape, stride, dtype and device match, which preserves ``data_ptr`` for captured CUDA graphs; - otherwise the parameter's ``.data`` is rebound to ``new_tensor``. A new + otherwise a smaller ``new_tensor`` is copied into the parameter's own storage + (so the loaded storage never shrinks), else ``.data`` is rebound to it. A new parameter is registered only if none exists under ``param_name``, or if the existing one is an online requantization source (released after conversion). @@ -131,8 +160,10 @@ def replace_parameter( and old_param.device == new_tensor.device ): old_param.data.copy_(new_tensor) + elif getattr(getattr(layer, "quant_method", None), "uses_meta_device", False): + old_param.data = new_tensor # online quantization frees the source else: - old_param.data = new_tensor + old_param.data = _in_own_storage(old_param, new_tensor) for attr_name, attr in new_tensor_attrs.items(): setattr(old_param, attr_name, attr) return @@ -149,6 +180,74 @@ def replace_parameter( setattr(layer, param_name, new_param) +def move_parameter(layer: torch.nn.Module, src_name: str, dst_name: str) -> None: + """Register the ``src_name`` parameter object under ``dst_name`` as well. + + Both names then refer to one object and so to one storage: the loader writes + through the checkpoint name, kernels read the runtime name, and a later + `replace_parameter` on either name updates both. ``dst_name`` becomes the + name `kernel_parameters` reports for the object.""" + param = getattr(layer, src_name) + existing = layer._parameters.get(dst_name) + if existing is not None and existing is not param: + raise ValueError(f"{dst_name!r} is already registered as another parameter") + layer.register_parameter(dst_name, param) + param.runtime_name = dst_name + + +def retire_parameter(layer: torch.nn.Module, name: str) -> None: + """Mark a loaded parameter that no kernel reads after post-processing. + + It stays registered with its loaded data (so it can be loaded again); + `kernel_parameters` and `get_live_param` skip it.""" + getattr(layer, name).retired = True + + +def set_kernel_parameter( + layer: torch.nn.Module, name: str, tensor: torch.Tensor +) -> None: + """`replace_parameter`, for a tensor kernels read under ``name``: the parameter + is live even if it was retired before, and ``name`` becomes its runtime name + if it is registered under several names (see `move_parameter`).""" + replace_parameter(layer, name, tensor) + param = getattr(layer, name) + param.__dict__.pop("retired", None) + if hasattr(param, "runtime_name"): + param.runtime_name = name + + +def _is_live(param: torch.Tensor | None) -> bool: + return ( + param is not None and not getattr(param, "retired", False) and param.numel() > 0 + ) + + +def get_live_param(layer: torch.nn.Module, name: str) -> torch.Tensor | None: + """``getattr(layer, name, None)``, treating retired or empty params as absent.""" + tensor = getattr(layer, name, None) + if isinstance(tensor, torch.nn.Parameter) and not _is_live(tensor): + return None + return tensor + + +def kernel_parameters(layer: torch.nn.Module) -> dict[str, torch.nn.Parameter]: + """The layer's own parameters as kernels see them: live parameters only, each + object once, under its runtime name (see `move_parameter`).""" + names: dict[int, str] = {} + params: dict[str, torch.nn.Parameter] = {} + for name, param in layer._parameters.items(): + if not _is_live(param): + continue + first = names.setdefault(id(param), name) + if first == name: + params[name] = param + elif getattr(param, "runtime_name", None) == name: + del params[first] + params[name] = param + names[id(param)] = name + return params + + def get_packed_modules_mapping(model: torch.nn.Module) -> dict[str, list[str]]: parent_map = getattr(model, "packed_modules_mapping", None) parent_map = copy.deepcopy(parent_map) if parent_map is not None else {} diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py index 021bef5b9c64..e2d3109ab7b5 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_sparse.py +++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py @@ -15,7 +15,7 @@ from vllm import _custom_ops as ops from vllm.forward_context import get_forward_context -from vllm.model_executor.utils import replace_parameter +from vllm.model_executor.utils import replace_parameter, retire_parameter from vllm.models.deepseek_v4.attention import ( DeepseekV4Attention, DeepseekV4Indexer, @@ -317,18 +317,17 @@ def _capture_and_pack(layer: torch.nn.Module) -> None: dequant = dequant.view(self.n_local_groups, self.o_lora_rank, k_dim) self._wo_a_packed = torch.ops._C.convert_weight_packed(dequant.contiguous()) orig_pwal(layer) - # orig_pwal VNNI-repacks layer.weight/scale for ordinary-linear - # use that _o_proj never needs -- free them, same pattern as - # dispatch_cpu_unquantized_gemm's remove_weight=True. - replace_parameter(layer, "weight", layer.weight.new_empty(0)) + # _o_proj reads only the packed copy: it becomes the loaded weight's + # storage, and the (small) scale keeps its data. + replace_parameter(layer, "weight", self._wo_a_packed) + self._wo_a_packed = layer.weight.data scale_attr = ( "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale" ) if hasattr(layer, scale_attr): - scale = getattr(layer, scale_attr) - replace_parameter(layer, scale_attr, scale.new_empty(0)) + retire_parameter(layer, scale_attr) quant_method.process_weights_after_loading = ( # type: ignore[method-assign] _capture_and_pack From af8f27d5d77dfcb4d0505185900ad1c23e2016cd Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 28 Oct 2026 16:14:55 -0500 Subject: [PATCH 081/406] [CI][Bugfix] Fit the MLA backend bf16 atol to a bf16 SDPA reference (#64259) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- tests/v1/attention/test_mla_backends.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index ddafc37ba647..24be854d4ab2 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -36,7 +36,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed from vllm.v1.attention.backend import CommonAttentionMetadata, MLAAttentionImpl from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla from vllm.v1.attention.backends.mla import flashmla as flashmla_module @@ -1710,6 +1710,7 @@ def _run_backend_correctness( multiple GPUs. This tests that backends work correctly with different head counts. """ + set_random_seed(42) # Filter backends to those that support the requested kv_cache_dtype backends_to_test = [ b @@ -2085,7 +2086,7 @@ def _run_backend_correctness( # 4. Run vLLM backends and compare rtol = 1e-2 atol = { - "auto": 1e-2, + "auto": 2e-2, "fp8": 1.5e-1, "fp8_e4m3": 1.5e-1, }[kv_cache_dtype] From 9f0893f1adab0359e91eb6ce0fd0673c53660393 Mon Sep 17 00:00:00 2001 From: Aarushi Jain Date: Wed, 28 Oct 2026 16:30:26 -0500 Subject: [PATCH 082/406] [CI][Bugfix] Fix flaky port race in test_coordination_store_held_only_by_the_online_dp_master (#64348) Signed-off-by: aarushjain29 Co-authored-by: Andreas Karatzas --- tests/v1/engine/test_core_engine_actor_manager.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/v1/engine/test_core_engine_actor_manager.py b/tests/v1/engine/test_core_engine_actor_manager.py index f8e2373d4c58..711db295d5db 100644 --- a/tests/v1/engine/test_core_engine_actor_manager.py +++ b/tests/v1/engine/test_core_engine_actor_manager.py @@ -15,7 +15,7 @@ import zmq from torch.distributed import TCPStore -from vllm.utils.network_utils import get_open_port, make_zmq_socket, split_zmq_path +from vllm.utils.network_utils import make_zmq_socket, split_zmq_path from vllm.v1.engine.core import EngineCoreActorMixin from vllm.v1.engine.core_client import BackgroundResources from vllm.v1.engine.utils import ( @@ -162,8 +162,9 @@ def __init__(self, **kwargs) -> None: parallel_config = vllm_config.parallel_config for name, value in overrides.items(): setattr(parallel_config, name, value) - # A fresh port per case so a lingering handshake socket cannot fail the next. - parallel_config.data_parallel_rpc_port = get_open_port() + # Port 0 makes the handshake socket take a free port at bind time; a port + # reserved up front can be handed to the coordination store meanwhile. + parallel_config.data_parallel_rpc_port = 0 parallel_config._coord_store_port = 0 monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineProcManager", FakeManager) monkeypatch.setattr("vllm.v1.engine.utils.CoreEngineActorManager", FakeManager) From 7b87062308855434fa0747852ccebbe950ee9e4f Mon Sep 17 00:00:00 2001 From: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:13:49 -0400 Subject: [PATCH 083/406] [Bugfix][Spec Decode] Honor standalone draft load config (#61699) Signed-off-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Co-authored-by: Benjamin Chislett Co-authored-by: vllm-agent --- requirements/test/rocm.txt | 2 ++ tests/v1/worker/test_gpu_model_runner_v2.py | 6 +++++- vllm/v1/spec_decode/draft_model.py | 2 ++ vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py | 5 +++++ vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py | 5 +++++ 5 files changed, 19 insertions(+), 1 deletion(-) diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index e3af4cbc3946..7328ac9495d6 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -654,6 +654,8 @@ numpy==2.2.6 # tritonclient # vocos # xgrammar +nvidia-cuda-cccl==13.3.4.3.1 + # via nvidia-nvshmem-cu13 open-clip-torch==2.32.0 # via -r requirements/test/rocm.in openai==2.31.0 diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 6a043f83c78e..1ffe87aadef8 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -12,7 +12,7 @@ import vllm.distributed.aux_output_connector.worker as aux_output_worker import vllm.v1.worker.gpu.model_runner as model_runner_module -from vllm.config import CacheConfig, LoRAConfig, ParallelConfig +from vllm.config import CacheConfig, LoadConfig, LoRAConfig, ParallelConfig from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.v1.kv_cache_interface import ( CircularBufferSpec, @@ -525,6 +525,10 @@ class Config: get_model = Mock(return_value=object()) monkeypatch.setattr(ar_module, "get_model", get_model) monkeypatch.setattr(ar_module, "set_model_tag", lambda _: contextlib.nullcontext()) + monkeypatch.setattr(ar_module, "get_draft_load_config", lambda _: LoadConfig()) + monkeypatch.setattr( + ar_module, "get_pp_safe_draft_load_config", lambda load_config: load_config + ) speculator.load_draft_model(object(), set()) diff --git a/vllm/v1/spec_decode/draft_model.py b/vllm/v1/spec_decode/draft_model.py index 17be2732e67c..0714fcab917d 100644 --- a/vllm/v1/spec_decode/draft_model.py +++ b/vllm/v1/spec_decode/draft_model.py @@ -9,6 +9,7 @@ from vllm.config.utils import replace from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model +from vllm.model_executor.model_loader.utils import get_draft_load_config from vllm.tokenizers.registry import get_tokenizer from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer from vllm.v1.spec_decode.vocab_mapping import VocabMapping @@ -108,6 +109,7 @@ def _get_model(self) -> nn.Module: model = get_model( vllm_config=draft_vllm_config, prefix="draft_model", + load_config=get_draft_load_config(draft_vllm_config), ) return model diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py index 8b4d8d7c3dae..8ede3b30afd6 100644 --- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py +++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py @@ -10,9 +10,11 @@ from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import set_forward_context from vllm.model_executor.model_loader import get_model +from vllm.model_executor.model_loader.utils import get_draft_load_config from vllm.v1.worker.gpu.dp_utils import DPSyncState from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config class ExtractHiddenStatesSpeculator(DraftModelSpeculator): @@ -68,6 +70,9 @@ def load_draft_model( return get_model( vllm_config=self.vllm_config, model_config=self.draft_model_config, + load_config=get_pp_safe_draft_load_config( + get_draft_load_config(self.vllm_config) + ), ) def load_model(self, target_model: nn.Module) -> None: diff --git a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py index 1df19a3f9693..dc8ff567d0c4 100644 --- a/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/standalone_ar/speculator.py @@ -11,6 +11,7 @@ from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model +from vllm.model_executor.model_loader.utils import get_draft_load_config from vllm.triton_utils import tl, triton from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.spec_decode.utils import next_power_of_2 @@ -21,6 +22,7 @@ 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.spec_decode.speculator import DraftModelSpeculator +from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) @@ -127,6 +129,9 @@ def load_draft_model( return get_model( vllm_config=draft_vllm_config, prefix="draft_model", + load_config=get_pp_safe_draft_load_config( + get_draft_load_config(draft_vllm_config) + ), ) @torch.inference_mode() From d482020608af39f5543a5f7802a3228e17e11f8e Mon Sep 17 00:00:00 2001 From: Maroon Ayoub Date: Wed, 28 Oct 2026 15:29:12 -0700 Subject: [PATCH 084/406] [Bugfix][Core] Defer async KV loads on HiSparse decoders until in-flight steps finish (#63827) Signed-off-by: Maroon Ayoub Signed-off-by: Matthew Bonanni Signed-off-by: Lucas Wilkinson Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 Co-authored-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson --- tests/v1/core/test_deferred_block_free.py | 105 +++++++++++++++++++++- vllm/v1/core/sched/scheduler.py | 65 ++++++++++++-- 2 files changed, 164 insertions(+), 6 deletions(-) diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py index 44cdf567f7c9..aa2244e069ed 100644 --- a/tests/v1/core/test_deferred_block_free.py +++ b/tests/v1/core/test_deferred_block_free.py @@ -13,13 +13,14 @@ import os import time -from unittest.mock import PropertyMock, patch +from unittest.mock import MagicMock, PropertyMock, patch import pytest from vllm.config import VllmConfig from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.metrics.stats import PrefixCacheStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import RequestStatus @@ -544,3 +545,105 @@ def test_cow_retentions_deferred_until_copy_step_processed(): assert not scheduler.deferred_frees out2 = scheduler.schedule() assert [r.req_id for r in out2.scheduled_new_reqs] == ["late"] + + +def _setup_async_load(defer_async_loads: bool): + """Schedule a running request's prefill and one in-flight decode, then + admit a request whose KV an async connector load fills. + + Returns (scheduler, importer, outputs, update_state_after_alloc mock). + """ + scheduler = create_scheduler( + model=MODEL, + async_scheduling=True, + use_kv_connector=mock_kv(matched_tokens=0, is_async=False), + ) + scheduler.defer_block_free = True + scheduler.defer_async_loads = defer_async_loads + _, out0, out1 = _setup_request_with_inflight_step(scheduler) + + importer = create_requests( + num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=5, req_ids=["imp"] + )[0] + scheduler.add_request(importer) + connector = scheduler.connector + assert connector is not None + with ( + patch.object( + connector, + "get_num_new_matched_tokens", + return_value=(2 * scheduler.block_size, True), + ), + patch.object(connector, "update_state_after_alloc") as registered, + ): + out2 = scheduler.schedule() + assert importer.status == RequestStatus.WAITING_FOR_REMOTE_KVS + return scheduler, importer, [out0, out1, out2], registered + + +def test_deferred_async_load_registers_after_inflight_step(): + """HiSparse keeps reading a released page until the pool reuses it, and an + async connector load isn't ordered against the in-flight step, so the load + is registered only once that step is processed. + """ + scheduler, importer, outputs, registered = _setup_async_load(defer_async_loads=True) + out0, out1, _ = outputs + registered.assert_not_called() + + connector = scheduler.connector + with patch.object(connector, "update_state_after_alloc") as registered: + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + outputs.append(scheduler.schedule()) + registered.assert_not_called() + + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + outputs.append(scheduler.schedule()) + registered.assert_called_once() + assert registered.call_args.args[0] is importer + + +def test_deferred_async_load_records_hit_sources_when_started(): + """The connector attributes cache-hit tiers only once the load is + registered, so a deferred load records its prefix cache stats then.""" + scheduler, importer, outputs, _ = _setup_async_load(defer_async_loads=True) + out0, out1, _ = outputs + connector = scheduler.connector + scheduler.connector_prefix_cache_stats = PrefixCacheStats() + sources = MagicMock() + with ( + patch.object(connector, "update_state_after_alloc"), + patch.object( + connector, "get_external_cache_hit_sources", return_value=sources + ) as hit_sources, + patch.object(PrefixCacheStats, "record") as record, + ): + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + scheduler.update_from_output(out1, _make_model_runner_output(out1)) + hit_sources.assert_called_once() + assert hit_sources.call_args.args[0] is importer + assert record.call_args.kwargs["hits_by_source"] is sources + + +def test_async_load_registers_immediately_without_deferral(): + scheduler, importer, _, registered = _setup_async_load(defer_async_loads=False) + registered.assert_called_once() + + +def test_abort_deferred_async_load_frees_after_inflight_step(): + """No load was posted for a deferred request, so nothing will report it + done; its blocks are freed once the step that may still read them is + processed. + """ + scheduler, importer, outputs, _ = _setup_async_load(defer_async_loads=True) + pool = scheduler.kv_cache_manager.block_pool + num_free_held = pool.get_num_free_blocks() + + scheduler.finish_requests(importer.request_id, RequestStatus.FINISHED_ABORTED) + assert importer.request_id not in scheduler.requests + assert scheduler.deferred_frees + assert pool.get_num_free_blocks() == num_free_held + + for output in outputs: + scheduler.update_from_output(output, _make_model_runner_output(output)) + assert not scheduler.deferred_frees + assert pool.get_num_free_blocks() > num_free_held diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 5d97ef2dd91e..0db5e1f6bfc1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -167,6 +167,7 @@ def __init__( self.connector_prefix_cache_stats: PrefixCacheStats | None = None self.recompute_kv_load_failures = True self.defer_block_free = False + self.defer_async_loads = False # Whether a preempted request's in-flight output must be dropped; see # KVConnectorBase_V1.requires_kv_delivery. self.requires_kv_delivery = False @@ -192,6 +193,12 @@ def __init__( multiple_inflight_batches = self.vllm_config.max_concurrent_batches > 1 if multiple_inflight_batches and kv_transfer_config.is_kv_consumer: self.defer_block_free = True + # HiSparse keeps reading a released page until the pool reuses it, + # so a load into a reused page must also wait for in-flight steps. + self.defer_async_loads = ( + self.defer_block_free + and kv_cache_config.hisparse_host_num_blocks is not None + ) self.requires_kv_delivery = self.connector.requires_kv_delivery @@ -463,6 +470,10 @@ def __init__( # FIFO of (fence_seq, blocks): blocks become safe to free once # processed_step_seq >= fence_seq. self.deferred_frees: deque[tuple[int, list[KVCacheBlock]]] = deque() + # Request id -> (fence_seq, num_external_tokens, connector prefix cache + # queries, hits): async KV loads started once processed_step_seq >= + # fence_seq. + self.deferred_loads: dict[str, tuple[int, int, int, int]] = {} self.perf_metrics: ModelMetrics | None = None if self.log_stats and vllm_config.observability_config.enable_mfu_metrics: @@ -1454,14 +1465,29 @@ def skip_request(from_queue: RequestQueue) -> None: # This information is used to determine if a load is # needed for this request. if self.connector is not None: - self.connector.update_state_after_alloc( - request, - self.kv_cache_manager.get_blocks(request_id), - num_external_computed_tokens, + defer_load = ( + load_kv_async + and self.defer_async_loads + and self.sched_step_seq > self.processed_step_seq ) + if defer_load: + request.last_sched_seq = self.sched_step_seq + self.deferred_loads[request_id] = ( + self.sched_step_seq, + num_external_computed_tokens, + connector_prefix_cache_queries, + connector_prefix_cache_hits, + ) + else: + self.connector.update_state_after_alloc( + request, + self.kv_cache_manager.get_blocks(request_id), + num_external_computed_tokens, + ) if ( self.connector_prefix_cache_stats is not None and connector_prefix_cache_queries != 0 + and not defer_load ): preempted = request.num_preemptions > 0 self.connector_prefix_cache_stats.record( @@ -2254,6 +2280,7 @@ def update_from_output( if self.defer_block_free and scheduler_output.total_num_scheduled_tokens > 0: self.processed_step_seq += 1 self._drain_deferred_frees() + self._drain_deferred_loads() perf_stats: PerfStats | None = None if self.perf_metrics and self.perf_metrics.is_enabled(): @@ -2986,8 +3013,10 @@ def finish_requests( for request in valid_requests: delay_free_blocks = False if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: + # A deferred load never started, so it will not report done. delay_free_blocks = ( - request.request_id not in self.finished_recving_kv_req_ids + self.deferred_loads.pop(request.request_id, None) is None + and request.request_id not in self.finished_recving_kv_req_ids ) self.finished_recving_kv_req_ids.discard(request.request_id) self.failed_recving_kv_req_ids.discard(request.request_id) @@ -3087,6 +3116,32 @@ def _drain_deferred_frees(self): # Free in reverse order so that the tail blocks are evicted first. self.kv_cache_manager.block_pool.free_blocks(reversed(blocks)) + def _drain_deferred_loads(self) -> None: + """Start deferred async loads whose fence step has completed.""" + for request_id, (fence, num_tokens, queries, hits) in list( + self.deferred_loads.items() + ): + if fence > self.processed_step_seq: + break + del self.deferred_loads[request_id] + assert self.connector is not None + request = self.requests[request_id] + self.connector.update_state_after_alloc( + request, self.kv_cache_manager.get_blocks(request_id), num_tokens + ) + if self.connector_prefix_cache_stats is not None and queries != 0: + preempted = request.num_preemptions > 0 + self.connector_prefix_cache_stats.record( + num_tokens=queries, + num_hits=hits, + preempted=preempted, + hits_by_source=( + self.connector.get_external_cache_hit_sources(request, hits) + if hits and not preempted + else None + ), + ) + def get_num_unfinished_requests(self) -> int: if self._pause_state == PauseState.PAUSED_ALL: return 0 From f4bf6a93b6694361d296fc9b3ccdad0975299213 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:39:01 +0000 Subject: [PATCH 085/406] [CI] Bump pre-commit hooks (except ruff) (#64349) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .pre-commit-config.yaml | 17 +++++++++-------- csrc/cpu/cpu_types_scalar.hpp | 2 +- csrc/quantization/w8a8/fp8/common.cu | 6 +++--- .../w8a8/fp8/per_token_group_quant.cu | 4 ++-- docs/contributing/README.md | 7 +++++++ pyproject.toml | 5 +++++ vllm/envs.py | 2 +- .../layers/quantization/modelopt.py | 4 ++-- vllm/utils/gc_utils.py | 4 ++-- 9 files changed, 32 insertions(+), 19 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a2709f09142f..b42d38d93a31 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,42 +5,43 @@ default_stages: - pre-commit # Run locally - manual # Run in CI exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*' +# Revs are frozen to commit SHAs, bump them with `pre-commit autoupdate --freeze [--repo ]` repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.0 + rev: f9351c924055bf6c7b4a4670237d3ce141e9f57c # frozen: v0.14.0 hooks: - id: ruff-check args: [--output-format, github, --fix] - id: ruff-format - repo: https://github.com/crate-ci/typos - rev: v1.43.5 + rev: 9023c1f779400877cafdd98052113fc777e2cfdc # frozen: v1.51.1 hooks: - id: typos args: [--force-exclude] - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v21.1.2 + rev: 98ad4ddb9c057de2f5b9e93e44e6db959807d68c # frozen: v23.1.3 hooks: - id: clang-format exclude: 'csrc/moe/topk_softmax_kernels.cu|vllm/third_party/.*' types_or: [c++, cuda] args: [--style=file, --verbose] - repo: https://github.com/DavidAnson/markdownlint-cli2 - rev: v0.21.0 + rev: 341e2a331ba578029e66c2e01183e7a6b5ca9ab4 # frozen: v0.24.0 hooks: - id: markdownlint-cli2 language_version: 24.22.1 # Node 26 (current lts) needs libatomic, which the CI runners lack args: [--config, tools/pre_commit/.markdownlint.yaml, --fix] - repo: https://github.com/rhysd/actionlint - rev: v1.7.7 + rev: 914e7df21a07ef503a81201c76d2b11c789d3fca # frozen: v1.7.12 hooks: - id: actionlint - repo: https://github.com/shellcheck-py/shellcheck-py - rev: 745eface02aef23e168a8afb6b5737818efbea95 # v0.11.0.1 + rev: a6564f5d57444e123433a5c9af733e3975c77229 # frozen: v0.11.0.1-1 hooks: - id: shellcheck args: ["-s", "bash", "--format=gcc", "--rcfile=tools/pre_commit/.shellcheckrc"] - repo: https://github.com/astral-sh/uv-pre-commit - rev: 0.11.1 + rev: 46787482b9fb2003c56e0f7592e312ef381829d9 # frozen: 0.13.6 hooks: - id: pip-compile args: [ @@ -159,7 +160,7 @@ repos: ] files: ^requirements/docs\.(in|txt)$ - repo: https://github.com/pre-commit/pre-commit-hooks - rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c + rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c # frozen: v6.0.0 hooks: - id: check-json files: ^\.buildkite/.*\.json$ diff --git a/csrc/cpu/cpu_types_scalar.hpp b/csrc/cpu/cpu_types_scalar.hpp index d94cdbc076e8..76b3aec5d45e 100644 --- a/csrc/cpu/cpu_types_scalar.hpp +++ b/csrc/cpu/cpu_types_scalar.hpp @@ -502,7 +502,7 @@ inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); }); } -inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) { +inline FP16Vec8::FP16Vec8(const FP32Vec8& v) { unroll_loop( [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); }); } diff --git a/csrc/quantization/w8a8/fp8/common.cu b/csrc/quantization/w8a8/fp8/common.cu index 471499786977..837948fdd2a7 100644 --- a/csrc/quantization/w8a8/fp8/common.cu +++ b/csrc/quantization/w8a8/fp8/common.cu @@ -95,7 +95,7 @@ __global__ void scaled_fp8_quant_kernel_strided_group_shape( int size, float inv_scale) { vectorize_with_alignment( in, out, size, tid, blockDim.x, - [=] __device__(fp8_type & dst, const scalar_t& src) { + [=] __device__(fp8_type& dst, const scalar_t& src) { dst = scaled_fp8_conversion(static_cast(src), inv_scale); }); @@ -194,7 +194,7 @@ __global__ void scaled_fp8_quant_kernel_strided_dynamic( const float reciprocal_scale = 1.0f / (*scale); vectorize_with_alignment<16>( token_in, token_out, hidden_size, tid, blockDim.x, - [=] __device__(fp8_type & dst, const scalar_t& src) { + [=] __device__(fp8_type& dst, const scalar_t& src) { dst = scaled_fp8_conversion(static_cast(src), reciprocal_scale); }); @@ -238,7 +238,7 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided( // 2) quantize vectorize_with_alignment<16>( token_in, token_out, hidden_size, tid, blockDim.x, - [=] __device__(fp8_type & dst, const scalar_t& src) { + [=] __device__(fp8_type& dst, const scalar_t& src) { dst = scaled_fp8_conversion(static_cast(src), token_scale); }); diff --git a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu index 40d3f92efb39..162a2d879f67 100644 --- a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -49,7 +49,7 @@ __device__ __forceinline__ float ComputeGroupScale( constexpr int vec_size = 16 / sizeof(T); // copy global -> shared & compute absmax - auto scalar_op_cache = [&] __device__(T & dst, const T& src) { + auto scalar_op_cache = [&] __device__(T& dst, const T& src) { float abs_v = fabsf(static_cast(src)); local_absmax = fmaxf(local_absmax, abs_v); dst = src; @@ -85,7 +85,7 @@ __device__ __forceinline__ void QuantizeGroup( const float inv_y = 1.0f / y_s; // quantize shared -> global 8-bit - auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) { + auto scalar_op_quant = [&] __device__(DST_DTYPE& dst, const T& src) { float q = fminf(fmaxf(static_cast(src) * inv_y, min_8bit), max_8bit); dst = DST_DTYPE(q); }; diff --git a/docs/contributing/README.md b/docs/contributing/README.md index d665116e47b5..c97b5bd6c17a 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -103,6 +103,13 @@ integration, use the configuration files in `tools/pre_commit/`: pre-commit run -a # runs on all files (short for --all-files) ``` + Hook revisions in `.pre-commit-config.yaml` are pinned to commit SHAs. To bump them, use: + + ```bash + pre-commit autoupdate --freeze # bumps all hooks + pre-commit autoupdate --freeze --repo # bumps one hook + ``` + ### Documentation MkDocs is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file, [mkdocs.yaml](../../mkdocs.yaml). diff --git a/pyproject.toml b/pyproject.toml index 46c836043ca2..138569bd9dda 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -226,6 +226,7 @@ FoPE = "FoPE" k_ot = "k_ot" view_seperator = "view_seperator" inverse_std_variences = "inverse_std_variences" +stride_bse = "stride_bse" [tool.typos.default.extend-words] certifi = "certifi" @@ -258,6 +259,10 @@ nd = "nd" eles = "eles" datas = "datas" ser = "ser" +# numpy flags.writeable +writeable = "writeable" +# HF CLIP weight name +layrnorm = "layrnorm" ure = "ure" VALU = "VALU" # Walsh-Hadamard Transform diff --git a/vllm/envs.py b/vllm/envs.py index a7bde6b336fa..145cc8b0e8e5 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2119,7 +2119,7 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_NCCL_INCLUDE_PATH": lambda: os.environ.get("VLLM_NCCL_INCLUDE_PATH", None), # GC debug config # - VLLM_GC_DEBUG=0: disable GC debugger - # - VLLM_GC_DEBUG=1: enable GC debugger with gc.collect elpased times + # - VLLM_GC_DEBUG=1: enable GC debugger with gc.collect elapsed times # - VLLM_GC_DEBUG='{"top_objects":5}': enable GC debugger with # top 5 collected objects "VLLM_GC_DEBUG": lambda: os.getenv("VLLM_GC_DEBUG", ""), diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 72b190e2b663..bc85d867233d 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -204,7 +204,7 @@ def is_layer_excluded(self, prefix: str) -> bool: return True # TODO: This special hard coded logic is not needed for quantized checkpoints - # generated by ModelOpt >= 0.39.0 where they are handled natually by the + # generated by ModelOpt >= 0.39.0 where they are handled naturally by the # exclude_modules config. But need to keep them for loading quantized # checkpoints generated by older versions. Then check substring matching # for patterns not caught by exact match @@ -240,7 +240,7 @@ def get_quant_method( return None # TODO: This special hard coded logic is not needed for quantized checkpoints - # generated by ModelOpt >= 0.39.0 where they are handled natually by the + # generated by ModelOpt >= 0.39.0 where they are handled naturally by the # exclude_modules config. But need to keep them for loading quantized # checkpoints generated by older versions. Then check substring matching # for patterns not caught by exact match diff --git a/vllm/utils/gc_utils.py b/vllm/utils/gc_utils.py index be6ab4a23fa1..5a1fb73f1f70 100644 --- a/vllm/utils/gc_utils.py +++ b/vllm/utils/gc_utils.py @@ -73,11 +73,11 @@ def handle(self, phase: str, info: dict[str, int]) -> None: elif phase == "stop": # After GC finished, Record GC elapsed time and # optionally top collected objects - elpased_ms = (time.monotonic_ns() - self.start_time_ns) / 1e6 + elapsed_ms = (time.monotonic_ns() - self.start_time_ns) / 1e6 logger.info( "GC took %.3fms to complete. " "Collected %s objects (out of %d) in GC generation %d.%s", - elpased_ms, + elapsed_ms, str(info.get("collected", "?")), self.num_objects, generation, From 6ba5dbeeb0c11f56798c71b8855b656e6cf07b2f Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:39:30 +0000 Subject: [PATCH 086/406] [Model] Support MTP drafts in the Transformers modeling backend (#61270) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .../test_speculative_draft_hf_overrides.py | 46 +++++- tests/models/registry.py | 12 ++ tests/models/transformers/test_backend.py | 121 ++++++++++++++- tests/test_config.py | 1 + vllm/config/speculative.py | 52 ++++++- vllm/model_executor/models/registry.py | 3 + .../models/transformers/__init__.py | 8 + .../models/transformers/base.py | 51 ++++--- .../model_executor/models/transformers/mtp.py | 139 ++++++++++++++++++ 9 files changed, 396 insertions(+), 37 deletions(-) create mode 100644 vllm/model_executor/models/transformers/mtp.py diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index 39ebd7924bf2..01ea1193953b 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -15,8 +15,9 @@ from unittest.mock import MagicMock, patch import pytest -from transformers import Glm5NextConfig, PreTrainedConfig +from transformers import AutoConfig, Glm5NextConfig, PreTrainedConfig +from vllm.config.model import ModelConfig from vllm.config.parallel import ParallelConfig from vllm.config.speculative import SpeculativeConfig @@ -38,13 +39,16 @@ def test_dict_overrides_are_not_forwarded_to_draft(): composed = SpeculativeConfig.compose_draft_hf_overrides( {"max_position_embeddings": 1234} ) - assert composed is SpeculativeConfig.hf_config_override + hf_config = composed(_make_hf_config(model_type="deepseek_v3")) + assert hf_config.architectures == ["DeepSeekMTPModel"] + assert getattr(hf_config, "max_position_embeddings", None) != 1234 @pytest.mark.cpu_test def test_none_overrides_fall_back_to_arch_mapping(): composed = SpeculativeConfig.compose_draft_hf_overrides(None) - assert composed is SpeculativeConfig.hf_config_override + hf_config = composed(_make_hf_config(model_type="deepseek_v3")) + assert hf_config.architectures == ["DeepSeekMTPModel"] def _make_speculative_config( @@ -57,6 +61,7 @@ def _make_speculative_config( draft.registry.inspect_model_cls.return_value = (None, architecture) target = MagicMock(max_model_len=128, quantization=None, hf_overrides={}) target.hf_config.model_type = hf_kwargs["model_type"] + target.using_transformers_backend.return_value = False with patch("vllm.config.speculative.ModelConfig", return_value=draft): return SpeculativeConfig( @@ -245,6 +250,7 @@ def _make_mtp_speculative_config( quantization=None, hf_overrides={}, ) + target_model_config.using_transformers_backend.return_value = False with patch("vllm.config.speculative.ModelConfig", return_value=draft_model_config): return SpeculativeConfig( @@ -288,3 +294,37 @@ def test_glm5_next_mtp_override_promotes_index_share(): assert draft_config.architectures == ["Glm5NextMTPModel"] assert draft_config.n_predict == 1 assert draft_config.index_share_for_mtp_iteration is True + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + ("model_type", "model_impl", "architecture"), + [ + ("glm4_moe", "transformers", "TransformersMoEForMTP"), + ("glm4_moe", "auto", "Glm4MoeMTPModel"), + # Does not follow the Transformers MTP standard (no `num_mtp_layers`) + ("qwen3_next", "transformers", "Qwen3NextMTP"), + ], +) +def test_mtp_draft_architecture_follows_target_model_impl( + tmp_path, model_type: str, model_impl: str, architecture: str +): + """The MTP draft of a Transformers backend target uses the Transformers backend + too if the model follows the Transformers MTP standard, while other targets keep + their native MTP draft.""" + config = AutoConfig.for_model(model_type, num_nextn_predict_layers=1) + config.architectures = [config.__class__.__name__.replace("Config", "ForCausalLM")] + config.save_pretrained(tmp_path) + target = ModelConfig(str(tmp_path), model_impl=model_impl, skip_tokenizer_init=True) + + speculative_config = SpeculativeConfig( + method="mtp", + num_speculative_tokens=1, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + + assert speculative_config.method == "mtp" + draft = speculative_config.draft_model_config + assert draft.architecture == architecture + assert draft.hf_config.n_predict == 1 diff --git a/tests/models/registry.py b/tests/models/registry.py index 8743775a3368..2036c33df95b 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1909,6 +1909,18 @@ def check_available_online( ), "TransformersMultiModalForCausalLM": _HfExamplesInfo("BAAI/Emu3-Chat-hf"), "TransformersMoEForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924"), + "TransformersForMTP": _HfExamplesInfo( + "zai-org/GLM-4.5", + speculative_model="zai-org/GLM-4.5", + speculative_method="mtp", + # TODO: Use a dense model once one opts into the Transformers MTP standard + is_available_online=False, + ), + "TransformersMoEForMTP": _HfExamplesInfo( + "zai-org/GLM-4.5", + speculative_model="zai-org/GLM-4.5", + speculative_method="mtp", + ), "TransformersMultiModalMoEForCausalLM": _HfExamplesInfo( "Qwen/Qwen3-VL-30B-A3B-Instruct" ), diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index ef76ecd754d4..c48d753b0ca5 100644 --- a/tests/models/transformers/test_backend.py +++ b/tests/models/transformers/test_backend.py @@ -13,9 +13,22 @@ import pytest import torch import torch.nn as nn -from transformers import AutoConfig, AutoModel, PreTrainedConfig +from safetensors.torch import load_file, save_file +from transformers import ( + AutoConfig, + AutoModel, + AutoModelForCausalLM, + PreTrainedConfig, +) -from vllm.config import ModelConfig, VllmConfig +from vllm.config import ( + ModelConfig, + ParallelConfig, + SpeculativeConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.model_executor import model_loader from vllm.model_executor.models.interfaces import SupportsMultiModal from vllm.model_executor.models.transformers.base import Base from vllm.model_executor.models.transformers.fusers import AttentionFuser @@ -708,17 +721,13 @@ def test_attention_layer_index_is_the_modules_own(): assert all(f.validate(a, None) for f, a in zip(fusers, attentions)) -def test_attention_dispatch_is_required(monkeypatch: pytest.MonkeyPatch): +def test_attention_dispatch_is_required(): """Models selected by the permissive registry fail clearly without dispatch.""" - monkeypatch.setattr( - "vllm.model_executor.models.transformers.base.get_pp_indices", - lambda *_: (0, 1), - ) model = SimpleNamespace( text_config=SimpleNamespace(num_hidden_layers=1), - pp_group=SimpleNamespace(rank_in_group=0, world_size=1), attention_fusers={}, _get_attn_cls=lambda: None, + _pp_indices=lambda: (0, 1), ) with pytest.raises( @@ -787,3 +796,99 @@ def test_get_mrope_input_positions_omits_unsupported_grid_kwargs(): assert positions.shape == (4, 3) assert delta == 0 + + +# The smallest sizes the fused MoE kernels accept +MTP_HIDDEN_SIZE = 128 +MTP_MOE_KWARGS = dict( + vocab_size=128, + hidden_size=MTP_HIDDEN_SIZE, + intermediate_size=256, + moe_intermediate_size=128, + n_routed_experts=4, + num_experts_per_tok=2, +) +# Real layer counts, since the MTP weight patterns hard-code them +MTP_MODELS = { + "glm4_moe": dict( + num_hidden_layers=46, num_attention_heads=2, num_key_value_heads=1, head_dim=64 + ), + "deepseek_v3": dict( + num_hidden_layers=61, + num_attention_heads=16, + q_lora_rank=64, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + first_k_dense_replace=1, + n_group=1, + topk_group=1, + ), +} + + +@pytest.mark.parametrize("model_type", MTP_MODELS) +def test_mtp_draft_loads_mtp_weights(model_type, tmp_path, tp_init, monkeypatch): + """An MTP draft for a model following the Transformers MTP standard loads each + MTP layer appended to the target's checkpoint into `MtpModel`'s layout.""" + kwargs = MTP_MODELS[model_type] | MTP_MOE_KWARGS + num_hidden_layers = kwargs.pop("num_hidden_layers") + num_mtp_layers = 2 + hidden_size = MTP_HIDDEN_SIZE + config = AutoConfig.for_model( + model_type, num_hidden_layers=num_hidden_layers + num_mtp_layers, **kwargs + ) + model_cls = AutoModelForCausalLM._model_mapping[type(config)] + config.architectures = [model_cls.__name__] + # These have one MTP layer, so register the second like Step-3.7 registers its own + monkeypatch.setattr( + model_cls, + "_keys_to_ignore_on_load_unexpected", + [rf"model\.layers\.{num_hidden_layers + 1}.*"] + + model_cls._keys_to_ignore_on_load_unexpected, + ) + # Save the MTP layers as the last decoder layers, then add their MTP modules + model_cls(config).save_pretrained(tmp_path) + weights = load_file(tmp_path / "model.safetensors") + for i in range(num_hidden_layers, num_hidden_layers + num_mtp_layers): + for name, shape in { + "enorm.weight": (hidden_size,), + "hnorm.weight": (hidden_size,), + "eh_proj.weight": (hidden_size, 2 * hidden_size), + "shared_head.norm.weight": (hidden_size,), + }.items(): + weights[f"model.layers.{i}.{name}"] = torch.randn(shape) + save_file(weights, tmp_path / "model.safetensors", metadata={"format": "pt"}) + config.num_hidden_layers = num_hidden_layers + config.num_nextn_predict_layers = num_mtp_layers + config.save_pretrained(tmp_path) + + target = ModelConfig(str(tmp_path), model_impl="transformers", dtype="bfloat16") + speculative_config = SpeculativeConfig( + method="mtp", + num_speculative_tokens=num_mtp_layers, + target_model_config=target, + target_parallel_config=ParallelConfig(), + ) + draft = speculative_config.draft_model_config + assert draft.architecture == "TransformersMoEForMTP" + vllm_config = VllmConfig(model_config=target, speculative_config=speculative_config) + with set_current_vllm_config(vllm_config): + model = model_loader.get_model(vllm_config=vllm_config, model_config=draft) + + assert len(model.model.layers) == num_mtp_layers + for k, layer in enumerate(model.model.layers): + mtp = f"model.layers.{num_hidden_layers + k}" + for ckpt_name, param_name in { + "enorm.weight": "enorm.weight", + "hnorm.weight": "hnorm.weight", + "eh_proj.weight": "eh_proj.weight", + "shared_head.norm.weight": "post_norm.weight", + "input_layernorm.weight": "mtp_block.input_layernorm.weight", + "self_attn.o_proj.weight": "mtp_block.self_attn.o_proj.weight", + }.items(): + torch.testing.assert_close( + layer.get_parameter(param_name).cpu(), + weights[f"{mtp}.{ckpt_name}"].to(torch.bfloat16), + ) diff --git a/tests/test_config.py b/tests/test_config.py index 97c20a19fbf7..c7b3d6e4720f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4081,6 +4081,7 @@ def test_mtp_draft_uses_model_weights_not_local_cache(mock_model_config_cls): target_config.hf_text_config.model_type = "deepseek_v3" target_config.quantization = None target_config.max_model_len = 4096 + target_config.using_transformers_backend.return_value = False SpeculativeConfig( method="mtp", diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 24b519af3d13..fd2e3a397e13 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -709,7 +709,15 @@ def compute_hash(self) -> str: return hash_str @staticmethod - def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig: + def hf_config_override( + hf_config: PreTrainedConfig, + transformers_mtp_overrides: dict[str, Any] | None = None, + ) -> PreTrainedConfig: + if transformers_mtp_overrides is not None: + # The Transformers modeling backend builds the MTP draft from the + # target's config, so it needs none of the rewrites below + hf_config.update(transformers_mtp_overrides) + return hf_config initial_architecture = hf_config.architectures[0] if initial_architecture == "RadioMBartPrefixForCausalLM" and getattr( getattr(hf_config, "decoder", None), "causal_vision_prefix", False @@ -1184,13 +1192,17 @@ def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig: def _apply_composed_hf_override( target_hf_overrides: Callable[[PreTrainedConfig], PreTrainedConfig], hf_config: PreTrainedConfig, + transformers_mtp_overrides: dict[str, Any] | None = None, ) -> PreTrainedConfig: - hf_config = SpeculativeConfig.hf_config_override(hf_config) + hf_config = SpeculativeConfig.hf_config_override( + hf_config, transformers_mtp_overrides + ) return target_hf_overrides(hf_config) @staticmethod def compose_draft_hf_overrides( target_hf_overrides: HfOverrides | None, + transformers_mtp_overrides: dict[str, Any] | None = None, ) -> Callable[[PreTrainedConfig], PreTrainedConfig]: """Build the ``hf_overrides`` for the draft ``ModelConfig``. @@ -1208,10 +1220,15 @@ def compose_draft_hf_overrides( method instead. """ if not callable(target_hf_overrides): - return SpeculativeConfig.hf_config_override + return functools.partial( + SpeculativeConfig.hf_config_override, + transformers_mtp_overrides=transformers_mtp_overrides, + ) return functools.partial( - SpeculativeConfig._apply_composed_hf_override, target_hf_overrides + SpeculativeConfig._apply_composed_hf_override, + target_hf_overrides, + transformers_mtp_overrides=transformers_mtp_overrides, ) @staticmethod @@ -1380,11 +1397,31 @@ def __post_init__(self): self.prompt_lookup_min = 0 if self.model is not None: + target = self.target_model_config + using_transformers_mtp = ( + self.method == "mtp" + and target is not None + and target.using_transformers_backend() + and getattr(target.hf_text_config, "num_mtp_layers", None) + is not None + ) + transformers_mtp_overrides = None + if using_transformers_mtp: + if target.hf_config is not target.hf_text_config: + raise NotImplementedError( + "The Transformers modeling backend does not support " + "MTP for multimodal models yet." + ) + moe = "MoE" if target.is_moe else "" + transformers_mtp_overrides = { + "architectures": [f"Transformers{moe}ForMTP"], + "n_predict": target.hf_text_config.num_mtp_layers, + } # Compose any callable hf_overrides set on the target so the # draft config receives the same transform (e.g. the test # shrink). Dict overrides stay target-only. draft_hf_overrides = SpeculativeConfig.compose_draft_hf_overrides( - self.target_model_config.hf_overrides + self.target_model_config.hf_overrides, transformers_mtp_overrides ) self.draft_model_config = ModelConfig( model=self.model, @@ -1449,8 +1486,9 @@ def __post_init__(self): ) ): self.method = "dspark" - elif self.draft_model_config.hf_config.model_type in get_args( - MTPModelTypes + elif using_transformers_mtp or ( + self.draft_model_config.hf_config.model_type + in get_args(MTPModelTypes) ): self.method = "mtp" if ( diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 8e4f7bba7a23..d4af5ef52751 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -788,6 +788,9 @@ # Text generation models "TransformersForCausalLM": ("transformers", "TransformersForCausalLM"), "TransformersMoEForCausalLM": ("transformers", "TransformersMoEForCausalLM"), + # Multi-token prediction draft models + "TransformersForMTP": ("transformers", "TransformersForMTP"), + "TransformersMoEForMTP": ("transformers", "TransformersMoEForMTP"), # Multimodal models "TransformersMultiModalForCausalLM": ( "transformers", diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index e19de6d4e891..2e50ed4e1f6a 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -25,6 +25,7 @@ from vllm.model_executor.models.transformers.causal import CausalMixin from vllm.model_executor.models.transformers.legacy import LegacyMixin from vllm.model_executor.models.transformers.moe import MoEMixin +from vllm.model_executor.models.transformers.mtp import MTPMixin from vllm.model_executor.models.transformers.multimodal import ( MultiModalDummyInputsBuilder, MultiModalMixin, @@ -132,6 +133,13 @@ class TransformersForCausalLM(CausalMixin): ... class TransformersMoEForCausalLM(MoEMixin, CausalMixin): ... +# Multi-token prediction draft models +class TransformersForMTP(MTPMixin): ... + + +class TransformersMoEForMTP(MoEMixin, MTPMixin): ... + + # Multimodal models @MULTIMODAL_REGISTRY.register_processor( MultiModalProcessor, diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 43898125ec7a..9985a1b5115e 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -32,6 +32,7 @@ from transformers import AutoModel, PreTrainedConfig from transformers.conversion_mapping import ( WeightRenaming, + WeightTransform, get_checkpoint_conversion_mapping, get_model_conversion_mapping, ) @@ -176,8 +177,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): self._mark_model_components(vllm_config), init_on_device_without_buffers("meta"), ): - from_config_kwargs = self._from_config_kwargs - self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs) + self.model: PreTrainedModel = self._create_model() # Create weight name to module qualname mapper self._create_hf_to_vllm_mapper() @@ -228,6 +228,10 @@ def _from_config_kwargs(self) -> dict[str, Any]: trust_remote_code=self.model_config.trust_remote_code, ) + def _create_model(self) -> "PreTrainedModel": + """Create the Transformers model that `self.model` wraps.""" + return AutoModel.from_config(**self._from_config_kwargs) + def _find_encoder_classes( self, model: "PreTrainedModel" ) -> dict[str, type["PreTrainedModel"]]: @@ -316,14 +320,7 @@ def _create_hf_to_vllm_mapper(self): orig_to_new_renaming: list[WeightRenaming] = [] orig_to_new_regex: dict[re.Pattern, str | None] = {} - mappings = get_model_conversion_mapping(self.model) - # Ensure layers that are part of the model, but not AutoModel, get their mapping - instantiated = {type(m).__name__ for m in self.model.modules()} - for arch in self.config.architectures or []: - if arch not in instantiated: - mappings.extend(get_checkpoint_conversion_mapping(arch) or []) - - for mapping in mappings: + for mapping in self._weight_conversions(): # Handle weights which have been renamed in Transformers if isinstance(mapping, WeightRenaming): orig_to_new_renaming.append(mapping) @@ -365,6 +362,20 @@ def _create_hf_to_vllm_mapper(self): # Apply mapping to quantization config if needed self._maybe_apply_model_mapping() + def _weight_conversions(self) -> list[WeightTransform]: + """The Transformers weight conversions that apply to `self.model`.""" + mappings = get_model_conversion_mapping(self.model) + # Ensure layers that are part of the model, but not AutoModel, get their mapping + instantiated = {type(m).__name__ for m in self.model.modules()} + for arch in self.config.architectures or []: + if arch not in instantiated: + mappings.extend(get_checkpoint_conversion_mapping(arch) or []) + return mappings + + def _tp_plan(self) -> dict: + """The tensor parallel plan of `self.model`, relative to `self.model`.""" + return self.model.tp_plan or {} + def _get_tie_word_embeddings(self): """Check if the model has tied word embeddings.""" # Models created with Transformers v4 and v5 will store this in different places @@ -423,11 +434,7 @@ def pipeline_parallel(self): attrsetter(name)(module, PPMissingLayer()) # Module list - start_layer, end_layer = get_pp_indices( - self.text_config.num_hidden_layers, - self.pp_group.rank_in_group, - self.pp_group.world_size, - ) + start_layer, end_layer = self._pp_indices() layers_name = names[module_list_idx] # attrgetter in case the module is nested (e.g. "text_model.layers") layers = attrgetter(layers_name)(module) @@ -481,7 +488,7 @@ def recursive_replace(self): - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm` - `nn.LayerNorm` with vLLM's `LayerNorm` """ - tp_plan = self.model.tp_plan or {} + tp_plan = self._tp_plan() if not tp_plan and self.tp_group.world_size > 1: tip = get_feature_request_tip( @@ -591,6 +598,14 @@ def _recursive_replace(module: nn.Module, prefix: str): self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked) + def _pp_indices(self) -> tuple[int, int]: + """The range of decoder layers on this pipeline parallel rank.""" + return get_pp_indices( + self.text_config.num_hidden_layers, + self.pp_group.rank_in_group, + self.pp_group.world_size, + ) + def _create_attention_instances(self, prefix: str = ""): """Create `Attention` instances to inform KV cache allocation.""" text_config = self.text_config @@ -610,9 +625,7 @@ def _create_attention_instances(self, prefix: str = ""): logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None) layer_types = getattr(text_config, "layer_types", None) - pp_rank = self.pp_group.rank_in_group - pp_size = self.pp_group.world_size - start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size) + start, end = self._pp_indices() for i in range(start, end): if i not in self.attention_fusers: diff --git a/vllm/model_executor/models/transformers/mtp.py b/vllm/model_executor/models/transformers/mtp.py new file mode 100644 index 000000000000..05a5d5b458fb --- /dev/null +++ b/vllm/model_executor/models/transformers/mtp.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2026 The vLLM team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Transformers modeling backend mixin for multi-token prediction draft models.""" + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +import torch +from transformers import AutoModelForCausalLM +from transformers.conversion_mapping import WeightTransform + +from vllm.model_executor.models.transformers.causal import CausalMixin +from vllm.sequence import IntermediateTensors + +if TYPE_CHECKING: + from transformers import PreTrainedModel + + +class MTPMixin(CausalMixin): + """Serve the MTP layers of a model which opts into the Transformers MTP standard + (i.e. its config sets `num_mtp_layers`) as a draft model. + + `self.model` is the Transformers `MtpModel`, whose embedding and LM head the + proposer replaces with the target's. + """ + + def _patch_config(self): + # `vllm_config.model_config` is the target's, which must not be patched + assert self.vllm_config.speculative_config is not None + self.model_config = self.vllm_config.speculative_config.draft_model_config + self.config = self.model_config.hf_config + self.text_config = self.config.get_text_config() + super()._patch_config() + + def _create_model(self) -> "PreTrainedModel": + self.check_version("5.20.0", "MTP") + from transformers.modeling_layers import MtpModel + + main_model = AutoModelForCausalLM.from_config(**self._from_config_kwargs) + num_mtp_layers = getattr(self.text_config, "num_mtp_layers", None) + mtp_model = MtpModel(main_model, num_mtp_layers) + self.mtp_weights = MtpModel.get_mtp_weights_regex(main_model) + self.mtp_weight_conversions = mtp_model.get_weight_conversions(main_model) + + # The proposer shares the target's LM head with the draft + del mtp_model.shared_head + # With PP, the last rank has no target embedding to share with the draft + if self.pp_group.world_size > 1: + mtp_model._keys_to_ignore_on_load_unexpected = [ + r"layers\..*embed_tokens\.weight" if "embed_tokens" in k else k + for k in mtp_model._keys_to_ignore_on_load_unexpected + ] + self.text_config = mtp_model.config + return mtp_model + + def _weight_conversions(self) -> list[WeightTransform]: + return self.mtp_weight_conversions + + def _tp_plan(self) -> dict: + tp_plan = self.text_config.base_model_tp_plan or {} + return { + k.replace("layers.*.", "layers.*.mtp_block.", 1): v + for k, v in tp_plan.items() + if k.startswith("layers.*.") + } + + # The draft runs whole on the last PP rank + def pipeline_parallel(self): + pass + + def _pp_indices(self) -> tuple[int, int]: + return 0, self.text_config.num_hidden_layers + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + embed_prefix = None + if self.pp_group.world_size > 1: + embed = self.model.get_input_embeddings() + embed_prefix = next(n for n, m in self.named_modules() if m is embed) + "." + + def is_draft_weight(name: str) -> bool: + if self.mtp_weights.search(name): + return True + if embed_prefix is None: + return False + mapped = self.hf_to_vllm_mapper.map_name(name) + return mapped is not None and mapped.startswith(embed_prefix) + + loaded = super().load_weights(w for w in weights if is_draft_weight(w[0])) + if embed_prefix is not None: + self.has_own_embed_tokens = any(n.startswith(embed_prefix) for n in loaded) + return loaded + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + **kwargs, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + # Add batch dimension before entering Transformers modules + inputs_embeds = inputs_embeds[None, ...] + hidden_states = hidden_states[None, ...] + + hidden_states = self.model.forward_layer( + spec_step_idx % len(self.model.layers), + hidden_states, + positions[None, ...], + mtp_cache=None, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + # Remove batch dimension after exiting Transformers modules + return hidden_states[0, ...] + + def compute_logits( + self, hidden_states: torch.Tensor, spec_step_idx: int = 0 + ) -> torch.Tensor | None: + # The MTP layers share the LM head + return super().compute_logits(hidden_states) From d7f68d4c7d31004e257b7b2ff142b4edd83e0d8c Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Wed, 28 Oct 2026 18:50:07 -0400 Subject: [PATCH 087/406] [Perf][MRV2] Skip idle sampler staged writes (#62748) Signed-off-by: mgoin Signed-off-by: Nick Hill Co-authored-by: Claude Opus 5.5 Co-authored-by: Nick Hill Co-authored-by: vllm-agent --- docs/features/custom_logitsprocs.md | 6 +++--- vllm/v1/worker/gpu/model_runner.py | 4 ++-- vllm/v1/worker/gpu/sample/logits_processor/interface.py | 6 +++--- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/features/custom_logitsprocs.md b/docs/features/custom_logitsprocs.md index 8c1c25939928..53bb258d6078 100644 --- a/docs/features/custom_logitsprocs.md +++ b/docs/features/custom_logitsprocs.md @@ -33,7 +33,7 @@ Custom logits processors must subclass `vllm.v1.worker.gpu.sample.logits_process * Return whether this processor modifies logits for the request. The sampler skips the logits-processing pipeline for batches in which no request needs it * `apply_staged_writes(self) -> None` (optional): - * Flush host-side writes staged by `add_request()` to the device; called once per step before the forward pass + * Flush host-side writes staged by `add_request()` to the device; called before the forward pass of every step that admits new requests, so only state set in `add_request()` is flushed. Read data that changes every step from `ctx` in `apply()` * `validate_params(cls, sampling_params) -> None` (optional classmethod): * Raise `ValueError` for invalid per-request arguments (especially custom arguments); runs at request admission, so invalid arguments fail the request instead of reaching the sampler @@ -70,7 +70,7 @@ The contrived example below implements a custom logits processor which masks out def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState): # Per-slot target; -1 means disabled. Staged on the host and - # flushed to the device once per step in apply_staged_writes(). + # flushed to the device in apply_staged_writes(). self.target_token = torch.full( (req_states.max_num_reqs,), -1, dtype=torch.int64 ) @@ -264,6 +264,6 @@ The examples below show how a user would pass a custom argument (`target_token`) 2. **The conditions under which the logits processor is or is not enabled on a per-request basis.** Unless your intention is for the custom logits processor to act on all requests all the time, you should write your logits processor in such a way that it is possible to disable the logits processor for a given request, i.e. by defaulting an argument to `None` or by passing in a specific do-nothing argument value i.e. `0.0`. Return `False` from `add_request()` for such requests so that the sampler can skip the logits-processing pipeline for batches in which no request needs it -* Stage host-side writes in `add_request()` and flush them once per step in `apply_staged_writes()`, instead of writing device tensors directly; this is the discipline every built-in sampler state follows +* Stage host-side writes in `add_request()` and flush them in `apply_staged_writes()`, instead of writing device tensors directly; this is the discipline every built-in sampler state follows * Since there is no removal hook, `add_request()` must fully (re)initialize the slot's state; do not rely on state left behind by the slot's previous occupant diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 74ae2c3c8624..c181dac5c2a0 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1298,8 +1298,8 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: if scheduler_output.scheduled_new_reqs: self.req_states.apply_staged_writes() self.model_state.apply_staged_writes() - if self.sampler is not None: - self.sampler.apply_staged_writes() + if self.sampler is not None: + self.sampler.apply_staged_writes() def update_requests(self, scheduler_output: SchedulerOutput) -> None: # Add new blocks and update num_computed_tokens for the existing requests. diff --git a/vllm/v1/worker/gpu/sample/logits_processor/interface.py b/vllm/v1/worker/gpu/sample/logits_processor/interface.py index 9934492ec227..d1a8d7750d58 100644 --- a/vllm/v1/worker/gpu/sample/logits_processor/interface.py +++ b/vllm/v1/worker/gpu/sample/logits_processor/interface.py @@ -136,9 +136,9 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: def apply_staged_writes(self) -> None: # noqa: B027 """Flush any host-side writes staged by ``add_request()`` to the device. - Called once per step before the forward pass, after the model runner - has flushed ``req_states``, so a processor that stages writes here can - read the request's tokens back on device. + Called before the forward pass of every step that admits new requests, + after the model runner has flushed ``req_states``, so a processor that + stages writes here can read the request's tokens back on device. """ @abstractmethod From a3ce9ceefe7f4f93fd36c05d5f7314a79b87f959 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 28 Oct 2026 15:53:14 -0700 Subject: [PATCH 088/406] [CI] Re-enable disabled jobs whose installs are already locked (#64010) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .../intel_xpu_ci/test-intel.yaml | 90 +- .buildkite/test-amd.yaml | 844 ++++++++---------- .buildkite/test_areas/misc.yaml | 344 ++++--- tools/pre_commit/test_tethering_allowlist.txt | 17 - 4 files changed, 606 insertions(+), 689 deletions(-) diff --git a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml index 96ac93049694..ed98d4408512 100644 --- a/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml +++ b/.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml @@ -15,29 +15,28 @@ steps: limit: 2 - exit_status: -10 # Agent was lost limit: 2 - # Disabled until tblib transitive dependencies are locked. - # - label: ":intel: (BMG) Example Test" - # key: xpu-example-test - # depends_on: - # - image-build-xpu - # timeout_in_minutes: 50 - # optional: true - # device: intel_gpu - # agent_tags: - # label: production - # gpu: 2+ - # mem: 24+ - # no_plugin: true - # env: - # REGISTRY: "public.ecr.aws/q9t5s3a7" - # REPO: "vllm-ci-test-repo" - # source_file_dependencies: - # - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml - # - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh - # commands: - # - >- - # bash .buildkite/scripts/hardware_ci/run-intel-test.sh - # 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example' + - label: ":intel: (BMG) Example Test" + key: xpu-example-test + depends_on: + - image-build-xpu + timeout_in_minutes: 50 + optional: true + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + source_file_dependencies: + - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml + - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example' - label: ":intel: (BMG) V1 test" key: xpu-v1-test depends_on: @@ -84,29 +83,28 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh w8a8-fp8-linear' - # Disabled until the av installation is locked. - # - label: ":intel: (BMG) Server test" - # key: xpu-server-test - # depends_on: - # - image-build-xpu - # timeout_in_minutes: 45 - # optional: true - # device: intel_gpu - # agent_tags: - # label: production - # gpu: 1+ - # mem: 16+ - # no_plugin: true - # env: - # REGISTRY: "public.ecr.aws/q9t5s3a7" - # REPO: "vllm-ci-test-repo" - # source_file_dependencies: - # - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml - # - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh - # commands: - # - >- - # bash .buildkite/scripts/hardware_ci/run-intel-test.sh - # 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server' + - label: ":intel: (BMG) Server test" + key: xpu-server-test + depends_on: + - image-build-xpu + timeout_in_minutes: 45 + optional: true + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + source_file_dependencies: + - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml + - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server' - label: ":intel: (BMG) Quantization test" key: xpu-quantization-test depends_on: diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index cfa5a1575888..ac593ad3a83a 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -274,29 +274,27 @@ steps: #------------------------------------------------------------ mi250 · v1 -------------------------------------------------------------# -# Disabled until the pytest plugins installation is locked. -# - label: ":amd: (MI250) Batch Invariance" # TBD -# timeout_in_minutes: 60 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] -# dind: false -# agent_pool: mi250_1 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# - vllm/_aiter_ops.py -# - vllm/platforms/rocm.py -# commands: -# - export VLLM_GPU_MEMORY_UTILIZATION=0.8 -# - export VLLM_TEST_CLEAN_GPU_MEMORY=1 -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pip install pytest-timeout pytest-forked -# - pytest -v -s v1/determinism/test_batch_invariance.py -# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA +- label: ":amd: (MI250) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] + dind: false + agent_pool: mi250_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_GPU_MEMORY_UTILIZATION=0.8 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/determinism/test_batch_invariance.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - label: ":amd: (MI250) CUDAGraph" # TBD timeout_in_minutes: 55 @@ -887,25 +885,23 @@ steps: - pytest -v -s -x lora/test_mixtral.py - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)' -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) Elastic EP Scaling" # TBD -# timeout_in_minutes: 55 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/ -# - vllm/engine/ -# - vllm/v1/executor/ -# - vllm/compilation/ -# - tests/distributed/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - pytest -v -s distributed/test_elastic_ep.py +- label: ":amd: (MI300) Elastic EP Scaling" # TBD + timeout_in_minutes: 55 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/ + - vllm/engine/ + - vllm/v1/executor/ + - vllm/compilation/ + - tests/distributed/ + - vllm/platforms/rocm.py + commands: + - pytest -v -s distributed/test_elastic_ep.py - label: ":amd: (MI300) RayExecutorV2" # TBD timeout_in_minutes: 75 @@ -2020,33 +2016,31 @@ steps: #------------------------------------------------------------ mi300 · v1 -------------------------------------------------------------# -# Disabled until the pytest plugins installation is locked. -# - label: ":amd: (MI300) Batch Invariance" # TBD -# timeout_in_minutes: 60 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_1 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/lora/ops/triton_ops -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# - vllm/_aiter_ops.py -# - vllm/platforms/rocm.py -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pip install pytest-timeout pytest-forked -# - pytest -v -s v1/determinism/test_batch_invariance.py -# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py -# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py -# - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py -# - pytest -v -s v1/determinism/test_matmul_batch_invariant.py -# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA -# - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN +- label: ":amd: (MI300) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN - label: ":amd: (MI300) Batch Invariance (Distributed)" # TBD timeout_in_minutes: 45 @@ -2204,39 +2198,33 @@ steps: - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -# Disabled until the tracing packages installation is locked. -# - label: ":amd: (MI300) Metrics, Tracing" # TBD -# timeout_in_minutes: 30 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_2 -# optional: true -# num_gpus: 2 -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/ -# - tests/v1/tracing -# - tests/tracing/ -# - vllm/config/ -# - vllm/distributed/ -# - vllm/engine/ -# - vllm/inputs/ -# - vllm/model_executor/ -# - vllm/multimodal/ -# - vllm/platforms/ -# - vllm/sampling_params.py -# - vllm/tracing/ -# - vllm/transformers_utils/ -# - vllm/utils/ -# - vllm/v1/ -# commands: -# - "pip install \ -# 'opentelemetry-sdk>=1.26.0' \ -# 'opentelemetry-api>=1.26.0' \ -# 'opentelemetry-exporter-otlp>=1.26.0' \ -# 'opentelemetry-semantic-conventions-ai>=0.4.1'" -# - pytest -v -s v1/tracing -# - pytest -v -s tracing +- label: ":amd: (MI300) Metrics, Tracing" # TBD + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + optional: true + num_gpus: 2 + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/v1/tracing + - tests/tracing/ + - vllm/config/ + - vllm/distributed/ + - vllm/engine/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/tracing/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + commands: + - pytest -v -s v1/tracing + - pytest -v -s tracing - label: ":amd: (MI300) V1 E2E" # TBD timeout_in_minutes: 30 @@ -2325,40 +2313,36 @@ steps: - export BASELINE_FILE='/tmp/mooncake-ec-e2e/baseline.json' - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) NixlConnector PD edge case" # TBD -# timeout_in_minutes: 65 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_2 -# num_gpus: 2 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - vllm/v1/core/sched/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh +- label: ":amd: (MI300) NixlConnector PD edge case" # TBD + timeout_in_minutes: 65 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD -# timeout_in_minutes: 60 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: ":amd: (MI300) Distributed DP Extended" # TBD timeout_in_minutes: 75 @@ -2384,133 +2368,119 @@ steps: - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s distributed/test_utils.py -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) Distributed NixlConnector PD accuracy" # TBD -# timeout_in_minutes: 45 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: ":amd: (MI300) Distributed NixlConnector PD accuracy" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" # TBD -# timeout_in_minutes: 55 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - tests/v1/kv_connector/nixl_push_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh +- label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" # TBD + timeout_in_minutes: 55 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - tests/v1/kv_connector/nixl_push_integration/ + - vllm/platforms/rocm.py + commands: + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" # TBD -# timeout_in_minutes: 30 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" # TBD + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD -# timeout_in_minutes: 80 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD + timeout_in_minutes: 80 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) Hybrid SSM NixlConnector PD prefix cache" # TBD -# timeout_in_minutes: 25 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_2 -# num_gpus: 2 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - vllm/v1/core/sched/ -# - vllm/v1/core/kv_cache_coordinator.py -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +- label: ":amd: (MI300) Hybrid SSM NixlConnector PD prefix cache" # TBD + timeout_in_minutes: 25 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/core/sched/ + - vllm/v1/core/kv_cache_coordinator.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD accuracy" # TBD -# timeout_in_minutes: 30 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_2 -# num_gpus: 2 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py -# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py -# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh +- label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD accuracy" # TBD + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD edge cases" # TBD -# timeout_in_minutes: 25 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] -# dind: false -# agent_pool: mi300_2 -# num_gpus: 2 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py -# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py -# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh +- label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD edge cases" # TBD + timeout_in_minutes: 25 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh ######################################################################################################################################### # # @@ -2928,132 +2898,117 @@ steps: - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]' -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD -# timeout_in_minutes: 45 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] -# dind: false -# agent_pool: mi355_2 -# num_gpus: 2 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - vllm/v1/worker/kv_connector_model_runner_mixin.py -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh +- label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/v1/worker/kv_connector_model_runner_mixin.py + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI355) Distributed NixlConnector PD accuracy" # TBD -# timeout_in_minutes: 60 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] -# dind: false -# agent_pool: mi355_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: ":amd: (MI355) Distributed NixlConnector PD accuracy" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy" -# timeout_in_minutes: 60 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] -# dind: false -# agent_pool: mi355_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - tests/v1/kv_connector/rocm_pd_accuracy_utils.py -# - tests/utils.py -# - vllm/v1/attention/backends/rocm_aiter_unified_attn.py -# - vllm/platforms/rocm.py -# - requirements/common.txt -# - requirements/kv_connectors_rocm.txt -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py +- label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy" + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - tests/v1/kv_connector/rocm_pd_accuracy_utils.py + - tests/utils.py + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/platforms/rocm.py + - requirements/common.txt + commands: + - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy" -# timeout_in_minutes: 30 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] -# dind: false -# agent_pool: mi355_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/ -# - tests/v1/kv_connector/mooncake_integration/ -# - tests/v1/kv_connector/rocm_pd_accuracy_utils.py -# - tests/utils.py -# - examples/disaggregated/mooncake_connector/ -# - vllm/v1/attention/backends/rocm_aiter_unified_attn.py -# - vllm/platforms/rocm.py -# - requirements/common.txt -# - requirements/kv_connectors_rocm.txt -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py +- label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy" + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/ + - tests/v1/kv_connector/mooncake_integration/ + - tests/v1/kv_connector/rocm_pd_accuracy_utils.py + - tests/utils.py + - examples/disaggregated/mooncake_connector/ + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/platforms/rocm.py + - requirements/common.txt + commands: + - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" # TBD -# timeout_in_minutes: 40 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] -# dind: false -# agent_pool: mi355_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" # TBD + timeout_in_minutes: 40 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/platforms/rocm.py + commands: + - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" -# timeout_in_minutes: 105 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] -# dind: false -# agent_pool: mi355_4 -# num_gpus: 4 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/platforms/rocm.py -# - vllm/models/kimi_k3/ -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - tests/v1/kv_connector/nixl_integration/ -# - requirements/kv_connectors_rocm.txt -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - export ENABLE_HMA_FLAG='1' -# - export DP_EP='1' -# - export GPU_MEMORY_UTILIZATION='0.9' -# - export VLLM_SSM_CONV_STATE_LAYOUT='DS' -# - export PREFILLER_TP_SIZE='2' -# - export DECODER_TP_SIZE='2' -# - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' -# - export PREFILL_BLOCK_SIZE='2048' -# - export DECODE_BLOCK_SIZE='2048' -# - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh +- label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" + timeout_in_minutes: 105 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/platforms/rocm.py + - vllm/models/kimi_k3/ + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - tests/v1/kv_connector/nixl_integration/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export ENABLE_HMA_FLAG='1' + - export DP_EP='1' + - export GPU_MEMORY_UTILIZATION='0.9' + - export VLLM_SSM_CONV_STATE_LAYOUT='DS' + - export PREFILLER_TP_SIZE='2' + - export DECODER_TP_SIZE='2' + - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' + - export PREFILL_BLOCK_SIZE='2048' + - export DECODE_BLOCK_SIZE='2048' + - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' + - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh - label: ":amd: (MI355) Distributed Features" # TBD timeout_in_minutes: 90 @@ -3167,45 +3122,42 @@ steps: not +quant_fp8 and not +rms_norm and DeepSeek-Coder and ROCM_AITER_MLA' -# Disabled until the KV connectors installation is locked. -# - label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" # TBD -# timeout_in_minutes: 85 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] -# dind: false -# agent_pool: mi355_8 -# num_gpus: 8 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - requirements/kv_connectors_rocm.txt -# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ -# - vllm/model_executor/layers/fused_moe/ -# - vllm/model_executor/layers/quantization/ -# - vllm/models/deepseek_v4/ -# - vllm/v1/attention/ -# - tests/v1/kv_connector/nixl_integration/ -# - vllm/_aiter_ops.py -# - vllm/platforms/rocm.py -# commands: -# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt -# # A mirror does not inherit the parent env, so it is repeated here in -# # the parent's order. The layouts must stay empty: unset, the script -# # defaults to HND, which the V4 indexer rejects. -# - >- -# VLLM_ROCM_USE_AITER=1 -# ENABLE_HMA_FLAG=1 -# DP_EP=1 -# GPU_MEMORY_UTILIZATION=0.85 -# PREFILLER_TP_SIZE=4 -# DECODER_TP_SIZE=4 -# PREFILL_BLOCK_SIZE=256 -# DECODE_BLOCK_SIZE=256 -# MODEL_NAMES=deepseek-ai/DeepSeek-V4-Flash -# VLLM_ENGINE_READY_TIMEOUT_S=1800 -# VLLM_SERVE_EXTRA_ARGS=--trust-remote-code,--kv-cache-dtype,fp8 -# PREFILLER_KV_LAYOUT= -# DECODER_KV_LAYOUT= -# bash v1/kv_connector/nixl_integration/run_accuracy_test.sh +- label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" # TBD + timeout_in_minutes: 85 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_8 + num_gpus: 8 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/models/deepseek_v4/ + - vllm/v1/attention/ + - tests/v1/kv_connector/nixl_integration/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + # A mirror does not inherit the parent env, so it is repeated here in + # the parent's order. The layouts must stay empty: unset, the script + # defaults to HND, which the V4 indexer rejects. + - >- + VLLM_ROCM_USE_AITER=1 + ENABLE_HMA_FLAG=1 + DP_EP=1 + GPU_MEMORY_UTILIZATION=0.85 + PREFILLER_TP_SIZE=4 + DECODER_TP_SIZE=4 + PREFILL_BLOCK_SIZE=256 + DECODE_BLOCK_SIZE=256 + MODEL_NAMES=deepseek-ai/DeepSeek-V4-Flash + VLLM_ENGINE_READY_TIMEOUT_S=1800 + VLLM_SERVE_EXTRA_ARGS=--trust-remote-code,--kv-cache-dtype,fp8 + PREFILLER_KV_LAYOUT= + DECODER_KV_LAYOUT= + bash v1/kv_connector/nixl_integration/run_accuracy_test.sh - label: ":amd: (MI355) Distributed Torchrun + Examples" # TBD timeout_in_minutes: 55 @@ -5276,37 +5228,35 @@ steps: #------------------------------------------------------------ mi355 · v1 -------------------------------------------------------------# -# Disabled until the pytest plugins installation is locked. -# - label: ":amd: (MI355) Batch Invariance" # TBD -# timeout_in_minutes: 60 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] -# dind: false -# agent_pool: mi355_1 -# optional: true -# working_dir: "/vllm-workspace/tests" -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/lora/ops/triton_ops -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# - vllm/_aiter_ops.py -# - vllm/platforms/rocm.py -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pip install pytest-timeout pytest-forked -# - pytest -v -s v1/determinism/test_batch_invariance.py -# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py -# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py -# - pytest -v -s v1/determinism/test_batch_invariance_vlm.py -# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA -# - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN -# - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py -# - pytest -v -s v1/determinism/test_matmul_batch_invariant.py -# - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py -# - pytest -v -s v1/determinism/test_online_batch_invariance.py -# - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py +- label: ":amd: (MI355) Batch Invariance" # TBD + timeout_in_minutes: 60 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py - label: ":amd: (MI355) Spec Decode AL DFlash2 Nightly" timeout_in_minutes: 60 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 0cbfe1d0a883..6e5b6b55bb9a 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -493,44 +493,38 @@ steps: depends_on: - image-build-amd -# Disabled until the tracing packages installation is locked. -# - label: ":nvidia: (L4) Metrics, Tracing" -# key: metrics-tracing-2-gpus -# timeout_in_minutes: 25 -# device: l4 -# num_devices: 2 -# source_file_dependencies: -# - vllm/config/ -# - vllm/distributed/ -# - "!vllm/distributed/kv_transfer/" -# - vllm/engine/ -# - vllm/inputs/ -# - vllm/model_executor/ -# - vllm/multimodal/ -# - vllm/platforms/ -# - vllm/sampling_params.py -# - vllm/tracing/ -# - vllm/transformers_utils/ -# - vllm/utils/ -# - vllm/v1/ -# - tests/v1/tracing -# - tests/tracing/ -# commands: -# - "pip install \ -# 'opentelemetry-sdk>=1.26.0' \ -# 'opentelemetry-api>=1.26.0' \ -# 'opentelemetry-exporter-otlp>=1.26.0' \ -# 'opentelemetry-semantic-conventions-ai>=0.4.1'" -# - pytest -v -s v1/tracing -# - pytest -v -s tracing -# mirror: -# amd: -# label: ":amd: (MI300) Metrics, Tracing" -# dind: false -# device: mi300_2 -# timeout_in_minutes: 30 -# depends_on: -# - image-build-amd +- label: ":nvidia: (L4) Metrics, Tracing" + key: metrics-tracing-2-gpus + timeout_in_minutes: 25 + device: l4 + num_devices: 2 + source_file_dependencies: + - vllm/config/ + - vllm/distributed/ + - "!vllm/distributed/kv_transfer/" + - vllm/engine/ + - vllm/inputs/ + - vllm/model_executor/ + - vllm/multimodal/ + - vllm/platforms/ + - vllm/sampling_params.py + - vllm/tracing/ + - vllm/transformers_utils/ + - vllm/utils/ + - vllm/v1/ + - tests/v1/tracing + - tests/tracing/ + commands: + - pytest -v -s v1/tracing + - pytest -v -s tracing + mirror: + amd: + label: ":amd: (MI300) Metrics, Tracing" + dind: false + device: mi300_2 + timeout_in_minutes: 30 + depends_on: + - image-build-amd - label: ":nvidia: (H200 MIG 18GB) [Nightly] Python-only Installation" key: python-only-installation @@ -845,149 +839,141 @@ steps: commands: - pytest -v -s tool_parsers -# Disabled until the pytest plugins installation is locked. -# - label: ":nvidia: (A100) Batch Invariance" -# key: batch-invariance-a100 -# timeout_in_minutes: 60 -# device: a100 -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pip install pytest-timeout pytest-forked -# - pytest -v -s v1/determinism/test_batch_invariance.py -# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA -# mirror: -# amd: -# label: ":amd: (MI250) Batch Invariance" -# dind: false -# device: mi250_1 -# timeout_in_minutes: 60 -# env: -# VLLM_GPU_MEMORY_UTILIZATION: "0.8" -# VLLM_TEST_CLEAN_GPU_MEMORY: "1" -# depends_on: -# - image-build-amd -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# - vllm/_aiter_ops.py -# - vllm/platforms/rocm.py +- label: ":nvidia: (A100) Batch Invariance" + key: batch-invariance-a100 + timeout_in_minutes: 60 + device: a100 + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/determinism/test_batch_invariance.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + mirror: + amd: + label: ":amd: (MI250) Batch Invariance" + dind: false + device: mi250_1 + timeout_in_minutes: 60 + env: + VLLM_GPU_MEMORY_UTILIZATION: "0.8" + VLLM_TEST_CLEAN_GPU_MEMORY: "1" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py -# Disabled until the pytest plugins installation is locked. -# - label: ":nvidia: (H100) Batch Invariance" -# key: batch-invariance-h100 -# automatic_shard: true -# timeout_in_minutes: 60 -# device: h100 -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/lora/ops/triton_ops -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pip install pytest-timeout pytest-forked -# - pytest -v -s v1/determinism/test_batch_invariance.py -# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py -# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py -# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA -# - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN -# mirror: -# amd: -# label: ":amd: (MI300) Batch Invariance" -# dind: false -# automatic_shard: true -# device: mi300_1 -# timeout_in_minutes: 60 -# depends_on: -# - image-build-amd -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/lora/ops/triton_ops -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# - vllm/_aiter_ops.py -# - vllm/platforms/rocm.py -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pip install pytest-timeout pytest-forked -# - pytest -v -s v1/determinism/test_batch_invariance.py -# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py -# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py -# - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py -# - pytest -v -s v1/determinism/test_matmul_batch_invariant.py -# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA -# - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN +- label: ":nvidia: (H100) Batch Invariance" + key: batch-invariance-h100 + automatic_shard: true + timeout_in_minutes: 60 + device: h100 + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN + mirror: + amd: + label: ":amd: (MI300) Batch Invariance" + dind: false + automatic_shard: true + device: mi300_1 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN -# Disabled until the pytest plugins installation is locked. -# - label: ":nvidia: (B200) Batch Invariance" -# key: batch-invariance-b200 -# automatic_shard: true -# timeout_in_minutes: 45 -# device: b200-k8s -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/lora/ops/triton_ops -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pip install pytest-timeout pytest-forked -# - pytest -v -s v1/determinism/test_batch_invariance.py -# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py -# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py -# - pytest -v -s v1/determinism/test_batch_invariance_vlm.py -# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA -# - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN -# - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py -# - pytest -v -s v1/determinism/test_matmul_batch_invariant.py -# - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py -# - pytest -v -s v1/determinism/test_online_batch_invariance.py -# mirror: -# amd: -# label: ":amd: (MI355) Batch Invariance" -# dind: false -# device: mi355_1 -# timeout_in_minutes: 60 -# depends_on: -# - image-build-amd -# source_file_dependencies: -# - vllm/v1/attention -# - vllm/model_executor/layers -# - vllm/lora/ops/triton_ops -# - vllm/model_executor/kernels/linear -# - vllm/model_executor/determinism -# - tests/v1/determinism/ -# - vllm/_aiter_ops.py -# - vllm/platforms/rocm.py -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pip install pytest-timeout pytest-forked -# - pytest -v -s v1/determinism/test_batch_invariance.py -# - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py -# - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py -# - pytest -v -s v1/determinism/test_batch_invariance_vlm.py -# - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA -# - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN -# - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py -# - pytest -v -s v1/determinism/test_matmul_batch_invariant.py -# - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py -# - pytest -v -s v1/determinism/test_online_batch_invariance.py -# - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py +- label: ":nvidia: (B200) Batch Invariance" + key: batch-invariance-b200 + automatic_shard: true + timeout_in_minutes: 45 + device: b200-k8s + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + mirror: + amd: + label: ":amd: (MI355) Batch Invariance" + dind: false + device: mi355_1 + timeout_in_minutes: 60 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/lora/ops/triton_ops + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s v1/determinism/test_batch_invariance.py + - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py + - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py + - pytest -v -s v1/determinism/test_batch_invariance_vlm.py + - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA + - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py + - pytest -v -s v1/determinism/test_online_batch_invariance.py + - pytest -v -s v1/determinism/test_mx_linear_batch_invariant.py - label: ":nvidia: (H100) [Nightly] Batch Invariance (Distributed)" key: batch-invariance-distributed-4xh100 diff --git a/tools/pre_commit/test_tethering_allowlist.txt b/tools/pre_commit/test_tethering_allowlist.txt index af3b990c8469..6a900d7f871c 100644 --- a/tools/pre_commit/test_tethering_allowlist.txt +++ b/tools/pre_commit/test_tethering_allowlist.txt @@ -104,7 +104,6 @@ tests/rocm/test_moe_weight_replay.py # CI jobs disabled pending locked torchao # Temporary gaps from disabling jobs with unpinned install environments. tests/device_allocator/test_alloc_conf.py # CI jobs disabled pending locked cmake installation -tests/distributed/test_elastic_ep.py # CI jobs disabled pending locked KV connectors installation tests/models/flint/test_decode_attention.py # CI jobs disabled pending locked gguf installation tests/models/flint/test_flint_moe.py # CI jobs disabled pending locked gguf installation tests/models/flint/test_flint_ops.py # CI jobs disabled pending locked gguf installation @@ -186,7 +185,6 @@ tests/tokenizers_/test_detokenize.py # CI jobs disabled pending locked cmake in tests/tokenizers_/test_hf.py # CI jobs disabled pending locked cmake installation tests/tokenizers_/test_mistral.py # CI jobs disabled pending locked cmake installation tests/tokenizers_/test_registry.py # CI jobs disabled pending locked cmake installation -tests/tracing/test_loading_tracing.py # CI jobs disabled pending locked tracing packages installation tests/transformers_utils/processors/test_pixtral.py # CI jobs disabled pending locked cmake installation tests/transformers_utils/processors/test_voxtral.py # CI jobs disabled pending locked cmake installation tests/transformers_utils/test_bailing_moe_v3_vl_config.py # CI jobs disabled pending locked cmake installation @@ -203,23 +201,8 @@ tests/transformers_utils/test_repo_utils.py # CI jobs disabled pending locked c tests/transformers_utils/test_speculators_dspark_config.py # CI jobs disabled pending locked cmake installation tests/transformers_utils/test_speculators_override.py # CI jobs disabled pending locked cmake installation tests/transformers_utils/test_utils.py # CI jobs disabled pending locked cmake installation -tests/v1/determinism/test_batch_invariance_vlm.py # CI jobs disabled pending locked pytest plugins installation -tests/v1/determinism/test_cutlass_batch_invariance.py # CI jobs disabled pending locked pytest plugins installation -tests/v1/determinism/test_lora_shrink_batch_invariant.py # CI jobs disabled pending locked pytest plugins installation -tests/v1/determinism/test_matmul_batch_invariant.py # CI jobs disabled pending locked pytest plugins installation -tests/v1/determinism/test_mx_linear_batch_invariant.py # CI jobs disabled pending locked pytest plugins installation -tests/v1/determinism/test_nvfp4_batch_invariant.py # CI jobs disabled pending locked pytest plugins installation -tests/v1/determinism/test_rms_norm_batch_invariant.py # CI jobs disabled pending locked pytest plugins installation tests/v1/fault_tolerance/test_engine_core_sentinel.py # CI jobs disabled pending locked KV connectors installation tests/v1/fault_tolerance/test_fault_tolerance_e2e.py # CI jobs disabled pending locked KV connectors installation tests/v1/kv_connector/mooncake_integration/test_accuracy.py # CI jobs disabled pending locked KV connectors installation -tests/v1/kv_connector/mooncake_integration/test_accuracy_rocm.py # CI jobs disabled pending locked KV connectors installation tests/v1/kv_connector/mooncake_integration/test_mooncake_imports.py # CI jobs disabled pending locked KV connectors installation tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py # CI jobs disabled pending locked KV connectors installation -tests/v1/kv_connector/nixl_integration/test_accuracy.py # CI jobs disabled pending locked KV connectors installation -tests/v1/kv_connector/nixl_integration/test_accuracy_rocm.py # CI jobs disabled pending locked KV connectors installation -tests/v1/kv_connector/nixl_integration/test_edge_cases.py # CI jobs disabled pending locked KV connectors installation -tests/v1/kv_connector/nixl_integration/test_mamba_prefix_cache.py # CI jobs disabled pending locked KV connectors installation -tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py # CI jobs disabled pending locked KV connectors installation -tests/v1/kv_connector/nixl_integration/test_nixl_imports.py # CI jobs disabled pending locked KV connectors installation -tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py # CI jobs disabled pending locked KV connectors installation From 65ff14056a1833785a843641b87bbc3f6330338e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 28 Oct 2026 15:55:48 -0700 Subject: [PATCH 089/406] [Bugfix][HiSparse] Fix compact KV row regression tests (#64238) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: vllm-agent --- tests/v1/attention/test_sparse_mla_backends.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 8a55a6771b32..ce65f2188862 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -2256,8 +2256,6 @@ def test_sparse_mla_index_groups_own_distinct_physical_buffers(): def test_hisparse_index_group_preserves_row_width( monkeypatch, kv_cache_dtype, row_width, kv_dtype ): - monkeypatch.setattr(index_group_module, "_create_side_stream", MagicMock()) - monkeypatch.setattr(index_group_module, "_create_event", MagicMock()) create_handle = MagicMock(side_effect=lambda *args, **kwargs: SimpleNamespace()) monkeypatch.setattr( index_group_module, "create_hisparse_cache_handle", create_handle @@ -3765,7 +3763,9 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows(row_width, kv_dtype device=device, ) - runtime = _make_hisparse_runtime(row_width=row_width, block_size=block_size) + runtime = _make_hisparse_runtime( + row_width=row_width, block_size=block_size, kv_dtype=kv_dtype + ) runtime.bind_source_cache(host_cache) staged = runtime.gather_prefill_cache( host_cache, plan, resident_cache=resident_cache @@ -3807,7 +3807,9 @@ def test_hisparse_staged_prefill_cache_backs_min_attention_rows(): ) host_cache = torch.zeros(4, block_size, row_width).pin_memory() - staged = HiSparseRuntime.gather_prefill_cache(None, host_cache, plan) + runtime = _make_hisparse_runtime(row_width=row_width, block_size=block_size) + runtime.bind_source_cache(host_cache) + staged = runtime.gather_prefill_cache(host_cache, plan) assert staged.shape[0] * block_size < 256 backing_rows = staged.untyped_storage().nbytes() // ( From ce775739e99debecf79ef147153738d7fa96e29c Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 28 Oct 2026 18:15:31 -0500 Subject: [PATCH 090/406] [CI][ROCm] Drop nvidia-cuda-cccl from the ROCm test lockfile again (#64377) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- requirements/test/rocm.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 7328ac9495d6..e3af4cbc3946 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -654,8 +654,6 @@ numpy==2.2.6 # tritonclient # vocos # xgrammar -nvidia-cuda-cccl==13.3.4.3.1 - # via nvidia-nvshmem-cu13 open-clip-torch==2.32.0 # via -r requirements/test/rocm.in openai==2.31.0 From 15cc0a8debd5f93ace4e1a3ea88975d7e43512ee Mon Sep 17 00:00:00 2001 From: Thang Nguyen <69278249+Thangnguyenvn98@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:54:44 -0400 Subject: [PATCH 091/406] [CI] Enroll Language Models (Standard) in automatic sharding (#64318) Signed-off-by: Thang Nguyen Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/models_language.yaml | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 909bb30e5e0f..887d079f7c0a 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) Language Models (Standard)" key: language-models-tests-standard + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: @@ -11,12 +12,14 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/language commands: - # Test standard language models, excluding a subset of slow tests + # Test standard language models, excluding a subset of slow tests. + # The exclusion sits in -k so the shard planner reads -m as picking tests. - pip freeze | grep -E 'torch' - - pytest -v -s models/language -m 'core_model and (not slow_test)' + - pytest -v -s models/language -m core_model -k 'not slow_test' mirror: amd: label: ":amd: (MI355 DPX) Language Models (Standard)" + automatic_shard: true dind: false device: mi355_dpx timeout_in_minutes: 60 From f79a3a7d48dad51bc12048433c1f04fc34d5942b Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Wed, 28 Oct 2026 20:18:23 -0400 Subject: [PATCH 092/406] [KV Connector] Mooncake store: skip padding in stored KV blocks (#63513) Signed-off-by: Wei Zhao Signed-off-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 Co-authored-by: Yifan Qiao Co-authored-by: OpenAI Codex --- .../unit/test_mooncake_store_hma_e2e.py | 299 ++++++++++++++++-- .../unit/test_mooncake_store_layout.py | 43 ++- .../unit/test_mooncake_store_worker.py | 18 +- .../kv_connector/v1/mooncake/store/data.py | 86 +++-- .../kv_connector/v1/mooncake/store/worker.py | 28 +- 5 files changed, 406 insertions(+), 68 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 0682f9f465f1..75eb088ee8da 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -3,11 +3,14 @@ """End-to-end save->lookup test for MooncakeStoreConnector on a hybrid (SWA + Full) attention config, using a dict-backed mock store.""" +import ctypes import sys import threading import types +from types import SimpleNamespace from unittest.mock import MagicMock, patch +import pytest import torch from tests.v1.attention.utils import dense_kv_cache_views @@ -31,6 +34,7 @@ KVCacheStoreRecvingThread, KVCacheStoreSendingThread, ) +from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.core.kv_cache_utils import BlockHash from vllm.v1.kv_cache_interface import ( CircularBufferSpec, @@ -39,9 +43,11 @@ KVCacheGroupSpec, KVCacheTensor, MambaSpec, + MLAAttentionSpec, SlidingWindowSpec, ) from vllm.v1.kv_cache_layout import KVCacheLayout +from vllm.v1.worker.utils import allocate_kv_cache class _DictStore: @@ -149,6 +155,33 @@ def _build_worker_with_dict_store(vllm_config, kv_cache_config, store): return worker +def _register_with_fake_threads(worker, kv_caches): + """Register kv_caches without starting the transfer threads.""" + + def _fake_thread_init(*args, **kwargs): + """Mock thread that sets all threading.Event args so waits don't block.""" + for v in list(args) + list(kwargs.values()): + if isinstance(v, threading.Event): + v.set() + m = MagicMock() + m.start = lambda: None + return m + + with ( + patch.object( + mooncake_store_worker, + "KVCacheStoreSendingThread", + side_effect=_fake_thread_init, + ), + patch.object( + mooncake_store_worker, + "KVCacheStoreRecvingThread", + side_effect=_fake_thread_init, + ), + ): + worker.register_kv_caches(kv_caches) + + def test_e2e_swa_plus_full_save_then_lookup_hits(): """E2E: build a SWA+Full hybrid worker, save all blocks via the sending thread (synchronously), then verify lookup returns the full hit length. @@ -201,28 +234,13 @@ def test_e2e_swa_plus_full_save_then_lookup_hits(): "L1": dense_kv_cache_views(raw_swa, swa, 4, 1, KVCacheLayout.LBNHC)[0], } - def _fake_thread_init(*args, **kwargs): - """Mock thread that sets all threading.Event args so waits don't block.""" - for v in list(args) + list(kwargs.values()): - if isinstance(v, threading.Event): - v.set() - m = MagicMock() - m.start = lambda: None - return m + _register_with_fake_threads(worker, kv_caches) - with ( - patch.object( - mooncake_store_worker, - "KVCacheStoreSendingThread", - side_effect=_fake_thread_init, - ), - patch.object( - mooncake_store_worker, - "KVCacheStoreRecvingThread", - side_effect=_fake_thread_init, - ), - ): - worker.register_kv_caches(kv_caches) + # Each group's store values hold only its own layer. + assert [db.kv_caches_base_addr for db in worker.token_dbs] == [ + [raw_full.data_ptr()], + [raw_swa.data_ptr()], + ] # Now build a real sending thread (no .start()) over the worker's token_dbs # and the dict-backed store, so _handle_request runs synchronously. @@ -283,6 +301,245 @@ def _fake_thread_init(*args, **kwargs): assert worker.lookup(num_tokens=64, block_hashes=hs).hit_length == 0 +def test_hybrid_mamba_values_hold_only_group_states(): + """Hybrid groups overlay one allocation. A Mamba group's values hold only + its own layers' states: not the layer slot only attention uses, and not + the padding of each state up to the attention page. Restoring the value + brings back the states.""" + attn = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=16, dtype=torch.float16 + ) + page = attn.page_size_bytes + mamba = MambaSpec( + block_size=16, + shapes=((4, 8), (2, 8)), + dtypes=(torch.float32, torch.float16), + page_size_padded=page, + mamba_cache_mode="align", + ) + assert mamba.state_content_size_bytes == 160 < page + num_layers, num_blocks = 3, 4 + + def tensor(layers): + return KVCacheTensor( + size=num_layers * num_blocks * page, + layers=layers, + layer_stride=num_blocks * page, + block_stride=page, + ) + + cfg = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[tensor(["A0", "A1", "A2"]), tensor(["M0", "M1"])], + kv_cache_groups=[ + KVCacheGroupSpec(["A0", "A1", "A2"], attn), + KVCacheGroupSpec(["M0", "M1"], mamba), + ], + ) + worker = _build_worker_with_dict_store( + _minimal_vllm_config(cache_block_size=16), cfg, _DictStore() + ) + kv_caches = allocate_kv_cache(cfg, torch.device("cpu"), KVCacheLayout.LBNHC) + # Like a Mamba backend, view each layer's page as its two states, here + # with a gap between them. + for name in ("M0", "M1"): + page_view = kv_caches[name] + kv_caches[name] = [ + page_view[..., :128].view(torch.float32), + page_view[..., 136:160].view(torch.float16), + ] + _register_with_fake_threads(worker, kv_caches) + + states = [state for name in ("M0", "M1") for state in kv_caches[name]] + state_lens = [state[0].numel() * state.element_size() for state in states] + attn_db, mamba_db = worker.token_dbs + attn_layout, mamba_layout = attn_db.store_layout, mamba_db.store_layout + assert attn_layout.block_len == [page] * num_layers + assert mamba_layout.kv_caches_base_addr == [s.data_ptr() for s in states] + assert mamba_layout.block_len == state_lens + + raw = torch.empty(0, dtype=torch.int8).set_(kv_caches["A0"].untyped_storage()) + raw.copy_(torch.randint(-128, 128, raw.shape, dtype=torch.int8)) + before = raw.clone() + addrs, sizes = mamba_db.prepare_value_for_block(2) + saved = [ctypes.string_at(addr, size) for addr, size in zip(addrs, sizes)] + # The value is each state of block 2 in turn. + assert saved == [ + state[2].contiguous().view(torch.int8).numpy().tobytes() for state in states + ] + + # Clobber block 2 of every layer slot, then load the saved value: only the + # Mamba states come back; the gap, the padding and slot A2 keep the clobber. + def block_2(addr, length): + start = addr - raw.data_ptr() + 2 * page + return slice(start, start + length) + + for name in ("A0", "A1", "A2"): + raw[block_2(kv_caches[name].data_ptr(), page)] = 7 + expected = raw.clone() + for state, length in zip(states, state_lens): + expected[block_2(state.data_ptr(), length)] = before[ + block_2(state.data_ptr(), length) + ] + for addr, data in zip(addrs, saved): + ctypes.memmove(addr, data, len(data)) + assert torch.equal(raw, expected) + + +_ATTN = FullAttentionSpec( + block_size=16, num_kv_heads=2, head_size=16, dtype=torch.float16 +) +_MAMBA = MambaSpec( + block_size=16, + shapes=((4, 8), (2, 8)), + dtypes=(torch.float32, torch.float16), + page_size_padded=_ATTN.page_size_bytes, + mamba_cache_mode="align", +) +# Model shapes as KV cache groups: each group is (layer names, spec). +_MODELS = { + "dense": [(["A0", "A1", "A2"], _ATTN)], + "mla": [ + ( + ["A0", "A1", "A2"], + MLAAttentionSpec( + block_size=16, num_kv_heads=1, head_size=40, dtype=torch.float16 + ), + ) + ], + "full_swa": [ + (["A0", "A1"], _ATTN), + ( + ["S0", "S1"], + SlidingWindowSpec( + block_size=16, + num_kv_heads=2, + head_size=16, + dtype=torch.float16, + sliding_window=32, + ), + ), + ], + # Like Kimi K3: Mamba groups with fewer layers, states padded to the page. + "attn_mamba": [ + (["A0", "A1", "A2"], _ATTN), + (["M0", "M1"], _MAMBA), + (["M2", "M3"], _MAMBA), + ], + # Pages of different sizes are packed side by side within a block. + "mixed_pages": [ + (["A0", "A1"], _ATTN), + ( + ["B0", "B1", "B2"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=8, dtype=torch.float16 + ), + ), + ], +} + + +def _byte_mask(t: torch.Tensor, nbytes: int) -> torch.Tensor: + """The bytes of the backing storage that ``t`` covers.""" + offsets = torch.tensor([t.storage_offset()]) + for size, stride in zip(t.shape, t.stride()): + offsets = (offsets[:, None] + torch.arange(size) * stride).flatten() + elem = t.element_size() + mask = torch.zeros(nbytes, dtype=torch.bool) + mask[(offsets[:, None] * elem + torch.arange(elem)).flatten()] = True + return mask + + +def _span_mask(spans, nbytes: int) -> torch.Tensor: + mask = torch.zeros(nbytes, dtype=torch.bool) + for start, length in spans: + mask[start : start + length] = True + return mask + + +@pytest.mark.parametrize("layout", list(KVCacheLayout), ids=lambda x: x.name) +@pytest.mark.parametrize("model", list(_MODELS)) +def test_store_values_round_trip_every_layout(model, layout): + """Through the planner's real layout, each group's value for a block holds + every byte of its layers' views of that block, reads and writes only that + block, and restores it exactly. Without padding the value is the group's + whole block, as before; with Mamba padding it is exactly the states.""" + num_blocks, block = 4, 2 + groups = [KVCacheGroupSpec(names, spec) for names, spec in _MODELS[model]] + planner_config = SimpleNamespace( + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: layout, + num_gpu_blocks_override=num_blocks, + prefix_cache_retention_interval=None, + ), + attention_config=SimpleNamespace(hisparse_config=None), + ) + try: + cfg = get_kv_cache_config_from_groups(planner_config, groups, 1 << 30) + kv_caches = allocate_kv_cache(cfg, torch.device("cpu"), layout) + except (ValueError, AssertionError): + pytest.skip(f"{layout.name} cannot express {model}") + worker = _build_worker_with_dict_store( + _minimal_vllm_config(cache_block_size=16), cfg, _DictStore() + ) + _register_with_fake_threads(worker, kv_caches) + + first = next(iter(kv_caches.values())) + raw = torch.empty(0, dtype=torch.int8).set_(first.untyped_storage()) + nbytes = raw.numel() + raw.copy_(torch.randint(-128, 128, raw.shape, dtype=torch.int8)) + before = raw.clone() + base = raw.data_ptr() + + # The block of every byte that some layer's view uses (-1: none, e.g. + # padding). A value must not reach another block's bytes. + byte_block = torch.full((nbytes,), -1) + for views in kv_caches.values(): + for view in views if isinstance(views, list) else [views]: + for b in range(num_blocks): + mask = _byte_mask(view[b], nbytes) + if ((byte_block[mask] != -1) & (byte_block[mask] != b)).any(): + pytest.skip(f"{layout.name} overlaps blocks of {model}") + byte_block[mask] = b + other_blocks = (byte_block != block) & (byte_block != -1) + + for group, db in zip(groups, worker.token_dbs, strict=True): + store_layout = db.store_layout + addrs, sizes = store_layout.prepare_value_for_block(block) + chunk_addrs, chunk_sizes, _ = store_layout.prepare_values( + [(0, 16), (16, 32)], [block, block], [0, 0] + ) + # Chunked saves and loads copy the same pieces as a single block. + assert chunk_addrs == [addrs, addrs] and chunk_sizes == [sizes, sizes] + + value = _span_mask([(a - base, n) for a, n in zip(addrs, sizes)], nbytes) + owned = torch.zeros(nbytes, dtype=torch.bool) + for name in group.layer_names: + views = kv_caches[name] + for view in views if isinstance(views, list) else [views]: + owned |= _byte_mask(view[block], nbytes) + assert not (value & other_blocks).any(), "value reaches other blocks" + assert not (owned & ~value).any(), "value misses bytes of the group" + # Each piece ends at the last byte the group uses in it. + assert all(owned[a - base + n - 1] for a, n in zip(addrs, sizes)) + if len(groups) == 1 and layout != KVCacheLayout.BHLNC: + # A single group stores whole blocks, as before. (A BHLNC head's + # block stride spans every head, so only the head's block is.) + assert store_layout.block_len == store_layout.block_stride + if not layout.is_block_outermost: + # Exactly the group's bytes: no padding, no other group's slots. + assert torch.equal(value, owned) and sum(sizes) == owned.sum() + + # Clobber everything but the other blocks, then load the saved value. + saved = [ctypes.string_at(a, n) for a, n in zip(addrs, sizes)] + raw[~other_blocks] = 7 + for addr, data in zip(addrs, saved): + ctypes.memmove(addr, data, len(data)) + assert torch.equal(raw[owned], before[owned]) + assert torch.equal(raw[other_blocks], before[other_blocks]) + raw.copy_(before) + + def test_recv_skips_swa_blocks_before_window(): """Producer stored every block for both groups; consumer must only fetch SWA blocks within the sliding window, not the head.""" diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_layout.py b/tests/v1/kv_connector/unit/test_mooncake_store_layout.py index 00ff3eea4d58..dce3dad7928b 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_layout.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_layout.py @@ -34,9 +34,9 @@ def _rank_local_reference( layout: RankLocalStoreLayout, start: int, end: int, block_ids: list[int] ) -> tuple[list[int], list[int], int]: block_id = block_ids[start // layout.block_size] - length = len(layout.block_len) + length = len(layout.block_stride) addrs = [ - base_addr + block_id * layout.block_len[index % length] + base_addr + block_id * layout.block_stride[index % length] for index, base_addr in enumerate(layout.kv_caches_base_addr) ] sizes = [ @@ -79,6 +79,45 @@ def test_rank_local_descriptors_handle_empty_and_invalid_chunks(): layout.prepare_values([(0, BLOCK_SIZE + 1)], [0, 1], [0]) +def test_rank_local_layout_skips_block_padding(): + """A state padded up to a larger page, like a Mamba state padded to the + attention page size, is stored without the padding.""" + num_layers, num_blocks, page, state = 2, 3, 64, 40 + # Layer-outermost buffer; each layer views its state at the start of + # every page. + raw = torch.zeros(num_layers * num_blocks * page, dtype=torch.int8) + views = [ + torch.as_strided( + raw, + (num_blocks, 1, 1, state), + (page, state, state, 1), + layer * num_blocks * page, + ) + for layer in range(num_layers) + ] + layout = RankLocalStoreLayout( + KeyMetadata("test-model", 0, 0, 0, 0), BLOCK_SIZE, BLOCK_SIZE + ) + layout.register_kv_caches(views, num_blocks) + + assert layout.kv_caches_base_addr == [view.data_ptr() for view in views] + assert layout.block_stride == [page] * num_layers + assert layout.block_len == [state] * num_layers + addrs, sizes, _ = layout.prepare_values( + [(0, BLOCK_SIZE), (BLOCK_SIZE, 3 * BLOCK_SIZE)], [2, 0, 1], [0, 0] + ) + assert addrs == [ + [view.data_ptr() + 2 * page for view in views], + [view.data_ptr() for view in views], + ] + # A two-block chunk copies the padding between its blocks. + assert sizes == [[state] * num_layers, [page + state] * num_layers] + assert layout.prepare_value_for_block(1) == ( + [view.data_ptr() + page for view in views], + [state] * num_layers, + ) + + @pytest.mark.parametrize( ("layout_cls", "store_format"), [ diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 09d27eb947a0..ee3961171b88 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -2340,10 +2340,10 @@ def test_recv_thread_splits_disk_offload_loads_by_budget(): layout = thread.token_databases[0].store_layout assert isinstance(layout, RankLocalStoreLayout) base_addr = layout.kv_caches_base_addr[0] - block_len = layout.block_len[0] - assert first_addrs == [[base_addr], [base_addr + block_len]] - assert second_addrs == [[base_addr + 2 * block_len]] - expected_size = block_len + block_stride = layout.block_stride[0] + assert first_addrs == [[base_addr], [base_addr + block_stride]] + assert second_addrs == [[base_addr + 2 * block_stride]] + expected_size = layout.block_len[0] assert first_sizes == [[expected_size], [expected_size]] assert second_sizes == [[expected_size]] @@ -4630,15 +4630,15 @@ def test_register_kv_caches_shared_storage(layout: KVCacheLayout): for cache in caches for head_idx in range(cache.shape[1]) ] - assert db.block_len == [head_cache.stride(0) * head_cache.element_size()] * ( + assert db.block_stride == [head_cache.stride(0) * head_cache.element_size()] * ( num_layers * caches[0].shape[1] ) elif layout.is_layer_compact: assert db.kv_caches_base_addr == [cache.data_ptr() for cache in caches] - assert db.block_len == [spec.page_size_bytes] * num_layers + assert db.block_stride == [spec.page_size_bytes] * num_layers else: assert db.kv_caches_base_addr == [raw.data_ptr()] - assert db.block_len == [num_layers * spec.page_size_bytes] + assert db.block_stride == [num_layers * spec.page_size_bytes] worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes) @@ -4707,7 +4707,7 @@ def test_register_kv_caches_uses_transfer_group_memory_domain(): db = worker.token_dbs[0] assert db.kv_caches_base_addr == [source.data_ptr()] - assert db.block_len == [page_size] + assert db.block_stride == [page_size] worker.store.register_buffer.assert_called_once_with( source.data_ptr(), source.nbytes ) @@ -4770,7 +4770,7 @@ def test_register_kv_caches_separate_head_groups(): ] db = worker.token_dbs[0] assert db.kv_caches_base_addr == expected_addrs - assert db.block_len == [head_block_bytes] * len(expected_addrs) + assert db.block_stride == [head_block_bytes] * len(expected_addrs) worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 4cd0402b7b4f..be4b5a030414 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -234,6 +234,10 @@ def __init__( super().__init__(metadata, block_size, hash_block_size) self._key_prefix = PoolKey.build_prefix(metadata) self.kv_caches_base_addr: list[int] = [] + self.block_stride: list[int] = [] + # Bytes at the start of each block that are stored: up to the last + # byte of the group's data, skipping e.g. the padding of a Mamba state + # up to the attention page size. self.block_len: list[int] = [] @property @@ -263,44 +267,68 @@ def set_kv_caches_base_addr(self, base_addrs: list[int]) -> None: self.kv_caches_base_addr = base_addrs def set_block_len(self, block_lens: list[int]) -> None: - self.block_len = block_lens + self.block_stride = block_lens + self.block_len = list(block_lens) def register_kv_caches( self, kv_caches: Sequence[torch.Tensor], num_blocks: int, ) -> None: - seen_region_ptrs: set[int] = set() + # Map each region's base address to its entry in the parallel lists. + region_index: dict[int, int] = {} base_addrs: list[int] = [] + block_strides: list[int] = [] block_lens: list[int] = [] + + def add_region(addr: int, block_stride: int, block_len: int) -> None: + assert 0 < block_len <= block_stride + index = region_index.get(addr) + if index is None: + region_index[addr] = len(base_addrs) + base_addrs.append(addr) + block_strides.append(block_stride) + block_lens.append(block_len) + else: + # Several views can share a region (e.g. the layers of a + # block-outermost buffer); keep the bytes up to the last one. + block_lens[index] = max(block_lens[index], block_len) + for cache in kv_caches: storage = cache.untyped_storage() storage_addr = storage.data_ptr() storage_len = storage.nbytes() + block_bytes = cache[0].numel() * cache.element_size() if not is_non_overlapping_and_dense(cache[0]): + # One region per head, copying only that head's block (in a + # block-outermost buffer the block stride spans every head). for head_idx in range(cache.shape[1]): head_cache = cache[:, head_idx] assert is_non_overlapping_and_dense(head_cache[0]) - region_addr = head_cache.data_ptr() - if region_addr in seen_region_ptrs: - continue - seen_region_ptrs.add(region_addr) - base_addrs.append(region_addr) - block_lens.append(head_cache.stride(0) * head_cache.element_size()) + add_region( + head_cache.data_ptr(), + head_cache.stride(0) * head_cache.element_size(), + head_cache[0].numel() * head_cache.element_size(), + ) elif cache.stride(0) * cache.element_size() * num_blocks == storage_len: - if storage_addr in seen_region_ptrs: - continue - seen_region_ptrs.add(storage_addr) - base_addrs.append(storage_addr) - block_lens.append(storage_len // num_blocks) + # Block-outermost buffer: each block holds every layer, so the + # whole storage is one region. In block 0 this layer's data + # starts at cache.data_ptr() - storage_addr and spans + # block_bytes, so the block is used up to that end; the + # region keeps the end of its last layer. + add_region( + storage_addr, + storage_len // num_blocks, + cache.data_ptr() - storage_addr + block_bytes, + ) else: - region_addr = cache.data_ptr() - if region_addr in seen_region_ptrs: - continue - seen_region_ptrs.add(region_addr) - base_addrs.append(region_addr) - block_lens.append(cache.stride(0) * cache.element_size()) + add_region( + cache.data_ptr(), + cache.stride(0) * cache.element_size(), + block_bytes, + ) self.kv_caches_base_addr = base_addrs + self.block_stride = block_strides self.block_len = block_lens def prepare_values( @@ -315,7 +343,11 @@ def prepare_values( shard_id == 0 for shard_id in shard_ids ) base = np.asarray(self.kv_caches_base_addr, dtype=np.int64) - length = len(self.block_len) + length = len(self.block_stride) + stride = np.asarray( + [self.block_stride[i % length] for i in range(base.shape[0])], + dtype=np.int64, + ) blen = np.asarray( [self.block_len[i % length] for i in range(base.shape[0])], dtype=np.int64, @@ -329,16 +361,18 @@ def prepare_values( dtype=np.int64, count=n, ) - addrs = base[None, :] + bids[:, None] * blen[None, :] + addrs = base[None, :] + bids[:, None] * stride[None, :] block_counts = (spans + self.block_size - 1) // self.block_size - sizes = blen[None, :] * block_counts[:, None] + # A chunk spanning several consecutive blocks also copies the + # padding between them. + sizes = stride[None, :] * (block_counts[:, None] - 1) + blen[None, :] return addrs.tolist(), sizes.tolist(), bids.tolist() def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]: - length = len(self.block_len) + length = len(self.block_stride) return ( [ - base_addr + block_id * self.block_len[index % length] + base_addr + block_id * self.block_stride[index % length] for index, base_addr in enumerate(self.kv_caches_base_addr) ], [ @@ -581,8 +615,8 @@ def kv_caches_base_addr(self) -> list[int]: return self._rank_local_layout().kv_caches_base_addr @property - def block_len(self) -> list[int]: - return self._rank_local_layout().block_len + def block_stride(self) -> list[int]: + return self._rank_local_layout().block_stride def _rank_local_layout(self) -> RankLocalStoreLayout: if not isinstance(self.store_layout, RankLocalStoreLayout): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 004b69286ea3..27bbd80a4f65 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1737,13 +1737,11 @@ def register_kv_caches( logger.warning("No KV caches to offload.") return - # Resolve each entry to a representative tensor for storage - # deduplication. For attention layers the value is already a tensor; - # for Mamba layers it is a list of tensors that all share the same - # underlying raw storage, so we take the first one. - def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: + # A layer's cache is one tensor, or a list of state tensors that share + # one raw storage (Mamba). + def _layer_views(v: torch.Tensor | list[torch.Tensor]) -> list[torch.Tensor]: assert isinstance(v, torch.Tensor | list) - return v if isinstance(v, torch.Tensor) else v[0] + return [v] if isinstance(v, torch.Tensor) else v assert self.cache_config.num_gpu_blocks is not None self.num_blocks = self.cache_config.num_gpu_blocks @@ -1757,7 +1755,7 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: } seen_storage_ptrs: set[int] = set() cache_tensors = [ - group_kernel_blocks(_repr_tensor(cache), self.num_blocks) + group_kernel_blocks(_layer_views(cache)[0], self.num_blocks) for layer_name, cache in kv_caches.items() if layer_name in store_layer_names ] @@ -1776,8 +1774,18 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: region_len, ret, ) - for db in self.token_dbs: - db.store_layout.register_kv_caches(cache_tensors, self.num_blocks) + for group, db in zip(self._kv_cache_groups, self.token_dbs, strict=True): + # In a hybrid model each KV cache group registers only its own + # layers. A group with no layers on this PP rank registers all. + group_caches = [ + group_kernel_blocks(view, self.num_blocks) + for layer_name in group.layer_names + if layer_name in kv_caches + for view in _layer_views(kv_caches[layer_name]) + ] + db.store_layout.register_kv_caches( + group_caches or cache_tensors, self.num_blocks + ) registered_buffers: dict[int, int] = {} num_host_segments = 0 @@ -1803,7 +1811,7 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: addrs: list[int] = [] block_lens: list[int] = [] for layer_name, value in group_caches: - cache = _repr_tensor(value) + cache = _layer_views(value)[0] is_host_resident = group.host_resident if is_host_resident: num_blocks = self._kv_cache_config.hisparse_host_num_blocks From 36d4ef654a83e227216b0c0b55cf4041d96b1675 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Wed, 28 Oct 2026 17:33:23 -0700 Subject: [PATCH 093/406] [CI] Import the merge-base commit's image cache in PR builds (#64387) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/image_build/image_build.sh | 30 +++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/.buildkite/image_build/image_build.sh b/.buildkite/image_build/image_build.sh index 6d22acd5b16c..53cb3bc044b1 100755 --- a/.buildkite/image_build/image_build.sh +++ b/.buildkite/image_build/image_build.sh @@ -203,6 +203,35 @@ resolve_parent_commit() { fi } +resolve_merge_base_cache_commit() { + # PR builds: the commit the PR branched from, whose build pushed + # vllm-ci-test-cache:. A cache source only (see ci.hcl), unlike + # the VLLM_MERGE_BASE_COMMIT build arg, which would change the cache key. + if [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]]; then + return 0 + fi + if [[ -n "${VLLM_CACHE_MERGE_BASE_COMMIT:-}" ]]; then + echo "Using provided VLLM_CACHE_MERGE_BASE_COMMIT: ${VLLM_CACHE_MERGE_BASE_COMMIT}" + return 0 + fi + local base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" + local base_ref="refs/remotes/vllm-cache-base/${base_branch}" + git fetch --quiet --no-tags --depth=200 origin \ + "+refs/heads/${base_branch}:${base_ref}" 2>/dev/null || true + VLLM_CACHE_MERGE_BASE_COMMIT=$(git merge-base HEAD "${base_ref}" 2>/dev/null || echo "") + if [[ -z "${VLLM_CACHE_MERGE_BASE_COMMIT}" ]]; then + git fetch --quiet --no-tags --deepen=1000 origin \ + "+refs/heads/${base_branch}:${base_ref}" 2>/dev/null || true + VLLM_CACHE_MERGE_BASE_COMMIT=$(git merge-base HEAD "${base_ref}" 2>/dev/null || echo "") + fi + if [[ -n "${VLLM_CACHE_MERGE_BASE_COMMIT}" ]]; then + echo "Computed merge base with ${base_branch} for cache fallback: ${VLLM_CACHE_MERGE_BASE_COMMIT}" + export VLLM_CACHE_MERGE_BASE_COMMIT + else + echo "Could not determine merge base with ${base_branch}; skipping that cache source" + fi +} + print_bake_config() { echo "--- :page_facing_up: Resolved bake configuration" # Write to a temp directory to avoid polluting the repo root (which is the @@ -353,6 +382,7 @@ setup_buildx_builder resolve_parent_commit export PARENT_COMMIT +resolve_merge_base_cache_commit print_bake_config From dc9278ee27e5d5175359150eae836e5bba15d8ed Mon Sep 17 00:00:00 2001 From: Park Soobin <101092510+subin9@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:10:57 +0900 Subject: [PATCH 094/406] [Bugfix][LoRA] Load the head of SEQ_CLS adapters saved without modules_to_save (#63286) Signed-off-by: subin9 <101092510+subin9@users.noreply.github.com> Co-authored-by: linitra24 --- vllm/lora/peft_helper.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index 721b7a039f00..459c36d2d8a5 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -31,6 +31,7 @@ class PEFTHelper: bias: Literal["none"] = field(default="none") modules_to_save: list[str] | None = field(default=None) + task_type: str | None = field(default=None) # True to use Rank-Stabilized LoRA (rsLoRA, see: https://arxiv.org/abs/2312.03732) use_rslora: bool = field(default=False) # True to use Weight-Decomposed Low-Rank Adaptation (DoRA, see: https://arxiv.org/abs/2402.09353) @@ -110,6 +111,12 @@ def __post_init__(self): f"LoRA rank for `{key}` in `rank_pattern` must be a positive " f"integer, got {rank}." ) + # PEFT adds the classification head to modules_to_save whenever it + # loads a SEQ_CLS or TOKEN_CLS adapter, even if the config omits it. + if self.task_type in ("SEQ_CLS", "TOKEN_CLS"): + self.modules_to_save = list( + dict.fromkeys([*(self.modules_to_save or []), "classifier", "score"]) + ) if self.use_rslora: logger.info_once("Loading LoRA weights trained with rsLoRA.") self.vllm_lora_scaling_factor = self._scaling(self.r, self.lora_alpha) From f5aa164b514819fea7dffa245457f04965cb9309 Mon Sep 17 00:00:00 2001 From: Park Soobin <101092510+subin9@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:12:01 +0900 Subject: [PATCH 095/406] [Bugfix][LoRA] Refuse adapters with LoRA on modules vLLM does not wrap (#63489) Signed-off-by: subin9 <101092510+subin9@users.noreply.github.com> Co-authored-by: linitra24 --- tests/lora/test_sequence_classification.py | 53 ++++++++++++++++++++++ vllm/lora/lora_model.py | 7 +++ vllm/lora/model_manager.py | 8 ++++ vllm/lora/worker_manager.py | 1 + 4 files changed, 69 insertions(+) diff --git a/tests/lora/test_sequence_classification.py b/tests/lora/test_sequence_classification.py index 773e5448c85f..3ed3cf3905c4 100644 --- a/tests/lora/test_sequence_classification.py +++ b/tests/lora/test_sequence_classification.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json + +import pytest import torch +from safetensors.torch import save_file from tests.utils import multi_gpu_test from vllm import LLM, PoolingParams @@ -10,6 +14,7 @@ MODEL_NAME = "Qwen/Qwen3-0.6B" NATIVE_MODEL_NAME = "Skywork/Skywork-Reward-V2-Qwen3-0.6B" +BERT_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L6-v2" PROMPTS = { "star_trek": "Does warp drive appear in Star Trek?", "new_zealand": "Wellington is the capital of New Zealand.", @@ -100,6 +105,54 @@ def test_native_classification_model_with_modules_to_save( cleanup_dist_env_and_memory() +def _save_dense_adapter(path, hidden_size: int, modules: list[str]) -> str: + path.mkdir() + config = {"r": 8, "lora_alpha": 16, "target_modules": ["dense"]} + (path / "adapter_config.json").write_text(json.dumps(config)) + tensors = {} + for module in modules: + tensors[f"base_model.model.{module}.lora_A.weight"] = torch.zeros( + 8, hidden_size + ) + tensors[f"base_model.model.{module}.lora_B.weight"] = torch.zeros( + hidden_size, 8 + ) + save_file(tensors, str(path / "adapter_model.safetensors")) + return str(path) + + +def test_lora_on_unsupported_module_is_rejected(tmp_path) -> None: + """BERT's pooler.dense is an nn.Linear that LoRA does not wrap. Its LoRA + weights must fail the load instead of being ignored, while the encoder's + dense layers with the same suffix still load.""" + llm = LLM( + model=BERT_MODEL_NAME, + runner="pooling", + dtype="float16", + enable_lora=True, + max_lora_rank=8, + enforce_eager=True, + max_model_len=512, + gpu_memory_utilization=0.3, + ) + hidden_size = llm.llm_engine.model_config.get_hidden_size() + encoder_dense = "bert.encoder.layer.0.attention.output.dense" + pooler_adapter = _save_dense_adapter( + tmp_path / "pooler", hidden_size, [encoder_dense, "bert.pooler.dense"] + ) + encoder_adapter = _save_dense_adapter( + tmp_path / "encoder", hidden_size, [encoder_dense] + ) + + with pytest.raises(Exception, match="bert.pooler.dense"): + llm.llm_engine.add_lora(LoRARequest("pooler", 1, pooler_adapter)) + assert llm.llm_engine.add_lora(LoRARequest("encoder", 2, encoder_adapter)) + assert set(llm.llm_engine.list_loras()) == {2} + + del llm + cleanup_dist_env_and_memory() + + def _test_batched_loras( qwen3_guard_star_trek_lora_files: str, qwen3_guard_new_zealand_lora_files: str, diff --git a/vllm/lora/lora_model.py b/vllm/lora/lora_model.py index f4520b12c75b..886d285d9ea7 100644 --- a/vllm/lora/lora_model.py +++ b/vllm/lora/lora_model.py @@ -207,6 +207,7 @@ def from_local_checkpoint( weights_mapper: WeightsMapper | None = None, skip_prefixes: list[str] | None = None, moe_ep_spec: MoEEPLoadSpec | None = None, + unsupported_lora_modules: set[str] | None = None, ) -> "LoRAModel": """Create a LoRAModel from a local checkpoint. @@ -231,6 +232,9 @@ def from_local_checkpoint( slicing metadata shared across all MoE layers. Non-local expert weights are skipped at read time instead of being loaded and discarded later. + unsupported_lora_modules: Full names of model layers that match + expected_lora_modules but are not wrapped by LoRA, such as + BERT's pooler.dense. LoRA weights on them are unexpected. Returns: Loaded LoRA Model. @@ -270,6 +274,9 @@ def check_unexpected_modules(modules: dict): ): unexpected_modules.append(module_name) + elif module_name in (unsupported_lora_modules or ()): + unexpected_modules.append(f"{module_name} (not a LoRA layer)") + if unexpected_modules: raise ValueError( f"While loading {lora_dir}, expected" diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index bf270d6d5667..bcefa043d4dc 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -482,6 +482,9 @@ def _parent_module(module_name: str) -> str: return module_name.rpartition(".")[0] wrapped_by_id: dict[int, BaseLayerWithLoRA] = {} + # Layers that share a suffix with a LoRA layer but are not wrapped, + # such as BERT's pooler.dense. Adapters with LoRA on them are refused. + self.unsupported_lora_modules: set[str] = set() for module_name, module in self.model.named_modules(remove_duplicate=False): if isinstance(module, PPMissingLayer): @@ -495,6 +498,11 @@ def _parent_module(module_name: str) -> str: module, self.supported_lora_modules, ): + if ( + module_name.rsplit(".", 1)[-1] in self.supported_lora_modules + and self._get_punica_wrapper(module_name) is not None + ): + self.unsupported_lora_modules.add(module_name) continue elif not self._match_target_modules(module_name) and not is_classifier_head: # Full classification heads bypass the A/B LoRA target filter. diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 7e3d485fabf8..a229718fe216 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -148,6 +148,7 @@ def _load_adapter(self, lora_request: LoRARequest) -> LoRAModel: weights_mapper=hf_to_vllm_mapper, skip_prefixes=lora_skip_prefixes, moe_ep_spec=self._adapter_manager.moe_ep_load_spec, + unsupported_lora_modules=self._adapter_manager.unsupported_lora_modules, ) # Stamp the on-disk MoE layout onto the loaded model so the # adapter manager can route 3D-format checkpoints through the From fed7de909f894fbee0c45130376aed2e2c9bf0d3 Mon Sep 17 00:00:00 2001 From: Moein Khazraee <33970824+mkhazraee@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:23:06 -0700 Subject: [PATCH 096/406] [KV Offload] Update KVCR to v0.2.0 (#64224) Signed-off-by: Moein Khazraee Co-authored-by: vllm-agent --- .buildkite/scripts/install-kv-offload.sh | 2 +- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 5 +---- vllm/v1/kv_offload/tiering/kvcr/manager.py | 5 ++--- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh index eb8bab9e393f..034d8e76b1c3 100644 --- a/.buildkite/scripts/install-kv-offload.sh +++ b/.buildkite/scripts/install-kv-offload.sh @@ -10,7 +10,7 @@ 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.1" + "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@8c6fb79be4dfd9d7bb797c7dc4656790273bc786" # Keep only the NIXL wheel matching the CI image's CUDA runtime. NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p') diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index beca574d8881..946e5cfabba2 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -7,7 +7,6 @@ from typing import Any from unittest.mock import Mock -import msgspec import numpy as np import pytest @@ -367,9 +366,7 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch, backpressure): tier.on_new_request(ctx) - assert kvcr.submit_hint_calls == [ - ((), {"request_id": "req", "hints": msgspec.to_builtins(kv_hint)}) - ] + assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": kv_hint})] bindings = kvcr.constructor_bindings assert bindings is not None diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index e24531120b5d..7f4669c962dd 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -11,7 +11,6 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar, cast -import msgspec from kvcr import ( DURATION_METRIC, KVCR, @@ -20,6 +19,7 @@ TRANSFER_BLOCKS_METRIC, TRANSFER_BYTES_METRIC, KVCRBindings, + ZmqPeerControlChannel, ) from kvcr.config import ( G3Options, @@ -30,7 +30,6 @@ LocalDramOptions, RemoteFWDramOptions, ) -from kvcr.control_channels import ZmqPeerControlChannel from kvcr.policy import ( FIFOPolicy, G3FIFOPolicy, @@ -618,7 +617,7 @@ def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: if (hint := req_context.kv_hints) is not None: self._kvcr.submit_hint( request_id=req_context.req_id, - hints=msgspec.to_builtins(hint), + hints=hint, ) return RequestOffloadingContext() From d295942d0d9dfd25e1daf09cb1156c7a5c43c9d6 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 29 Oct 2026 09:25:36 +0800 Subject: [PATCH 097/406] [Spec Decode] Support trace replay with speculative decoding (#63275) Signed-off-by: aoshen524 Co-authored-by: aoshen524 Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- docs/serving/online_serving/trace_replay.md | 2 +- tests/v1/sample/test_logprobs.py | 82 ++++++++++++++ tests/v1/sample/test_trace_replay_params.py | 11 +- .../test_rejection_sampler_utils.py | 1 + .../test_gpu_rejection_sampler_chunking.py | 1 + tests/v1/worker/test_gpu_trace_replay.py | 94 ++++++++++++++++ vllm/sampling_params.py | 4 - vllm/v1/worker/gpu/sample/trace_replay.py | 100 ++++++++++++++++++ .../gpu/spec_decode/rejection_sampler.py | 4 + 9 files changed, 288 insertions(+), 11 deletions(-) diff --git a/docs/serving/online_serving/trace_replay.md b/docs/serving/online_serving/trace_replay.md index 86ed13644c7a..9acd09b57ad6 100644 --- a/docs/serving/online_serving/trace_replay.md +++ b/docs/serving/online_serving/trace_replay.md @@ -44,6 +44,7 @@ When `trace_decode_token_ids` is set: - The trace is truncated if it does not fit within `max_model_len` given the prompt length. - All stop conditions (EOS, stop strings, stop token IDs) are disabled. - Generation produces exactly the trace tokens, then stops. +- With speculative decoding, a draft token is accepted only if it equals the trace token at its position, so a step emits the matching drafts plus one trace token, up to the end of the trace. The returned logprobs are those of the trace tokens, but each verification step scores several positions in one forward pass, so they are bitwise equal to the logprobs without speculative decoding only with batch invariance (`VLLM_BATCH_INVARIANT=1`). ## Limitations @@ -53,7 +54,6 @@ When `trace_decode_token_ids` is set: | --- | --- | | `n > 1` | Trace replay produces a single deterministic sequence | | `prompt_logprobs` | Expanded logit layout conflicts with trace kernel indexing | -| Speculative decoding | Multi-token speculation conflicts with single-token trace stepping | | Structured outputs | Grammar constraints conflict with forced token injection | | `repetition_detection` | Would terminate the request on repeated patterns in the trace | | `thinking_token_budget` | Logit masking corrupts logprobs for trace tokens | diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index f7f51ac47d88..2f13694bccad 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -1628,3 +1628,85 @@ def test_token_logprobs_large_batch_int64_row_offset(): assert torch.allclose(logprobs[last, 0], ref, atol=1e-2), ( f"logprob {logprobs[last, 0].item()} != ref {ref.item()}" ) + + +def test_trace_replay_with_spec_decode(monkeypatch): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + + prompt = "Hello world " * 50 + llm_kwargs = dict( + max_model_len=512, + max_logprobs=5, + seed=42, + enable_trace_replay=True, + enable_chunked_prefill=True, + max_num_batched_tokens=32, + enable_prefix_caching=False, + gpu_memory_utilization=0.25, + **GPU_DETERMINISM_KWARGS, + ) + + def replay(runner, traces): + outputs = runner.llm.generate( + [prompt] * len(traces), + [ + SamplingParams( + trace_decode_token_ids=trace, max_tokens=len(trace), logprobs=3 + ) + for trace in traces + ], + ) + results = [] + for output in outputs: + completion = output.outputs[0] + assert completion.logprobs is not None + results.append((list(completion.token_ids), completion.logprobs)) + return results + + with VllmRunner("Qwen/Qwen3-0.6B", **llm_kwargs) as ref_runner: + sampled = ref_runner.llm.generate( + [prompt, prompt], + [ + SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True), + SamplingParams(temperature=1.0, seed=7, max_tokens=32, ignore_eos=True), + ], + ) + greedy_trace = list(sampled[0].outputs[0].token_ids) + sampled_trace = list(sampled[1].outputs[0].token_ids) + diverged_trace = greedy_trace[:10] + sampled_trace[10:20] + greedy_trace[20:] + traces = [greedy_trace, sampled_trace, diverged_trace] + ref_results = replay(ref_runner, traces) + + with VllmRunner( + "Qwen/Qwen3-0.6B", + speculative_config={ + "method": "ngram", + "prompt_lookup_max": 5, + "prompt_lookup_min": 3, + "num_speculative_tokens": 3, + }, + disable_log_stats=False, + **llm_kwargs, + ) as spec_runner: + spec_results = replay(spec_runner, traces) + num_accepted = sum( + m.value + for m in spec_runner.llm.get_metrics() + if m.name == "vllm:spec_decode_num_accepted_tokens" + ) + assert num_accepted > 0 + + for trace, (ref_ids, ref_lps), (spec_ids, spec_lps) in zip( + traces, ref_results, spec_results + ): + assert ref_ids == trace + assert spec_ids == trace + assert len(ref_lps) == len(trace) + assert len(spec_lps) == len(trace) + for token_id, ref_lp, spec_lp in zip(trace, ref_lps, spec_lps): + assert math.isclose( + ref_lp[token_id].logprob, + spec_lp[token_id].logprob, + rel_tol=5e-2, + abs_tol=2.5e-1, + ) diff --git a/tests/v1/sample/test_trace_replay_params.py b/tests/v1/sample/test_trace_replay_params.py index 43f8266dede5..36e0749e556d 100644 --- a/tests/v1/sample/test_trace_replay_params.py +++ b/tests/v1/sample/test_trace_replay_params.py @@ -102,13 +102,12 @@ def test_validate_trace_replay_noop_when_unset(): ) -def test_trace_decode_token_ids_rejects_speculative_decoding(): +def test_trace_decode_token_ids_allows_speculative_decoding(): params = SamplingParams(trace_decode_token_ids=[1]) - with pytest.raises(ValueError, match="not supported with speculative decoding"): - params._validate_trace_replay( - _make_model_config(vocab_size=100), - speculative_config=object(), - ) + params._validate_trace_replay( + _make_model_config(vocab_size=100), + speculative_config=object(), + ) def test_trace_decode_token_ids_rejects_structured_outputs(): diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index fd282f434624..e51a6421ee2a 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -1098,6 +1098,7 @@ def test_verify_rejects_unproposed_drafts(): seeds=SimpleNamespace(gpu=inputs["seed"]), ), use_fp64_gumbel=False, + trace_replay_state=None, ) rejection_sampler.num_speculative_steps = K rejection_sampler.synthetic_conditional_rates = None diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py index 13f70139d373..a8d9d3d76739 100644 --- a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py +++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py @@ -187,6 +187,7 @@ def test_replay_on_off_preserves_rejection_sampling_and_rng(monkeypatch): ), use_fp64_gumbel=False, apply_sampling_params=lambda logits, *_: logits, + trace_replay_state=None, ) rejection_sampler = _make_rejection_sampler(sampler, num_speculative_steps) pos = torch.arange(12, dtype=torch.int32, device=device) diff --git a/tests/v1/worker/test_gpu_trace_replay.py b/tests/v1/worker/test_gpu_trace_replay.py index 646842108744..d82ea0b86a9a 100644 --- a/tests/v1/worker/test_gpu_trace_replay.py +++ b/tests/v1/worker/test_gpu_trace_replay.py @@ -24,6 +24,7 @@ from vllm.v1.worker.gpu.sample.trace_replay import ( TraceReplayState, apply_trace_tokens, + apply_trace_tokens_to_verified, ) from vllm.v1.worker.gpu.states import RequestState @@ -195,3 +196,96 @@ def test_slot_reuse_clears_trace(): _set_lens(state, [3, 0], [3, 0]) state.apply_trace(sampled, idx_mapping) assert sampled.tolist() == [888] + + +class TestApplyTraceToVerified: + @pytest.mark.parametrize( + "drafts,trace,step,expected_tokens,expected_num_sampled", + [ + ([10, 11, 12], [10, 11, 12, 13, 14], 0, [10, 11, 12, 13], 4), + ([10, 99, 12], [10, 11, 12, 13, 14], 0, [10, 11], 2), + ([99, 11, 12], [10, 11, 12, 13, 14], 0, [10], 1), + ([-1, -1, -1], [10, 11, 12, 13, 14], 0, [10], 1), + ([11, 12, 13], [10, 11, 12, 13, 14], 1, [11, 12, 13, 14], 4), + ([12, 13, 14], [10, 11, 12, 13, 14], 2, [12, 13, 14], 3), + ([13, 14, 15], [10, 11, 12, 13, 14], 3, [13, 14], 2), + ([], [10, 11, 12], 1, [11], 1), + ], + ids=[ + "all-accepted", + "reject-middle", + "reject-first", + "placeholder-drafts", + "offset-step", + "trace-ends-at-bonus", + "trace-ends-mid-step", + "no-drafts", + ], + ) + def test_accepts_drafts_matching_trace( + self, drafts, trace, step, expected_tokens, expected_num_sampled + ): + num_steps = 3 + prompt_len = 6 + trace_token_ids = torch.zeros( + 1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE + ) + trace_token_ids[0, : len(trace)] = _i32(trace) + sampled = torch.full((1, num_steps + 1), -7, dtype=torch.int64, device=DEVICE) + num_sampled = _i32([1]) + draft_sampled = _i64([5, *drafts]) + cu_num_logits = _i32([0, len(drafts) + 1]) + + apply_trace_tokens_to_verified( + sampled, + num_sampled, + draft_sampled, + cu_num_logits, + _i32([0]), + trace_token_ids, + _i32([len(trace)]), + _i32([prompt_len + step]), + _i32([prompt_len]), + ) + + assert num_sampled.tolist() == [expected_num_sampled] + assert sampled[0, :expected_num_sampled].tolist() == expected_tokens + + def test_batch_with_mixed_requests(self): + trace_token_ids = torch.zeros( + 3, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE + ) + trace_token_ids[0, :4] = _i32([100, 101, 102, 103]) + trace_token_ids[2, :3] = _i32([300, 301, 302]) + trace_len = _i32([4, 0, 3]) + prompt_len = _i32([4, 4, 4]) + total_len = _i32([5, 9, 7]) + sampled = torch.tensor( + [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], + dtype=torch.int64, + device=DEVICE, + ) + num_sampled = _i32([1, 2, 3, 1]) + draft_sampled = _i64([0, 101, 102, 0, 5, 6, 0, 0, 0, 7]) + cu_num_logits = _i32([0, 3, 6, 9, 10]) + idx_mapping = _i32([0, 1, -1, 2]) + + apply_trace_tokens_to_verified( + sampled, + num_sampled, + draft_sampled, + cu_num_logits, + idx_mapping, + trace_token_ids, + trace_len, + total_len, + prompt_len, + ) + + assert num_sampled.tolist() == [3, 2, 3, 1] + assert sampled.tolist() == [ + [101, 102, 103], + [4, 5, 6], + [7, 8, 9], + [10, 11, 12], + ] diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 823840312864..73dd2d06d4b9 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -1125,10 +1125,6 @@ def _validate_trace_replay( raise ValueError( "trace_decode_token_ids is not supported with prompt_logprob_token_ids." ) - if speculative_config is not None: - raise ValueError( - "trace_decode_token_ids is not supported with speculative decoding." - ) if self.structured_outputs is not None: raise ValueError( "trace_decode_token_ids is not supported with structured outputs." diff --git a/vllm/v1/worker/gpu/sample/trace_replay.py b/vllm/v1/worker/gpu/sample/trace_replay.py index 9b58d4cf2923..7ecec14393ba 100644 --- a/vllm/v1/worker/gpu/sample/trace_replay.py +++ b/vllm/v1/worker/gpu/sample/trace_replay.py @@ -57,6 +57,26 @@ def apply_trace( self.req_states.prompt_len.gpu, ) + def apply_trace_to_verified( + self, + sampled: torch.Tensor, + num_sampled: torch.Tensor, + draft_sampled: torch.Tensor, + cu_num_logits: torch.Tensor, + idx_mapping: torch.Tensor, + ) -> None: + apply_trace_tokens_to_verified( + sampled, + num_sampled, + draft_sampled, + cu_num_logits, + idx_mapping, + self.trace_token_ids.gpu, + self.trace_len.gpu, + self.req_states.total_len.gpu, + self.req_states.prompt_len.gpu, + ) + @triton.jit def _trace_replay_kernel( @@ -111,3 +131,83 @@ def apply_trace_tokens( total_len, prompt_len, ) + + +@triton.jit +def _trace_replay_verified_kernel( + sampled_ptr, + sampled_stride, + num_sampled_ptr, + draft_sampled_ptr, + cu_num_logits_ptr, + idx_mapping_ptr, + trace_token_ids_ptr, + trace_token_ids_stride, + trace_len_ptr, + total_len_ptr, + prompt_len_ptr, +): + req_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + if req_state_idx < 0: + return + + trace_len = tl.load(trace_len_ptr + req_state_idx) + if trace_len <= 0: + return + + step = tl.load(total_len_ptr + req_state_idx) - tl.load( + prompt_len_ptr + req_state_idx + ) + if step < 0 or step >= trace_len: + return + + start = tl.load(cu_num_logits_ptr + req_idx) + num_draft_tokens = tl.load(cu_num_logits_ptr + req_idx + 1) - start - 1 + num_replayed = tl.minimum(num_draft_tokens + 1, trace_len - step) + trace_row = ( + trace_token_ids_ptr + req_state_idx.to(tl.int64) * trace_token_ids_stride + step + ) + num_accepted = tl.zeros((), tl.int32) + matching = tl.full((), 1, tl.int1) + for i in range(num_replayed): + if matching: + token_id = tl.load(trace_row + i).to(tl.int64) + tl.store(sampled_ptr + req_idx * sampled_stride + i, token_id) + num_accepted = i + 1 + draft_token_id = tl.load( + draft_sampled_ptr + start + i + 1, + mask=i < num_draft_tokens, + other=-1, + ).to(tl.int64) + matching = draft_token_id == token_id + tl.store(num_sampled_ptr + req_idx, num_accepted) + + +def apply_trace_tokens_to_verified( + sampled: torch.Tensor, + num_sampled: torch.Tensor, + draft_sampled: torch.Tensor, + cu_num_logits: torch.Tensor, + idx_mapping: torch.Tensor, + trace_token_ids: torch.Tensor, + trace_len: torch.Tensor, + total_len: torch.Tensor, + prompt_len: torch.Tensor, +) -> None: + """Replace verified tokens with the trace, accepting each draft only if it + equals the trace token at its position.""" + num_reqs = idx_mapping.shape[0] + _trace_replay_verified_kernel[(num_reqs,)]( + sampled, + sampled.stride(0), + num_sampled, + draft_sampled, + cu_num_logits, + idx_mapping, + trace_token_ids, + trace_token_ids.stride(0), + trace_len, + total_len, + prompt_len, + ) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index fd9808063c4a..cce3dfbba319 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -298,6 +298,10 @@ def _verify( draft_sampled, expanded_idx_mapping, expanded_local_pos ), ) + if self.sampler.trace_replay_state is not None: + self.sampler.trace_replay_state.apply_trace_to_verified( + sampled, num_sampled, verify_draft_sampled, cu_num_logits, idx_mapping + ) return processed_logits, sampled, num_sampled def _verify_in_chunks( From 38b1775405f1307aa611e1e7a19d16abb6d5db49 Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Wed, 28 Oct 2026 18:26:31 -0700 Subject: [PATCH 098/406] [XPU][UT] reuse test_phi3v.py::test_models_text on XPU (#64228) Signed-off-by: mayuyuace Co-authored-by: vllm-agent --- tests/models/multimodal/pooling/test_phi3v.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/models/multimodal/pooling/test_phi3v.py b/tests/models/multimodal/pooling/test_phi3v.py index 35df42127eca..ba017f065a13 100644 --- a/tests/models/multimodal/pooling/test_phi3v.py +++ b/tests/models/multimodal/pooling/test_phi3v.py @@ -11,7 +11,6 @@ from vllm.assets.image import VLM_IMAGES_DIR from vllm.config import ModelConfig from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.platforms import current_platform from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner from ....utils import large_gpu_test @@ -103,15 +102,6 @@ def _run_test( @pytest.mark.core_model -@pytest.mark.skipif( - current_platform.is_xpu(), - reason=( - "vllm-xpu-kernels paged_decode_default.conf is missing the " - "8,96,64,false,true,false tuple needed by VLM2Vec-Full at " - "block_size=64. Re-enable once vllm-xpu-kernels ships the updated " - "config. See github.com/vllm-project/vllm-xpu-kernels/pull/587." - ), -) @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) def test_models_text( From b7f0dda503ef69a9784b1fb8a6068a32ebc80ae9 Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Wed, 28 Oct 2026 19:32:41 -0600 Subject: [PATCH 099/406] [ROCm][Kimi-K3] Allow DSpark adaptive verification with DCP on ROCM_AITER_MLA (#62990) Signed-off-by: Oxana Korzh Co-authored-by: Claude Opus 5.5 --- .../test_rocm_aiter_mla_dcp_cprr_numerics.py | 171 ++++++++++++++++++ .../test_rocm_aiter_mla_mtp_split.py | 113 +++++++++++- .../attention/backends/mla/rocm_aiter_mla.py | 70 ++++++- 3 files changed, 341 insertions(+), 13 deletions(-) diff --git a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py index 21626e9a675e..4c48508c7159 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py +++ b/tests/v1/attention/test_rocm_aiter_mla_dcp_cprr_numerics.py @@ -547,3 +547,174 @@ def test_cprr_kernel_is_selected_for_the_k3_shape(): ) assert route is m._DCPDecodeRoute.CPRR assert m._asm_dcp_verify_heads(96) in m._NATIVE_CPRR_HEADS + + +VARLEN_QO = (1, 5, 3, 2, 4, 1, 5, 2, 0, 0) +VARLEN_CTX = (256, 3, 100, 0, 37, 1, 251, 2, 0, 0) +PAD_TOKENS = 6 + + +def _varlen_reference(q_f32, kv_f32, seq_lens, qo_lens, positions_per_req): + heads = q_f32.shape[1] + o = torch.zeros(sum(qo_lens), heads, KV_LORA_RANK, device=q_f32.device) + lse = torch.full((sum(qo_lens), heads), float("-inf"), device=q_f32.device) + token = 0 + for req, (seq_len, qo_len) in enumerate(zip(seq_lens, qo_lens)): + positions = positions_per_req[req] + kv_sel = kv_f32[req].index_select(0, positions) + for t in range(qo_len): + valid = positions <= seq_len - qo_len + t + if bool(valid.any()): + k = kv_sel[valid] + s = (q_f32[token] @ k.T) * SCALE + lse[token] = torch.logsumexp(s, dim=-1) + o[token] = torch.softmax(s, dim=-1) @ k[:, :KV_LORA_RANK] + token += 1 + return o, lse + + +@requires_cprr +@pytest.mark.parametrize( + ("asm", "expected"), + [ + (True, rocm_aiter_mla._DCPDecodeRoute.CPRR), + (False, rocm_aiter_mla._DCPDecodeRoute.SEGMENTED), + ], +) +@torch.inference_mode() +def test_adaptive_varlen_shards_match_reference(monkeypatch, asm, expected): + import aiter + + from vllm.config import set_current_vllm_config + from vllm.v1.attention.backend import CommonAttentionMetadata + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder + from vllm.v1.attention.ops.dcp import ( + CPTritonContext, + correct_attn_out, + mask_dcp_empty_shards_, + ) + from vllm.v1.kv_cache_interface import MLAAttentionSpec + + if not asm: + monkeypatch.setattr( + rocm_aiter_mla, "_asm_dcp_verify_configured", lambda *a, **k: False + ) + heads_per_rank = 12 + device = torch.device("cuda", 0) + torch.accelerator.set_device_index(device.index) + torch.manual_seed(0) + _ensure_workspace_manager(device) + _install_fake_groups(DCP, 0) + vllm_config = _make_vllm_config(heads_per_rank) + vllm_config.speculative_config.enable_adaptive_verification = True + + nr = len(VARLEN_QO) + qo_lens = list(VARLEN_QO) + seq_lens = [c + q for c, q in zip(VARLEN_CTX, VARLEN_QO)] + num_tokens = sum(qo_lens) + max_len = max(seq_lens) + gathered_heads = heads_per_rank * DCP + real_reqs = sum(q > 0 for q in qo_lens) + base, extra = divmod(num_tokens, real_reqs) + host_qo = [base + (i < extra) for i in range(real_reqs)] + [0] * (nr - real_reqs) + + def qsl(lens): + return torch.tensor([0] + lens, dtype=torch.int32).cumsum(0, dtype=torch.int32) + + with set_current_vllm_config(vllm_config): + layer_name = "model.layers.0.self_attn.attn" + vllm_config.compilation_config.static_forward_context[layer_name] = ( + _make_ctx_layer(vllm_config, monkeypatch) + ) + spec = MLAAttentionSpec( + block_size=1, + num_kv_heads=1, + head_size=HEAD_SIZE, + dtype=vllm_config.model_config.dtype, + non_causal_multi_token_decode=False, + ) + kv_flat = torch.randn( + nr * max_len, 1, HEAD_SIZE, dtype=torch.float32, device=device + ).to(aiter.dtypes.fp8) + q_flat = ( + torch.randn( + num_tokens + PAD_TOKENS, + gathered_heads, + HEAD_SIZE, + dtype=torch.float32, + device=device, + ) + .mul_(0.5) + .to(aiter.dtypes.fp8) + ) + kv_f32 = kv_flat.float().view(nr, max_len, HEAD_SIZE) + q_f32 = q_flat.float() + ref_o, _ = _varlen_reference( + q_f32, + kv_f32, + seq_lens, + qo_lens, + [torch.arange(s, device=device) for s in seq_lens], + ) + + raw_o, raw_lse = [], [] + for r in range(DCP): + positions = [ + torch.arange(r, max(r, s), DCP, device=device) for s in seq_lens + ] + local = torch.tensor([p.numel() for p in positions], dtype=torch.int32) + block_table = torch.zeros( + nr, + vllm_config.model_config.max_model_len // DCP, + dtype=torch.int32, + device=device, + ) + for req, p in enumerate(positions): + block_table[req, : p.numel()] = (req * max_len + p).to(torch.int32) + common = CommonAttentionMetadata( + query_start_loc=qsl(qo_lens).to(device), + query_start_loc_cpu=qsl(host_qo), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), + num_reqs=nr, + num_actual_tokens=num_tokens, + max_query_len=QLEN, + max_seq_len=max_len, + block_table_tensor=block_table, + slot_mapping=torch.arange(num_tokens, dtype=torch.int64, device=device), + causal=True, + dcp_local_seq_lens=local.to(device), + dcp_local_seq_lens_cpu_upper_bound=local, + ) + _install_fake_groups(DCP, r) + builder = AiterMLAMetadataBuilder(spec, [layer_name], vllm_config, device) + md = builder.build(0, common) + assert md.decode is not None + assert md.decode.dcp_route is expected + impl = _build_impl(builder.num_heads, DCP, r) + o, lse = impl.forward_mqa(q_flat, kv_flat, md, _Layer(device)) + torch.accelerator.synchronize() + + lse = lse.float().view(-1, gathered_heads).contiguous() + mask_dcp_empty_shards_(lse, md.decode.seq_lens, common.query_start_loc) + raw_o.append(o.view(-1, gathered_heads, KV_LORA_RANK)) + raw_lse.append(lse) + + po, plse = _varlen_reference(q_f32, kv_f32, seq_lens, qo_lens, positions) + seen = torch.isfinite(plse) + assert not bool((~seen & (lse[:num_tokens] > -1e30)).any()) + shard_o = raw_o[-1][:num_tokens].float()[seen] + rel = (shard_o - po[seen]).abs().max() / po.abs().max() + assert rel < SHARD_RTOL, f"rank {r} rel err {rel:.3e}" + lse_err = (lse[:num_tokens][seen] - plse[seen]).abs().max() + assert lse_err < SHARD_RTOL, f"rank {r} lse err {lse_err:.3e}" + + lses = torch.stack(raw_lse) + merged = sum( + correct_attn_out(raw_o[r].contiguous(), lses, r, CPTritonContext())[ + 0 + ].float() + for r in range(DCP) + ) + assert bool(torch.isfinite(merged).all()) + merge_rel = (merged[:num_tokens] - ref_o).abs().max() / ref_o.abs().max() + assert merge_rel < MERGE_RTOL, f"merged rel err {merge_rel:.3e}" 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 8ed68986b267..563c7744c027 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -1462,7 +1462,11 @@ def test_a_two_token_bf16_block_is_allowed(monkeypatch): def _adaptive_config( - *, adaptive: bool = True, dcp: int = 1, num_speculative_tokens: int = 7 + *, + adaptive: bool = True, + dcp: int = 1, + num_speculative_tokens: int = 7, + cp_interleave: int = 1, ) -> SimpleNamespace: return SimpleNamespace( speculative_config=SimpleNamespace( @@ -1470,18 +1474,117 @@ def _adaptive_config( enable_adaptive_verification=adaptive, parallel_drafting=False, ), - parallel_config=SimpleNamespace(decode_context_parallel_size=dcp), + parallel_config=SimpleNamespace( + decode_context_parallel_size=dcp, + cp_kv_cache_interleave_size=cp_interleave, + ), use_v2_model_runner=True, ) -def test_varlen_cudagraph_bound_follows_adaptive_verification(): +def test_varlen_cudagraph_bound_follows_adaptive_verification(monkeypatch): + monkeypatch.setattr(rocm_aiter_mla, "_asm_dcp_verify_configured", lambda *a: True) + for route in ("_segmented_dcp_verify_supported", "_triton_dcp_verify_supported"): + monkeypatch.setattr(rocm_aiter_mla, route, lambda *a: False) bound = AiterMLAMetadataBuilder.get_varlen_cudagraph_max_query_len assert bound(_adaptive_config(), SimpleNamespace()) == 8 assert bound(_adaptive_config(adaptive=False), SimpleNamespace()) is None - # DCP routes plan per-row windows from host lengths. - assert bound(_adaptive_config(dcp=2), SimpleNamespace()) is None + assert bound(_adaptive_config(dcp=8), SimpleNamespace()) == 8 + assert bound(_adaptive_config(dcp=8, cp_interleave=2), SimpleNamespace()) is None + assert ( + bound(_adaptive_config(dcp=8, num_speculative_tokens=1), SimpleNamespace()) + is None + ) # Kimi-K3 DSpark draft layers share the target's group; causality is # decided per batch, so the merged non-causal spec keeps the bound. shared_draft_group = SimpleNamespace(non_causal_multi_token_decode=True) assert bound(_adaptive_config(), shared_draft_group) == 8 + + +def test_backend_rejects_unsupported_adaptive_dcp(monkeypatch): + import vllm.config + + monkeypatch.setattr( + vllm.config, + "get_current_vllm_config_or_none", + lambda: _adaptive_config(dcp=8, cp_interleave=2), + ) + + def validate(**flags): + return rocm_aiter_mla.AiterMLABackend.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=None, + use_mla=True, + has_sink=False, + use_sparse=False, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=None, + attn_type="decoder", + **flags, + ) + + with pytest.raises(ValueError, match="decode context parallelism"): + validate(use_adaptive_verification=True, use_dcp=True) + validate(use_dcp=True) + + +class _RecordingExpandKernel(_NoOpTritonKernel): + def __call__(self, *args, CLAMP_TAIL=False, **kwargs): + self.clamp_tail = CLAMP_TAIL + + +def test_dcp_varlen_right_aligns_rows_and_keeps_lengths(monkeypatch): + monkeypatch.setitem( + sys.modules, "aiter", SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()) + ) + expand = _RecordingExpandKernel() + monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand) + width, num_reqs = 5, 4 + builder = _builder( + mtp_decode_qlen=width, + dcp_world_size=8, + max_decode_rows=num_reqs * width, + num_heads=12, + asm_dcp_verify=True, + supports_segmented_dcp_verify=True, + ) + rows = torch.arange(num_reqs * width, dtype=torch.int64) + builder._varlen_decode_qlen = width + builder._varlen_right_align = AiterMLAMetadataBuilder.__dict__[ + "_varlen_right_align" + ].__get__(builder) + builder._varlen_row_req = rows // width + builder._varlen_row_pos = rows % width + builder._varlen_row_ids = rows + builder._varlen_q_rows = torch.zeros_like(rows) + builder._varlen_o_rows = torch.zeros(num_reqs * width + 1, dtype=torch.int64) + builder._fill_varlen_rows = AiterMLAMetadataBuilder._fill_varlen_rows.__get__( + builder + ) + + metadata = AiterMLAMetadataBuilder._build_decode( + builder, + block_table_tensor=torch.arange(num_reqs * 8, dtype=torch.int32).view(-1, 8), + seq_lens_device=torch.tensor([5, 3, 2, 0], dtype=torch.int32), + max_seq_len=40, + query_start_loc_cpu=torch.tensor([0, 3, 5, 7, 9], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, 1, 6, 9, 9], dtype=torch.int32), + num_decode_tokens=9, + max_query_len=width, + dcp_tot_seq_lens_device=torch.tensor([40, 21, 9, 0], dtype=torch.int32), + ) + + assert metadata.dcp_route is rocm_aiter_mla._DCPDecodeRoute.CPRR + assert metadata.seq_lens.tolist() == [5, 3, 2, width] + assert metadata.g_kv_indptr is not None + assert metadata.g_kv_indptr.tolist()[: num_reqs + 1] == [0, 40, 61, 70, 110] + assert metadata.qo_indptr is not None + assert metadata.qo_indptr.tolist() == [0, 5, 10, 15, 20] + assert expand.clamp_tail is False + assert metadata.varlen_q_rows is not None and metadata.varlen_o_rows is not None + real_rows = [4, 5, 6, 7, 8, 9, 12, 13, 14] + assert metadata.varlen_o_rows[:9].tolist() == real_rows + assert metadata.varlen_q_rows[real_rows].tolist() == list(range(9)) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 1cfb2ae215fd..963eee2d7d2e 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -486,13 +486,40 @@ def _aiter_mla_small_head_mode() -> str: return mode +def _adaptive_dcp_verify_unsupported_reason(vllm_config: VllmConfig) -> str | None: + parallel_config = vllm_config.parallel_config + dcp_world_size = parallel_config.decode_context_parallel_size + if dcp_world_size == 1: + return None + cp_interleave = parallel_config.cp_kv_cache_interleave_size + if cp_interleave != 1: + return ( + "adaptive verification with decode context parallelism requires " + f"cp_kv_cache_interleave_size 1, got {cp_interleave}" + ) + asm_selected = ( + _asm_dcp_verify_configured(dcp_world_size, cp_interleave, True) + and max_decode_query_len(vllm_config) >= _MIN_CPRR_QLEN + ) + if not ( + asm_selected + or _segmented_dcp_verify_supported(dcp_world_size, cp_interleave) + or _triton_dcp_verify_supported(dcp_world_size, cp_interleave) + ): + return ( + "adaptive verification with decode context parallelism needs the " + "round-robin asm decode, segmented MLA or the gfx942 Triton verify" + ) + return None + + def _adaptive_varlen_verify_enabled(vllm_config: VllmConfig) -> bool: speculative_config = vllm_config.speculative_config return ( speculative_config is not None and bool(getattr(speculative_config, "enable_adaptive_verification", False)) - and vllm_config.parallel_config.decode_context_parallel_size == 1 and max_decode_query_len(vllm_config) > 1 + and _adaptive_dcp_verify_unsupported_reason(vllm_config) is None ) @@ -587,6 +614,20 @@ def validate_configuration( and kv_cache_dtype in (None, "auto", "float16") ): invalid_reasons.append("non-causal fp16 MLA decode not supported") + if use_adaptive_verification and use_dcp: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + reason = ( + _adaptive_dcp_verify_unsupported_reason(vllm_config) + if vllm_config is not None + else None + ) + if reason is not None: + raise ValueError( + f"{reason}. Disable enable_adaptive_verification in the " + "speculative config or set decode_context_parallel_size=1." + ) return invalid_reasons @staticmethod @@ -764,6 +805,10 @@ def _uniform_padded_mtp_qo_len( return uniform_qo_len + @property + def _varlen_right_align(self) -> bool: + return self.dcp_world_size > 1 + @classmethod def get_varlen_cudagraph_max_query_len( cls: type["AiterMLAMetadataBuilder"], @@ -922,7 +967,9 @@ def __init__( max_num_reqs = vllm_config.scheduler_config.max_num_seqs max_num_pages = max_num_reqs * max_num_pages_per_req speculative_config = vllm_config.speculative_config - if _adaptive_varlen_verify_enabled(vllm_config): + if _adaptive_varlen_verify_enabled(vllm_config) and not ( + self._varlen_right_align + ): assert speculative_config is not None max_num_pages += max_num_reqs * ( speculative_config.num_speculative_tokens or 0 @@ -1564,8 +1611,10 @@ def _build_decode( ) varlen_q_rows = varlen_o_rows = None - if varlen_decode: + varlen_left_align = varlen_decode and not self._varlen_right_align + if varlen_left_align: seq_lens_for_kernel = seq_lens_device + (max_qo_len - qo_lens_device) + if varlen_decode: varlen_q_rows, varlen_o_rows = self._fill_varlen_rows( query_start_loc_device, qo_lens_device, num_reqs ) @@ -1630,7 +1679,7 @@ def _build_decode( # device synchronisation. Programs whose chunk lies past # num_tokens mask out entirely. max_tokens_per_req = block_table_tensor.shape[1] * self.kernel_block_size - if varlen_decode: + if varlen_left_align: max_tokens_per_req += max_qo_len - 1 num_chunks = max(1, cdiv(max_tokens_per_req, 1024)) _expand_page_indices_kernel[(num_reqs, num_chunks)]( @@ -1642,7 +1691,7 @@ def _build_decode( seq_lens_device, KERNEL_BLOCK_SIZE=self.kernel_block_size, BLOCK_SIZE=1024, - CLAMP_TAIL=varlen_decode, + CLAMP_TAIL=varlen_left_align, ) paged_kv_indices = self.paged_kv_indices @@ -1886,7 +1935,9 @@ def _fill_varlen_rows( row_pos = self._varlen_row_pos[:num_rows] row_qlen = qo_lens_device.to(torch.int64).index_select(0, row_req) row_start = query_start_loc_device.to(torch.int64).index_select(0, row_req) - is_real = row_pos < row_qlen + if self._varlen_right_align: + row_pos = row_pos - (width - row_qlen) + is_real = (row_pos >= 0) & (row_pos < row_qlen) packed_row = row_start + row_pos q_rows = self._varlen_q_rows[:num_rows] q_rows.copy_(torch.where(is_real, packed_row, 0)) @@ -2994,8 +3045,11 @@ def forward_mqa( num_tokens = q.shape[0] q = q.index_select(0, q_rows) o, lse = self._forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer) - assert lse is None - return o.index_select(0, decode.varlen_o_rows[:num_tokens]), None + o_rows = decode.varlen_o_rows[:num_tokens] + assert lse is None or self.dcp_world_size > 1 + if lse is not None: + lse = lse.index_select(0, o_rows) + return o.index_select(0, o_rows), lse def _forward_mqa( self, From 8d3f6f7e788a3bcbd73cae851e7717d222eb8bb9 Mon Sep 17 00:00:00 2001 From: Drago1234 <32420036+drago1234@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:49:52 +0800 Subject: [PATCH 100/406] [Bugfix] Parenthesize penalties Triton chained `or` for Ascend backends (#64013) Signed-off-by: drago1234 <593849302dzq@gmail.com> Co-authored-by: vllm-agent --- vllm/v1/worker/gpu/sample/penalties.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py index 4f0f003f05a8..232b86e0c634 100644 --- a/vllm/v1/worker/gpu/sample/penalties.py +++ b/vllm/v1/worker/gpu/sample/penalties.py @@ -130,7 +130,8 @@ def _penalties_kernel( use_rep_penalty = rep_penalty != 1.0 use_freq_penalty = freq_penalty != 0.0 use_pres_penalty = pres_penalty != 0.0 - use_penalty = use_rep_penalty or use_freq_penalty or use_pres_penalty + # Parenthesize for Triton backends that reject chained `A or B or C`. + use_penalty = (use_rep_penalty or use_freq_penalty) or use_pres_penalty if not use_penalty: # Early return to avoid loading logits. return From 45fd60d4d556d4a9eb7894287eb67dd28f02dc2b Mon Sep 17 00:00:00 2001 From: Mathew Odden <1471252+mrodden@users.noreply.github.com> Date: Wed, 28 Oct 2026 20:57:19 -0500 Subject: [PATCH 101/406] [ROCm][Qwen4Exp] Enable fp8 KV and indexer caches on the AMD QSA path (#61529) Signed-off-by: Mathew Odden --- .buildkite/test_areas/models_basic.yaml | 19 +++ tests/models/qwen4_exp/test_qsa_amd.py | 49 ++++++- .../models/qwen4_exp/test_qsa_prepare_amd.py | 112 ++++++++++++--- tests/models/qwen4_exp/test_qsa_reference.py | 6 +- vllm/models/qwen4_exp/amd/indexer_qsa.py | 33 ++++- vllm/models/qwen4_exp/amd/ops/qsa.py | 19 ++- vllm/models/qwen4_exp/amd/ops/qsa_prepare.py | 2 +- vllm/models/qwen4_exp/amd/qsa.py | 129 ++++++++++++++++-- vllm/models/qwen4_exp/common/qsa_cache.py | 10 ++ vllm/models/qwen4_exp/nvidia/qsa.py | 12 +- 10 files changed, 338 insertions(+), 53 deletions(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index b5c64ec7b3f0..6135a46c6165 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -349,6 +349,25 @@ steps: # PLE and QSA import conflicting HC registrations in a shared process. - pytest -v -s models/qwen4_exp/test_ple.py +- label: ":amd: (MI300) Qwen4 Exp" + key: qwen4-exp-unit-tests-mi300 + timeout_in_minutes: 40 + dind: false + device: mi300_1 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/models/qwen4_exp/ + - vllm/config/compilation.py + - vllm/v1/attention/backends/short_conv_attn.py + - vllm/model_executor/layers/mamba/checkpoint.py + - tests/models/qwen4_exp/ + commands: + - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_prepare_amd.py + # PLE and QSA import conflicting HC registrations in a shared process. + - pytest -v -s models/qwen4_exp/test_ple.py + - label: ":computer: (CPU) Qwen4 Exp" key: qwen4-exp-unit-tests-cpu depends_on: diff --git a/tests/models/qwen4_exp/test_qsa_amd.py b/tests/models/qwen4_exp/test_qsa_amd.py index 58f33b6cd53c..8684f01236d0 100644 --- a/tests/models/qwen4_exp/test_qsa_amd.py +++ b/tests/models/qwen4_exp/test_qsa_amd.py @@ -132,6 +132,7 @@ def _make_attn(*, use_fused: bool) -> SimpleNamespace: return SimpleNamespace( use_fused_qsa_prepare=use_fused, kv_cache=torch.zeros(1, 2, 16, 8), + kv_cache_dtype="auto", q_norm=GemmaRMSNorm(4, eps=1e-6), k_norm=GemmaRMSNorm(4, eps=1e-6), _k_scale_float=1.0, @@ -154,6 +155,7 @@ def _make_indexer( indexer.index_kv_heads = 1 indexer.index_head_dim = head_dim indexer.compress_ratio = 2 + indexer.indexer_dtype = torch.bfloat16 indexer.use_fused_pre_indexer = use_fused indexer.skip_topk = True indexer.index_qk_proj = lambda hidden_states: ( @@ -397,8 +399,9 @@ def expand_qsa_block_indices(*args) -> None: @requires_qsa_kernels +@pytest.mark.parametrize("use_fp8", [False, True]) @pytest.mark.parametrize( - ("num_rows", "num_query_heads", "num_kv_heads", "page_size"), + "num_rows,num_query_heads,num_kv_heads,page_size", [ pytest.param(1, 24, 2, 1792, id="tp1_split64"), pytest.param(16, 12, 1, 1792, id="tp2_split32"), @@ -412,6 +415,7 @@ def test_qsa_sparse_paged_attention_matches_reference( num_query_heads: int, num_kv_heads: int, page_size: int, + use_fp8: bool, ) -> None: torch.manual_seed(2) head_dim = 256 @@ -433,7 +437,18 @@ def test_qsa_sparse_paged_attention_matches_reference( device="cuda", dtype=torch.bfloat16, ) - k_cache, v_cache = kv_cache.split(head_dim, dim=-1) + if use_fp8: + fp8_dtype = current_platform.fp8_dtype() + # The kernel's contract is that the cache already holds values in + # the attention dtype's numeric space; quantize the stored cache and + # have the reference consume the stored (quantized) values upcast + # to float, not the original bf16 — the fp8 cast is lossy. + kv_stored = kv_cache.to(fp8_dtype) + k_cache, v_cache = kv_stored.split(head_dim, dim=-1) + ref_k, ref_v = k_cache.float(), v_cache.float() + else: + k_cache, v_cache = kv_cache.split(head_dim, dim=-1) + ref_k, ref_v = k_cache, v_cache block_table = ( torch.randperm(num_cache_blocks, device="cuda") .reshape(num_requests, num_pages_per_request) @@ -486,8 +501,8 @@ def test_qsa_sparse_paged_attention_matches_reference( ) expected = _qsa_sparse_paged_attention_reference( q, - k_cache, - v_cache, + ref_k, + ref_v, logical_indices, block_table, token_to_req, @@ -495,3 +510,29 @@ def test_qsa_sparse_paged_attention_matches_reference( ) torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +@requires_qsa_kernels +@pytest.mark.parametrize("kv_cache_dtype", ["fp8", "fp8_e4m3"]) +def test_qsa_impl_constructs_with_fp8_kv(kv_cache_dtype: str) -> None: + """The impl constructor must accept an FP8 KV cache dtype. + + The QSA Triton kernel dequantizes fp8 itself and never runs flash-attn + over the cache, so the parent constructor's flash-attn quantized-KV + probe (which raises on ROCm) must not reject it: the impl hands the + parent "auto" for that probe and restores the real dtype afterwards, + mirroring the NVIDIA implementation. Guards both the construction + succeeding and the dtype surviving the parent round-trip. + """ + from vllm.models.qwen4_exp.amd.qsa import Qwen4ExpQSAFlashAttentionImpl + + impl = Qwen4ExpQSAFlashAttentionImpl( + num_heads=24, + head_size=128, + scale=0.088, + num_kv_heads=2, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype=kv_cache_dtype, + ) + assert impl.kv_cache_dtype == kv_cache_dtype diff --git a/tests/models/qwen4_exp/test_qsa_prepare_amd.py b/tests/models/qwen4_exp/test_qsa_prepare_amd.py index 718929720955..6d5335a6e3dc 100644 --- a/tests/models/qwen4_exp/test_qsa_prepare_amd.py +++ b/tests/models/qwen4_exp/test_qsa_prepare_amd.py @@ -77,15 +77,17 @@ def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> N assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all()) -def _make_main_inputs(num_tokens: int) -> dict: - """Random main-attention arguments with a paged BF16 cache in vLLM's layout.""" +def _make_main_inputs(num_tokens: int, main_fp8: bool = False) -> dict: + """Random main-attention arguments with a paged cache in vLLM's layout.""" num_blocks = num_tokens // MAIN_PAGE + 2 + # FP8 configs store e4m3 codes in the platform's native flavor + # (FNUZ on MI300/MI325, OCP elsewhere), matching the production views. cache = torch.zeros( num_blocks, MAIN_HK, MAIN_PAGE, 2 * MAIN_D, - dtype=torch.bfloat16, + dtype=current_platform.fp8_dtype() if main_fp8 else torch.bfloat16, device="cuda", ).transpose(1, 2) slots = torch.randperm(num_blocks * MAIN_PAGE, device="cuda")[:num_tokens] @@ -103,8 +105,8 @@ def _make_main_inputs(num_tokens: int) -> dict: main_eps=EPS, main_kv_cache=cache, main_slot_mapping=slots, - main_k_scale=1.0, - main_v_scale=1.0, + main_k_scale=0.5 if main_fp8 else 1.0, + main_v_scale=2.0 if main_fp8 else 1.0, # The AMD op hands the gate back through a buffer it does not own. main_gate_out=torch.empty( num_tokens, MAIN_HQ, MAIN_D, dtype=torch.bfloat16, device="cuda" @@ -129,14 +131,86 @@ def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None: ).view(num_tokens, -1) q, k = rope(positions, q, k) kv_cache = main["main_kv_cache"] - cache = torch.zeros_like(kv_cache) - key_cache, value_cache = cache.split(MAIN_D, dim=-1) + if kv_cache.dtype in (torch.float8_e4m3fnuz, torch.float8_e4m3fn): + torch.testing.assert_close(q_out.flatten(1), q, rtol=RTOL, atol=ATOL) + assert torch.equal(gate_out, gate) + + # C++ reference: both the fused Triton kernel and the C++ + # reshape_and_cache_flash write the platform-native fp8 flavor + # (OCP on gfx950, FNUZ on gfx942), so the C++ op is the unfused + # reference on both arches. The reference buffers are separate + # allocations so the reference cannot alias the fused cache. + slot = main["main_slot_mapping"] + valid = slot >= 0 + blocks = (slot[valid] // MAIN_PAGE).long() + pages = (slot[valid] % MAIN_PAGE).long() + stored = kv_cache[blocks, pages] + + key_ref = torch.zeros( + kv_cache.shape[:-1] + (MAIN_D,), + dtype=torch.uint8, + device=kv_cache.device, + ) + value_ref = torch.zeros( + kv_cache.shape[:-1] + (MAIN_D,), + dtype=torch.uint8, + device=kv_cache.device, + ) + k_scale_t = torch.tensor(float(main["main_k_scale"]), device="cuda") + v_scale_t = torch.tensor(float(main["main_v_scale"]), device="cuda") + reshape_and_cache_flash( + k.view(-1, MAIN_HK, MAIN_D), + v.view(-1, MAIN_HK, MAIN_D), + key_ref, + value_ref, + main["main_slot_mapping"], + "fp8", + k_scale_t, + v_scale_t, + ) + # View the uint8 reference codes as the cache's fp8 dtype so + # _assert_codes_match's float tier compares fp8 values, not raw + # uint8 codes (advanced indexing produces contiguous copies, so + # the dtype view is valid). + ref_k_stored = key_ref[blocks, pages].view(kv_cache.dtype) + ref_v_stored = value_ref[blocks, pages].view(kv_cache.dtype) + + # Compare in code space (the uint8 codes): the fused kernel and + # the C++ op round in a slightly different order, so a value on a + # step boundary may flip to the adjacent code. Allowing 2 codes + # covers the boundary flips; a real bug — like a scale that was + # never applied — moves normal-sized numbers by many codes. + def _assert_codes_match(actual, expected): + code_diff = ( + actual.view(torch.uint8).int() - expected.view(torch.uint8).int() + ).abs() + abs_diff = (actual.float() - expected.float()).abs() + assert bool( + ( + (code_diff <= 2) + | (abs_diff <= RTOL * expected.float().abs() + 2 * ATOL) + ).all() + ) + + _assert_codes_match(stored[..., :MAIN_D], ref_k_stored) + _assert_codes_match(stored[..., MAIN_D:], ref_v_stored) + return + key_ref = torch.zeros( + kv_cache.shape[:-1] + (MAIN_D,), + dtype=kv_cache.dtype, + device=kv_cache.device, + ) + value_ref = torch.zeros( + kv_cache.shape[:-1] + (MAIN_D,), + dtype=kv_cache.dtype, + device=kv_cache.device, + ) scale = torch.tensor(1.0, device="cuda") reshape_and_cache_flash( k.view(-1, MAIN_HK, MAIN_D), v.view(-1, MAIN_HK, MAIN_D), - key_cache, - value_cache, + key_ref, + value_ref, main["main_slot_mapping"], "auto", scale, @@ -144,12 +218,16 @@ def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None: ) torch.testing.assert_close(q_out.flatten(1), q, rtol=RTOL, atol=ATOL) assert torch.equal(gate_out, gate) - torch.testing.assert_close(kv_cache, cache, rtol=RTOL, atol=ATOL) + torch.testing.assert_close(kv_cache[..., :MAIN_D], key_ref, rtol=RTOL, atol=ATOL) + torch.testing.assert_close(kv_cache[..., MAIN_D:], value_ref, rtol=RTOL, atol=ATOL) @requires_qsa_kernels @pytest.mark.usefixtures("default_vllm_config") -@pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("main_fp8", [False, True]) +@pytest.mark.parametrize( + "indexer_dtype", [torch.bfloat16, pytest.param(None, id="fp8_platform")] +) @pytest.mark.parametrize( "mrope,is_2d_positions,cache_rope_positions,state_size,seq_lens,query_lens,history_lens", [ @@ -173,6 +251,7 @@ def _check_main_outputs(main: dict, q_out, gate_out, rope, positions) -> None: ], ) def test_qsa_fused_prepare_matches_unfused( + main_fp8, indexer_dtype, mrope, is_2d_positions, @@ -184,9 +263,8 @@ def test_qsa_fused_prepare_matches_unfused( ) -> None: from vllm.model_executor.layers.rotary_embedding import get_rope - if indexer_dtype == torch.float8_e4m3fn: - pytest.skip("the AMD QSA indexer cache is bf16-only") - + if indexer_dtype is None: + indexer_dtype = current_platform.fp8_dtype() # The FP8 comparison below is a one-ulp bound on data this test generates, # so an unseeded run decides for itself whether it reproduces a rounding # difference. Seed it, or the kernel's bit-exactness is only sampled. @@ -338,7 +416,7 @@ def test_qsa_fused_prepare_matches_unfused( k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2 fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device) - main = _make_main_inputs(num_tokens) + main = _make_main_inputs(num_tokens, main_fp8) main_q_out, main_gate_out = qsa_prepare( projected_qk[:, : HQ * D], projected_qk[:, HQ * D :], @@ -400,7 +478,7 @@ def test_qsa_fused_prepare_matches_unfused( if rope_positions is not None: qsa_store_cache_rows(rope_positions, raw_slots, position_rows) - if indexer_dtype == torch.float8_e4m3fn: + if indexer_dtype != torch.bfloat16: # Both paths round the same ~bf16 intermediates to e4m3. Bitwise # equality (the dsv4 indexer test's bar) does not hold here: the 1D # reference rope (forward_cuda) differs from _norm_rope by up to 1 @@ -410,7 +488,7 @@ def test_qsa_fused_prepare_matches_unfused( else: torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL) assert torch.equal(fused_raw.view(torch.int16), unfused_raw.view(torch.int16)) - if indexer_dtype == torch.float8_e4m3fn: + if indexer_dtype != torch.bfloat16: assert_fp8_within_one_ulp(fused_compressed, unfused_compressed) else: torch.testing.assert_close( diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index 8cea1f83d420..e2ba804eb617 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -8,14 +8,16 @@ import torch from vllm.models.qwen4_exp.common import qsa_cache -from vllm.models.qwen4_exp.common.qsa_cache import QSAMetadataBuilder +from vllm.models.qwen4_exp.common.qsa_cache import ( + QSAMetadataBuilder, + qsa_kv_cache_dtype, +) from vllm.models.qwen4_exp.nvidia import indexer_qsa from vllm.models.qwen4_exp.nvidia import ( model as _qwen4_exp_model, # noqa: F401 ) from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops from vllm.models.qwen4_exp.nvidia.ops import qsa_indexer as qsa_indexer_ops -from vllm.models.qwen4_exp.nvidia.qsa import qsa_kv_cache_dtype from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON from vllm.v1.worker.utils import clear_layer_kv_caches diff --git a/vllm/models/qwen4_exp/amd/indexer_qsa.py b/vllm/models/qwen4_exp/amd/indexer_qsa.py index 09cb39acf919..092fb0138916 100644 --- a/vllm/models/qwen4_exp/amd/indexer_qsa.py +++ b/vllm/models/qwen4_exp/amd/indexer_qsa.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope +from vllm.platforms import current_platform from ..common.qsa_cache import ( QSACompressedKeyCache, @@ -131,6 +132,19 @@ def __init__( cache_config = vllm_config.cache_config cache_prefix = f"{prefix}." if prefix else "" + self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype( + default="bf16" + ) + if self.indexer_kv_dtype == "fp8": + indexer_dtype = current_platform.fp8_dtype() + elif self.indexer_kv_dtype == "bf16": + indexer_dtype = torch.bfloat16 + else: + raise NotImplementedError( + f"indexer_kv_dtype={self.indexer_kv_dtype!r} is not supported " + "by the Qwen4Exp QSA indexer (only 'bf16' or 'fp8')." + ) + self.indexer_dtype = indexer_dtype self.raw_key_cache = QSAKeyStateCache( head_size=self.index_head_dim, dtype=torch.bfloat16, @@ -142,7 +156,7 @@ def __init__( ) self.compressed_key_cache = QSACompressedKeyCache( head_size=self.index_head_dim, - dtype=torch.bfloat16, + dtype=indexer_dtype, compress_ratio=self.compress_ratio, prefix=f"{cache_prefix}compressed_key_cache", cache_config=cache_config, @@ -180,6 +194,7 @@ def project_qk( q.reshape(-1, self.index_head_dim), ).reshape_as(q) q = apply_qsa_rope(self.rotary_emb, positions, q) + q = q.to(self.indexer_dtype) return q, token_k.reshape(-1, 1, self.index_head_dim) def normalize_compressed_keys( @@ -291,7 +306,19 @@ def _fused_prepare( """ num_tokens = raw_metadata.num_actual_tokens raw_key_cache = self.raw_key_cache - q = projected_q.new_empty(num_tokens, self.index_n_heads, self.index_head_dim) + q = projected_q.new_empty( + num_tokens, + self.index_n_heads, + self.index_head_dim, + dtype=self.indexer_dtype, + ) + main_kv_cache = attn.kv_cache.transpose(1, 2) + if attn.kv_cache_dtype in ("fp8", "fp8_e4m3"): + # Reinterpret the uint8 cache in the platform's native fp8 flavor, + # matching the C++ cache-update op and the attention backend's + # read-side views. + fp8_dtype = current_platform.fp8_dtype() + main_kv_cache = main_kv_cache.view(fp8_dtype) main_outputs = qsa_prepare( projected_q, raw_keys, @@ -320,7 +347,7 @@ def _fused_prepare( main_q_norm_weight=attn.q_norm.weight, main_k_norm_weight=attn.k_norm.weight, main_eps=attn.q_norm.variance_epsilon, - main_kv_cache=attn.kv_cache.transpose(1, 2), + main_kv_cache=main_kv_cache, main_slot_mapping=slot_mapping[:num_tokens], main_k_scale=attn._k_scale_float, main_v_scale=attn._v_scale_float, diff --git a/vllm/models/qwen4_exp/amd/ops/qsa.py b/vllm/models/qwen4_exp/amd/ops/qsa.py index 6bdb1fe592df..1ad1f8323ac8 100644 --- a/vllm/models/qwen4_exp/amd/ops/qsa.py +++ b/vllm/models/qwen4_exp/amd/ops/qsa.py @@ -220,6 +220,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + IS_FP8: tl.constexpr, ) -> None: row = tl.program_id(0) kv_head = tl.program_id(1) @@ -292,6 +293,14 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[:, None], other=0.0, ) + if IS_FP8: + # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (an fp8 + # QK dot measured slower and less accurate on the NVIDIA port). + keys = keys.to(query.dtype) + # Dequant V to fp16 for the PV dot: P <= 1 (online softmax) so + # fp16 has the range, its wider mantissa is more accurate, and + # the fp8->fp16 upcast with an fp16 PV dot is faster. + values = values.to(tl.float16) scores = tl.dot(query, keys) # Scaling scores avoids re-quantizing a scaled query to BF16. scores *= softmax_scale_log2 @@ -825,7 +834,7 @@ def qsa_sparse_paged_attention( token_to_req: torch.Tensor, out: torch.Tensor | None = None, ) -> torch.Tensor: - """Run sparse GQA directly over paged BF16 K/V caches.""" + """Run sparse GQA directly over paged bf16 or fp8 K/V caches.""" if not q.is_cuda or not HAS_TRITON: raise RuntimeError("paged QSA sparse attention requires a GPU and Triton") if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: @@ -842,7 +851,12 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse attention requires valid grouped-query heads") head_dim = q.shape[2] assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 - assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16 + is_fp8 = k_cache.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + assert q.dtype == torch.bfloat16 + assert k_cache.dtype == v_cache.dtype + assert is_fp8 or k_cache.dtype == torch.bfloat16, ( + "QSA sparse attention requires BF16 or FP8-e4m3 K/V" + ) assert logical_indices.dtype == block_table.dtype == torch.int32 assert token_to_req.dtype == torch.int32 assert q.device == k_cache.device == v_cache.device @@ -938,6 +952,7 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, + IS_FP8=is_fp8, num_warps=partial_warps, num_stages=partial_stages, ) diff --git a/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py b/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py index bc2fee3b3a47..80fe084031ad 100644 --- a/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py +++ b/vllm/models/qwen4_exp/amd/ops/qsa_prepare.py @@ -124,7 +124,7 @@ def _to_dst_dtype(x, dst, scale): """Round to BF16 like the unfused path, then scale for an FP8 destination.""" out_ty = dst.dtype.element_ty x = x.to(tl.bfloat16) - if out_ty == tl.float8e4nv: + if out_ty.is_fp8(): x = x.to(tl.float32) / scale return x.to(out_ty) diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index 6a246dd31ec8..cf9f0fbbc28f 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/vllm/models/qwen4_exp/amd/qsa.py @@ -25,6 +25,7 @@ from vllm.model_executor.models.qwen3_next import Qwen3NextAttention from vllm.model_executor.models.qwen3_vl import mrope_positions_factor from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, @@ -58,7 +59,7 @@ is_workspace_manager_initialized, ) -from ..common.qsa_cache import QSAForwardMetadata +from ..common.qsa_cache import QSAForwardMetadata, qsa_kv_cache_dtype from . import model from .indexer_qsa import QSAIndexer @@ -112,7 +113,13 @@ def __init__( "octave_k3v3", "octave_k3v3_compact", ) -_QSA_CACHE_DTYPES: tuple[CacheDType, ...] = ("auto", "bfloat16", *_OCTAVE_CACHE_DTYPES) +_QSA_CACHE_DTYPES: tuple[CacheDType, ...] = ( + "auto", + "bfloat16", + "fp8", + "fp8_e4m3", + *_OCTAVE_CACHE_DTYPES, +) class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): @@ -124,8 +131,47 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): """ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + # fp8/fp8_e4m3: e4m3 bytes in a uint8 cache, written by qsa_prepare with + # the layer's per-tensor scales and cast to the query dtype for the QK + # dot (fp16 for PV) inside the QSA Triton kernel. flash-attn never runs + # over this cache, so its fp8 probe does not apply (see + # supports_combination below). + # octave_*: packed 3-bit slots handled by the Octave kernels. supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = list(_QSA_CACHE_DTYPES) + @classmethod + def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: + return kv_cache_dtype is None or kv_cache_dtype in cls.supported_kv_cache_dtypes + + @classmethod + def supports_combination( + 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, + device_capability: DeviceCapability, + ) -> str | None: + # The QSA Triton kernel reads the quantized KV itself (fp8 casts + # inside the kernel, Octave decode kernels); flash-attn never runs + # over this cache, so bypass only the parent's flash-attn + # kv-cache-dtype probe and keep every other combination check. + return super().supports_combination( + head_size, + dtype, + None, + block_size, + use_mla, + has_sink, + use_sparse, + use_mm_prefix, + device_capability, + ) + @classmethod def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: from vllm.v1.attention.backends.rocm_octave_attn import ( @@ -166,8 +212,44 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): supports_dcp: bool = False supports_pcp: bool = False - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None = None, + attn_type: AttentionType = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + sinks: torch.Tensor | None = None, + ) -> None: + # The parent constructor probes flash-attn for quantized-KV support + # and raises where it is unavailable (ROCm), but QSA dequantizes fp8 + # inside its own Triton kernel and never runs flash-attn over the + # cache. Hand the parent "auto" for that probe and restore the real + # dtype afterwards, mirroring the NVIDIA implementation: the parent + # only uses it there, and the cache-update path reads the attribute + # at call time. + real_kv_cache_dtype = kv_cache_dtype + if kv_cache_dtype in ("fp8", "fp8_e4m3"): + kv_cache_dtype = "auto" + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + sinks, + ) + self.kv_cache_dtype = real_kv_cache_dtype if not is_flash_attn_varlen_func_available(): raise NotImplementedError("Qwen4Exp QSA requires FlashAttention") if self.dcp_world_size != 1: @@ -176,7 +258,7 @@ def __init__(self, *args, **kwargs) -> None: ) if self.kv_cache_dtype not in _QSA_CACHE_DTYPES: raise NotImplementedError( - "Qwen4Exp QSA requires a BF16 or Octave main KV cache" + "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache" ) self.supports_quant_query_input = False self.octave_fmt: rocm_octave.OctaveFormat | None = None @@ -332,8 +414,24 @@ def forward_qsa( key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) key_cache = canonicalize_singleton_dim_strides(key_cache) value_cache = canonicalize_singleton_dim_strides(value_cache) - if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16: - raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V") + if self.kv_cache_dtype in ("fp8", "fp8_e4m3"): + assert key_cache.dtype == torch.uint8, ( + "fp8 KV cache must be uint8 storage; a config mismatch would " + "silently double the view's shape" + ) + fp8_dtype = current_platform.fp8_dtype() + key_cache = key_cache.view(fp8_dtype) + value_cache = value_cache.view(fp8_dtype) + if query.dtype != torch.bfloat16: + raise NotImplementedError("Qwen4Exp QSA requires BF16 queries") + if key_cache.dtype != torch.bfloat16 and ( + layer._k_scale_float != 1.0 or layer._v_scale_float != 1.0 + ): + raise NotImplementedError( + "Qwen4Exp QSA with an FP8 KV cache requires unit scales on " + "the AMD path (the NVIDIA port's host-side dequant fold is " + "not yet ported here). Use bf16 KV caches or unit scales." + ) from .ops.qsa import qsa_sparse_paged_attention @@ -371,9 +469,10 @@ def __init__( raise ValueError("Qwen4Exp QSA requires a paged KV cache") if model_config.dtype != torch.bfloat16: raise NotImplementedError("Qwen4Exp QSA currently requires BF16") - if cache_config.cache_dtype not in _QSA_CACHE_DTYPES: + cache_dtype = qsa_kv_cache_dtype(cache_config, prefix) + if cache_dtype not in _QSA_CACHE_DTYPES: raise NotImplementedError( - "Qwen4Exp QSA requires a BF16 or Octave main KV cache" + "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache" ) if getattr(quant_config, "kv_cache_scheme", None) is not None: raise NotImplementedError("Qwen4Exp QSA does not support KV quantization") @@ -452,15 +551,17 @@ def __init__( self.layer_name = f"{prefix}.attn" self.attn_type = AttentionType.DECODER - self.kv_cache_dtype = cache_config.cache_dtype + self.kv_cache_dtype = cache_dtype self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( self.kv_cache_dtype, model_config ) if ( - self.kv_cache_dtype not in _OCTAVE_CACHE_DTYPES + self.kv_cache_dtype not in _QSA_CACHE_DTYPES and self.kv_cache_torch_dtype != torch.bfloat16 ): - raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage") + raise NotImplementedError( + "Qwen4Exp QSA requires a BF16, FP8-e4m3, or Octave main KV cache" + ) self.kv_sharing_target_layer_name = None self.kv_cache = torch.tensor([]) set_default_quant_scales(self, register_buffer=True) @@ -486,8 +587,8 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.indexer", ) - # The fused prepare writes raw BF16 K/V; an Octave cache is packed by - # do_kv_cache_update instead. + # The fused prepare writes bf16 or platform-fp8 K/V; an Octave cache + # is packed by do_kv_cache_update instead. self.use_fused_qsa_prepare = ( self.attn_output_gate and self.indexer.use_fused_pre_indexer diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index 5eba092eec25..8b7acee4a246 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -883,4 +883,14 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: "canonical_qsa_rope_positions", "circular_qsa_slot_mapping", "compressed_qsa_slot_mapping", + "qsa_kv_cache_dtype", ] + + +def qsa_kv_cache_dtype(cache_config: CacheConfig, prefix: str) -> CacheDType | str: + """The layer's KV cache dtype, honoring ``--kv-cache-dtype-skip-layers``.""" + from vllm.model_executor.models.utils import extract_layer_index + + if str(extract_layer_index(prefix)) in cache_config.kv_cache_dtype_skip_layers: + return "auto" + return cache_config.cache_dtype diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 36e4864570ee..0e9d87b995e5 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -14,7 +14,7 @@ import vllm.envs as envs from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import VllmConfig, get_current_vllm_config_or_none -from vllm.config.cache import CacheConfig, CacheDType +from vllm.config.cache import CacheDType from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention.attention import ( @@ -34,7 +34,6 @@ from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, - extract_layer_index, ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability @@ -61,7 +60,7 @@ get_kv_quant_mode, ) -from ..common.qsa_cache import QSAForwardMetadata +from ..common.qsa_cache import QSAForwardMetadata, qsa_kv_cache_dtype from . import model from .indexer_qsa import QSAIndexer from .ops.qsa_indexer import _TOPK_WORKSPACE_BYTES @@ -133,13 +132,6 @@ def remap_shards(): return super().load_weights(remap_shards()) -def qsa_kv_cache_dtype(cache_config: CacheConfig, prefix: str) -> CacheDType | str: - """The layer's KV cache dtype, honoring ``--kv-cache-dtype-skip-layers``.""" - if str(extract_layer_index(prefix)) in cache_config.kv_cache_dtype_skip_layers: - return "auto" - return cache_config.cache_dtype - - class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder): """Flash metadata supporting uniform decode and target-verify graphs.""" From b42f25b00c3d13b565ec60e81928cd21efefce32 Mon Sep 17 00:00:00 2001 From: nv-nikitaj Date: Wed, 28 Oct 2026 19:03:01 -0700 Subject: [PATCH 102/406] [Bugfix][NVFP4] Fix NaN logits and garbage output from uninitialized fused act+quant scales (#63382) Signed-off-by: Nikita Jain Co-authored-by: vllm-agent Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Thien Tran --- .../activation_nvfp4_quant_fusion_kernels.cu | 45 +++++++++++++------ .../quantization/test_silu_mul_nvfp4_quant.py | 41 +++++++++++++++++ 2 files changed, 73 insertions(+), 13 deletions(-) diff --git a/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu index 2064775111fb..085d65430bb0 100644 --- a/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu +++ b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu @@ -38,7 +38,7 @@ namespace vllm { template __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, - int32_t num_packed_cols, + int32_t num_packed_cols, int32_t num_padded_cols, Type const* __restrict__ in, float const* __restrict__ SFScale, uint32_t* __restrict__ out, @@ -52,6 +52,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) // Precompute SF layout parameter (constant for entire kernel). int32_t const numKTiles = (numCols + 63) / 64; + int const sf_m = round_up(numRows, 128); + // Get the global scaling factor, which will be applied to the SF. // Note SFScale is the same as next GEMM's alpha, which is // (448.f / (Alpha_A / 6.f)). @@ -60,9 +62,10 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x; int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD; - // Input tensor row/col loops. - for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) { - if (colIdx < num_packed_cols) { + // Iterate over all rows and cols including padded ones - + // ensures we visit every single scale factor address to initialize it. + for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) { + if (colIdx < num_padded_cols) { PackedVec in_vec; PackedVec in_vec2; int64_t inOffset = @@ -70,7 +73,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) int64_t inOffset2 = rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) + numCols / CVT_FP4_ELTS_PER_THREAD + colIdx; - bool valid = (rowIdx < numRows) && (elem_idx < numCols); + // If we are outside valid rows/columns, feed zeros + bool valid = (rowIdx < numRows) && (colIdx < num_packed_cols); if constexpr (CVT_FP4_PACK16) { ld256_cg_or_zero(reinterpret_cast(in_vec), &reinterpret_cast(in)[inOffset * 8], @@ -95,9 +99,17 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) CVT_FP4_NUM_THREADS_PER_SF>( rowIdx, colIdx, numKTiles, SFout); - auto out_val = - cvt_warp_fp16_to_fp4( - out_silu_mul, SFScaleVal, sf_out); + fp4_packed_t out_val{}; + if (valid) { + out_val = + cvt_warp_fp16_to_fp4( + out_silu_mul, SFScaleVal, sf_out); + } else if (sf_out != nullptr) { + // The swizzled scale tensor may be allocated with torch.empty. + // Explicitly initialize padded rows and columns in this kernel so + // downstream NVFP4 GEMMs never observe stale scale factors. + *sf_out = 0; + } if (valid) { if constexpr (CVT_FP4_PACK16) { @@ -108,7 +120,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) } else { int64_t outOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; - out[outOffset] = out_val; + reinterpret_cast(out)[outOffset] = out_val; } } } @@ -146,9 +158,16 @@ void silu_and_mul_nvfp4_quant_sm1xxa( int num_packed_cols = int(n / CVT_FP4_ELTS_PER_THREAD); - int grid_y = vllm::div_round_up(num_packed_cols, static_cast(block.x)); - int grid_x = std::min( - int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); + // The swizzled scale tensor is padded to 128 rows x 4 scale columns; the + // grid covers the padding so the kernel can initialize it. + int sf_n_int = int(vllm::round_up(n / CVT_FP4_SF_VEC_SIZE, 4) / 4); + int num_padded_cols = + sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; + + int grid_y = vllm::div_round_up(num_padded_cols, static_cast(block.x)); + int grid_x = + std::min(vllm::computeEffectiveRows(m), + std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); dim3 grid(grid_x, grid_y); VLLM_STABLE_DISPATCH_HALF_TYPES( @@ -156,7 +175,7 @@ void silu_and_mul_nvfp4_quant_sm1xxa( using cuda_type = vllm::CUDATypeConverter::Type; auto input_ptr = static_cast(input.data_ptr()); vllm::silu_mul_cvt_fp16_to_fp4<<>>( - m, n, num_packed_cols, input_ptr, input_sf_ptr, + m, n, num_packed_cols, num_padded_cols, input_ptr, input_sf_ptr, reinterpret_cast(output_ptr), reinterpret_cast(sf_out)); }); diff --git a/tests/kernels/quantization/test_silu_mul_nvfp4_quant.py b/tests/kernels/quantization/test_silu_mul_nvfp4_quant.py index 02ab4f009fa8..8c0f931a2c98 100644 --- a/tests/kernels/quantization/test_silu_mul_nvfp4_quant.py +++ b/tests/kernels/quantization/test_silu_mul_nvfp4_quant.py @@ -76,3 +76,44 @@ def test_silu_mul_nvfp4_quant( torch.testing.assert_close( ref_output_dequant, fused_output_dequant, atol=atol, rtol=rtol ) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("shape", [(90, 160), (1, 160), (130, 96), (128, 256)]) +@torch.inference_mode() +def test_silu_mul_nvfp4_quant_scale_padding_initialized( + default_vllm_config, + dtype: torch.dtype, + shape: tuple[int, int], +) -> None: + """The kernel must overwrite dirty scale padding (rows and columns) + without relying on a zero-initialized allocation.""" + set_random_seed(42) + device = "cuda:0" + torch.set_default_device(device) + + m, two_d = shape + d = two_d // 2 + round_up = lambda x, y: (x + y - 1) // y * y + rounded_m = round_up(m, 128) + rounded_scale_n = round_up(d // BLOCK_SIZE, 4) + + x = torch.randn(shape, dtype=dtype) + global_scale = torch.tensor(1.0, dtype=torch.float32) + + out = torch.empty((m, d // 2), dtype=FP4_DTYPE) + out_scale = torch.full( + (rounded_m, rounded_scale_n // 4), 0x7F7F7F7F, dtype=torch.int32 + ) + torch.ops._C.silu_and_mul_nvfp4_quant(out, out_scale, x, global_scale) + + tmp = out_scale.view(FP8_DTYPE).reshape( + 1, rounded_m // 128, rounded_scale_n // 4, 32, 4, 4 + ) + full_scale = tmp.permute(0, 1, 4, 3, 2, 5).reshape(rounded_m, rounded_scale_n) + full_scale = full_scale.to(torch.float32) + + assert torch.count_nonzero(full_scale[m:, :]) == 0 + assert torch.count_nonzero(full_scale[:m, d // BLOCK_SIZE :]) == 0 + # Valid entries must have been written (no leftover 0x7F == NaN fill). + assert not torch.isnan(full_scale[:m, : d // BLOCK_SIZE]).any() From a6c3635e3660f6343c14e2158c7ab44c62be88b9 Mon Sep 17 00:00:00 2001 From: "yhcheong (vllmellm)" Date: Thu, 29 Oct 2026 10:07:41 +0800 Subject: [PATCH 103/406] [ROCm][Model] Enable ROCm AITER unified attention for MiMo-V2.6 (#64064) Signed-off-by: vllmellm --- .buildkite/test-amd.yaml | 3 +- .../benchmark_reshape_and_cache_flash.py | 116 ++++++++-- csrc/cache_kernels.cu | 113 +++++++++- tests/kernels/attention/test_cache.py | 186 ++++++++++++++++ .../models/test_mimo_v2_attention_backend.py | 205 ++++++++++++++++++ vllm/model_executor/models/mimo_v2.py | 26 ++- vllm/v1/attention/backend.py | 5 + .../attention/backends/flash_attn_diffkv.py | 4 + .../backends/rocm_aiter_unified_attn.py | 6 + .../attention/backends/triton_attn_diffkv.py | 4 + 10 files changed, 637 insertions(+), 31 deletions(-) create mode 100644 tests/models/test_mimo_v2_attention_backend.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index ac593ad3a83a..993bee49f186 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -471,7 +471,8 @@ steps: 1) pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' && pytest models/language -v -s -m 'distributed(num_gpus=2)' && - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' ;; + VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' && + pytest -v -s models/test_mimo_v2_attention_backend.py ;; 2) pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py && diff --git a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py index 84d9cc707c40..229f5385a4c2 100644 --- a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py +++ b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py @@ -22,10 +22,45 @@ @torch.inference_mode() +def _make_packed_kv_caches( + num_blocks: int, + block_size: int, + num_heads: int, + head_size: int, + v_head_size: int, + kv_cache_dtype: str, + dtype: torch.dtype, + kv_cache_layout: str, + device: str, +): + """K and V as views of one [.., head_size + v_head_size] buffer. + + The asymmetric kernel needs both sides to share a head stride, which is what + a packed K|V buffer gives (and what backends such as ROCm AITER unified + attention produce). + """ + if kv_cache_dtype not in ("auto", "bfloat16", "float16", "fp8"): + raise ValueError(f"asymmetric benchmark does not support {kv_cache_dtype}.") + content = head_size + v_head_size + if kv_cache_layout == "NHD": + physical = (num_blocks, block_size, num_heads, content) + permute = None + else: + physical = (num_blocks, num_heads, block_size, content) + permute = (0, 2, 1, 3) + + cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype + packed = torch.zeros(physical, dtype=cache_dtype, device=device) + view = packed if permute is None else packed.permute(permute) + key_cache, value_cache = view.split([head_size, v_head_size], dim=-1) + return key_cache, value_cache + + def run_benchmark( num_tokens: int, num_heads: int, head_size: int, + v_head_size: int | None, block_size: int, num_blocks: int, dtype: torch.dtype, @@ -37,8 +72,14 @@ def run_benchmark( device: str = "cuda", ) -> float: """Return latency (seconds) for given num_tokens.""" - if kv_cache_dtype == "fp8" and head_size % 16: - raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.") + if v_head_size is None: + v_head_size = head_size + asymmetric = v_head_size != head_size + + if kv_cache_dtype == "fp8" and (head_size % 16 or v_head_size % 16): + raise ValueError( + "fp8 kv-cache requires head_size and v_head_size to be multiples of 16." + ) if implementation not in ("cuda", "triton"): raise ValueError( @@ -47,13 +88,23 @@ def run_benchmark( ) if implementation == "triton" and kv_cache_layout == "HND": return float("nan") # Triton does not support HND layout yet. + if asymmetric and implementation != "cuda": + raise ValueError( + "asymmetric head sizes (v_head_size != head_size) are only supported " + "by the 'cuda' implementation." + ) set_random_seed(42) torch.set_default_device(device) # create random key / value tensors [T, H, D]. key = torch.randn(num_tokens, num_heads, head_size, dtype=dtype, device=device) - value = torch.randn_like(key) + if asymmetric: + value = torch.randn( + num_tokens, num_heads, v_head_size, dtype=dtype, device=device + ) + else: + value = torch.randn_like(key) # prepare the slot mapping. # each token is assigned a unique slot in the KV-cache. @@ -63,20 +114,36 @@ def run_benchmark( slot_mapping_lst = random.sample(range(num_slots), num_tokens) slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device) - key_caches, value_caches = create_kv_caches_with_random_flash( - num_blocks, - block_size, - 1, # num_layers - num_heads, - head_size, - kv_cache_dtype, - dtype, - device=device, - cache_layout=kv_cache_layout, - ) - key_cache, value_cache = key_caches[0], value_caches[0] - # to free unused memory - del key_caches, value_caches + if asymmetric: + # K and V are views of one buffer, so both sides share the head stride; + # that is the only layout the asymmetric path accepts (see + # reshape_and_cache_flash). + key_cache, value_cache = _make_packed_kv_caches( + num_blocks, + block_size, + num_heads, + head_size, + v_head_size, + kv_cache_dtype, + dtype, + kv_cache_layout, + device, + ) + else: + key_caches, value_caches = create_kv_caches_with_random_flash( + num_blocks, + block_size, + 1, # num_layers + num_heads, + head_size, + kv_cache_dtype, + dtype, + device=device, + cache_layout=kv_cache_layout, + ) + key_cache, value_cache = key_caches[0], value_caches[0] + # to free unused memory + del key_caches, value_caches # compute per-kernel scaling factors for fp8 conversion (if used). k_scale = (key.amax() / 64.0).to(torch.float32) @@ -142,6 +209,7 @@ def main(args): num_tokens=n_tok, num_heads=args.num_heads, head_size=args.head_size, + v_head_size=args.v_head_size, block_size=args.block_size, num_blocks=args.num_blocks, dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype], @@ -154,9 +222,14 @@ def main(args): ) rows.append([n_tok, layout, f"{lat * 1e6:.3f}"]) + v_note = ( + "" + if args.v_head_size is None + else f", head_size={args.head_size} v_head_size={args.v_head_size}" + ) print( f"Benchmark results for implementation {args.implementation}" - f" (measuring with {args.mode}):" + f" (measuring with {args.mode}{v_note}):" ) print(tabulate(rows, headers=["num_tokens", "layout", "latency (µs)"])) @@ -171,6 +244,13 @@ def main(args): choices=[64, 80, 96, 112, 120, 128, 192, 256], default=128, ) + parser.add_argument( + "--v-head-size", + type=int, + default=None, + help="V head dim; defaults to --head-size. Setting it differently " + "benchmarks the asymmetric K/V path (cuda implementation only).", + ) parser.add_argument("--block-size", type=int, choices=[16, 32], default=16) parser.add_argument("--num-blocks", type=int, default=128 * 512) diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index 3541d1050128..bd4a080c25ff 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -420,6 +420,73 @@ __global__ void reshape_and_cache_flash_kernel( } } +// Asymmetric K/V head widths. K and V are views of one packed buffer, so they +// share a head stride and the head-contiguous path below cannot apply; this +// kernel therefore walks heads explicitly and takes V's width as its one extra +// argument. It is kept separate so the symmetric kernel's codegen is unchanged. +template +__global__ void reshape_and_cache_flash_kernel_asym( + const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] + const scalar_t* __restrict__ value, // [num_tokens, num_heads, v_head_size] + // NHD or HND: [num_blocks, block_size, num_heads, head_size], addressed + // through the strides passed by the launcher + cache_t* __restrict__ key_cache, + cache_t* __restrict__ value_cache, // same, with v_head_size + const int64_t* __restrict__ slot_mapping, // [num_tokens] + const int64_t block_stride, const int64_t page_stride, + const int64_t head_stride, const int64_t key_stride, + const int64_t value_stride, const int num_heads, const int head_size, + const int v_head_size, const int block_size, const float* k_scale, + const float* v_scale, const int kv_scale_stride) { + const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; + // NOTE: slot_idx can be -1 if the token is padded + if (slot_idx < 0) { + return; + } + const int64_t block_idx = slot_idx / block_size; + const int64_t block_offset = slot_idx % block_size; + + const scalar_t* __restrict__ key_src = key + token_idx * key_stride; + const scalar_t* __restrict__ value_src = value + token_idx * value_stride; + + cache_t* __restrict__ key_dst = + key_cache + block_idx * block_stride + block_offset * page_stride; + cache_t* __restrict__ value_dst = + value_cache + block_idx * block_stride + block_offset * page_stride; + + constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4; + + const int lane = threadIdx.x & 31; // 0..31 within warp + const int warp_id = threadIdx.x >> 5; // warp index within block + const int warps_per_block = blockDim.x >> 5; + + for (int head = warp_id; head < num_heads; head += warps_per_block) { + const scalar_t* __restrict__ k_src_h = key_src + head * head_size; + const scalar_t* __restrict__ v_src_h = value_src + head * v_head_size; + + cache_t* __restrict__ k_dst_h = + key_dst + static_cast(head) * head_stride; + cache_t* __restrict__ v_dst_h = + value_dst + static_cast(head) * head_stride; + + float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) + ? 0.f + : k_scale[head * kv_scale_stride]; + float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) + ? 0.f + : v_scale[head * kv_scale_stride]; + + CopyWithScaleOp k_op{k_scale_val}; + CopyWithScaleOp v_op{v_scale_val}; + + vectorize_with_alignment(k_src_h, k_dst_h, head_size, lane, 32, + k_op); + vectorize_with_alignment(v_src_h, v_dst_h, v_head_size, lane, 32, + v_op); + } +} + template __global__ void concat_and_cache_mla_kernel( @@ -827,6 +894,20 @@ void reshape_and_cache( // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. +#define CALL_RESHAPE_AND_CACHE_FLASH_ASYM(KV_T, CACHE_T, KV_DTYPE) \ + vllm::reshape_and_cache_flash_kernel_asym \ + <<>>( \ + reinterpret_cast(key.data_ptr()), \ + reinterpret_cast(value.data_ptr()), \ + reinterpret_cast(key_cache.data_ptr()), \ + reinterpret_cast(value_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, page_stride, \ + head_stride, key_stride, value_stride, num_heads, head_size, \ + v_head_size, block_size, \ + reinterpret_cast(k_scale.data_ptr()), \ + reinterpret_cast(v_scale.data_ptr()), \ + kv_scale_stride); + #define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \ vllm::reshape_and_cache_flash_kernel \ <<>>( \ @@ -864,6 +945,7 @@ void reshape_and_cache_flash( int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); + int v_head_size = value.size(2); const torch::stable::accelerator::DeviceGuard device_guard( key.get_device_index()); @@ -897,8 +979,8 @@ void reshape_and_cache_flash( int64_t block_stride = key_cache.stride(0); int64_t page_stride = key_cache.stride(1); int64_t head_stride = key_cache.stride(2); - STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); + STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()), "k_scale and v_scale must have the same shape"); STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads, @@ -906,10 +988,31 @@ void reshape_and_cache_flash( int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0; dim3 grid(num_tokens); - dim3 block(std::min(num_heads * head_size, 512)); - - DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, - CALL_RESHAPE_AND_CACHE_FLASH); + // Only models with v_head_size != head_size take the kernel below; everything + // else keeps using the kernel above, so its codegen is untouched. + if (v_head_size == head_size) { + dim3 block(std::min(num_heads * head_size, 512)); + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, + CALL_RESHAPE_AND_CACHE_FLASH); + } else { + // V is a view of K's buffer, so the two must agree on geometry; these + // checks stay out of the symmetric path so its host-side work is unchanged. + STD_TORCH_CHECK(value_cache.stride(1) == page_stride, + "key_cache and value_cache must share the page stride"); + STD_TORCH_CHECK(key_cache.size(0) == value_cache.size(0) && + key_cache.size(1) == value_cache.size(1), + "key_cache and value_cache must share block geometry"); + STD_TORCH_CHECK(key.size(1) == value.size(1), + "key and value must have the same number of heads"); + STD_TORCH_CHECK(value_cache.stride(2) == head_stride, + "key_cache and value_cache must share a head stride"); + // Only the strided walk applies here, and V's width is the single extra + // argument. One warp per head, and at least one warp so the head loop + // always advances. + dim3 block(32 * std::min(num_heads, 16)); + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, + CALL_RESHAPE_AND_CACHE_FLASH_ASYM); + } } // KV_T is the data type of key and value tensors. diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 116835e93223..fded1554d92f 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -508,6 +508,192 @@ def quantize_and_dequantize(cache_dtype: str) -> torch.Tensor: assert four_over_six_mse < default_mse +# MiMo-V2.6 Flash: K 192 / V 128, 4 KV heads, block sizes 16 and 64 in use. +ASYM_HEAD_SIZE = 192 +ASYM_HEAD_SIZE_V = 128 +ASYM_NUM_KV_HEADS = 4 +ASYM_BLOCK_SIZES = [16, 64] +ASYM_NUM_TOKENS = [6, 512] +# "unified": one page holds K's head group then V's, addressed the way ROCm +# AITER unified attention hands the kernel (transpose(1, 2).split(head_size)). +# "grouped": the same page content laid out per token, which gives the other +# (page-contiguous, head-strided) stride pattern. No in-tree backend produces it +# today; it keeps both stride shapes covered for asymmetric head dims. +ASYM_LAYOUTS = ["unified", "grouped"] + + +def test_reshape_and_cache_flash_asymmetric_tiny_head() -> None: + """A single small head still fills the cache correctly. + + num_heads * max(head_size, v_head_size) is below one warp here, so a launch + that sized the block from the element count would schedule zero warps and + the per-head loop would never advance. + """ + set_random_seed(0) + device = "cuda:0" + torch.set_default_device(device) + torch.accelerator.set_device_index(device) + + num_tokens, num_heads = 4, 1 + head_size, v_head_size = 16, 8 + block_size, num_blocks = 16, 2 + content = head_size + v_head_size + + packed = torch.zeros( + num_blocks, num_heads, block_size, content, dtype=torch.bfloat16 + ) + key_cache, value_cache = packed.transpose(1, 2).split( + [head_size, v_head_size], dim=-1 + ) + slot_mapping = torch.tensor([0, 5, 16, 31], dtype=torch.long) + key = torch.randn(num_tokens, num_heads, head_size, dtype=torch.bfloat16) + value = torch.randn(num_tokens, num_heads, v_head_size, dtype=torch.bfloat16) + k_scale = torch.tensor(1.0, dtype=torch.float32) + v_scale = torch.tensor(1.0, dtype=torch.float32) + + ops.reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + "auto", + k_scale, + v_scale, + ) + + expected_key = torch.zeros_like(key_cache) + expected_value = torch.zeros_like(value_cache) + for token, slot in enumerate(slot_mapping.tolist()): + block, offset = slot // block_size, slot % block_size + expected_key[block, offset] = key[token] + expected_value[block, offset] = value[token] + torch.testing.assert_close(key_cache, expected_key) + torch.testing.assert_close(value_cache, expected_value) + + +@pytest.mark.parametrize("kv_cache_layout", ASYM_LAYOUTS) +@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES) +@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) +@pytest.mark.parametrize("block_size", ASYM_BLOCK_SIZES) +@pytest.mark.parametrize("num_tokens", ASYM_NUM_TOKENS) +@torch.inference_mode() +def test_reshape_and_cache_flash_asymmetric_head_sizes( + kv_cache_layout: str, + kv_scale_type: str, + kv_cache_dtype: str, + block_size: int, + num_tokens: int, +) -> None: + """V is written with its own head width and head stride. + + Both layouts keep K and V in one layer buffer, so only per-head geometry + differs between the two sides. + """ + set_random_seed(0) + device = "cuda:0" + torch.set_default_device(device) + torch.accelerator.set_device_index(device) + + num_heads = ASYM_NUM_KV_HEADS + num_blocks = num_tokens // block_size + 2 + content = ASYM_HEAD_SIZE + ASYM_HEAD_SIZE_V + # A kernel that copies K's width into V runs past the page, so leave enough + # slack to turn that into a failed assertion instead of an out-of-bounds + # write that takes the whole test process down. + guard = num_heads * content + cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else torch.bfloat16 + # The last slot is always written: a too-wide copy would run past the end. + num_slots = block_size * num_blocks + slot_mapping = torch.randperm(num_slots - 1, device=device)[: num_tokens - 1] + slot_mapping = torch.cat( + [slot_mapping, torch.tensor([num_slots - 1], device=device)] + ).to(torch.long) + + key = torch.randn( + num_tokens, num_heads, ASYM_HEAD_SIZE, dtype=torch.bfloat16, device=device + ) + value = torch.randn( + num_tokens, num_heads, ASYM_HEAD_SIZE_V, dtype=torch.bfloat16, device=device + ) + + # K and V are views of one packed layer buffer, as in vLLM. + content = ASYM_HEAD_SIZE + ASYM_HEAD_SIZE_V + if kv_cache_layout == "unified": + # [num_blocks, num_heads, block_size, content], split per side. + num_bytes = num_blocks * num_heads * block_size * content + buf = torch.zeros(num_bytes + guard, dtype=cache_dtype, device=device) + packed = buf[:num_bytes].view(num_blocks, num_heads, block_size, content) + key_cache, value_cache = packed.transpose(1, 2).split( + [ASYM_HEAD_SIZE, ASYM_HEAD_SIZE_V], dim=-1 + ) + else: + # [num_blocks, block_size, num_heads, content], K group then V group. + num_bytes = num_blocks * block_size * num_heads * content + buf = torch.zeros(num_bytes + guard, dtype=cache_dtype, device=device) + packed = buf[:num_bytes].view(num_blocks, block_size, num_heads, content) + key_cache, value_cache = packed.split( + [ASYM_HEAD_SIZE, ASYM_HEAD_SIZE_V], dim=-1 + ) + + if kv_cache_dtype == "fp8": + # Compare in fp16 after dequantizing, as the symmetric case does. + expected_key = torch.zeros(key_cache.shape, dtype=torch.float16, device=device) + expected_value = torch.zeros( + value_cache.shape, dtype=torch.float16, device=device + ) + else: + expected_key = torch.zeros_like(key_cache) + expected_value = torch.zeros_like(value_cache) + for token, slot in enumerate(slot_mapping.tolist()): + block, offset = slot // block_size, slot % block_size + expected_key[block, offset] = key[token] + expected_value[block, offset] = value[token] + + scale_shape = (1,) if kv_scale_type == "tensor" else (num_heads,) + k_scale = torch.ones(scale_shape, dtype=torch.float32, device=device) + v_scale = torch.ones(scale_shape, dtype=torch.float32, device=device) + + ops.reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + torch.accelerator.synchronize() + + if kv_cache_dtype == "fp8": + + def _dequant(cache, scale): + # gfx94x stores fp8 as e4m3fnuz; reinterpret with the platform dtype + # so the bytes are decoded the way the kernel wrote them. + values = cache.view(current_platform.fp8_dtype()).to(torch.float16) + scale = scale.to(torch.float16) + if scale.numel() == 1: + return values * scale + # per-attn-head scales broadcast over the head dimension + return values * scale.view(1, 1, -1, 1) + + result_key = _dequant(key_cache, k_scale) + result_value = _dequant(value_cache, v_scale) + torch.testing.assert_close(result_key, expected_key, atol=0.001, rtol=0.1) + torch.testing.assert_close(result_value, expected_value, atol=0.001, rtol=0.1) + else: + torch.testing.assert_close(key_cache, expected_key) + torch.testing.assert_close(value_cache, expected_value) + # Untouched slots stay zero, and nothing is written behind the cache. + assert torch.count_nonzero(key_cache) == num_tokens * num_heads * ASYM_HEAD_SIZE + assert ( + torch.count_nonzero(value_cache) + == num_tokens * num_heads * ASYM_HEAD_SIZE_V + ) + assert torch.count_nonzero(buf[num_bytes:]) == 0 + + @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) @pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS) diff --git a/tests/models/test_mimo_v2_attention_backend.py b/tests/models/test_mimo_v2_attention_backend.py new file mode 100644 index 000000000000..14279b6ee4aa --- /dev/null +++ b/tests/models/test_mimo_v2_attention_backend.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiMo-V2.6 has head_dim 192 but v_head_dim 128, so only backends that keep +V's width out of the K path can serve it. An explicit ``--attention-backend`` +from that set has to be honored; any other request has to fall back to the +DiffKV auto-pick with a warning rather than mount a backend that cannot serve +the shape. +""" + +import pytest +import torch + +from vllm.config import AttentionConfig, VllmConfig, set_current_vllm_config +from vllm.platforms import current_platform +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +HEAD_DIM = 192 +V_HEAD_DIM = 128 + +DIFFKV_BACKENDS = ( + AttentionBackendEnum.FLASH_ATTN_DIFFKV, + AttentionBackendEnum.TRITON_ATTN_DIFFKV, +) + + +@pytest.fixture(autouse=True) +def bf16_default_dtype(): + # Backend selection is dtype sensitive; use the dtype the model ships in. + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + yield + torch.set_default_dtype(default_dtype) + + +@pytest.fixture(autouse=True) +def restore_diffkv_head_size_v(): + """``set_head_size_v`` writes backend class state. + + Leaving it modified would make these tests order dependent, so snapshot and + restore it on every backend that can be selected here. + """ + from vllm.v1.attention.backends.flash_attn_diffkv import ( + FlashAttentionDiffKVBackend, + ) + from vllm.v1.attention.backends.triton_attn_diffkv import ( + TritonAttentionDiffKVBackend, + ) + + backends = (FlashAttentionDiffKVBackend, TritonAttentionDiffKVBackend) + saved = [backend.head_size_v for backend in backends] + yield + for backend, head_size_v in zip(backends, saved): + backend.head_size_v = head_size_v + + +def _build_attention( + backend: AttentionBackendEnum | None, v_head_dim: int = V_HEAD_DIM +): + from vllm.model_executor.models.mimo_v2 import MiMoV2Attention + + config = VllmConfig(attention_config=AttentionConfig(backend=backend)) + with set_current_vllm_config(config): + layer = MiMoV2Attention( + hidden_size=512, + num_heads=4, + num_kv_heads=2, + head_dim=HEAD_DIM, + v_head_dim=v_head_dim, + prefix="model.layers.0.self_attn", + ) + return layer + + +def test_default_backend_is_diffkv(dist_init) -> None: + """No explicit backend: the DiffKV auto-pick is unchanged.""" + layer = _build_attention(None) + assert layer.attn.backend in DIFFKV_BACKENDS + # V's width has to reach the layer: Attention defaults head_size_v to + # head_size (192), which would corrupt the output reshape and KV spec. + assert layer.attn.head_size_v == V_HEAD_DIM + + +@pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCM_AITER_UNIFIED_ATTN is ROCm only" +) +def test_explicit_unified_attention_is_honored(dist_init) -> None: + from vllm._aiter_ops import IS_AITER_FOUND + + if not IS_AITER_FOUND: + pytest.skip("aiter is not installed") + + layer = _build_attention(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN) + assert layer.attn.backend is AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN + # The enum alone would not catch the bug this guards: the impl has to be + # unified attention's, not the DiffKV one the auto-pick installs. + assert type(layer.attn.impl).__name__ == "RocmAiterUnifiedAttentionImpl" + + +def test_explicit_diffkv_backend_is_honored(dist_init) -> None: + """The accept path also works for the DiffKV family, on any platform.""" + layer = _build_attention(AttentionBackendEnum.TRITON_ATTN_DIFFKV) + assert layer.attn.backend is AttentionBackendEnum.TRITON_ATTN_DIFFKV + assert layer.attn.head_size_v == V_HEAD_DIM + + +def test_explicit_symmetric_backend_falls_back_to_diffkv(dist_init) -> None: + """A backend that cannot take head_size_v != head_size must not be mounted.""" + layer = _build_attention(AttentionBackendEnum.TRITON_ATTN) + assert layer.attn.backend in DIFFKV_BACKENDS + + +def test_layer_v_head_dim_reaches_the_backend(dist_init) -> None: + """The backend must be told V's width before its impl is built. + + The DiffKV backends default to ``head_size_v = 128``, which is this model's + own V width, so a missing ``set_head_size_v`` call is invisible at the + model's geometry. That attribute is what the metadata builder and forward + read, so assert it carries the layer's width instead of the default. + """ + from vllm.v1.attention.backends.triton_attn_diffkv import ( + TritonAttentionDiffKVBackend, + ) + + _build_attention(AttentionBackendEnum.TRITON_ATTN_DIFFKV, v_head_dim=64) + assert TritonAttentionDiffKVBackend.head_size_v == 64 + + +def test_only_backends_that_can_serve_v_head_size_are_admitted() -> None: + """`supports_head_size_v` is what the model's gate consults.""" + from vllm.v1.attention.backends.flash_attn_diffkv import ( + FlashAttentionDiffKVBackend, + ) + from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( + RocmAiterUnifiedAttentionBackend, + ) + from vllm.v1.attention.backends.triton_attn_diffkv import ( + TritonAttentionDiffKVBackend, + ) + + assert FlashAttentionDiffKVBackend.supports_head_size_v(HEAD_DIM, V_HEAD_DIM) + assert TritonAttentionDiffKVBackend.supports_head_size_v(HEAD_DIM, V_HEAD_DIM) + # The packed cache is split with `split(head_size)`, so V cannot be wider. + assert RocmAiterUnifiedAttentionBackend.supports_head_size_v(HEAD_DIM, V_HEAD_DIM) + assert not RocmAiterUnifiedAttentionBackend.supports_head_size_v( + V_HEAD_DIM, HEAD_DIM + ) + + +def test_cache_write_matches_pytorch_reference(dist_init) -> None: + """Fill the packed cache through the impl, then compare with torch. + + The layer hands the kernel two views of one layer buffer, and V has its own + width, so the check is what lands in that buffer: K in the first + ``head_dim`` slots of a row, V in the remaining ``v_head_dim``. Same shape + as the aiter FA check in tests/kernels/attention/test_rocm_aiter_fa.py, for + the asymmetric case those backends cannot serve. + """ + from vllm._aiter_ops import IS_AITER_FOUND + + if not IS_AITER_FOUND: + pytest.skip("aiter is not installed") + + from tests.kernels.allclose_default import get_default_atol, get_default_rtol + + layer = _build_attention(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN) + impl = layer.attn.impl + num_kv_heads = layer.num_kv_heads + num_tokens, block_size, num_blocks = 6, 16, 2 + + # One packed layer buffer per token row: [K head_dim | V v_head_dim]. The + # spare trailing block absorbs an over-wide write from a regressed kernel, + # so that fails the comparison instead of writing out of bounds. + packed = torch.zeros( + num_blocks + 1, + num_kv_heads, + block_size, + HEAD_DIM + V_HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + key = torch.randn( + num_tokens, num_kv_heads, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + value = torch.randn( + num_tokens, num_kv_heads, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + # Slot 31 is the last slot of the cache. + slot_mapping = torch.tensor([0, 5, 16, 31, 7, 17], dtype=torch.long, device="cuda") + + # The impl reads only the scales off `layer`; they are not the subject here. + scales = torch.nn.Module() + scales._k_scale = torch.ones((), dtype=torch.float32, device="cuda") + scales._v_scale = torch.ones((), dtype=torch.float32, device="cuda") + impl.do_kv_cache_update(scales, key, value, packed, slot_mapping) + + expected = torch.zeros_like(packed) + for token, slot in enumerate(slot_mapping.tolist()): + block, offset = divmod(slot, block_size) + # packed buffer is [blocks, kv_heads, block_tokens, content] + expected[block, :, offset, :HEAD_DIM] = key[token] + expected[block, :, offset, HEAD_DIM:] = value[token] + + torch.testing.assert_close( + packed, expected, atol=get_default_atol(packed), rtol=get_default_rtol(packed) + ) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index d35e4faf22f9..28f89fb99e8e 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -317,15 +317,24 @@ def __init__( sliding_window = sliding_window_size if sliding_window_size > -1 else None - # Use DiffKV backend when V has a different head dim than K. - # Auto-pick FA-DiffKV when FA3/4 is usable on this device, else fall - # back to TRITON_ATTN_DIFFKV. Users can force a choice via - # `--attention-backend `. + # V has a different head dim than K. Honor an explicit backend that + # declares it can serve that geometry; anything else keeps the auto-pick + # below, with a warning. if self.v_head_dim != self.head_dim: requested = get_current_vllm_config().attention_config.backend - if requested is not None and requested.name.endswith("_DIFFKV"): + if requested is not None and requested.get_class().supports_head_size_v( + self.head_dim, self.v_head_dim + ): backend_enum = requested else: + if requested is not None: + logger.warning_once( + "%s cannot serve head_size_v=%d != head_size=%d; " + "using the DiffKV auto-pick instead.", + requested.name, + self.v_head_dim, + self.head_dim, + ) fa_backend = AttentionBackendEnum.FLASH_ATTN_DIFFKV.get_class() assert hasattr(fa_backend, "is_supported_on_current_device") if fa_backend.is_supported_on_current_device( @@ -337,8 +346,11 @@ def __init__( else: backend_enum = AttentionBackendEnum.TRITON_ATTN_DIFFKV attn_backend = backend_enum.get_class() - assert hasattr(attn_backend, "set_head_size_v") - attn_backend.set_head_size_v(self.v_head_dim) + if backend_enum.name.endswith("_DIFFKV"): + # DiffKV takes V's head size as a backend class attribute; + # unified attention reads it from the layer's KV spec. + assert hasattr(attn_backend, "set_head_size_v") + attn_backend.set_head_size_v(self.v_head_dim) logger.info_once("Using %s for attention.", attn_backend.get_name()) else: attn_backend = None diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 761a582b20ca..5f53f7699298 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -106,6 +106,11 @@ def supports_head_size(cls, head_size: int) -> bool: supported_head_sizes = cls.get_supported_head_sizes() return (not supported_head_sizes) or head_size in supported_head_sizes + @classmethod + def supports_head_size_v(cls, head_size: int, head_size_v: int) -> bool: + """Whether the backend can serve a V head size that differs from K's.""" + return head_size_v == head_size + @classmethod def supports_dtype(cls, dtype: torch.dtype) -> bool: return dtype in cls.supported_dtypes diff --git a/vllm/v1/attention/backends/flash_attn_diffkv.py b/vllm/v1/attention/backends/flash_attn_diffkv.py index 758e2c563c60..e17488db81cf 100644 --- a/vllm/v1/attention/backends/flash_attn_diffkv.py +++ b/vllm/v1/attention/backends/flash_attn_diffkv.py @@ -35,6 +35,10 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend): # Default to 128 for this backend head_size_v: int = 128 + @classmethod + def supports_head_size_v(cls, head_size: int, head_size_v: int) -> bool: + return True + @classmethod def set_head_size_v(cls, head_size_v: int) -> None: cls.head_size_v = head_size_v diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 26aea0095635..3944d17cdaa3 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -87,6 +87,12 @@ def supports_block_size(cls, block_size: int | None) -> bool: def supports_head_size(cls, head_size: int) -> bool: return head_size >= 32 + @classmethod + def supports_head_size_v(cls, head_size: int, head_size_v: int) -> bool: + # The cache is split with `split(head_size)` over a 2*head_size page, so + # V can be narrower than K but not wider. + return head_size_v <= head_size + @classmethod def supports_mm_prefix(cls) -> bool: # Not implemented diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py index 2fff3226d67d..c4c9d083be12 100644 --- a/vllm/v1/attention/backends/triton_attn_diffkv.py +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -69,6 +69,10 @@ class TritonAttentionDiffKVBackend(TritonAttentionBackend): # V head dim — set per layer via ``set_head_size_v`` before instantiation. head_size_v: int = 128 + @classmethod + def supports_head_size_v(cls, head_size: int, head_size_v: int) -> bool: + return True + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", "bfloat16", From d6fbd947824959eabc6165751f8b9cd4aa83f546 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Thu, 29 Oct 2026 10:30:34 +0800 Subject: [PATCH 104/406] [Models] Deprecate `torch.compile` for Qwen3.5 models (#51012) Signed-off-by: Isotr0py Signed-off-by: Thien Tran Co-authored-by: vllm-agent Co-authored-by: Thien Tran Co-authored-by: Kimi Code --- .buildkite/hardware_tests/cpu.yaml | 1 + .buildkite/test-amd.yaml | 6 + .buildkite/test_areas/kernels.yaml | 3 + .buildkite/test_areas/lm_eval.yaml | 2 + .buildkite/test_areas/misc.yaml | 2 + .buildkite/test_areas/models_basic.yaml | 2 + .buildkite/test_areas/models_language.yaml | 1 + .buildkite/test_areas/spec_decode.yaml | 6 + tests/conftest.py | 6 +- tests/kernels/core/test_fused_pre_fc_norms.py | 66 ++ .../core/test_fused_sigmoid_gate_mul.py | 75 ++ tests/kernels/fla/test_fused_gdn_post_conv.py | 45 + .../fla/test_fused_recurrent_packed_decode.py | 54 + .../mamba/test_gdn_output_projection.py | 61 ++ .../test_runai_model_streamer_loader.py | 2 +- tests/models/test_qwen3_5.py | 47 + tests/v1/attention/utils.py | 29 +- vllm/config/vllm.py | 15 + vllm/kernels/triton/fused_pre_fc_norms.py | 73 ++ vllm/kernels/triton/fused_sigmoid_gate_mul.py | 86 ++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 34 +- vllm/model_executor/models/qwen3_next.py | 10 +- vllm/model_executor/models/registry.py | 15 +- vllm/models/qwen3_5/__init__.py | 48 + vllm/models/qwen3_5/nvidia/__init__.py | 2 + vllm/models/qwen3_5/nvidia/model.py | 939 ++++++++++++++++++ vllm/models/qwen3_5/nvidia/mtp.py | 295 ++++++ vllm/models/qwen3_5/nvidia/ops/__init__.py | 2 + .../nvidia/ops/fused_gdn_gated_norm.py | 123 +++ 29 files changed, 2029 insertions(+), 21 deletions(-) create mode 100644 tests/kernels/core/test_fused_pre_fc_norms.py create mode 100644 tests/kernels/core/test_fused_sigmoid_gate_mul.py create mode 100644 tests/models/test_qwen3_5.py create mode 100644 vllm/kernels/triton/fused_pre_fc_norms.py create mode 100644 vllm/kernels/triton/fused_sigmoid_gate_mul.py create mode 100644 vllm/models/qwen3_5/__init__.py create mode 100644 vllm/models/qwen3_5/nvidia/__init__.py create mode 100644 vllm/models/qwen3_5/nvidia/model.py create mode 100644 vllm/models/qwen3_5/nvidia/mtp.py create mode 100644 vllm/models/qwen3_5/nvidia/ops/__init__.py create mode 100644 vllm/models/qwen3_5/nvidia/ops/fused_gdn_gated_norm.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index c085a6adb706..d4d90eeaa6eb 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -79,6 +79,7 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/attention/backends/cpu_attn.py - vllm/config/speculative.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/llama_eagle3.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 993bee49f186..8857f8515110 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2154,7 +2154,9 @@ steps: - vllm/model_executor/model_loader/ - vllm/model_executor/models/extract_hidden_states.py - vllm/model_executor/models/llama.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/registry.py - vllm/transformers_utils/configs/extract_hidden_states.py @@ -3574,6 +3576,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen2.py @@ -5025,6 +5028,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen2.py @@ -5279,7 +5283,9 @@ steps: - vllm/v1/attention/ - vllm/model_executor/models/qwen3_dflash.py - vllm/model_executor/models/qwen3_dflash2.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/model_loader/ - vllm/model_executor/layers/quantization/ - vllm/model_executor/kernels/linear/ diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 18c69b08f690..a5629f2f8016 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -38,6 +38,7 @@ steps: source_file_dependencies: - csrc/ - tests/kernels/core + - vllm/kernels/triton/ - tests/kernels/test_concat_mla_q.py - tests/kernels/test_fused_qk_norm_rope_gate.py commands: @@ -60,6 +61,7 @@ steps: source_file_dependencies: - csrc/ - tests/kernels/core + - vllm/kernels/triton/ - tests/kernels/test_top_k_per_row.py - tests/kernels/test_concat_mla_q.py - tests/kernels/test_rocm_fp32_router_gemm.py @@ -555,6 +557,7 @@ steps: - tests/kernels/mamba - vllm/model_executor/layers/mamba/checkpoint.py - vllm/model_executor/layers/mamba/gdn + - vllm/models/qwen3_5/ - vllm/model_executor/layers/mamba/ops - vllm/v1/attention/backends/gdn_attn.py commands: diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index b55515d0beea..376b17abc803 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -454,6 +454,7 @@ steps: optional: true num_devices: 2 source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen3_next.py @@ -471,6 +472,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen2.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 6e5b6b55bb9a..17a4a6cebdf8 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -401,7 +401,9 @@ steps: - vllm/model_executor/model_loader/ - vllm/model_executor/models/extract_hidden_states.py - vllm/model_executor/models/llama.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/registry.py - vllm/transformers_utils/configs/extract_hidden_states.py diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 6135a46c6165..38bb4674bfc7 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -397,6 +397,7 @@ steps: - tests/models/test_vision.py - tests/models/test_adapters.py - tests/models/test_qwen3_5_mtp_config.py + - tests/models/test_qwen3_5.py - tests/models/test_dory.py - tests/models/test_alice_ai.py - tests/models/test_alice_ai_config.py @@ -419,6 +420,7 @@ steps: models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py + models/test_qwen3_5.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 887d079f7c0a..d58a92db45a1 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -186,6 +186,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen2.py diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 411e37aa2db6..e9c39537246d 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -450,7 +450,9 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - vllm/model_executor/models/qwen3_dflash.py - vllm/model_executor/models/qwen3_dflash2.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/laguna_dflash.py - tests/v1/e2e/spec_decode/ commands: @@ -494,7 +496,9 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - vllm/model_executor/models/qwen3_dflash.py - vllm/model_executor/models/qwen3_dflash2.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "dflash2" @@ -519,7 +523,9 @@ steps: - vllm/v1/attention/ - vllm/model_executor/models/qwen3_dflash.py - vllm/model_executor/models/qwen3_dflash2.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/model_loader/ - vllm/model_executor/layers/quantization/ - vllm/model_executor/kernels/linear/ diff --git a/tests/conftest.py b/tests/conftest.py index a979968c1621..9ae47a776db4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -251,12 +251,16 @@ def dist_init(): @pytest.fixture -def default_vllm_config(): +def default_vllm_config(monkeypatch): """Set a default VllmConfig for tests that directly test CustomOps or pathways that use get_current_vllm_config() outside of a full engine context. """ from vllm.config import VllmConfig, set_current_vllm_config + # Creating a breakable-by-default model config sets + # VLLM_USE_BREAKABLE_CUDAGRAPH process-wide; shield the default config + # from such leaks (it would flip custom_ops to "all"). + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "0") config = VllmConfig() with set_current_vllm_config(config): yield config diff --git a/tests/kernels/core/test_fused_pre_fc_norms.py b/tests/kernels/core/test_fused_pre_fc_norms.py new file mode 100644 index 000000000000..5cffb6d1bb73 --- /dev/null +++ b/tests/kernels/core/test_fused_pre_fc_norms.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the fused dual-GemmaRMSNorm + concat kernel used by the +Qwen3.5 MTP drafter. + +The kernel must match the eager 3-op reference: two Gemma-style RMSNorms +(scale by ``1 + weight``, fp32 accumulation) followed by a last-dim concat. +""" + +import pytest +import torch + +from vllm.kernels.triton.fused_pre_fc_norms import fused_pre_fc_norms +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +EPS = 1e-6 +DTYPES = [torch.bfloat16, torch.float16, torch.float32] +NUM_TOKENS = [1, 129] +HIDDEN_SIZES = [2048, 4096] +# (rtol, atol) per dtype; the kernel accumulates in fp32 like the reference, +# so the gap is output rounding and reduction order only. +TOLS = { + torch.bfloat16: (1e-2, 1e-3), + torch.float16: (1e-3, 1e-4), + torch.float32: (1e-5, 1e-5), +} + +requires_cuda = pytest.mark.skipif( + not current_platform.is_cuda(), + reason="fused pre-fc norms Triton kernel requires a CUDA device", +) + + +def _gemma_rmsnorm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor: + dtype = x.dtype + x = x.float() + var = x.pow(2).mean(dim=-1, keepdim=True) + return (x * torch.rsqrt(var + eps) * (w.float() + 1.0)).to(dtype) + + +@requires_cuda +@torch.inference_mode() +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) +@pytest.mark.parametrize("dtype", DTYPES) +def test_fused_pre_fc_norms_correctness( + num_tokens: int, hidden_size: int, dtype: torch.dtype +): + """Fused kernel == cat([norm_e(embeds), norm_h(hidden)]) for Gemma-style + (1 + w) norms with non-zero weights.""" + set_random_seed(0) + embeds = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda") + hidden = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda") + emb_w = torch.randn(hidden_size, dtype=dtype, device="cuda") * 0.2 + hid_w = torch.randn(hidden_size, dtype=dtype, device="cuda") * 0.2 + + fused = fused_pre_fc_norms(embeds, hidden, emb_w, hid_w, EPS) + ref = torch.cat( + [_gemma_rmsnorm(embeds, emb_w, EPS), _gemma_rmsnorm(hidden, hid_w, EPS)], + dim=-1, + ) + + assert fused.shape == (num_tokens, 2 * hidden_size) + rtol, atol = TOLS[dtype] + torch.testing.assert_close(fused.float(), ref.float(), rtol=rtol, atol=atol) diff --git a/tests/kernels/core/test_fused_sigmoid_gate_mul.py b/tests/kernels/core/test_fused_sigmoid_gate_mul.py new file mode 100644 index 000000000000..2c2a51e534ff --- /dev/null +++ b/tests/kernels/core/test_fused_sigmoid_gate_mul.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.kernels.triton.fused_sigmoid_gate_mul import fused_sigmoid_gate_mul +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +DTYPES = [torch.bfloat16, torch.float16, torch.float32] +# Odd token counts, non-power-of-2 and large hidden, plus an empty batch. +SHAPES = [(1, 512), (3, 1792), (17, 4096), (128, 8192), (1023, 2048), (0, 1024)] +requires_cuda = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="fused_sigmoid_gate_mul Triton kernel requires a CUDA/ROCm device", +) + + +_TOLS = { + torch.float32: (1e-6, 1e-5), + torch.float16: (2e-3, 2e-3), + torch.bfloat16: (1e-2, 1e-2), +} + + +@requires_cuda +@pytest.mark.parametrize("shape", SHAPES) +@pytest.mark.parametrize("dtype", DTYPES) +@torch.inference_mode() +def test_fused_sigmoid_gate_mul_correctness(shape: tuple[int, int], dtype: torch.dtype): + """The fused kernel must match ``x * torch.sigmoid(gate)`` and must not + write in place.""" + set_random_seed(13) + num_tokens, hidden = shape + x = torch.randn(num_tokens, hidden, dtype=dtype, device="cuda") + # Scale the gate to exercise both tails of the sigmoid. + gate = torch.randn(num_tokens, hidden, dtype=dtype, device="cuda") * 4 + + result = fused_sigmoid_gate_mul(x, gate) + reference = x * torch.sigmoid(gate) + + if result.numel() > 0: + assert result.data_ptr() != x.data_ptr() + atol, rtol = _TOLS[dtype] + torch.testing.assert_close(result, reference, atol=atol, rtol=rtol) + + +@requires_cuda +@torch.inference_mode() +def test_fused_sigmoid_gate_mul_offset_row_view(): + """A row-offset (contiguous but misaligned-pointer) view must compile and + match eager -- this exercises the unaligned pointer specialization.""" + set_random_seed(13) + base_x = torch.randn(65, 1023, dtype=torch.bfloat16, device="cuda") + base_g = torch.randn(65, 1023, dtype=torch.bfloat16, device="cuda") + x, gate = base_x[1:], base_g[1:] + + result = fused_sigmoid_gate_mul(x, gate) + reference = x * torch.sigmoid(gate) + + torch.testing.assert_close(result, reference, atol=1e-2, rtol=1e-2) + + +@torch.inference_mode() +def test_fused_sigmoid_gate_mul_fallback_noncontiguous(): + """Non-CUDA or non-contiguous inputs take the eager fallback (CPU-safe).""" + x = torch.randn(64, 256, dtype=torch.float32).t() + gate = torch.randn(64, 256, dtype=torch.float32).t() + assert not x.is_contiguous() + + result = fused_sigmoid_gate_mul(x, gate) + expected = x * torch.sigmoid(gate) + + torch.testing.assert_close(result, expected, atol=0.0, rtol=0.0) diff --git a/tests/kernels/fla/test_fused_gdn_post_conv.py b/tests/kernels/fla/test_fused_gdn_post_conv.py index f6d8500c5a9c..da37739b861f 100644 --- a/tests/kernels/fla/test_fused_gdn_post_conv.py +++ b/tests/kernels/fla/test_fused_gdn_post_conv.py @@ -402,3 +402,48 @@ def test_fused_gdn_decode_post_conv_mtp_head_ratios( ) torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2) + + +@torch.inference_mode() +def test_fused_post_conv_prep_strided_ba_bitwise() -> None: + """Packed b/a chunk views (row stride 2*HV) must match contiguous inputs + bitwise.""" + torch.manual_seed(0) + dtype = torch.bfloat16 + device = "cuda" + num_tokens = 33 + H, HV, K, V = 2, 4, 128, 128 + qkv_dim = 2 * H * K + HV * V + + conv_output = torch.randn(num_tokens, qkv_dim, dtype=dtype, device=device) + ba = torch.randn(num_tokens, 2 * HV, dtype=dtype, device=device) + b_view, a_view = ba.chunk(2, dim=-1) + A_log = torch.randn(HV, dtype=torch.float32, device=device) + dt_bias = torch.randn(HV, dtype=torch.float32, device=device) + + ref = fused_post_conv_prep( + conv_output=conv_output, + a=a_view.contiguous(), + b=b_view.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + num_k_heads=H, + head_k_dim=K, + head_v_dim=V, + apply_l2norm=True, + output_g_exp=False, + ) + got = fused_post_conv_prep( + conv_output=conv_output, + a=a_view, + b=b_view, + A_log=A_log, + dt_bias=dt_bias, + num_k_heads=H, + head_k_dim=K, + head_v_dim=V, + apply_l2norm=True, + output_g_exp=False, + ) + for r, g in zip(ref, got, strict=True): + torch.testing.assert_close(g, r, rtol=0, atol=0) diff --git a/tests/kernels/fla/test_fused_recurrent_packed_decode.py b/tests/kernels/fla/test_fused_recurrent_packed_decode.py index 5ea66ced2390..e43734a49f16 100644 --- a/tests/kernels/fla/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/fla/test_fused_recurrent_packed_decode.py @@ -186,3 +186,57 @@ def test_packed_decode_supports_large_batch_head_grid(): ) assert torch.count_nonzero(out).item() == 0 + + +@pytest.mark.parametrize("num_tokens", [1, 9]) +@torch.inference_mode() +def test_packed_decode_strided_ba_bitwise(num_tokens: int) -> None: + """Packed b/a chunk views (row stride 2*HV) must match contiguous inputs + bitwise.""" + torch.manual_seed(0) + dtype = torch.bfloat16 + device = torch.device(DEVICE) + H, HV, K, V = 2, 4, 128, 128 + qkv_dim = 2 * H * K + HV * V + + qkvz = torch.randn(num_tokens, qkv_dim + HV * V, dtype=dtype, device=device) + mixed_qkv = qkvz[:, :qkv_dim] + ba = torch.randn(num_tokens, 2 * HV, dtype=dtype, device=device) + b_view, a_view = ba.chunk(2, dim=-1) + A_log = torch.randn(HV, dtype=torch.float32, device=device) + dt_bias = torch.randn(HV, dtype=torch.float32, device=device) + state = torch.randn(num_tokens + 1, HV, V, K, dtype=torch.float32, device=device) + indices = torch.arange(1, num_tokens + 1, dtype=torch.int32, device=device) + + out_strided = torch.empty(num_tokens, 1, HV, V, dtype=dtype, device=device) + state_strided = state.clone() + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv, + a=a_view, + b=b_view, + A_log=A_log, + dt_bias=dt_bias, + scale=K**-0.5, + initial_state=state_strided, + out=out_strided, + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=True, + ) + + out_contig = torch.empty(num_tokens, 1, HV, V, dtype=dtype, device=device) + state_contig = state.clone() + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv.contiguous(), + a=a_view.contiguous(), + b=b_view.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + scale=K**-0.5, + initial_state=state_contig, + out=out_contig, + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=True, + ) + + torch.testing.assert_close(out_strided, out_contig, rtol=0, atol=0) + torch.testing.assert_close(state_strided, state_contig, rtol=0, atol=0) diff --git a/tests/kernels/mamba/test_gdn_output_projection.py b/tests/kernels/mamba/test_gdn_output_projection.py index f9abfb4ac666..2b59b5062241 100644 --- a/tests/kernels/mamba/test_gdn_output_projection.py +++ b/tests/kernels/mamba/test_gdn_output_projection.py @@ -18,6 +18,7 @@ from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( QwenGatedDeltaNetAttention, ) +from vllm.platforms import current_platform @pytest.mark.parametrize("dtype", [torch.float32]) @@ -51,3 +52,63 @@ def test_output_projection_norms_per_head_and_flattens( assert z.shape == (num_tokens, num_heads, head_dim) assert out.shape == (num_tokens, hidden) assert out.dtype == dtype + + +@pytest.mark.parametrize("num_tokens", [1, 37]) +@pytest.mark.parametrize("norm_before_gate", [True, False]) +@pytest.mark.parametrize("activation", ["silu", "sigmoid"]) +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="the fused gated RMS norm Triton kernel requires CUDA", +) +@torch.inference_mode() +def test_fused_gated_rms_norm_strided_z_bitwise( + num_tokens: int, norm_before_gate: bool, activation: str +) -> None: + """The Qwen3.5 fused gated RMS norm reads z from its packed strided view + of the qkvz projection and matches ``rmsnorm_fn`` on the copied/reshaped + z bitwise.""" + from vllm.models.qwen3_5.nvidia.ops.fused_gdn_gated_norm import ( + fused_gdn_gated_rms_norm, + ) + from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( + rmsnorm_fn, + ) + + torch.manual_seed(0) + dtype = torch.bfloat16 + num_heads, head_dim = 4, 128 + qkvz_row = 2 * 2 * 128 + 2 * num_heads * head_dim + qkvz = torch.randn(num_tokens, qkvz_row, dtype=dtype, device="cuda") + z_view = qkvz[:, qkvz_row - num_heads * head_dim :] + x = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype, device="cuda") + weight = torch.randn(head_dim, dtype=torch.float32, device="cuda") + eps = 1e-6 + + ref = rmsnorm_fn( + x, + weight, + None, + z=z_view.reshape(num_tokens, num_heads, head_dim), + eps=eps, + group_size=None, + norm_before_gate=norm_before_gate, + activation=activation, + ) + + out = torch.empty(num_tokens, num_heads * head_dim, dtype=dtype, device="cuda") + fused_gdn_gated_rms_norm( + x=x.view(num_tokens, -1), + z=z_view, + weight=weight, + out=out, + eps=eps, + num_v_heads=num_heads, + head_v_dim=head_dim, + norm_before_gate=norm_before_gate, + activation=activation, + ) + + torch.testing.assert_close( + out.view(num_tokens, num_heads, head_dim), ref, rtol=0, atol=0 + ) diff --git a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py index a33f11b61aa4..f946b5c71450 100644 --- a/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py +++ b/tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py @@ -90,7 +90,7 @@ def test_runai_model_loader_download_files(vllm_runner): @create_new_process_for_each_test("spawn") def test_runai_mtp_real_checkpoint(vllm_runner, tmp_path, monkeypatch): """Skip main-only shards without changing a real MTP model's generation.""" - from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP + from vllm.models.qwen3_5 import Qwen3_5MTP from vllm.platforms import current_platform from vllm.transformers_utils.repo_utils import hf_api from vllm.v1.metrics.reader import Counter diff --git a/tests/models/test_qwen3_5.py b/tests/models/test_qwen3_5.py new file mode 100644 index 000000000000..333099635ae5 --- /dev/null +++ b/tests/models/test_qwen3_5.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +def test_qwen3_5_models_use_platform_specific_implementation(): + from vllm.models.qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5MTP + from vllm.platforms import current_platform + + # Only the NVIDIA variant drops torch.compile; ROCm and other platforms + # keep the previous torch.compile implementation in + # vllm/model_executor/models/qwen3_5{,_mtp}.py. + if current_platform.is_cuda(): + model_module = "vllm.models.qwen3_5.nvidia.model" + mtp_module = "vllm.models.qwen3_5.nvidia.mtp" + else: + model_module = "vllm.model_executor.models.qwen3_5" + mtp_module = "vllm.model_executor.models.qwen3_5_mtp" + assert Qwen3_5ForCausalLM.__module__ == model_module + assert Qwen3_5MTP.__module__ == mtp_module + + +def test_qwen3_5_nvidia_models_do_not_use_torch_compile_wrapper(): + # Import the NVIDIA modules directly so the assertion is independent of + # the platform this test runs on. + from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + from vllm.models.qwen3_5.nvidia import model as model_module + from vllm.models.qwen3_5.nvidia import mtp as mtp_module + + assert not issubclass(model_module.Qwen3_5Model, TorchCompileWithNoGuardsWrapper) + assert not issubclass( + mtp_module.Qwen3_5MultiTokenPredictor, TorchCompileWithNoGuardsWrapper + ) + assert not issubclass(mtp_module.Qwen3_5MTP, TorchCompileWithNoGuardsWrapper) + + +def test_qwen3_5_legacy_models_keep_torch_compile_wrapper(): + # The legacy implementation in vllm/model_executor/models/ serves ROCm and + # other non-NVIDIA platforms; it must keep the torch.compile wrapper. + from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + from vllm.model_executor.models import qwen3_5 as model_module + from vllm.model_executor.models import qwen3_5_mtp as mtp_module + + assert issubclass(model_module.Qwen3_5Model, TorchCompileWithNoGuardsWrapper) + assert issubclass( + mtp_module.Qwen3_5MultiTokenPredictor, TorchCompileWithNoGuardsWrapper + ) + assert issubclass(mtp_module.Qwen3_5MTP, TorchCompileWithNoGuardsWrapper) diff --git a/tests/v1/attention/utils.py b/tests/v1/attention/utils.py index 0b85e4010050..14265e490abf 100644 --- a/tests/v1/attention/utils.py +++ b/tests/v1/attention/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for attention-related v1 tests.""" +import os from dataclasses import dataclass import pytest @@ -232,15 +233,25 @@ def create_vllm_config( if hf_config_override: model_config.hf_config.update(hf_config_override) - return VllmConfig( - model_config=model_config, - cache_config=cache_config, - parallel_config=parallel_config, - scheduler_config=scheduler_config, - device_config=device_config, - load_config=load_config, - compilation_config=compilation_config, - ) + # VllmConfig() auto-enables breakable cudagraphs process-wide via + # VLLM_USE_BREAKABLE_CUDAGRAPH for breakable-by-default models; keep + # that from leaking into unrelated tests sharing this process. + breakable_env = os.environ.get("VLLM_USE_BREAKABLE_CUDAGRAPH") + try: + return VllmConfig( + model_config=model_config, + cache_config=cache_config, + parallel_config=parallel_config, + scheduler_config=scheduler_config, + device_config=device_config, + load_config=load_config, + compilation_config=compilation_config, + ) + finally: + if breakable_env is None: + os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None) + else: + os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = breakable_env @dataclass diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b2f33c272da0..da0a7008df5c 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -73,6 +73,17 @@ logger = init_logger(__name__) +_QWEN3_5_ARCHITECTURES = frozenset( + { + "Qwen3_5ForCausalLM", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MTP", + "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5MoeMTP", + } +) + DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( { "AliceAIForCausalLM", @@ -126,6 +137,10 @@ def default_breakable_cudagraph_architectures() -> frozenset[str]: # graphs are on. Enable this architecture so the published AMD # recipe can start. return frozenset({"DeepseekV41ForCausalLM"}) + if current_platform.is_cuda(): + # Qwen3.5 migrates to breakable CUDA graphs on NVIDIA only; other + # platforms keep the torch.compile path. + return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES | _QWEN3_5_ARCHITECTURES return DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES diff --git a/vllm/kernels/triton/fused_pre_fc_norms.py b/vllm/kernels/triton/fused_pre_fc_norms.py new file mode 100644 index 000000000000..453b06f6a45e --- /dev/null +++ b/vllm/kernels/triton/fused_pre_fc_norms.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _gemma_rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr): + x = x.to(tl.float32) + mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE + rrms = tl.rsqrt(mean_sq + eps) + w = w.to(tl.float32) + 1.0 + return (x * rrms) * w + + +@triton.jit +def _fused_pre_fc_norms_kernel( + embeds_ptr, + embeds_stride, + hidden_ptr, + hidden_stride, + emb_w_ptr, + hid_w_ptr, + eps, + out_ptr, + out_stride, + H: tl.constexpr, + BLOCK: tl.constexpr, +): + tok = tl.program_id(0) + off = tl.arange(0, BLOCK) + mask = off < H + + e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0) + ew = tl.load(emb_w_ptr + off, mask=mask) + e_normed = _gemma_rms_norm(e, ew, eps, H) + tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask) + + h = tl.load(hidden_ptr + tok * hidden_stride + off, mask=mask, other=0.0) + hw = tl.load(hid_w_ptr + off, mask=mask) + h_normed = _gemma_rms_norm(h, hw, eps, H) + tl.store(out_ptr + tok * out_stride + H + off, h_normed, mask=mask) + + +def fused_pre_fc_norms( + inputs_embeds: torch.Tensor, + hidden_states: torch.Tensor, + emb_weight: torch.Tensor, + hid_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Returns cat([norm_e(inputs_embeds), norm_h(hidden_states)]) -> [N, 2H]. + + Both norms are Gemma-style RMSNorm (scale by ``1 + weight``, fp32 + accumulation, multiply in fp32 before the cast back to the input dtype). + """ + n, h = inputs_embeds.shape + out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device) + _fused_pre_fc_norms_kernel[(n,)]( + inputs_embeds, + inputs_embeds.stride(0), + hidden_states, + hidden_states.stride(0), + emb_weight, + hid_weight, + eps, + out, + out.stride(0), + h, + triton.next_power_of_2(h), + ) + return out diff --git a/vllm/kernels/triton/fused_sigmoid_gate_mul.py b/vllm/kernels/triton/fused_sigmoid_gate_mul.py new file mode 100644 index 000000000000..233452ccc0a5 --- /dev/null +++ b/vllm/kernels/triton/fused_sigmoid_gate_mul.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused sigmoid-gate-multiply Triton kernel: ``out = x * sigmoid(gate)``.""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_sigmoid_gate_mul_kernel( + x_ptr, + gate_ptr, + out_ptr, + hidden, + x_stride_t, + gate_stride_t, + out_stride_t, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = offs < hidden + + x = tl.load(x_ptr + row * x_stride_t + offs, mask=mask, other=0.0) + g = tl.load(gate_ptr + row * gate_stride_t + offs, mask=mask, other=0.0) + out = x.to(tl.float32) * tl.sigmoid(g.to(tl.float32)) + tl.store( + out_ptr + row * out_stride_t + offs, + out.to(out_ptr.dtype.element_ty), + mask=mask, + ) + + +def fused_sigmoid_gate_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """Compute ``x * sigmoid(gate)`` in a single Triton kernel. + + Falls back to the eager ``x * torch.sigmoid(gate)`` sequence unless both + inputs are same-shape, same-dtype, contiguous 2D CUDA tensors. + + Args: + x: (num_tokens, hidden) tensor, e.g. attention output. + gate: (num_tokens, hidden) pre-sigmoid gate. + + Returns: + A fresh (num_tokens, hidden) tensor with x's dtype and device. + + """ + if ( + x.ndim != 2 + or x.shape != gate.shape + or x.dtype != gate.dtype + or not x.is_contiguous() + or not gate.is_contiguous() + or not x.is_cuda + ): + return x * torch.sigmoid(gate) + + num_tokens, hidden = x.shape + out = torch.empty_like(x) + if num_tokens == 0 or hidden == 0: + return out + + if hidden <= 1024: + block = triton.next_power_of_2(hidden) + num_warps = 4 + elif hidden <= 4096: + block = 1024 + num_warps = 4 + else: + block = 2048 + num_warps = 8 + + grid = (num_tokens, triton.cdiv(hidden, block)) + _fused_sigmoid_gate_mul_kernel[grid]( + x, + gate, + out, + hidden, + x.stride(0), + gate.stride(0), + out.stride(0), + BLOCK=block, + num_warps=num_warps, + ) + return out diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index b712268a9460..3c6c15b961a6 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -703,6 +703,10 @@ def __init__( self.enable_packed_recurrent_decode = ( envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE ) + # Off by default; the Qwen3.5 NVIDIA variant installs a fused gated + # RMS norm that reads z as a strided view of the packed qkvz + # projection, removing the z reshape copy (see Qwen3_5DecoderLayer). + self._fused_gdn_gated_rms_norm: Callable[..., None] | None = None self.gdn_decode_kernel = envs.VLLM_GDN_DECODE_KERNEL.strip().lower() if self.gdn_decode_kernel == "cuda" and current_platform.is_cuda_alike(): reason = self._fused_gdn_decode_unsupported_reason(vllm_config) @@ -1181,7 +1185,8 @@ def forward_cuda( qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size z_size = self.value_dim // self.tp_size mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1) - z = z.reshape(z.size(0), -1, self.head_v_dim) + if self._fused_gdn_gated_rms_norm is None: + z = z.reshape(z.size(0), -1, self.head_v_dim) b, a = self.split_ba(ba) # ============================================================ @@ -1195,14 +1200,37 @@ def forward_cuda( device=hidden_states.device, ) + # b/a may stay strided chunk views of ba: the packed decode, prefill + # (fused_post_conv_prep) and ROCm fallback kernels consume row strides + # directly, and fused_sigmoid_gating_delta_rule_update contiguous() + # internally where needed. torch.ops.vllm.qwen_gdn_attention_core( mixed_qkv, - b.contiguous(), - a.contiguous(), + b, + a, core_attn_out, layer_name=_encode_layer_name(self.prefix), ) + if ( + self._fused_gdn_gated_rms_norm is not None + and not self.gqa_interleaved_layout + ): + out2d = core_attn_out.view(num_tokens, -1) + self._fused_gdn_gated_rms_norm( + x=out2d, + z=z, + weight=self.norm.weight, + out=out2d, + eps=self.norm.eps, + num_v_heads=self.num_v_heads // self.tp_size, + head_v_dim=self.head_v_dim, + norm_before_gate=self.norm.norm_before_gate, + activation=self.norm.activation, + ) + output, _ = self.out_proj(out2d) + return output + # ============================================================ # Part 3: Output Projection # ============================================================ diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index f6df90594a69..ca0e7cddacf7 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -18,6 +18,7 @@ tensor_model_parallel_all_gather, tensor_model_parallel_reduce_scatter, ) +from vllm.kernels.triton.fused_sigmoid_gate_mul import fused_sigmoid_gate_mul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, @@ -393,6 +394,10 @@ def __init__( and (text_only or supports_mrope) ) + # Off by default; Qwen3.5 enables the fused sigmoid-gate-mul Triton + # kernel post-construction (see Qwen3_5DecoderLayer). + self.use_fused_output_gate = False + def _project_qkv_gate( self, qkv: torch.Tensor, @@ -464,7 +469,10 @@ def forward( q, k, v, gate = self._project_qkv_gate(qkv, positions) attn_output = self.attn(q, k, v) if gate is not None: - attn_output = attn_output * torch.sigmoid(gate) + if self.use_fused_output_gate: + attn_output = fused_sigmoid_gate_mul(attn_output, gate) + else: + attn_output = attn_output * torch.sigmoid(gate) output, _ = self.o_proj(attn_output) return output diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index d4af5ef52751..0092feacd9ba 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -201,8 +201,8 @@ "Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"), "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"), "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"), - "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"), - "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"), + "Qwen3_5ForCausalLM": ("vllm.models.qwen3_5", "Qwen3_5ForCausalLM"), + "Qwen3_5MoeForCausalLM": ("vllm.models.qwen3_5", "Qwen3_5MoeForCausalLM"), "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"), "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"), "Step1ForCausalLM": ("step1", "Step1ForCausalLM"), @@ -613,9 +613,12 @@ "qwen3_vl_moe", "Qwen3VLMoeForConditionalGeneration", ), - "Qwen3_5ForConditionalGeneration": ("qwen3_5", "Qwen3_5ForConditionalGeneration"), + "Qwen3_5ForConditionalGeneration": ( + "vllm.models.qwen3_5", + "Qwen3_5ForConditionalGeneration", + ), "Qwen3_5MoeForConditionalGeneration": ( - "qwen3_5", + "vllm.models.qwen3_5", "Qwen3_5MoeForConditionalGeneration", ), "Qwen4ExpForConditionalGeneration": ( @@ -739,8 +742,8 @@ "Step5MTP": ("vllm.models.step5", "Step5MTP"), "Qwen4ExpMTP": ("vllm.models.qwen4_exp", "Qwen4ExpMTP"), "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"), - "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"), - "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"), + "Qwen3_5MTP": ("vllm.models.qwen3_5", "Qwen3_5MTP"), + "Qwen3_5MoeMTP": ("vllm.models.qwen3_5", "Qwen3_5MoeMTP"), "InternS2MobiusMTP": ("interns2_mobius", "InternS2MobiusMTP"), "HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"), "HYV4MTPModel": ("vllm.models.hy_v4", "HYV4MTP"), diff --git a/vllm/models/qwen3_5/__init__.py b/vllm/models/qwen3_5/__init__.py new file mode 100644 index 000000000000..b6def8434743 --- /dev/null +++ b/vllm/models/qwen3_5/__init__.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen3.5 model — hardware-isolated entry point. + +Only the NVIDIA variant migrates to the new eager + breakable CUDA-graph +modelling code with explicit fused kernels. ROCm and all other platforms keep +the previous torch.compile implementation in ``vllm/model_executor/models/`` +so their behavior is unchanged. +""" + +from typing import TYPE_CHECKING + +from vllm.platforms import current_platform + +if TYPE_CHECKING or current_platform.is_cuda(): + from .nvidia.model import ( + Qwen3_5ForCausalLM, + Qwen3_5ForCausalLMBase, + Qwen3_5ForConditionalGeneration, + Qwen3_5Model, + Qwen3_5MoeForCausalLM, + Qwen3_5MoeForConditionalGeneration, + ) + from .nvidia.mtp import Qwen3_5MoeMTP, Qwen3_5MTP +else: + from vllm.model_executor.models.qwen3_5 import ( # type: ignore[assignment] + Qwen3_5ForCausalLM, + Qwen3_5ForCausalLMBase, + Qwen3_5ForConditionalGeneration, + Qwen3_5Model, + Qwen3_5MoeForCausalLM, + Qwen3_5MoeForConditionalGeneration, + ) + from vllm.model_executor.models.qwen3_5_mtp import ( # type: ignore[assignment] + Qwen3_5MoeMTP, + Qwen3_5MTP, + ) + +__all__ = [ + "Qwen3_5ForCausalLM", + "Qwen3_5ForCausalLMBase", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5Model", + "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5MoeMTP", + "Qwen3_5MTP", +] diff --git a/vllm/models/qwen3_5/nvidia/__init__.py b/vllm/models/qwen3_5/nvidia/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/qwen3_5/nvidia/__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/qwen3_5/nvidia/model.py b/vllm/models/qwen3_5/nvidia/model.py new file mode 100644 index 000000000000..b5e4df663416 --- /dev/null +++ b/vllm/models/qwen3_5/nvidia/model.py @@ -0,0 +1,939 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2025 The vLLM team. +# Copyright 2025 The Qwen Team. +# Copyright 2025 The HuggingFace Inc. team. +# All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only Qwen3.5 Series compatible with HuggingFace weights.""" + +from collections.abc import Iterable +from itertools import islice +from typing import ClassVar, cast + +import torch +from torch import nn +from transformers import ( + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, +) + +from vllm.config import MultiModalConfig, VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + flashinfer_trtllm_fused_allreduce_norm, + fused_allreduce_gemma_rms_norm, +) +from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + HasInnerState, + IsHybrid, + MixtureOfExperts, + MultiModalEmbeddings, + SupportsEagle3, + SupportsLoRA, + SupportsMRoPE, + SupportsPP, + _require_is_multimodal, +) +from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP +from vllm.model_executor.models.qwen3_5 import ( + Qwen3_5MoeProcessingInfo, + Qwen3_5ProcessingInfo, +) +from vllm.model_executor.models.qwen3_next import ( + Qwen3NextAttention, + Qwen3NextDecoderLayer, + Qwen3NextModel, + Qwen3NextSparseMoeBlock, + QwenNextMixtureOfExperts, +) +from vllm.model_executor.models.qwen3_vl import ( + Qwen3_VisionTransformer, + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + mrope_positions_factor, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + _merge_multimodal_embeddings, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, + sequence_parallel_chunk, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.sequence import IntermediateTensors +from vllm.tokenizers.registry import cached_tokenizer_from_config + +from .ops.fused_gdn_gated_norm import fused_gdn_gated_rms_norm + +logger = init_logger(__name__) + + +class Qwen3_5RMSNorm(nn.Module): + """Gemma-style RMS normalization backed by flashinfer kernels. + + When ``residual`` is given, the fused add + norm runs in place and the + updated ``(x, residual)`` pair is returned. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + set_weight_attrs(self.weight, {"dummy_weight_value": 0.0}) + self.variance_epsilon = eps + + def forward( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + from flashinfer.norm import gemma_fused_add_rmsnorm, gemma_rmsnorm + + if residual is None: + return gemma_rmsnorm(x, self.weight, self.variance_epsilon) + + # gemma_fused_add_rmsnorm mutates x and residual in place. + gemma_fused_add_rmsnorm(x, residual, self.weight, self.variance_epsilon) + return x, residual + + +class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer): + def __init__( + self, + vllm_config: VllmConfig, + layer_type: str, + prefix: str = "", + ) -> None: + super(Qwen3NextDecoderLayer, self).__init__() + + config = vllm_config.model_config.hf_text_config + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + parallel_config = vllm_config.parallel_config + quant_config = vllm_config.quant_config + + self.layer_type = layer_type + self.layer_idx = extract_layer_index(prefix) + is_moe_layer = config.model_type == "qwen3_5_moe_text" + self.use_attn_reduce_scatter_for_moe = ( + parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size == 1 + and is_moe_layer + ) + # Defer row-parallel all-reduces to the following norm + # (deepseek_v32 pattern); exclude PP and SP. + self.use_fused_allreduce_norm = ( + get_tensor_model_parallel_world_size() > 1 + and parallel_config.pipeline_parallel_size == 1 + and not self.use_attn_reduce_scatter_for_moe + and flashinfer_trtllm_fused_allreduce_norm is not None + ) + # Whether this layer's MLP/MoE leaves an un-reduced partial, i.e. + # whether the *next* norm (or the final norm) must do the all-reduce. + self.mlp_output_is_partial = self.use_fused_allreduce_norm + reduce_results = not ( + self.use_attn_reduce_scatter_for_moe or self.use_fused_allreduce_norm + ) + + if self.layer_type == "linear_attention": + self.linear_attn = QwenGatedDeltaNetAttention( + config=config, + vllm_config=vllm_config, + prefix=f"{prefix}.linear_attn", + gqa_interleaved_layout=False, + reduce_results=reduce_results, + ) + # Recover the inductor gated-norm fusion with a Triton kernel that + # also reads z strided from the packed qkvz projection. + self.linear_attn._fused_gdn_gated_rms_norm = fused_gdn_gated_rms_norm + elif self.layer_type == "full_attention": + self.self_attn = Qwen3NextAttention( + config, + model_config=model_config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + reduce_results=reduce_results, + mrope_positions_factor=mrope_positions_factor(vllm_config), + ) + self.self_attn.use_fused_output_gate = True + else: + raise ValueError(f"Invalid layer_type {self.layer_type}") + + if config.model_type == "qwen3_5_moe_text": + self.mlp = Qwen3NextSparseMoeBlock( + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + reduce_results=not self.use_fused_allreduce_norm, + ) + if self.use_fused_allreduce_norm: + # reduce_results=False is only honored on the plain-TP late-AR + # path (skip_final_all_reduce); EP/all2all, SP, zero experts, + # and replicated shared experts all produce or add a complete + # output, which a downstream fused norm would all-reduce again. + # NOTE: moe_config is shared across layers; the condition is + # layer-uniform, so this reset is idempotent. + moe_config = self.mlp.experts.moe_config + moe_partial_ok = ( + moe_config.skip_final_all_reduce + and moe_config.tp_size > 1 + and moe_config.ep_size == 1 + and moe_config.dp_size == 1 + and not self.mlp.replicate_shared_expert + ) + if moe_config.skip_final_all_reduce and not moe_partial_ok: + moe_config.skip_final_all_reduce = False + self.mlp_output_is_partial = moe_partial_ok + elif config.model_type == "qwen3_5_text": + self.mlp = Qwen3NextMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + reduce_results=not self.use_fused_allreduce_norm, + ) + else: + raise ValueError(f"Invalid model_type {config.model_type}") + + self.input_layernorm = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.layer_scale = getattr(config, "layer_scale", False) + if self.layer_scale: + self.attn_layer_scale = torch.nn.Parameter( + torch.zeros( + 1, + 1, + config.hidden_size, + ), + ) + self.ffn_layer_scale = torch.nn.Parameter( + torch.zeros( + 1, + 1, + config.hidden_size, + ), + ) + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + positions: torch.Tensor, + **kwargs: object, + ): + if not self.use_fused_allreduce_norm: + return super().forward(hidden_states, residual, positions, **kwargs) + + # layer_scale is a per-channel multiply, which commutes with the + # deferred all-reduce. + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + elif self.mlp_output_is_partial: + hidden_states, residual = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.input_layernorm + ) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + if self.layer_type == "linear_attention": + hidden_states = self.linear_attn(hidden_states=hidden_states) + elif self.layer_type == "full_attention": + hidden_states = self.self_attn( + hidden_states=hidden_states, + positions=positions, + ) + else: + raise ValueError("Invalid layer_type") + + if self.layer_scale: + if len(hidden_states.shape) == 2: + hidden_states = hidden_states * ( + self.attn_layer_scale.to(hidden_states.dtype)[0] + 1 + ) + else: + hidden_states = hidden_states * ( + self.attn_layer_scale.to(hidden_states.dtype) + 1 + ) + + hidden_states, residual = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + hidden_states = self.mlp(hidden_states) + + if self.layer_scale: + if len(hidden_states.shape) == 2: + hidden_states = hidden_states * ( + self.ffn_layer_scale.to(hidden_states.dtype)[0] + 1 + ) + else: + assert len(hidden_states.shape) == len(self.ffn_layer_scale.shape), ( + f"shape must be the same {len(hidden_states.shape)}, " + f"{len(self.ffn_layer_scale.shape)}" + ) + hidden_states = hidden_states * ( + self.ffn_layer_scale.to(hidden_states.dtype) + 1 + ) + + return hidden_states, residual + + +class Qwen3_5Model(nn.Module, EagleModelMixin): + # Qwen3.5 ships the GDN in_proj checkpoints separately (qwen3-next + # pre-fuses them); fuse them on top of the qwen3-next QKV/gate_up mapping. + hf_to_vllm_mapper = Qwen3NextModel.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_stacked={ + ".in_proj_qkv": (".in_proj_qkvz", (0, 1, 2)), + ".in_proj_z": (".in_proj_qkvz", 3), + ".in_proj_b": (".in_proj_ba", 0), + ".in_proj_a": (".in_proj_ba", 1), + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = ( + vllm_config.model_config.hf_text_config + ) + parallel_config = vllm_config.parallel_config + + eplb_config = parallel_config.eplb_config + self.num_redundant_experts = eplb_config.num_redundant_experts + + self.config = config + self.quant_config = vllm_config.quant_config + + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + ) + + def get_layer(prefix: str): + return Qwen3_5DecoderLayer( + vllm_config, + layer_type=config.layer_types[extract_layer_index(prefix)], + prefix=prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" + ) + # Mirror of the decoder layers' flag (uniform across layers): when + # True, every layer returns its MLP/MoE output as an un-reduced TP + # partial, so the final norm and eagle aux captures must all-reduce. + self.mlp_output_is_partial = getattr( + self.layers[self.start_layer], "mlp_output_is_partial", False + ) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + if get_pp_group().is_last_rank: + self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + self.aux_hidden_state_layers: tuple[int, ...] = () + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + @property + def use_sequence_parallel(self) -> bool: + return self.layers[self.start_layer].use_attn_reduce_scatter_for_moe + + def _maybe_add_hidden_state( + self, + aux_hidden_states: list[torch.Tensor], + layer_idx: int, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> list[torch.Tensor]: + if residual is None and layer_idx in self.aux_hidden_state_layers: + # The layer-0 capture aliases the embeddings, which become the + # first residual and are mutated in place by the next layer's + # fused-add norm; capture by value. + hidden_states = hidden_states.clone() + if ( + residual is not None + and self.mlp_output_is_partial + and layer_idx in self.aux_hidden_state_layers + ): + # hidden_states is an un-reduced TP partial; the aux capture needs + # the complete hidden state. + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + return super()._maybe_add_hidden_state( + aux_hidden_states, layer_idx, hidden_states, residual + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + full_num_tokens = positions.shape[-1] + if self.use_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + assert residual is None + + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + for layer_idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, residual + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + if self.mlp_output_is_partial: + # The last layer's MLP/MoE output is an un-reduced TP partial; + # fuse its all-reduce into the final norm. + hidden_states, _ = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.norm + ) + else: + hidden_states, _ = self.norm(hidden_states, residual) + if self.use_sequence_parallel: + if aux_hidden_states: + hidden_size = hidden_states.shape[-1] + hidden_states = torch.cat([hidden_states, *aux_hidden_states], dim=-1) + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + hidden_states, *aux_hidden_states = hidden_states.split( + hidden_size, dim=-1 + ) + else: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + 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 Qwen3_5ForCausalLMBase( + nn.Module, + HasInnerState, + IsHybrid, + SupportsEagle3, + SupportsLoRA, + SupportsMRoPE, + SupportsPP, +): + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": ["gate_proj", "up_proj"], + # GDN fused projections. + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + } + # Maps PEFT embed/lm_head LoRA targets onto vLLM embedding wrappers. + embedding_modules = { + "embed_tokens": "input_embeddings", + "lm_head": "output_embeddings", + } + + # Some community text-only checkpoints keep the extraneous + # `model.language_model.` prefix inherited from the VL training stack. + # Strip it so both prefixed and clean checkpoints load correctly. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"model.language_model.": "model.", "mtp.": None}, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + config = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + + scheduler_config = vllm_config.scheduler_config + self.quant_config = vllm_config.quant_config + + super().__init__() + self.config = config + self.scheduler_config = scheduler_config + self.model = Qwen3_5Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + + if get_pp_group().is_last_rank: + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + self.model.aux_hidden_state_layers = layers + + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: + num_layers = len(self.model.layers) + return (2, num_layers // 2, num_layers - 3) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ): + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + return hidden_states + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + num_spec, + ) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def compute_logits_local( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True) + + 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) + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec], + ) -> tuple[torch.Tensor, int]: + positions = torch.arange(len(input_tokens), dtype=torch.long) + return positions.unsqueeze(0).expand(3, -1), 0 + + +class Qwen3_5ForCausalLM(Qwen3_5ForCausalLMBase): + pass + + +class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + + # set MoE hyperparameters + self.set_moe_parameters() + + +######################################################## +# Qwen3_5-Dense +######################################################## + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen3_5ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid): + supports_multimodal_pruning = True + + hf_to_vllm_mapper = ( + Qwen3VLForConditionalGeneration.hf_to_vllm_mapper + | WeightsMapper(orig_to_new_prefix={"mtp.": None}) + ) + + packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | { + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): + # protocols have not __init__ method, so we need to use nn.Module.__init__ + nn.Module.__init__(self) + config: Qwen3_5Config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = cast( + MultiModalConfig, vllm_config.model_config.multimodal_config + ) + + self.config = config + self.model_config = vllm_config.model_config + self.multimodal_config = multimodal_config + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + self.video_pruning_rate = self.multimodal_config.video_pruning_rate + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + # attributes needed by EVS-related functions inherited from Qwen3-VL + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = Qwen3_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + input_norm=build_mm_input_norm(self.model_config), + prefix=maybe_prefix(prefix, "visual"), + ) + + with self._mark_language_model(vllm_config): + self.language_model = Qwen3_5ForCausalLM( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model") + ) + + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.language_model.embed_input_ids, + is_multimodal=is_multimodal, + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + is_multimodal = _require_is_multimodal(is_multimodal) + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + return inputs_embeds + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + """Run forward pass for Qwen3.5. + + With mrope enabled (default for Qwen3VL-class models), ``positions`` + has shape `(3, seq_len)`, otherwise `(seq_len,)`. + """ + if intermediate_tensors is not None: + inputs_embeds = None + + hidden_states = self.language_model.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + return hidden_states + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.language_model.compute_logits_local(hidden_states) + + 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) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: "VllmConfig", + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: "VllmConfig" + ) -> tuple[tuple[int, int], tuple[int, int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + num_spec, + ) + + @classmethod + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() + + +######################################################## +# Qwen3_5-MoE +######################################################## + + +class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts): + language_model: Qwen3_5MoeForCausalLM + num_local_physical_experts: int + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for layer in self.language_model.model.layers: + if isinstance(layer.mlp, Qwen3NextSparseMoeBlock): + moe = layer.mlp + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + def set_moe_parameters(self): + self.moe_layers = [] + example_moe = None + for layer in self.language_model.model.layers: + if isinstance(layer, Qwen3_5DecoderLayer) and isinstance( + layer.mlp, Qwen3NextSparseMoeBlock + ): + example_moe = layer.mlp + self.moe_layers.append(layer.mlp.experts) + + if example_moe is None: + raise RuntimeError( + "No Qwen3_5 layer found in the language_model.model.layers." + ) + + # Set MoE hyperparameters + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = 1 + self.num_shared_experts = 0 + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen3_5MoeProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class Qwen3_5MoeForConditionalGeneration( + Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts +): + # For MoE LoRA weights loading + is_3d_moe_weight: ClassVar[bool] = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): + # protocols have not __init__ method, so we need to use nn.Module.__init__ + nn.Module.__init__(self) + config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = cast( + MultiModalConfig, vllm_config.model_config.multimodal_config + ) + + self.config = config + self.model_config = vllm_config.model_config + self.multimodal_config = multimodal_config + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + self.video_pruning_rate = self.multimodal_config.video_pruning_rate + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + # attributes needed by EVS-related functions inherited from Qwen3-VL + self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = Qwen3_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + input_norm=build_mm_input_norm(self.model_config), + prefix=maybe_prefix(prefix, "visual"), + ) + + with self._mark_language_model(vllm_config): + self.language_model = Qwen3_5MoeForCausalLM( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model") + ) + + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + + # set MoE hyperparameters + self.set_moe_parameters() diff --git a/vllm/models/qwen3_5/nvidia/mtp.py b/vllm/models/qwen3_5/nvidia/mtp.py new file mode 100644 index 000000000000..d0f11a7b82f5 --- /dev/null +++ b/vllm/models/qwen3_5/nvidia/mtp.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Qwen3_5 MTP model.""" + +from collections.abc import Iterable + +import torch +from torch import nn +from transformers import Qwen3_5MoeTextConfig, Qwen3_5TextConfig + +from vllm.config import VllmConfig +from vllm.distributed import tensor_model_parallel_all_gather +from vllm.distributed.parallel_state import get_pp_group +from vllm.kernels.triton.fused_pre_fc_norms import fused_pre_fc_norms +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import ( + fused_allreduce_gemma_rms_norm, +) +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import ( + LocalArgmaxMixin, + MultiModalEmbeddings, + SupportsMultiModal, + SupportsPP, + _require_is_multimodal, +) +from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + _merge_multimodal_embeddings, + make_empty_intermediate_tensors_factory, + maybe_prefix, + sequence_parallel_chunk, +) +from vllm.sequence import IntermediateTensors + +from .model import Qwen3_5DecoderLayer, Qwen3_5Model, Qwen3_5RMSNorm + +logger = init_logger(__name__) + + +class Qwen3_5MultiTokenPredictor(nn.Module): + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + model_config = vllm_config.model_config + quant_config = vllm_config.quant_config + + config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = model_config.hf_text_config + + self.config = config + + self.vocab_size = config.vocab_size + + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1) + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + ) + + # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is + # missing from hf_quant_config.json exclude_modules. Force unquantized. + # Ref: https://github.com/vllm-project/vllm/pull/38650 + # Ref: https://github.com/NVIDIA/Model-Optimizer/pull/1124 + fc_quant = ( + None + if (quant_config and quant_config.get_name() == "modelopt_fp4") + else quant_config + ) + self.fc = ColumnParallelLinear( + self.config.hidden_size * 2, + self.config.hidden_size, + gather_output=True, + bias=False, + return_bias=False, + quant_config=fc_quant, + prefix=f"{prefix}.fc", + ) + + # GPTQ: quantized checkpoints may exclude MTP from quantization via + # quantization_config.dynamic with "-:pattern" entries. When detected, + # disable quantization for MTP layers so they use unquantized params. + original_quant = vllm_config.quant_config + if quant_config and quant_config.get_name() not in ("modelopt_fp4",): + hf_qc = getattr(model_config.hf_config, "quantization_config", None) + if isinstance(hf_qc, dict): + dynamic = hf_qc.get("dynamic", {}) + if any(k.startswith("-:") and "mtp" in k for k in dynamic): + vllm_config.quant_config = None + self.layers = torch.nn.ModuleList( + Qwen3_5DecoderLayer( + vllm_config, + layer_type="full_attention", + prefix=f"{prefix}.layers.{idx}", + ) + for idx in range(self.num_mtp_layers) + ) + vllm_config.quant_config = original_quant + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.pre_fc_norm_hidden = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_fc_norm_embedding = Qwen3_5RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + # Branch on the inputs, not the rank: the drafter is built entirely on + # the last PP stage, where `is_first_rank` is False. + if intermediate_tensors is None: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + assert hidden_states.shape[-1] == inputs_embeds.shape[-1] + hidden_states = fused_pre_fc_norms( + inputs_embeds, + hidden_states, + self.pre_fc_norm_embedding.weight, + self.pre_fc_norm_hidden.weight, + self.pre_fc_norm_embedding.variance_epsilon, + ) + hidden_states = self.fc(hidden_states) + residual = None + else: + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[current_step_idx] + if mtp_layer.use_attn_reduce_scatter_for_moe: + assert hidden_states.shape[0] == positions.shape[-1] + hidden_states = sequence_parallel_chunk(hidden_states) + assert residual is None + hidden_states, residual = mtp_layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + if mtp_layer.mlp_output_is_partial: + # The MTP layer's MLP/MoE output is an un-reduced TP partial; + # fuse its all-reduce into the final norm. + hidden_states, _ = fused_allreduce_gemma_rms_norm( + hidden_states, residual, self.norm + ) + else: + hidden_states, _ = self.norm(hidden_states, residual) + if mtp_layer.use_attn_reduce_scatter_for_moe: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[: positions.shape[-1]] + return hidden_states + + 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 Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP): + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + config = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + self.quant_config = vllm_config.quant_config + + super().__init__() + self.config = config + self.model = Qwen3_5MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp") + ) + + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if config.tie_word_embeddings: + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.model.embed_input_ids, + is_multimodal=is_multimodal, + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + is_multimodal = _require_is_multimodal(is_multimodal) + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + return inputs_embeds + + def forward( # type: ignore[override] + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ): + hidden_states = self.model( + input_ids, positions, hidden_states, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def is_unused_checkpoint_weight(self, name: str) -> bool: + return not name.startswith("mtp.") and not any( + key in name for key in ["embed_tokens", "lm_head"] + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + def remap_weight_names(weights): + for name, weight in weights: + if self.is_unused_checkpoint_weight(name): + continue + if name.startswith("mtp."): + name = name.replace("mtp.", "model.") + elif "embed_tokens" in name: + name = name.replace("language_model.", "") + yield name, weight + + loader = AutoWeightsLoader(self) + return loader.load_weights(remap_weight_names(weights)) + + +class Qwen3_5MoeMTP(Qwen3_5MTP, QwenNextMixtureOfExperts): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + self.set_moe_parameters() diff --git a/vllm/models/qwen3_5/nvidia/ops/__init__.py b/vllm/models/qwen3_5/nvidia/ops/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/qwen3_5/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/qwen3_5/nvidia/ops/fused_gdn_gated_norm.py b/vllm/models/qwen3_5/nvidia/ops/fused_gdn_gated_norm.py new file mode 100644 index 000000000000..f7de88323f10 --- /dev/null +++ b/vllm/models/qwen3_5/nvidia/ops/fused_gdn_gated_norm.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( + calc_rows_per_block, +) +from vllm.triton_utils import tl, triton + + +@triton.jit(do_not_specialize=["M"]) +def _fused_gdn_gated_rms_norm_kernel( + x_ptr, + z_ptr, + w_ptr, + out_ptr, + stride_x_tok, + stride_z_tok, + stride_out_tok, + M, # number of tokens + eps, + N: tl.constexpr, # head_v_dim + BLOCK_N: tl.constexpr, + ROWS_PER_BLOCK: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + ACTIVATION: tl.constexpr, +): + """FLA-equivalent per-head gated RMS norm with token-strided z.""" + row_start = tl.program_id(0) * ROWS_PER_BLOCK + head = tl.program_id(1) + + rows = row_start + tl.arange(0, ROWS_PER_BLOCK) + cols = tl.arange(0, BLOCK_N) + row_mask = rows[:, None] < M + col_mask = cols[None, :] < N + mask = row_mask & col_mask + + x_ptrs = x_ptr + rows[:, None] * stride_x_tok + head * N + cols[None, :] + z_ptrs = z_ptr + rows[:, None] * stride_z_tok + head * N + cols[None, :] + + x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) + + if not NORM_BEFORE_GATE: + z = tl.load(z_ptrs, mask=mask, other=0.0).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + x *= z * tl.sigmoid(z) + elif ACTIVATION == "sigmoid": + x *= tl.sigmoid(z) + + xbar = tl.where(mask, x, 0.0) + var = tl.sum(xbar * xbar, axis=1) / N + rstd = tl.rsqrt(var + eps) + + w = tl.load(w_ptr + cols, mask=cols < N, other=0.0).to(tl.float32) + y = x * rstd[:, None] * w[None, :] + + if NORM_BEFORE_GATE: + z = tl.load(z_ptrs, mask=mask, other=0.0).to(tl.float32) + if ACTIVATION == "swish" or ACTIVATION == "silu": + y *= z * tl.sigmoid(z) + elif ACTIVATION == "sigmoid": + y *= tl.sigmoid(z) + + out_ptrs = out_ptr + rows[:, None] * stride_out_tok + head * N + cols[None, :] + tl.store(out_ptrs, y.to(out_ptr.dtype.element_ty), mask=mask) + + +def fused_gdn_gated_rms_norm( + x: torch.Tensor, + z: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + eps: float, + num_v_heads: int, + head_v_dim: int, + norm_before_gate: bool, + activation: str, +) -> None: + """Gated RMS norm ``out = rms_norm(x) * act(z)`` for Qwen3.5 GDN. + + Args: + x: (num_tokens, num_v_heads * head_v_dim), contiguous. May alias out. + z: (num_tokens, num_v_heads * head_v_dim) gate values, typically a + strided view into the packed in_proj_qkvz output. Only the last + dim must be contiguous. + weight: (head_v_dim,) shared per-head norm weight. + out: (num_tokens, num_v_heads * head_v_dim), contiguous. + eps: RMS norm epsilon. + num_v_heads: number of V heads per TP rank. + head_v_dim: head_v_dim (per-row norm size). + norm_before_gate: if True, out = norm(x) * act(z), else + out = norm(x * act(z)). + activation: "swish"/"silu" or "sigmoid". + + """ + assert x.shape == z.shape == out.shape and x.dim() == 2 + assert x.stride(-1) == 1 and z.stride(-1) == 1 and out.stride(-1) == 1 + assert weight.shape == (head_v_dim,) and weight.stride(0) == 1 + assert x.shape[1] == num_v_heads * head_v_dim + + num_tokens = x.shape[0] + if num_tokens == 0: + return + block_n = triton.next_power_of_2(head_v_dim) + rows_per_block = calc_rows_per_block(num_tokens, x.device) + grid = (triton.cdiv(num_tokens, rows_per_block), num_v_heads) + _fused_gdn_gated_rms_norm_kernel[grid]( + x, + z, + weight, + out, + x.stride(0), + z.stride(0), + out.stride(0), + num_tokens, + eps, + N=head_v_dim, + BLOCK_N=block_n, + ROWS_PER_BLOCK=rows_per_block, + NORM_BEFORE_GATE=norm_before_gate, + ACTIVATION=activation, + num_warps=min(max(block_n // 256, 1), 8), + ) From bcbcd02e53bed1a9455a6c7a15e4184fdbb27dba Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Thu, 29 Oct 2026 10:40:11 +0800 Subject: [PATCH 105/406] [ROCm][Perf] Fuse MXFP8 activation quant into the 1x32 dot_scaled linear for M <= 16 (#63915) Signed-off-by: RuibinCheung Co-authored-by: vllm-agent --- .../quantization/test_rocm_mxfp8_linear.py | 57 ++++++ .../kernels/linear/mxfp8/rocm_native.py | 189 ++++++++++++++---- 2 files changed, 206 insertions(+), 40 deletions(-) diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py index d5e48f2aafa2..8c9f4a6d0055 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -259,6 +259,63 @@ def test_rocm_mxfp8_block32_split_k_first_use_in_capture(packed): assert _rel_err(out, expected) < 5e-3 +FUSED_QUANT_SHAPES = [(6144, 1024), (6144, 768), (6144, 384), (8192, 512), (4096, 640)] +FUSED_QUANT_TOKENS = [1, 2, 3, 5, 8, 16] +FUSED_QUANT_CASES = [ + (n, k, m) for n, k in FUSED_QUANT_SHAPES for m in FUSED_QUANT_TOKENS +] + [(6144, 2048, m) for m in (1, 2, 3, 4)] + + +def _spy_launch(monkeypatch) -> list[bool]: + """Record, per ``_launch``, whether the activation quant ran in-kernel.""" + from vllm.model_executor.kernels.linear.mxfp8 import rocm_native + + fused: list[bool] = [] + real_launch = rocm_native._launch + + def spy(x, x_scale, *args, **kwargs): + fused.append(x_scale is None) + return real_launch(x, x_scale, *args, **kwargs) + + monkeypatch.setattr(rocm_native, "_launch", spy) + return fused + + +def _mxfp8_weight(n: int, k: int): + weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.05 + return _mxfp8_e4m3_quantize_torch(weight) + + +@pytest.mark.parametrize("n,k,num_tokens", FUSED_QUANT_CASES, ids=lambda v: str(v)) +@torch.inference_mode() +def test_rocm_mxfp8_fused_act_quant_bitwise_eq_unfused(monkeypatch, n, k, num_tokens): + """Small-M bf16 GEMMs quantize the activation in-kernel, bit-exact with + the unfused path: ``mxfp8_e4m3_quantize`` and the ``_select_cfg`` tile.""" + from vllm.model_executor.kernels.linear.mxfp8 import rocm_native + + torch.manual_seed(num_tokens) + w, w_scale = _mxfp8_weight(n, k) + x = torch.randn(num_tokens, k, device="cuda", dtype=torch.bfloat16) + x[0, :32] = 0.0 # an all-zero block clamps its scale byte to 0 + if num_tokens > 1: + x[1] *= 3e6 # output beyond the fp16 range, which bf16 must keep + + assert rocm_native._select_fused_quant_cfg(num_tokens, n, k) is not None + x_q, x_scale = mxfp8_e4m3_quantize(x) + expected = torch.empty((num_tokens, n), dtype=torch.bfloat16, device="cuda") + rocm_native._launch( + x_q, x_scale, w, w_scale, expected, *rocm_native._select_cfg(num_tokens, n, k) + ) + + fused = _spy_launch(monkeypatch) + out = rocm_native._mxfp8_dot_scaled_linear(x, w, w_scale) + assert fused == [True] + assert torch.isfinite(out).all() + assert torch.equal(out, expected), ( + f"max abs diff {(out.float() - expected.float()).abs().max().item()}" + ) + + @pytest.mark.parametrize("num_rows", [1, 64, 65]) @torch.inference_mode() def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows): diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 134567555cd7..93246d9160e3 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -3,12 +3,12 @@ """Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``. Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16); -activations are MXFP8-quantized per token, or arrive pre-quantized as a -``QuantizedActivation`` with row-major (not swizzled) E8M0 scales. Uses the -CDNA4 hardware microscaling 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. +activations are MXFP8-quantized per token (inside the GEMM for small M), or +arrive pre-quantized as a ``QuantizedActivation`` with row-major (not swizzled) +E8M0 scales. Uses the CDNA4 hardware microscaling 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 @@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, MXFP8_SCALE_DTYPE, + MXFP8_VALUE_DTYPE, dequant_mxfp8_to_bf16, mxfp8_e4m3_quantize, ) @@ -35,6 +36,25 @@ from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig from .rocm_block32_gemm import BLOCK_ROWS, rocm_mxfp8_block32_gemm +_FP8_MAX = float(torch.finfo(MXFP8_VALUE_DTYPE).max) +_TINY = float(torch.finfo(torch.float32).tiny) + + +@triton.jit +def _quant_mxfp8_tile(x, FP8_MAX: tl.constexpr, TINY: tl.constexpr): + """MXFP8-quantize a [BLOCK_M, BLOCK_K] tile: bit-exact with + ``mxfp8_e4m3_quantize`` (see ``_mxfp8_quant_triton_kernel`` for why the + scale is applied as a reciprocal).""" + BM: tl.constexpr = x.shape[0] + BK: tl.constexpr = x.shape[1] + xb = tl.reshape(x.to(tl.float32), (BM, BK // 32, 32)) + amax = tl.maximum(tl.max(tl.abs(xb), axis=2), TINY) + sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0 + sb = tl.minimum(tl.maximum(sb, 0.0), 254.0) + rescale = tl.exp2(127.0 - sb) + xq = tl.reshape((xb * rescale[:, :, None]).to(tl.float8e4nv), (BM, BK)) + return xq, sb.to(tl.uint8) + @triton.jit def _mxfp8_linear_kernel( @@ -45,7 +65,7 @@ def _mxfp8_linear_kernel( out_ptr, M, N, - K, + K: tl.constexpr, stride_xm, stride_xk, stride_xsm, @@ -59,31 +79,59 @@ def _mxfp8_linear_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + QUANT_X: tl.constexpr, + QUANT_ROWS: tl.constexpr, + FP8_MAX: tl.constexpr, + TINY: tl.constexpr, ): pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + # Rows of x that are loaded: all BLOCK_M, or with QUANT_X the first + # QUANT_ROWS (a power of 2 >= the rows this tile stores). + offs_xm = pid_m * BLOCK_M + tl.arange(0, QUANT_ROWS) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) offs_sk = tl.arange(0, BLOCK_K // 32) m_mask = offs_m < M + xm_mask = offs_xm < M n_mask = offs_n < N - x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk + x_ptrs = x_ptr + offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk xs_ptrs = xs_ptr + offs_m[:, None] * stride_xsm + offs_sk[None, :] * stride_xsk w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :] * stride_wsk acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for _ in range(0, tl.cdiv(K, BLOCK_K)): - x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0) - w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0) - xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0) - ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0) + if QUANT_X: + # x is the unquantized activation; every N-tile quantizes its own + # copy, which is cheaper than a separate quant launch at small M. + xf = tl.load(x_ptrs, mask=xm_mask[:, None], other=0.0) + w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0) + ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0) + x, xs = _quant_mxfp8_tile(xf, FP8_MAX, TINY) + if QUANT_ROWS < BLOCK_M: + # BLOCK_M is the smallest MFMA tile; repeating the quantized + # rows fills it, and rows >= M are never stored. + REPS: tl.constexpr = BLOCK_M // QUANT_ROWS + x = tl.reshape( + tl.broadcast_to(x[None, :, :], (REPS, QUANT_ROWS, BLOCK_K)), + (BLOCK_M, BLOCK_K), + ) + xs = tl.reshape( + tl.broadcast_to(xs[None, :, :], (REPS, QUANT_ROWS, BLOCK_K // 32)), + (BLOCK_M, BLOCK_K // 32), + ) + else: + x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0) + w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0) + xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0) + ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0) + xs_ptrs += (BLOCK_K // 32) * stride_xsk acc += tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3") x_ptrs += BLOCK_K * stride_xk w_ptrs += BLOCK_K * stride_wk - xs_ptrs += (BLOCK_K // 32) * stride_xsk ws_ptrs += (BLOCK_K // 32) * stride_wsk o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on @@ -92,11 +140,98 @@ def _mxfp8_linear_kernel( ) +def _launch( + x: torch.Tensor, + x_scale: torch.Tensor | None, + w: torch.Tensor, + w_scale: torch.Tensor, + out: torch.Tensor, + block_m: int, + block_n: int, + block_k: int, + num_warps: int, + num_stages: int, +) -> None: + """``out = x @ w.T``. ``x_scale`` None means ``x`` is the unquantized + activation and is MXFP8-quantized inside the kernel.""" + M, K = x.shape + N = w.shape[0] + quant_x = x_scale is None + xs = x if x_scale is None else x_scale # Unused by the kernel if quant_x. + # Quantizing only the rows that can be stored keeps the fused quant's cost + # (serial per N-tile) from scaling with BLOCK_M at M < BLOCK_M. At + # BLOCK_K=128 that cost is small and repeating the rows costs more. + if quant_x and block_k >= 256: + quant_rows = min(block_m, triton.next_power_of_2(M)) + else: + quant_rows = block_m + grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n)) + _mxfp8_linear_kernel[grid]( + x, + xs, + w, + w_scale, + out, + M, + N, + K, + x.stride(0), + x.stride(1), + xs.stride(0), + xs.stride(1), + w.stride(0), + w.stride(1), + w_scale.stride(0), + w_scale.stride(1), + out.stride(0), + out.stride(1), + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + QUANT_X=quant_x, + QUANT_ROWS=quant_rows, + FP8_MAX=_FP8_MAX, + TINY=_TINY, + num_warps=num_warps, + num_stages=num_stages, + ) + + +# Past this many rows, quantizing the activation once per N-tile inside the +# GEMM costs more than one separate quant launch. +_FUSED_QUANT_MAX_M = 16 + + +def _select_fused_quant_cfg(M: int, N: int, K: int): + """Tile config with the activation quant fused (M <= 16), or None to + quantize separately. + """ + if K == 2048 and M <= 4: + return 16, 32, 512, 4, 3 + if K > 1024: + return None + if K == 1024: + return 16, 32, 512, 4, 3 + return 16, 32, 256 if K % 256 == 0 else 128, 4, 3 + + def _mxfp8_dot_scaled_linear( x: torch.Tensor, # [M, K] bf16/fp16 w: torch.Tensor, # [N, K] fp8 e4m3 w_scale: torch.Tensor, # [N, K//32] uint8 (E8M0) ) -> torch.Tensor: + M, K = x.shape + N = w.shape[0] + if ( + 0 < M <= _FUSED_QUANT_MAX_M + and x.dtype == torch.bfloat16 + and not envs.VLLM_BATCH_INVARIANT + ): + cfg = _select_fused_quant_cfg(M, N, K) + if cfg is not None: + out = torch.empty((M, N), dtype=x.dtype, device=x.device) + _launch(x, None, w, w_scale, out, *cfg) + return out x_q, x_scale = mxfp8_e4m3_quantize(x) return _mxfp8_dot_scaled_gemm(x_q, x_scale, w, w_scale, x.dtype) @@ -111,33 +246,7 @@ def _mxfp8_dot_scaled_gemm( M, K = x_q.shape N = w.shape[0] out = torch.empty((M, N), dtype=out_dtype, device=x_q.device) - BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _select_cfg(M, N, K) - grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) - _mxfp8_linear_kernel[grid]( - x_q, - x_scale, - w, - w_scale, - out, - M, - N, - K, - x_q.stride(0), - x_q.stride(1), - x_scale.stride(0), - x_scale.stride(1), - w.stride(0), - w.stride(1), - w_scale.stride(0), - w_scale.stride(1), - out.stride(0), - out.stride(1), - BLOCK_M=BLOCK_M, - BLOCK_N=BLOCK_N, - BLOCK_K=BLOCK_K, - num_warps=num_warps, - num_stages=num_stages, - ) + _launch(x_q, x_scale, w, w_scale, out, *_select_cfg(M, N, K)) return out From 0c81a0371f586cc50a0daad6d216fbac21cf1a2d Mon Sep 17 00:00:00 2001 From: Jon Belof Date: Wed, 28 Oct 2026 19:41:42 -0700 Subject: [PATCH 106/406] [KV Connector][NIXL] Add a ucx_backend_params key for UCX backend init params (#62905) Signed-off-by: Jon Belof --- docs/features/nixl_connector_usage.md | 13 ++++ .../kv_connector/unit/test_nixl_connector.py | 68 +++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 18 ++++- 3 files changed, 98 insertions(+), 1 deletion(-) diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index 4402095b2ec1..00f5860af4c2 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -63,6 +63,19 @@ vllm serve \ !!! note Backend availability depends on how NIXL was built and what plugins are present in your environment. Refer to the [NIXL repository](https://github.com/ai-dynamo/nixl) for available backends and build instructions. +#### UCX backend parameters + +`kv_connector_extra_config.ucx_backend_params` passes extra init parameters to NIXL's UCX backend. Values are strings. For example, NIXL creates UCX with `ucx_error_handling_mode=peer` by default; UCX then uses only transports that can detect peer failure, which can exclude intra-node GPU transports such as `rocm_ipc`. To trade that error handling for those transports: + +```bash +vllm serve \ + --kv-transfer-config '{ + "kv_connector":"NixlConnector", + "kv_role":"kv_both", + "kv_connector_extra_config":{"ucx_backend_params":{"ucx_error_handling_mode":"none"}} + }' +``` + ## Basic Usage (on the same host) ### Producer (Prefiller) Configuration diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 2889a8bf36b8..c26e3a098b0c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -5132,3 +5132,71 @@ def test_explicit_kv_role_no_deprecation_warning(default_vllm_config, dist_init) mock_logger.warning_once.assert_not_called(), (f"kv_role={role!r} should not emit deprecation warning"), ) + + +class _BackendRecordingNixlWrapper(FakeNixlWrapper): + def __init__(self, agent_name: str, config=None, *args, **kwargs): + super().__init__(agent_name, *args, **kwargs) + self.config = config + self.created_backends: list[tuple[str, dict[str, str]]] = [] + + def create_backend(self, backend: str, init_params: dict[str, str]) -> None: + self.created_backends.append((backend, dict(init_params))) + + +@pytest.mark.parametrize( + "extra_config, agent_backends, created_backends", + [ + # No UCX params: the agent config creates the backends, as before. + ({}, ["UCX"], []), + ( + {"ucx_backend_params": {"ucx_error_handling_mode": "none"}}, + [], + [("UCX", {"num_threads": "4", "ucx_error_handling_mode": "none"})], + ), + ( + { + "backends": ["UCX", "LIBFABRIC"], + "ucx_backend_params": {"ucx_error_handling_mode": "none"}, + }, + [], + [("UCX", {"ucx_error_handling_mode": "none"}), ("LIBFABRIC", {})], + ), + # UCX not selected: the params have nothing to apply to. + ( + { + "backends": ["LIBFABRIC"], + "ucx_backend_params": {"ucx_error_handling_mode": "none"}, + }, + ["LIBFABRIC"], + [], + ), + ], +) +def test_ucx_backend_params(dist_init, extra_config, agent_backends, created_backends): + """kv_connector_extra_config.ucx_backend_params reaches UCX's create_backend.""" + + def agent_config(backends=None, num_threads=0, capture_telemetry=False): + return SimpleNamespace( + backends=["UCX"] if backends is None else backends, + num_threads=num_threads, + ) + + vllm_config = create_vllm_config(kv_connector_extra_config=extra_config) + with ( + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + _BackendRecordingNixlWrapper, + ), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.nixl_agent_config", + agent_config, + ), + set_current_vllm_config(vllm_config), + ): + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + wrapper = connector.connector_worker.nixl_wrapper + assert wrapper.config.backends == agent_backends + assert wrapper.created_backends == created_backends 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 ffd75e2d3e18..189c07fc4237 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 @@ -756,8 +756,15 @@ def __init__( else nixl_agent_config(num_threads=num_threads, capture_telemetry=True) ) + # Extra UCX backend init params, e.g. {"ucx_error_handling_mode": "none"}. + params = self.kv_transfer_config.get_from_extra_config("ucx_backend_params", {}) + ucx_backend_params = {k: str(v) for k, v in params.items()} sleep_mode = vllm_config.model_config.enable_sleep_mode and config is not None - if sleep_mode: + # With UCX params, create every backend here, in the configured order. + explicit_backends = (config is not None and not sleep_mode) and ( + bool(ucx_backend_params) and "UCX" in self.nixl_backends + ) + if sleep_mode or explicit_backends: config = nixl_agent_config(backends=[], capture_telemetry=True) self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config) if sleep_mode: @@ -767,8 +774,17 @@ def __init__( { "num_threads": str(num_threads), "engine_config": "RCACHE_ENABLE=n,TLS=^cuda_ipc", + **ucx_backend_params, }, ) + elif explicit_backends: + for backend in self.nixl_backends: + nixl_params: dict[str, str] = {} + if backend == "UCX": + if not non_ucx_backends: + nixl_params["num_threads"] = str(num_threads) + nixl_params.update(ucx_backend_params) + self.nixl_wrapper.create_backend(backend, nixl_params) # Map of engine_id -> {(pp_rank, tp_rank): agent_name, ...}. # non-PP remote uses pp_rank 0, i.e. (0, tp_rank). self._remote_agents: dict[EngineId, dict[tuple[int, int], str]] = defaultdict( From a94b768d964e5f463737857606aadac5690061fa Mon Sep 17 00:00:00 2001 From: hellozmz <407190054@qq.com> Date: Thu, 29 Oct 2026 12:03:41 +0800 Subject: [PATCH 107/406] [Kernel][MiniMax-M3] Enable SM90 split-K in the Triton indexer (#64037) Signed-off-by: hellozmz Co-authored-by: hellozmz Co-authored-by: vllm-agent --- tests/kernels/attention/test_minimax_m3.py | 19 ++++++++++++------- .../minimax_m3/common/ops/index_topk.py | 16 ++++++++++------ 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 8f0f38be5d6c..d46224bd5fd8 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -627,13 +627,18 @@ def _reference_decode_index_score( # 3600: 29 blocks over 12 splits of 3, so the last busy split has 2 blocks. @pytest.mark.parametrize("prefix_len", [1024, 3600]) -@pytest.mark.parametrize("force_split_k", [False, True]) -def test_prefill_index_topk_correctness(monkeypatch, prefix_len, force_split_k): - if force_split_k: - # Take the SM12.0-only split-K launch on any CUDA GPU. - monkeypatch.setattr( - current_platform, "is_device_capability", lambda cap, *_: cap == (12, 0) - ) +@pytest.mark.parametrize("device_capability", [80, 90, 120]) +def test_prefill_index_topk_correctness(monkeypatch, prefix_len, device_capability): + monkeypatch.setattr( + current_platform, + "is_device_capability", + lambda cap, *_: device_capability == cap[0] * 10 + cap[1], + ) + monkeypatch.setattr( + current_platform, + "is_device_capability_family", + lambda cap, *_: device_capability // 10 == cap // 10, + ) topk = 6 init_blocks = 0 local_blocks = 1 diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index 2d612c12f254..bc24a9a3e285 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -703,15 +703,19 @@ def minimax_m3_index_score( ) BLOCK_SIZE_Q = 64 n_q_tiles = triton.cdiv(max_query_len, BLOCK_SIZE_Q) - SCORE_TARGET_GRID = 48 + # Splits write disjoint score columns, preserving the unsplit result. + if current_platform.is_cuda() and current_platform.is_device_capability((12, 0)): + score_target_grid = 48 + elif current_platform.is_cuda() and current_platform.is_device_capability_family( + 90 + ): + score_target_grid = 1056 + else: + score_target_grid = 1 split_k = max( 1, - min(max_block, SCORE_TARGET_GRID // max(1, n_q_tiles * batch * num_idx_heads)), + min(max_block, score_target_grid // max(1, n_q_tiles * batch * num_idx_heads)), ) - if not ( - current_platform.is_cuda() and current_platform.is_device_capability((12, 0)) - ): - split_k = 1 grid_score = (n_q_tiles, batch * num_idx_heads, split_k) _index_block_score_kernel[grid_score]( idx_q, From 18450e5ef778fa9a08001e51d1711c5e39d10df2 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Thu, 29 Oct 2026 06:03:48 +0200 Subject: [PATCH 108/406] [Fusion] Fuse Qwen MoE shared MLP activation quantization (#61107) Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Signed-off-by: Thien Tran Co-authored-by: Thien Tran Co-authored-by: Claude Opus 5.5 --- vllm/model_executor/models/qwen2_moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index 1c095d6f091c..46320b2c8eae 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -44,6 +44,7 @@ FusedMoEFactory, GateLinear, ) +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 ( MergedColumnParallelLinear, @@ -114,7 +115,7 @@ def __init__( def forward(self, x): gate_up, _ = self.gate_up_proj(x) - out = self.act_fn(gate_up) + out = maybe_fused_act_quant(self.act_fn, gate_up, self.down_proj) out, _ = self.down_proj(out) if self.expert_gate is not None: From 66875d3127534183b5ea723964fc7ec388e89722 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:08:42 -0700 Subject: [PATCH 109/406] [Spec Decode] Batch-sharded draft logits refactor (#64356) Signed-off-by: Giancarlo Delfin Co-authored-by: vllm-agent Co-authored-by: Woosuk Kwon --- .../amd/test_dsv41_gfx942_fused_markov.py | 13 +- tests/v1/spec_decode/test_dspark_topk.py | 9 +- tests/v1/spec_decode/test_lilicorr.py | 23 +- .../test_rejection_sampler_utils.py | 272 +++++++++++++++--- tests/v1/worker/test_gpu_batch_shard.py | 182 +++++++++++- .../worker/test_gpu_model_runner_v2_eplb.py | 1 + tests/v1/worker/test_gpu_pcp_manager.py | 9 +- .../test_gpu_rejection_sampler_chunking.py | 11 +- .../worker/test_gpu_rejection_sampler_i64.py | 9 +- tests/watermarking/test_watermarking.py | 25 +- .../deepseek_v41/amd/ops/fused_markov.py | 6 +- vllm/v1/watermarking/spec_decode.py | 5 +- vllm/v1/worker/gpu/model_runner.py | 22 +- vllm/v1/worker/gpu/pcp_manager.py | 7 +- vllm/v1/worker/gpu/sample/batch_shard.py | 62 +++- vllm/v1/worker/gpu/sample/gumbel.py | 25 +- .../gpu/spec_decode/dflash2/speculator.py | 43 +-- .../v1/worker/gpu/spec_decode/draft_logits.py | 201 +++++++++++++ .../gpu/spec_decode/dspark/speculator.py | 7 +- .../gpu/spec_decode/rejection_sampler.py | 23 +- .../spec_decode/rejection_sampler_utils.py | 57 ++-- vllm/v1/worker/gpu/spec_decode/speculator.py | 75 ++--- 22 files changed, 845 insertions(+), 242 deletions(-) create mode 100644 vllm/v1/worker/gpu/spec_decode/draft_logits.py diff --git a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py index 42bac7e4e7fc..892987dfd7f3 100644 --- a/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py +++ b/tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py @@ -80,6 +80,7 @@ def map_draft_to_target(self, ids): def _speculator(temp: float, seed: int): + from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator dev = torch.device("cuda") @@ -98,9 +99,9 @@ def _speculator(temp: float, seed: int): s._d2t_scatter_index = None s._draft_scatter_buf = None s.use_fp64_gumbel = False - s.draft_logits = torch.zeros(MAX_REQS, N_SPEC, VOCAB, dtype=bf16, device=dev) - s.draft_logits_num_shards = 1 - s.draft_logits_shard_idx = 0 + s.draft_logits = DraftLogitsCache( + MAX_REQS, N_SPEC, VOCAB, dtype=bf16, fill=0.0, device=dev + ) s.temperature = torch.full((MAX_REQS,), temp, dtype=torch.float32, device=dev) s.seeds = torch.randint(0, 2**62, (MAX_REQS,), generator=g, device=dev) s._step_cols = torch.arange(N_SPEC, dtype=torch.int32, device=dev) @@ -132,17 +133,17 @@ def test_fused_markov_matches_vllm_loop(monkeypatch, num_reqs, temp): _set_gate(monkeypatch, on=False) DSparkSpeculator._sample_sequential(s, num_reqs, hidden) torch.accelerator.synchronize() - ref_tok, ref_cache = s.draft_tokens.clone(), s.draft_logits.clone() + ref_tok, ref_cache = s.draft_tokens.clone(), s.draft_logits.values.clone() s.draft_tokens.zero_() - s.draft_logits.zero_() + s.draft_logits.values.zero_() _set_gate(monkeypatch, on=True) try: assert fused_markov.sample_sequential(s, num_reqs, hidden) torch.accelerator.synchronize() finally: _set_gate(monkeypatch, on=False) - got_tok, got_cache = s.draft_tokens.clone(), s.draft_logits.clone() + got_tok, got_cache = s.draft_tokens.clone(), s.draft_logits.values.clone() # rc and gc are [num_reqs, N_SPEC, VOCAB]: the cached logits of each # request and draft step. The differences are counted for each request diff --git a/tests/v1/spec_decode/test_dspark_topk.py b/tests/v1/spec_decode/test_dspark_topk.py index fe5f5a389bea..941b81236c20 100644 --- a/tests/v1/spec_decode/test_dspark_topk.py +++ b/tests/v1/spec_decode/test_dspark_topk.py @@ -16,6 +16,7 @@ from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead from vllm.platforms import current_platform from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator @@ -86,9 +87,9 @@ def test_dspark_revisited_position_uses_step_seed(use_fp64): spec.temperature = torch.ones(batch, device=device) spec.seeds = torch.arange(batch, dtype=torch.int64, device=device) spec._step_cols = torch.arange(2, dtype=torch.int32, device=device) - spec.draft_logits = torch.empty(batch, 2, 2, device=device) - spec.draft_logits_num_shards = 1 - spec.draft_logits_shard_idx = 0 + spec.draft_logits = DraftLogitsCache( + batch, 2, 2, dtype=torch.float32, fill=0.0, device=device + ) logits = torch.tensor([0.7, 0.3], device=device).log().expand(batch, -1) mapping = torch.arange(batch, device=device) positions = torch.full((batch,), 22, dtype=torch.int64, device=device) @@ -105,7 +106,7 @@ def test_dspark_revisited_position_uses_step_seed(use_fp64): positions - 1, apply_temperature=True, is_drafting=True, - logits_cache=spec.draft_logits, + logits_cache=spec.draft_logits.values, logits_cache_col=spec._step_cols[0], use_fp64=use_fp64, ) diff --git a/tests/v1/spec_decode/test_lilicorr.py b/tests/v1/spec_decode/test_lilicorr.py index f6d7834ebe6e..f9d655c52248 100644 --- a/tests/v1/spec_decode/test_lilicorr.py +++ b/tests/v1/spec_decode/test_lilicorr.py @@ -24,6 +24,7 @@ from vllm.transformers_utils.configs.eagle import EAGLEConfig from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import CandidateSampler +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.lilicorr.speculator import LiLiCorrSpeculator @@ -368,8 +369,6 @@ def init_base(self, config, device): spec.seeds = torch.arange(3) spec.draft_tokens = torch.empty(3, 2, dtype=torch.long) spec.draft_logits = None - spec.draft_logits_num_shards = 1 - spec.draft_logits_shard_idx = 0 spec.use_fp64_gumbel = False spec.enable_adaptive_verification = True spec.input_buffers = SimpleNamespace(input_ids=torch.arange(9)) @@ -456,7 +455,9 @@ def test_candidate_walk_preserves_conditional_scores_and_padding(probabilistic): scores = torch.randn(batch, steps, top_k, top_k, device=device) tokens = torch.full((batch, steps), -1, dtype=torch.long, device=device) cache = ( - torch.full((batch, steps, vocab), -float("inf"), device=device) + DraftLogitsCache( + batch, steps, vocab, dtype=torch.float32, fill=-float("inf"), device=device + ) if probabilistic else None ) @@ -498,13 +499,17 @@ def run(): assert index == realized.argmax().item() if cache is not None: torch.testing.assert_close( - cache[state, step, candidates[row, step]], realized + cache.values[state, step, candidates[row, step]], realized + ) + assert ( + cache.values[state, step, old_candidates[row, step]] + .isneginf() + .all() ) - assert cache[state, step, old_candidates[row, step]].isneginf().all() previous = index assert (tokens[2] == -1).all() if cache is not None: - assert cache[1].isneginf().all() + assert cache.values[1].isneginf().all() @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -522,7 +527,9 @@ def test_candidate_walk_revisited_position_matches_step_seed(use_fp64): seeds = mapping.to(torch.int64) positions = torch.arange(21, 21 + steps, device=device).repeat(batch) tokens = torch.empty(batch, steps, dtype=torch.int64, device=device) - cache = torch.empty(batch, steps, top_k, device=device) + cache = DraftLogitsCache( + batch, steps, top_k, dtype=torch.float32, fill=0.0, device=device + ) first_round = None for start in (21, 22): positions.view(batch, steps).copy_( @@ -549,7 +556,7 @@ def test_candidate_walk_revisited_position_matches_step_seed(use_fp64): positions.view(batch, steps)[:, step] - 1, apply_temperature=True, is_drafting=True, - logits_cache=cache, + logits_cache=cache.values, logits_cache_col=torch.tensor(step, device=device), use_fp64=use_fp64, ) diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index e51a6421ee2a..a00f41929c56 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -12,6 +12,7 @@ from vllm.v1.watermarking.spec_decode import watermarked_rejection_sample from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.rejection_sampler import ( RejectionSampler, gather_draft_sampled, @@ -28,6 +29,40 @@ pytest.skip("CUDA required for rejection sampler tests", allow_module_level=True) +def _draft_logits( + values: torch.Tensor | None, row_mapping: torch.Tensor | None = None +) -> DraftLogitsCache | None: + """Wrap the draft logits tensor a test builds in a DraftLogitsCache.""" + if values is None: + return None + draft_logits = DraftLogitsCache( + 1, 1, 1, dtype=values.dtype, fill=0.0, device=values.device + ) + draft_logits.values = values + draft_logits.row_mapping = row_mapping + return draft_logits + + +def _rejection_sample( + *, + draft_logits: torch.Tensor | None, + draft_logits_row_mapping: torch.Tensor | None = None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + return rejection_sample( + draft_logits_cache=_draft_logits(draft_logits, draft_logits_row_mapping), + **kwargs, + ) + + +def _watermarked_rejection_sample( + *, draft_logits: torch.Tensor | None, **kwargs +) -> tuple[torch.Tensor, torch.Tensor]: + return watermarked_rejection_sample( + draft_logits_cache=_draft_logits(draft_logits), **kwargs + ) + + def _build_rejection_sample_inputs( target_logits_1d: torch.Tensor, draft_logits_1d: torch.Tensor, @@ -138,7 +173,7 @@ def test_fly_requires_entropy_and_native_lookahead( [0, *torch.tensor(counts).cumsum(0).tolist()], dtype=torch.int32, device=device ) - sampled, num_sampled = rejection_sample( + sampled, num_sampled = _rejection_sample( **inputs, num_speculative_steps=4, fly_window_size=2, @@ -196,10 +231,10 @@ def test_fly_pending_window_preserves_recovery_position( if step < rejected_step: reference_logits[:, step] = native_logits[step] reference_inputs["target_logits"] = reference_logits.flatten(0, 1) - expected, expected_counts = rejection_sample( + expected, expected_counts = _rejection_sample( **reference_inputs, num_speculative_steps=num_steps ) - sampled, counts = rejection_sample( + sampled, counts = _rejection_sample( **inputs, num_speculative_steps=num_steps, fly_window_size=2 ) assert counts.tolist() == [rejected_step + 1] * num_trials @@ -302,7 +337,7 @@ def test_watermarked_bonus_uses_target_key(): [[10, 11], [11, draft_token.item()]], dtype=torch.int64, device="cuda" ) - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=1, contexts=contexts, @@ -330,7 +365,7 @@ def test_watermarked_recovery_uses_target_key(): inputs["draft_sampled"][1] = 0 contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda") - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=1, contexts=contexts, @@ -358,9 +393,11 @@ def test_repeated_draft_context_does_not_change_acceptance(): temperature=1.0, num_trials=num_trials, ) - expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1) + expected, expected_num_sampled = _rejection_sample( + **inputs, num_speculative_steps=1 + ) - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=1, contexts=torch.tensor([[10, 11], [11, 12]], device="cuda").repeat( @@ -403,14 +440,14 @@ def test_repeated_context_skip_mask_is_read_per_row_in_a_block(): contexts = torch.tensor( [[10, 11], [11, 3], [3, 0], [0, 1]], dtype=torch.int64, device="cuda" ) - expected, expected_num_sampled = rejection_sample( + expected, expected_num_sampled = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps ) assert expected_num_sampled.item() == 2 assert expected[0, 0].item() == 3 watermarker = GumbelWatermarker(key=42, context_width=2) - skipped, skipped_num_sampled = watermarked_rejection_sample( + skipped, skipped_num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=num_speculative_steps, contexts=contexts, @@ -422,7 +459,7 @@ def test_repeated_context_skip_mask_is_read_per_row_in_a_block(): assert skipped_num_sampled.item() == 2 torch.testing.assert_close(skipped[:, :2], expected[:, :2], rtol=0, atol=0) - watermarked, watermarked_num_sampled = watermarked_rejection_sample( + watermarked, watermarked_num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=num_speculative_steps, contexts=contexts, @@ -454,9 +491,11 @@ def test_repeated_recovery_context_uses_ordinary_residual_sample(): num_trials=1, ) inputs["draft_sampled"][1] = 0 - expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1) + expected, expected_num_sampled = _rejection_sample( + **inputs, num_speculative_steps=1 + ) - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=1, contexts=torch.tensor([[10, 11], [11, 0]], device="cuda"), @@ -478,9 +517,11 @@ def test_repeated_bonus_context_uses_ordinary_target_sample(): temperature=1.0, num_trials=1, ) - expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1) + expected, expected_num_sampled = _rejection_sample( + **inputs, num_speculative_steps=1 + ) - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=1, contexts=torch.tensor([[10, 11], [11, 12]], device="cuda"), @@ -503,9 +544,11 @@ def test_repeated_context_does_not_change_unwatermarked_or_greedy_rows(): num_trials=2, ) inputs["temperature"][1] = 0 - expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1) + expected, expected_num_sampled = _rejection_sample( + **inputs, num_speculative_steps=1 + ) - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=1, contexts=torch.tensor([[10, 11], [11, 12], [20, 21], [21, 22]], device="cuda"), @@ -536,7 +579,7 @@ def test_watermarked_recovery_supports_smaller_draft_vocabulary(): contexts = torch.tensor([[10, 11], [11, 0]], dtype=torch.int64, device="cuda") watermarker = GumbelWatermarker(key=42, context_width=2) - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=1, contexts=contexts, @@ -592,7 +635,7 @@ def test_stochastic_rejection_sample( num_trials=num_trials, ) - sampled, num_sampled = rejection_sample( + sampled, num_sampled = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps ) @@ -684,7 +727,7 @@ def test_gumbel_drafted_rejection_sample_is_unbiased(num_speculative_steps: int) inputs, draft_logits_1d, NARROW_NUM_TRIALS, num_speculative_steps ) - sampled, num_sampled = rejection_sample( + sampled, num_sampled = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps ) @@ -744,7 +787,7 @@ def test_multiround_rejection_sample_is_unbiased( logits_cache=logits_cache, logits_cache_col=torch.tensor(step, device=device), ) - sampled, counts = rejection_sample( + sampled, counts = _rejection_sample( target_logits=target_probs.log().expand(n * (k + 1), -1).contiguous(), draft_logits=logits_cache, draft_sampled=draft_sampled.flatten(), @@ -792,7 +835,7 @@ def test_greedy_rejection_sample(num_speculative_steps: int): num_trials=num_trials, ) - sampled, num_sampled = rejection_sample( + sampled, num_sampled = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps ) @@ -855,7 +898,7 @@ def test_synthetic_rejection_sample( conditional_rates, dtype=torch.float32, device=device ) - _, num_sampled = rejection_sample( + _, num_sampled = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps, synthetic_conditional_rates=synthetic_conditional_rates, @@ -903,7 +946,7 @@ def test_all_nan_target_logits_in_range(temperature: float): num_trials=num_trials, ) - sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K) + sampled, num_sampled = _rejection_sample(**inputs, num_speculative_steps=K) assert ((num_sampled >= 1) & (num_sampled <= K + 1)).all() steps = torch.arange(K + 1, device=device).unsqueeze(0) @@ -934,7 +977,7 @@ def test_placeholder_draft_token_rejected(): ) inputs["draft_sampled"].view(num_trials, K + 1)[:, 1:] = -1 - sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K) + sampled, num_sampled = _rejection_sample(**inputs, num_speculative_steps=K) assert torch.equal(num_sampled, torch.ones_like(num_sampled)) recovered = sampled[:, 0] @@ -971,7 +1014,7 @@ def test_block_verification_placeholder_truncates_block( inputs["draft_logits"] = None inputs["draft_sampled"].view(num_trials, K + 1)[:, K + 1 - num_placeholders :] = -1 - sampled, num_sampled = rejection_sample( + sampled, num_sampled = _rejection_sample( **inputs, num_speculative_steps=K, use_block_verification=True ) @@ -1013,7 +1056,7 @@ def test_greedy_placeholder_emits_target_argmax(): draft_sampled[:, 1] = target_argmax draft_sampled[:, 2:] = -1 - sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K) + sampled, num_sampled = _rejection_sample(**inputs, num_speculative_steps=K) assert torch.equal(num_sampled, torch.full_like(num_sampled, 2)) steps = torch.arange(K + 1, device=device).unsqueeze(0) @@ -1049,7 +1092,7 @@ def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool): ) inputs["draft_sampled"].view(num_trials, K + 1)[:, 2] = -1 - _, num_sampled = rejection_sample( + _, num_sampled = _rejection_sample( **inputs, num_speculative_steps=K, use_block_verification=use_block_verification, @@ -1106,8 +1149,14 @@ def test_verify_rejects_unproposed_drafts(): rejection_sampler.fly_window_size = 0 rejection_sampler.fly_entropy_threshold = 0.3 rejection_sampler.fly_entropy_top_k = 3 - rejection_sampler.draft_logits_num_shards = 1 rejection_sampler.watermark_key = None + draft_logits = DraftLogitsCache( + *inputs["draft_logits"].shape, + dtype=inputs["draft_logits"].dtype, + fill=0.0, + device=device, + ) + draft_logits.values.copy_(inputs["draft_logits"]) def verify() -> tuple[torch.Tensor, torch.Tensor]: draft_sampled, pos = gather_draft_sampled( @@ -1120,7 +1169,7 @@ def verify() -> tuple[torch.Tensor, torch.Tensor]: ) _, sampled, num_sampled = rejection_sampler._verify( inputs["target_logits"], - inputs["draft_logits"], + draft_logits, draft_sampled, pos, inputs["cu_num_logits"], @@ -1184,7 +1233,7 @@ def test_block_verification_rejection_sample( if not has_draft_logits: inputs["draft_logits"] = None - sampled, num_sampled = rejection_sample( + sampled, num_sampled = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps, use_block_verification=True, @@ -1224,10 +1273,10 @@ def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int) num_trials=num_trials, ) - _, num_sampled_standard = rejection_sample( + _, num_sampled_standard = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps ) - _, num_sampled_block = rejection_sample( + _, num_sampled_block = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps, use_block_verification=True, @@ -1273,10 +1322,10 @@ def test_chunked_requests_match_full_batch( inputs["draft_logits"] = None if fly_window_size: - native, native_counts = rejection_sample( + native, native_counts = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps ) - gated, gated_counts = rejection_sample( + gated, gated_counts = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps, fly_window_size=fly_window_size, @@ -1289,7 +1338,7 @@ def test_chunked_requests_match_full_batch( ) assert torch.equal(gated[valid], native[valid]) - sampled, num_sampled = rejection_sample( + sampled, num_sampled = _rejection_sample( **inputs, num_speculative_steps=num_speculative_steps, fly_window_size=fly_window_size, @@ -1313,7 +1362,7 @@ def test_chunked_requests_match_full_batch( chunk_inputs["cu_num_logits"] = inputs["cu_num_logits"][start : end + 1] - lo chunk_inputs["idx_mapping"] = inputs["idx_mapping"][start:end] - chunk_sampled, chunk_num_sampled = rejection_sample( + chunk_sampled, chunk_num_sampled = _rejection_sample( **chunk_inputs, num_speculative_steps=num_speculative_steps, fly_window_size=fly_window_size, @@ -1381,7 +1430,7 @@ def test_adaptive_draft_lengths_match_and_bonus_oracle( num_speculative_steps=K, use_block_verification=use_block_verification, ) - sampled, num_sampled = rejection_sample(**inputs) + sampled, num_sampled = _rejection_sample(**inputs) for req_idx, num_matches in greedy_num_matches.items(): lo = int(cu_num_logits[req_idx]) @@ -1405,7 +1454,7 @@ def test_adaptive_draft_lengths_match_and_bonus_oracle( req_inputs["cu_num_logits"] = cu_num_logits[req_idx : req_idx + 2] - lo req_inputs["idx_mapping"] = idx_mapping[req_idx : req_idx + 1] - req_sampled, req_num_sampled = rejection_sample(**req_inputs) + req_sampled, req_num_sampled = _rejection_sample(**req_inputs) n = int(req_num_sampled[0]) assert n <= num_drafts[req_idx] + 1 @@ -1507,7 +1556,7 @@ def test_watermarked_recovery_matches_philox_gumbel_sample( ) watermarker = GumbelWatermarker(key=key, context_width=context_width) - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=2, contexts=contexts, @@ -1545,10 +1594,10 @@ def test_watermarked_recovery_leaves_disabled_rows_on_the_stock_draw(): watermarking[::2] = True inputs["temperature"][:8] = 0.0 - stock_sampled, stock_num_sampled = rejection_sample( + stock_sampled, stock_num_sampled = _rejection_sample( **inputs, num_speculative_steps=3 ) - sampled, num_sampled = watermarked_rejection_sample( + sampled, num_sampled = _watermarked_rejection_sample( **inputs, num_speculative_steps=3, contexts=contexts, @@ -1590,8 +1639,8 @@ def test_watermarked_recovery_is_unchanged_by_fp64_gumbel(): watermarking=torch.ones(64, dtype=torch.bool, device="cuda"), watermarker=watermarker, ) - sampled32, num_sampled32 = watermarked_rejection_sample(**inputs, **kwargs) - sampled64, num_sampled64 = watermarked_rejection_sample( + sampled32, num_sampled32 = _watermarked_rejection_sample(**inputs, **kwargs) + sampled64, num_sampled64 = _watermarked_rejection_sample( **inputs, **kwargs, use_fp64=True ) rows = torch.arange(64, device="cuda") @@ -1602,3 +1651,140 @@ def test_watermarked_recovery_is_unchanged_by_fp64_gumbel(): rtol=0, atol=0, ) + + +def _position_owned_rows( + slots: torch.Tensor, num_shards: int, shard_idx: int, max_num_reqs: int +) -> tuple[torch.Tensor, torch.Tensor]: + """The batch positions a shard owns and its DraftLogitsCache.row_mapping table: + the request at batch position i is cached by shard i % num_shards at row + i // num_shards.""" + owned = torch.arange(shard_idx, slots.shape[0], num_shards, device=slots.device) + rows = torch.full((max_num_reqs,), -1, dtype=torch.int32, device=slots.device) + rows[slots[owned].long()] = torch.arange( + owned.shape[0], dtype=torch.int32, device=slots.device + ) + return owned, rows + + +@pytest.mark.parametrize("num_shards", [2, 4]) +@pytest.mark.parametrize("use_block_verification", [False, True]) +def test_sharded_draft_logits_verify_like_replicated( + num_shards: int, use_block_verification: bool +): + """A TP rank that caches only its owned requests' draft logits, at the + rows of draft_logits_row_mapping, verifies those requests exactly like the + replicated cache does.""" + torch.manual_seed(0) + device = "cuda" + num_reqs = 12 + num_speculative_steps = 3 + vocab_size = 1000 + num_rows = num_speculative_steps + 1 + inputs = _build_rejection_sample_inputs( + torch.zeros(vocab_size, device=device), + torch.zeros(vocab_size, device=device), + num_speculative_steps, + temperature=1.0, + num_trials=num_reqs, + ) + # Per-request distributions, close enough that some drafts are accepted. + draft_logits = torch.randn( + num_reqs, num_speculative_steps, vocab_size, device=device + ) + target_logits = torch.randn(num_reqs, num_rows, vocab_size, device=device) + target_logits[:, :num_speculative_steps] = draft_logits + 0.5 * torch.randn_like( + draft_logits + ) + inputs["target_logits"] = target_logits.view(-1, vocab_size) + draft_tokens = torch.multinomial( + torch.softmax(draft_logits.view(-1, vocab_size), dim=-1), 1 + ).view(num_reqs, num_speculative_steps) + inputs["draft_sampled"].view(num_reqs, num_rows)[:, 1:] = draft_tokens + # Shuffle request slots so cache rows aren't request slots. + slots = torch.randperm(num_reqs, device=device).to(torch.int32) + inputs["idx_mapping"] = slots + inputs["expanded_idx_mapping"] = slots.repeat_interleave(num_rows) + inputs["draft_logits"] = torch.empty_like(draft_logits) + inputs["draft_logits"][slots.long()] = draft_logits + + sampled, num_sampled = _rejection_sample( + **inputs, + num_speculative_steps=num_speculative_steps, + use_block_verification=use_block_verification, + ) + assert 1 < num_sampled.float().mean() < num_rows + + for shard_idx in range(num_shards): + owned, draft_logits_row_mapping = _position_owned_rows( + slots, num_shards, shard_idx, num_reqs + ) + rows = owned.view(-1, 1) * num_rows + torch.arange(num_rows, device=device) + rows = rows.view(-1) + local_inputs = dict(inputs) + for name in ( + "target_logits", + "draft_sampled", + "pos", + "expanded_idx_mapping", + "expanded_local_pos", + ): + local_inputs[name] = inputs[name][rows] + local_inputs["idx_mapping"] = slots[owned] + local_inputs["cu_num_logits"] = inputs["cu_num_logits"][: len(owned) + 1] + local_inputs["draft_logits"] = draft_logits[owned] + + local_sampled, local_num_sampled = _rejection_sample( + **local_inputs, + num_speculative_steps=num_speculative_steps, + use_block_verification=use_block_verification, + draft_logits_row_mapping=draft_logits_row_mapping, + ) + assert torch.equal(local_num_sampled, num_sampled[owned]) + valid = torch.arange(num_rows, device=device) < local_num_sampled[:, None] + assert torch.equal(local_sampled[valid], sampled[owned][valid]) + + +@pytest.mark.parametrize("num_shards", [2, 4]) +def test_sharded_draft_logits_cache_writes(num_shards: int): + """gumbel_sample caches only the owned rows, at the rows of + logits_cache_row_mapping, without changing what it samples.""" + torch.manual_seed(0) + device = "cuda" + max_num_reqs = 16 + num_cols = 3 + num_tokens = max_num_reqs * num_cols + slots = torch.randperm(max_num_reqs, device=device).to(torch.int32) + idx_mapping = slots.repeat_interleave(num_cols) + # A CUDA graph padding row. + idx_mapping[-num_cols:] = -1 + cols = torch.arange(num_cols, device=device, dtype=torch.int32).repeat(max_num_reqs) + logits = torch.randn(num_tokens, VOCAB_SIZE, device=device) + kwargs = dict( + expanded_idx_mapping=idx_mapping, + temperature=torch.ones(max_num_reqs, device=device), + seed=torch.arange(max_num_reqs, device=device), + pos=torch.arange(num_tokens, device=device), + apply_temperature=True, + is_drafting=True, + logits_cache_col=cols, + ) + cache = torch.zeros(max_num_reqs, num_cols, VOCAB_SIZE, device=device) + sampled = gumbel_sample(logits, logits_cache=cache, **kwargs) + + for shard_idx in range(num_shards): + owned, logits_cache_row_mapping = _position_owned_rows( + slots, num_shards, shard_idx, max_num_reqs + ) + local_cache = torch.zeros(len(owned), num_cols, VOCAB_SIZE, device=device) + local_sampled = gumbel_sample( + logits, + logits_cache=local_cache, + logits_cache_row_mapping=logits_cache_row_mapping, + **kwargs, + ) + assert torch.equal(local_sampled, sampled) + # The padding request caches nothing. + is_real = owned < max_num_reqs - 1 + expected = cache[slots[owned].long()] * is_real.view(-1, 1, 1) + assert torch.equal(local_cache, expected) diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py index 4ee76818df4f..72d3affed239 100644 --- a/tests/v1/worker/test_gpu_batch_shard.py +++ b/tests/v1/worker/test_gpu_batch_shard.py @@ -23,6 +23,7 @@ BatchSharder, _shard_grammar_output, ) +from vllm.v1.worker.gpu.spec_decode import draft_logits as draft_logits_module DEVICE = current_platform.device_type @@ -43,13 +44,10 @@ def _make_batch( num_reqs: int, max_num_reqs: int, max_spec: int, - slot_pool: np.ndarray | None = None, ) -> InputBatch: """Consistent batch with random unique request slots, 1 + k_i logits rows per request at the tail of each query segment.""" - if slot_pool is None: - slot_pool = np.arange(max_num_reqs) - slots = rng.choice(slot_pool, size=num_reqs, replace=False) + slots = rng.choice(max_num_reqs, size=num_reqs, replace=False) slots = slots.astype(np.int32) k = rng.integers(0, max_spec + 1, size=num_reqs) num_logits_per_req = (1 + k).astype(np.int32) @@ -116,8 +114,21 @@ def _shard_all_ranks( max_num_reqs: int, tp_size: int, adaptive_verification_enabled: bool = False, + owners_np: np.ndarray | None = None, ): - """Simulate every TP rank by faking get_tp_group() per BatchSharder.""" + """Simulate every TP rank by faking get_tp_group() per BatchSharder. + + ``owners_np`` is a sharded DraftLogitsCache owner table indexed by request + slot. Without it, requests are owned by batch position. + """ + draft_logits = None + if owners_np is not None: + owners_np = owners_np.astype(np.int32) + owners = torch.from_numpy(owners_np).to(DEVICE) + # BatchSharder reads only the owner table of a sharded cache. + draft_logits = SimpleNamespace( + owners_of=lambda idx_mapping_np: (owners_np[idx_mapping_np], owners) + ) results = [] for rank in range(tp_size): fake_group = SimpleNamespace(rank_in_group=rank, world_size=tp_size) @@ -127,6 +138,7 @@ def _shard_all_ranks( max_num_logits_per_req=8, device=torch.device(DEVICE), adaptive_verification_enabled=adaptive_verification_enabled, + draft_logits=draft_logits, ) local_batch, sorted_logits_indices, _, metadata = sharder.shard_sampler_inputs( batch, None @@ -215,7 +227,7 @@ def test_local_batch_partition(tp_size: int, seed: int): for rank, (local, _, metadata) in enumerate(results): owned = _owned_batch_indices(local) assert local.num_reqs == len(owned) == metadata.num_local_reqs - assert (batch.idx_mapping_np[owned] % tp_size == rank).all() + assert (owned % tp_size == rank).all() # Stable sort: batch order is preserved within an owner. assert (np.diff(owned) > 0).all() all_owned.extend(owned.tolist()) @@ -398,15 +410,10 @@ def test_gather_src_indices_round_trip(tp_size: int): batch order, including when some rank owns zero requests.""" rng = np.random.default_rng(3) max_num_reqs = 40 - all_slots = np.arange(max_num_reqs) - batch = _make_batch( - rng, - num_reqs=17, - max_num_reqs=max_num_reqs, - max_spec=2, - slot_pool=all_slots[all_slots % tp_size != 0], # rank 0 owns nothing - ) - results = _shard_all_ranks(batch, max_num_reqs, tp_size) + batch = _make_batch(rng, num_reqs=17, max_num_reqs=max_num_reqs, max_spec=2) + # An owner table under which rank 0 owns nothing. + owners_np = 1 + np.arange(max_num_reqs) % (tp_size - 1) + results = _shard_all_ranks(batch, max_num_reqs, tp_size, owners_np=owners_np) metadata0 = results[0][2] assert metadata0.num_local_reqs == 0 @@ -670,3 +677,148 @@ def test_finish_requests_frees_slots_in_sorted_order(): GPUModelRunner.finish_requests(fake_runner, scheduler_output) assert removed == sorted(removed) assert set(removed) == {"req-a", "req-b", "req-c", "req-d", "req-e"} + + +def _make_sharded_draft_logits(tp_size: int, max_num_reqs: int) -> list: + """One sharded DraftLogitsCache per simulated TP rank.""" + caches = [] + for rank in range(tp_size): + fake_group = SimpleNamespace(rank_in_group=rank, world_size=tp_size) + with mock.patch.object( + draft_logits_module, "get_tp_group", return_value=fake_group + ): + caches.append( + draft_logits_module.DraftLogitsCache( + max_num_reqs, + 1, + 1, + dtype=torch.float32, + fill=0.0, + device=torch.device(DEVICE), + sharded=True, + ) + ) + return caches + + +def _draft_batch(slots: list[int], num_logits_per_req: int) -> SimpleNamespace: + """The fields of an InputBatch that DraftLogitsCache reads.""" + slots_np = np.array(slots, dtype=np.int32) + return SimpleNamespace( + idx_mapping=torch.from_numpy(slots_np).to(DEVICE), + idx_mapping_np=slots_np, + expanded_idx_mapping=torch.from_numpy( + np.repeat(slots_np, num_logits_per_req) + ).to(DEVICE), + expanded_local_pos=torch.arange(num_logits_per_req, device=DEVICE).repeat( + len(slots) + ), + ) + + +@requires_gpu +@pytest.mark.parametrize("tp_size", [2, 4]) +def test_sharded_draft_logits_balance_and_stale_rows(tp_size: int): + """Owners follow batch position, every rank caches at most + cdiv(num_reqs, tp_size) requests, and a request skipped while another + request took its row has stale drafts.""" + max_num_reqs = 64 + num_logits_per_req = 3 + owners = _make_sharded_draft_logits(tp_size, max_num_reqs) + # Skewed slots: under slot % tp_size, rank 0 would own every request. + step0 = [s * tp_size for s in range(2 * tp_size)] + batch0 = _draft_batch(step0, num_logits_per_req) + for o in owners: + o.assign(batch0.idx_mapping, batch0.idx_mapping_np) + + for rank, o in enumerate(owners): + assert o.shard_mapping_np[step0].tolist() == [ + i % tp_size for i in range(len(step0)) + ] + rows = _np(o.row_mapping)[step0] + owned = [i for i in range(len(step0)) if i % tp_size == rank] + assert rows[owned].tolist() == [i // tp_size for i in owned] + assert (np.delete(rows, owned) == -1).all() + assert _np(o.inverse_row_mapping)[: len(owned)].tolist() == [ + step0[i] for i in owned + ] + + # Every drafted request verifies against its own row. + for rank, o in enumerate(owners): + owned = [s for s in step0 if o.shard_mapping_np[s] == rank] + stale = _np(o.stale_drafts(_draft_batch(owned, num_logits_per_req))) + assert not stale.any() + + # Step 1 skips the request at position 1 and puts a new request there. + skipped, new = step0[1], 63 + step1 = [step0[0], new] + step0[2:] + batch1 = _draft_batch(step1, num_logits_per_req) + for o in owners: + o.assign(batch1.idx_mapping, batch1.idx_mapping_np) + + # The skipped request keeps its owner, whose row now holds `new`. + owner_of_skipped = 1 % tp_size + o = owners[owner_of_skipped] + assert o.shard_mapping_np[skipped] == owner_of_skipped + stale = _np(o.stale_drafts(_draft_batch([skipped, new], num_logits_per_req))) + stale = stale.reshape(2, num_logits_per_req) + # Only draft rows (local position > 0) of the skipped request are stale. + assert stale[0].tolist() == [False] + [True] * (num_logits_per_req - 1) + assert not stale[1].any() + + +@pytest.mark.parametrize("pcp_size", [1, 2]) +def test_draft_logits_stay_replicated_under_pcp(pcp_size: int): + """PCPManager.sync_drafts indexes the draft logits cache by request state, + so PCP keeps the cache replicated even under batch-sharded sampling.""" + from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator + + speculator = SimpleNamespace( + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size) + ), + max_num_reqs=8, + num_speculative_steps=2, + vocab_size=4, + device=torch.device("cpu"), + draft_watermarker=None, + draft_logits_spec=lambda _: (torch.float32, 0.0), + supports_batch_sharded_draft_logits=lambda: True, + ) + fake_group = SimpleNamespace(rank_in_group=0, world_size=2) + with ( + mock.patch( + "vllm.v1.worker.gpu.spec_decode.speculator.batch_sharded_sampling_enabled", + return_value=True, + ), + mock.patch.object(draft_logits_module, "get_tp_group", return_value=fake_group), + ): + DraftModelSpeculator._init_draft_logits(speculator, target_model=None) + + draft_logits = speculator.draft_logits + assert (draft_logits.row_mapping is not None) == (pcp_size == 1) + assert draft_logits.values.shape[0] == (4 if pcp_size == 1 else 8) + + +def test_owners_of_reads_owner_table_or_batch_position(): + """A sharded cache owns each request at the rank its owner table records; + a replicated cache leaves requests round-robin over batch positions.""" + fake_group = SimpleNamespace(rank_in_group=1, world_size=4) + idx_mapping_np = np.array([6, 1, 3, 4, 0], dtype=np.int32) + with mock.patch.object( + draft_logits_module, "get_tp_group", return_value=fake_group + ): + sharded, replicated = ( + draft_logits_module.DraftLogitsCache( + 8, 1, 1, dtype=torch.float32, fill=0.0, device="cpu", sharded=sharded + ) + for sharded in (True, False) + ) + sharded_owners_np, sharded_owners = sharded.owners_of(idx_mapping_np) + replicated_owners_np, replicated_owners = replicated.owners_of(idx_mapping_np) + + # Before any assign(), a sharded cache owns each request by its slot. + assert sharded_owners_np.tolist() == (idx_mapping_np % 4).tolist() + assert sharded_owners is sharded.shard_mapping + assert replicated_owners_np.tolist() == [0, 1, 2, 3, 0] + assert replicated_owners is None diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index c9dadd2e9f80..4cc000c7989a 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -204,6 +204,7 @@ class FakeDSparkSpeculator: def __init__(self): self.model = draft_model self.eplb_state = None + self.draft_logits = None def load_model(self, target_model): return None diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index a4e50a36e512..6596cd43529a 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -17,6 +17,7 @@ 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 +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache def _copy_to_cpu(value, device=None, dtype=None, out=None): @@ -165,13 +166,15 @@ def test_sync_drafts_writes_logits_to_request_slots(monkeypatch): ) monkeypatch.setattr(pcp_manager_module, "get_pcp_group", lambda: group) draft_tokens = torch.full((1, 3), -1) - draft_logits = torch.full((4, 3, 8), -1.0) + draft_logits = DraftLogitsCache( + 4, 3, 8, dtype=torch.float32, fill=-1.0, device="cpu" + ) manager.sync_drafts(draft_tokens, torch.tensor([2]), draft_logits) torch.testing.assert_close(draft_tokens, root_tokens) - torch.testing.assert_close(draft_logits[2], root_logits[0]) - assert (draft_logits[[0, 1, 3]] == -1).all() + torch.testing.assert_close(draft_logits.values[2], root_logits[0]) + assert (draft_logits.values[[0, 1, 3]] == -1).all() @pytest.mark.parametrize( diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py index a8d9d3d76739..932961f0d64e 100644 --- a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py +++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py @@ -11,6 +11,7 @@ from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode from vllm.platforms import current_platform from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.rejection_sampler import ( RejectionSampler, _iter_request_chunks, @@ -56,7 +57,6 @@ def _make_rejection_sampler( rejection_sampler.fly_window_size = 0 rejection_sampler.fly_entropy_threshold = 0.3 rejection_sampler.fly_entropy_top_k = 3 - rejection_sampler.draft_logits_num_shards = 1 rejection_sampler.watermark_key = None return rejection_sampler @@ -167,8 +167,13 @@ def test_replay_on_off_preserves_rejection_sampling_and_rng(monkeypatch): processed_logits = torch.zeros((12, vocab_size), device=device) processed_logits[0, 0] = -float("inf") processed_logits[rows_per_request + 1, 1] = -float("inf") - draft_logits = torch.zeros( - (num_reqs, num_speculative_steps, vocab_size), device=device + draft_logits = DraftLogitsCache( + num_reqs, + num_speculative_steps, + vocab_size, + dtype=torch.float32, + fill=0.0, + device=device, ) draft_sampled = torch.zeros( (num_reqs, rows_per_request), dtype=torch.int64, device=device diff --git a/tests/v1/worker/test_gpu_rejection_sampler_i64.py b/tests/v1/worker/test_gpu_rejection_sampler_i64.py index 769444ed64c6..5a571d03149b 100644 --- a/tests/v1/worker/test_gpu_rejection_sampler_i64.py +++ b/tests/v1/worker/test_gpu_rejection_sampler_i64.py @@ -15,6 +15,7 @@ import torch from vllm.platforms import current_platform +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample VOCAB_SIZE = 155264 @@ -42,14 +43,16 @@ def _run( num_logits, VOCAB_SIZE, dtype=torch.bfloat16, device=device ) full_target[roi_row : roi_row + LOGITS_PER_REQ] = target_logits - full_draft = torch.randn( + full_draft = DraftLogitsCache( max_num_reqs, NUM_SPECULATIVE_STEPS, VOCAB_SIZE, dtype=torch.bfloat16, + fill=0.0, device=device, ) - full_draft[req_state_idx] = draft_logits + full_draft.values.normal_() + full_draft.values[req_state_idx] = draft_logits # Draft token for step i of request r lives at row r*LOGITS_PER_REQ + i + 1. draft_sampled = torch.zeros(num_logits, dtype=torch.int64, device=device) @@ -80,7 +83,7 @@ def _run( sampled, num_sampled = rejection_sample( target_logits=full_target, - draft_logits=full_draft, + draft_logits_cache=full_draft, draft_sampled=draft_sampled, cu_num_logits=cu_num_logits, pos=pos, diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index dd1e90926d7a..cf76f772d017 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -37,6 +37,7 @@ repeated_context_mask, ) from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator @@ -1387,7 +1388,9 @@ def sample(logits, contexts, random_sampler=None, skip_mask=None): temperature=torch.ones(2), seeds=torch.zeros(2, dtype=torch.int64), draft_step=torch.tensor(0), - draft_logits=torch.zeros(2, 1, 8), + draft_logits=DraftLogitsCache( + 2, 1, 8, dtype=torch.float32, fill=0.0, device="cpu" + ), ) assert torch.equal(sampled, torch.tensor([7, 4])) @@ -1755,7 +1758,9 @@ def ordinary_sample(logits, idx_mapping, temperature, seed, pos, **kwargs): speculator.acceptance_estimator = None batch_slots = torch.tensor([2, 0, 1]) logits = torch.arange(num_reqs * num_steps * vocab_size, dtype=torch.float32) - draft_logits = torch.zeros(num_reqs, num_steps, vocab_size) + draft_logits = DraftLogitsCache( + num_reqs, num_steps, vocab_size, dtype=torch.float32, fill=0.0, device="cpu" + ) sampled = speculator.sample_draft( hidden_states=logits.view(-1, vocab_size), @@ -1769,7 +1774,9 @@ def ordinary_sample(logits, idx_mapping, temperature, seed, pos, **kwargs): ) assert torch.equal(sampled, torch.tensor([3, 4, 102, 103, 204, 205])) - assert torch.equal(draft_logits[batch_slots], logits.view(num_reqs, num_steps, -1)) + assert torch.equal( + draft_logits.values[batch_slots], logits.view(num_reqs, num_steps, -1) + ) @pytest.mark.skipif( @@ -1860,9 +1867,9 @@ def run(fused: bool): def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch): speculator = object.__new__(DSparkSpeculator) - speculator.draft_logits = torch.zeros(2, 1, 8) - speculator.draft_logits_num_shards = 1 - speculator.draft_logits_shard_idx = 0 + speculator.draft_logits = DraftLogitsCache( + 2, 1, 8, dtype=torch.float32, fill=0.0, device="cpu" + ) speculator._d2t_scatter_index = torch.tensor([1, 5]) speculator._draft_scatter_buf = torch.full((2, 8), float("-inf")) speculator.temperature = torch.ones(2) @@ -1906,9 +1913,9 @@ def sample( def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch): speculator = object.__new__(DSparkSpeculator) - speculator.draft_logits = torch.zeros(2, 1, 8) - speculator.draft_logits_num_shards = 1 - speculator.draft_logits_shard_idx = 0 + speculator.draft_logits = DraftLogitsCache( + 2, 1, 8, dtype=torch.float32, fill=0.0, device="cpu" + ) speculator._d2t_scatter_index = None speculator.temperature = torch.ones(2) speculator.seeds = torch.zeros(2, dtype=torch.int64) diff --git a/vllm/models/deepseek_v41/amd/ops/fused_markov.py b/vllm/models/deepseek_v41/amd/ops/fused_markov.py index 72408eeb96b3..aaf8e389019b 100644 --- a/vllm/models/deepseek_v41/amd/ops/fused_markov.py +++ b/vllm/models/deepseek_v41/amd/ops/fused_markov.py @@ -225,8 +225,8 @@ def _why_not(spec, num_reqs: int) -> str | None: return "Markov weights not bf16" if w1.shape[1] != w2.shape[1] or not (w1.is_contiguous() and w2.is_contiguous()): return "Markov weight layout" - if spec.draft_logits.dtype != torch.bfloat16: - return f"a {spec.draft_logits.dtype} draft logits cache" + if spec.draft_logits.values.dtype != torch.bfloat16: + return f"a {spec.draft_logits.values.dtype} draft logits cache" if head.markov_w2.quant_method.__class__.__name__ != "UnquantizedEmbeddingMethod": return "a quantized Markov head" return None @@ -260,7 +260,7 @@ def sample_sequential(spec, num_reqs: int, head_hidden: torch.Tensor) -> bool: idx_map = spec.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) sample_pos = spec.sample_pos[:num_sample].view(num_reqs, n_spec) assert idx_map.stride(1) == 1 and sample_pos.stride(1) == 1 - cache = spec.draft_logits + cache = spec.draft_logits.values tokens = spec.draft_tokens use_fp64 = bool(spec.use_fp64_gumbel) bufs = _scratch(base_logits.device, vocab, use_fp64) diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py index 6c432db959f3..05b8c255c9e3 100644 --- a/vllm/v1/watermarking/spec_decode.py +++ b/vllm/v1/watermarking/spec_decode.py @@ -14,6 +14,7 @@ Watermarker, ) from vllm.v1.worker.gpu.sample.watermark import draft_watermarking_mask +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample @@ -293,7 +294,7 @@ def speculative_target_watermark_key( def watermarked_rejection_sample( target_logits: torch.Tensor, - draft_logits: torch.Tensor | None, + draft_logits_cache: DraftLogitsCache | None, draft_sampled: torch.Tensor, cu_num_logits: torch.Tensor, pos: torch.Tensor, @@ -313,7 +314,7 @@ def watermarked_rejection_sample( assert contexts.shape == (target_logits.shape[0], watermarker.context_width) return rejection_sample( target_logits, - draft_logits, + draft_logits_cache, draft_sampled, cu_num_logits, pos, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c181dac5c2a0..eea9cbe0d042 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -454,6 +454,11 @@ def load_model( + self.model_state.num_new_sampled_tokens_per_step ) + draft_logits = ( + self.speculator.draft_logits + if isinstance(self.speculator, DraftModelSpeculator) + else None + ) if self.parallel_config.enable_batch_sharded_sampling: if batch_sharded_sampling_enabled(self.vllm_config, self.model): self.batch_sharder = BatchSharder( @@ -464,6 +469,7 @@ def load_model( self.speculative_config is not None and self.speculative_config.enable_adaptive_verification ), + draft_logits=draft_logits, ) logger.info("Batch-sharded sampling enabled.") else: @@ -521,13 +527,8 @@ def load_model( ) custom = self.model_state.custom_sampler(self.sampler) - draft_logits_num_shards = ( - self.speculator.draft_logits_num_shards - if isinstance(self.speculator, DraftModelSpeculator) - else 1 - ) if custom: - assert draft_logits_num_shards == 1, ( + assert draft_logits is None or draft_logits.row_mapping is None, ( "A custom rejection sampler cannot read batch-sharded draft logits" ) self.vllm_config._check_supports_watermarking(custom_sampler=True) @@ -540,7 +541,6 @@ def load_model( watermark_key=speculative_target_watermark_key( self.vllm_config.watermark_config ), - draft_logits_num_shards=draft_logits_num_shards, ) self.prompt_logprobs_worker = PromptLogprobsWorker( self.max_num_reqs, @@ -2324,16 +2324,12 @@ def sample_tokens( 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 + self.speculator.prepare_propose( + input_batch, num_sampled, num_rejected, self.sampler ) spec_hidden_states = get_drafter_hidden_states( self.model, draft_hidden_states ) - if isinstance(self.sampler, GPUWatermarkSampler): - self.speculator.prepare_watermarking( - self.sampler, input_batch.idx_mapping - ) with use_workspace_lane(self._draft_workspace_lane): draft_tokens = self.speculator.propose( input_batch, diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index a906993ec2bf..e2aa3e4ab14b 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache logger = init_logger(__name__) @@ -883,7 +884,7 @@ def sync_drafts( self, draft_tokens: torch.Tensor, idx_mapping: torch.Tensor, - draft_logits: torch.Tensor | None, + draft_logits: "DraftLogitsCache | None", ) -> None: """Broadcast PCP rank 0's drafts and draft logits. @@ -894,9 +895,9 @@ def sync_drafts( pcp_group.broadcast(draft_tokens, src=0) if draft_logits is not None: # Draft logits are indexed by persistent request slots. - logits = draft_logits[idx_mapping] + logits = draft_logits.values[idx_mapping] pcp_group.broadcast(logits, src=0) - draft_logits[idx_mapping] = logits + draft_logits.values[idx_mapping] = logits def restore_for_sampling( self, hidden_states: torch.Tensor, aux_hidden_states: list[torch.Tensor] | None diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py index daaf8806ff89..9c2da8d018c7 100644 --- a/vllm/v1/worker/gpu/sample/batch_shard.py +++ b/vllm/v1/worker/gpu/sample/batch_shard.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project 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 InputBatch from vllm.v1.worker.gpu.sample.output import SamplerOutput +if TYPE_CHECKING: + from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache + @dataclass class BatchShardMetadata: @@ -61,18 +65,32 @@ def batch_sharded_sampling_enabled( @triton.jit -def shard_owner(req_state_idx, NUM_SHARDS: tl.constexpr): - """TP rank that samples a request and holds its draft logits. +def shard_owner(req_idx, NUM_SHARDS: tl.constexpr): + """TP rank that samples the request at a batch position. - Mirrored on the host in `BatchSharder.shard_sampler_inputs`. + Mirrored on the host by `shard_owners`. """ - return req_state_idx % NUM_SHARDS + return req_idx % NUM_SHARDS + + +def shard_owners(num_reqs: int, num_shards: int) -> np.ndarray: + """[num_reqs] The `shard_owner` of each batch position.""" + return np.arange(num_reqs) % num_shards @triton.jit -def shard_row(req_state_idx, NUM_SHARDS: tl.constexpr): - """Row of a request in its owner's batch-sharded draft logits cache.""" - return req_state_idx // NUM_SHARDS +def shard_row(req_idx, NUM_SHARDS: tl.constexpr): + """Row of the request at a batch position in its `shard_owner`'s sharded + draft logits cache. + + Mirrored on the host by `shard_rows`. + """ + return req_idx // NUM_SHARDS + + +def shard_rows(num_reqs: int, num_shards: int) -> np.ndarray: + """[num_reqs] The `shard_row` of each batch position.""" + return np.arange(num_reqs) // num_shards @triton.jit( @@ -91,6 +109,7 @@ def _build_shard_plan_kernel( local_expanded_idx_mapping_ptr, local_expanded_local_pos_ptr, local_seq_lens_ptr, + owners_ptr, num_reqs, local_logits_start, max_num_reqs_per_rank, @@ -101,16 +120,22 @@ def _build_shard_plan_kernel( ): req_idx = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + req_idx) - owner = shard_owner(req_state_idx, TP_SIZE) num_logits = tl.load(cu_num_logits_ptr + req_idx + 1) - tl.load( cu_num_logits_ptr + req_idx ) req_block = tl.arange(0, PADDED_NUM_REQS) req_mask = req_block < num_reqs - owners = shard_owner( - tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0), TP_SIZE - ) + if owners_ptr is None: + owner = shard_owner(req_idx, TP_SIZE) + owners = shard_owner(req_block, TP_SIZE) + else: + owner = tl.load(owners_ptr + req_state_idx) + owners = tl.load( + owners_ptr + tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0), + mask=req_mask, + other=0, + ) num_logits_per_req = tl.load( cu_num_logits_ptr + req_block + 1, mask=req_mask, other=0 ) - tl.load(cu_num_logits_ptr + req_block, mask=req_mask, other=0) @@ -181,12 +206,14 @@ def __init__( max_num_logits_per_req: int, device: torch.device, adaptive_verification_enabled: bool, + draft_logits: "DraftLogitsCache | None" = None, ): tp_group = get_tp_group() self.tp_rank = tp_group.rank_in_group self.tp_size = tp_group.world_size self.device = device self.adaptive_verification_enabled = adaptive_verification_enabled + self.draft_logits = draft_logits self._padded_num_reqs = triton.next_power_of_2(max_num_reqs) self._padded_num_logits_per_req = triton.next_power_of_2(max_num_logits_per_req) self._num_warps = max(1, min(8, self._padded_num_reqs // 128)) @@ -208,11 +235,15 @@ def shard_sampler_inputs( tp_size = self.tp_size num_reqs = input_batch.idx_mapping_np.shape[0] - # Deterministically assign requests to ranks, round-robin over slot - # indices so ownership stays balanced when the slot allocator fills - # low slots first (partial occupancy). + # Deterministically assign requests to ranks: to the rank that cached + # their draft logits, or else round-robin over batch positions. # NOTE: This mirrors the assignment used in _build_shard_plan_kernel. - req_owner_np = input_batch.idx_mapping_np % tp_size + if self.draft_logits is None: + req_owner_np, req_owner = shard_owners(num_reqs, tp_size), None + else: + req_owner_np, req_owner = self.draft_logits.owners_of( + input_batch.idx_mapping_np + ) local_req_indices_np = np.flatnonzero(req_owner_np == tp_rank) local_idx_mapping_np = input_batch.idx_mapping_np[local_req_indices_np] num_local_reqs = local_req_indices_np.shape[0] @@ -282,6 +313,7 @@ def shard_sampler_inputs( local_expanded_idx_mapping, local_expanded_local_pos, local_seq_lens, + req_owner, num_reqs, local_logits_start, max_num_reqs_per_rank, diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 0c6503cbd6da..91ea3e1240f3 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -3,7 +3,7 @@ import torch from vllm.triton_utils import tl, triton -from vllm.v1.worker.gpu.sample.batch_shard import shard_owner, shard_row +from vllm.v1.worker.gpu.spec_decode.draft_logits import draft_logits_row # Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the # retained Philox helper for rejection sampling. @@ -229,13 +229,12 @@ def gumbel_block_argmax( logits_cache_col_ptr, logits_cache_source_ptr, logits_cache_source_stride, + logits_cache_row_mapping_ptr, vocab_size, IS_DRAFTING: tl.constexpr, APPLY_TEMPERATURE: tl.constexpr, USE_FP64: tl.constexpr, PER_TOKEN_COL: tl.constexpr = False, - CACHE_NUM_SHARDS: tl.constexpr = 1, - CACHE_SHARD_IDX: tl.constexpr = 0, ): req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64) is_valid_req = req_state_idx >= 0 @@ -255,16 +254,14 @@ def gumbel_block_argmax( logits_cache_source_ptr + token_idx * logits_cache_source_stride + block, mask=mask, ) + row = draft_logits_row(req_state_idx, logits_cache_row_mapping_ptr) tl.store( logits_cache_ptr - + shard_row(req_state_idx, CACHE_NUM_SHARDS) * logits_cache_stride_0 + + row * logits_cache_stride_0 + col * logits_cache_stride_1 + block, cached_logits, - mask=mask - & is_valid_req - # A batch-sharded cache holds only this rank's requests, packed by row. - & (shard_owner(req_state_idx, CACHE_NUM_SHARDS) == CACHE_SHARD_IDX), + mask=mask & (row >= 0), ) seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0) @@ -297,6 +294,7 @@ def _gumbel_sample_kernel( logits_cache_col_ptr, logits_cache_source_ptr, logits_cache_source_stride, + logits_cache_row_mapping_ptr, logits_ptr, logits_stride, expanded_idx_mapping_ptr, @@ -309,8 +307,6 @@ def _gumbel_sample_kernel( APPLY_TEMPERATURE: tl.constexpr, USE_FP64: tl.constexpr, PER_TOKEN_COL: tl.constexpr, - CACHE_NUM_SHARDS: tl.constexpr, - CACHE_SHARD_IDX: tl.constexpr, ): token_idx = tl.program_id(0).to(tl.int64) block_idx = tl.program_id(1) @@ -338,13 +334,12 @@ def _gumbel_sample_kernel( logits_cache_col_ptr, logits_cache_source_ptr, logits_cache_source_stride, + logits_cache_row_mapping_ptr, vocab_size, IS_DRAFTING=IS_DRAFTING, APPLY_TEMPERATURE=APPLY_TEMPERATURE, USE_FP64=USE_FP64, PER_TOKEN_COL=PER_TOKEN_COL, - CACHE_NUM_SHARDS=CACHE_NUM_SHARDS, - CACHE_SHARD_IDX=CACHE_SHARD_IDX, ) token_id = block_idx * BLOCK_SIZE + idx tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id) @@ -363,8 +358,7 @@ def gumbel_sample( logits_cache_col: torch.Tensor | None = None, # scalar or [num_tokens] use_fp64: bool = False, logits_cache_source: torch.Tensor | None = None, - logits_cache_num_shards: int = 1, - logits_cache_shard_idx: int = 0, + logits_cache_row_mapping: torch.Tensor | None = None, # [max_num_reqs] ) -> torch.Tensor: # Enforce contiguity on non-strided input tensors expanded_idx_mapping = expanded_idx_mapping.contiguous() @@ -410,6 +404,7 @@ def gumbel_sample( logits_cache_col, logits_cache_source, logits_cache_source.stride(0) if logits_cache_source is not None else 0, + logits_cache_row_mapping, logits, logits.stride(0), expanded_idx_mapping, @@ -422,8 +417,6 @@ def gumbel_sample( APPLY_TEMPERATURE=apply_temperature, USE_FP64=use_fp64, PER_TOKEN_COL=per_token_col, - CACHE_NUM_SHARDS=logits_cache_num_shards, - CACHE_SHARD_IDX=logits_cache_shard_idx, ) # NOTE(woosuk): Use int64 for later indexing. max_block_idx = local_max.argmax(dim=-1, keepdim=True) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index 198da76a06cb..e0d12efcabb8 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -8,9 +8,12 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.triton_utils import tl, triton -from vllm.v1.worker.gpu.sample.batch_shard import shard_owner, shard_row from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax, spec_step_seed from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator +from vllm.v1.worker.gpu.spec_decode.draft_logits import ( + DraftLogitsCache, + draft_logits_row, +) @triton.jit @@ -83,31 +86,24 @@ def _cache_draft_logits_kernel( candidate_ptr, scores_ptr, req_state_ptr, + draft_logits_row_mapping_ptr, draft_logits_stride_0, draft_logits_stride_1, num_steps: tl.constexpr, top_k: tl.constexpr, BLOCK_K: tl.constexpr, - CACHE_NUM_SHARDS: tl.constexpr, - CACHE_SHARD_IDX: tl.constexpr, ): flat = tl.program_id(0) req_state = tl.load(req_state_ptr + flat) step = flat % num_steps offsets = tl.arange(0, BLOCK_K) - mask = ( - (req_state >= 0) - & (offsets < top_k) - # A batch-sharded cache holds only this rank's requests, packed by row. - & (shard_owner(req_state, CACHE_NUM_SHARDS) == CACHE_SHARD_IDX) - ) + row = draft_logits_row(req_state, draft_logits_row_mapping_ptr) + mask = (row >= 0) & (offsets < top_k) candidate_base = flat * top_k - cache_base = (req_state * num_steps + step) * top_k + cache_base = (row * num_steps + step) * top_k old_token_ids = tl.load(cached_candidate_ptr + cache_base + offsets, mask=mask) logits_base = ( - draft_logits_ptr - + shard_row(req_state, CACHE_NUM_SHARDS) * draft_logits_stride_0 - + step * draft_logits_stride_1 + draft_logits_ptr + row * draft_logits_stride_0 + step * draft_logits_stride_1 ) tl.store(logits_base + old_token_ids, -float("inf"), mask=mask) token_ids = tl.load(candidate_ptr + candidate_base + offsets, mask=mask) @@ -141,10 +137,8 @@ def sample( temperature: torch.Tensor, seeds: torch.Tensor, draft_tokens: torch.Tensor, - draft_logits: torch.Tensor | None, + draft_logits: DraftLogitsCache | None, use_fp64: bool, - draft_logits_num_shards: int = 1, - draft_logits_shard_idx: int = 0, ) -> None: candidate_ids = candidate_ids.contiguous() sample_pos = sample_pos.contiguous() @@ -173,8 +167,6 @@ def sample( num_reqs * self.num_steps, idx_mapping, draft_logits, - draft_logits_num_shards, - draft_logits_shard_idx, ) def _cache_draft_logits( @@ -182,24 +174,21 @@ def _cache_draft_logits( candidate_ids: torch.Tensor, num_sample: int, idx_mapping: torch.Tensor, - draft_logits: torch.Tensor, - draft_logits_num_shards: int, - draft_logits_shard_idx: int, + draft_logits: DraftLogitsCache, ) -> None: block_k = triton.next_power_of_2(self.top_k) _cache_draft_logits_kernel[(num_sample,)]( - draft_logits, + draft_logits.values, self.cached_candidate_ids, candidate_ids, self.scores, idx_mapping, - draft_logits.stride(0), - draft_logits.stride(1), + draft_logits.row_mapping, + draft_logits.values.stride(0), + draft_logits.values.stride(1), num_steps=self.num_steps, top_k=self.top_k, BLOCK_K=block_k, - CACHE_NUM_SHARDS=draft_logits_num_shards, - CACHE_SHARD_IDX=draft_logits_shard_idx, num_warps=1, ) @@ -261,8 +250,6 @@ def _generate_draft( self.draft_tokens, self.draft_logits, self.use_fp64_gumbel, - self.draft_logits_num_shards, - self.draft_logits_shard_idx, ) if self.enable_adaptive_verification: self._maybe_predict_acceptance( diff --git a/vllm/v1/worker/gpu/spec_decode/draft_logits.py b/vllm/v1/worker/gpu/spec_decode/draft_logits.py new file mode 100644 index 000000000000..1bebcf69b75b --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/draft_logits.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np +import torch + +from vllm.distributed.parallel_state import get_tp_group +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.sample.batch_shard import ( + shard_owner, + shard_owners, + shard_row, + shard_rows, +) + + +@triton.jit +def draft_logits_row(req_state_idx, row_mapping_ptr): + """Row of a request in this rank's draft logits cache, or -1 if no row. + + ``row_mapping_ptr`` is None for a replicated cache, which is indexed by request + state. Otherwise it is `DraftLogitsCache.row_mapping`. + """ + row = req_state_idx + if row_mapping_ptr is not None: + row = tl.load( + row_mapping_ptr + req_state_idx, mask=req_state_idx >= 0, other=-1 + ) + return row.to(tl.int64) + + +@triton.jit(do_not_specialize=["num_reqs"]) +def _assign_rows_kernel( + idx_mapping_ptr, + shard_mapping_ptr, + row_mapping_ptr, + inverse_row_mapping_ptr, + num_reqs, + NUM_SHARDS: tl.constexpr, + SHARD_IDX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + req_block_idx = tl.program_id(0) + req_block = req_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = req_block < num_reqs + req_state_block = tl.load(idx_mapping_ptr + req_block, mask=mask, other=0) + owners = shard_owner(req_block, NUM_SHARDS) + rows = shard_row(req_block, NUM_SHARDS) + is_local = owners == SHARD_IDX + tl.store(shard_mapping_ptr + req_state_block, owners, mask=mask) + tl.store(row_mapping_ptr + req_state_block, tl.where(is_local, rows, -1), mask=mask) + tl.store(inverse_row_mapping_ptr + rows, req_state_block, mask=mask & is_local) + + +@triton.jit(do_not_specialize=["num_logits"]) +def _stale_drafts_kernel( + stale_ptr, + expanded_idx_mapping_ptr, + expanded_local_pos_ptr, + row_mapping_ptr, + inverse_row_mapping_ptr, + num_logits, + BLOCK_SIZE: tl.constexpr, +): + logit_block_idx = tl.program_id(0) + logit_block = logit_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = logit_block < num_logits + req_state_block = tl.load( + expanded_idx_mapping_ptr + logit_block, mask=mask, other=-1 + ) + local_pos = tl.load(expanded_local_pos_ptr + logit_block, mask=mask, other=0) + rows = draft_logits_row(req_state_block, row_mapping_ptr) + holders = tl.load(inverse_row_mapping_ptr + rows, mask=mask & (rows >= 0), other=-1) + # Position 0 is the bonus token, which has no draft. + stale = (local_pos > 0) & ((rows < 0) | (holders != req_state_block)) + tl.store(stale_ptr + logit_block, stale, mask=mask) + + +class DraftLogitsCache: + """Pre-temperature draft logits, cached by a draft step for the next + verification. + + A replicated cache holds every request on every TP rank, at row + req_state_idx. + + A sharded cache (batch-sharded sampling) holds on each TP rank only the + requests that rank verifies. Every draft step gives the request at batch + position i to rank i % tp, at row i // tp of its cache, so no rank caches + more than cdiv(num_reqs, tp) requests however the request slots are spread. + + Rows are reused every step. A request that the scheduler skips for a step + keeps its drafts, but its row may since have been given to another + request. `inverse_row_mapping` detects this, and `stale_drafts` marks those + drafts invalid, so verification rejects them and samples from the target. + """ + + def __init__( + self, + max_num_reqs: int, + num_steps: int, + vocab_size: int, + *, + dtype: torch.dtype, + fill: float, + device: torch.device | str, + sharded: bool = False, + ): + num_rows = max_num_reqs + # [max_num_reqs] Each request's row in values, -1 if another rank + # caches it. None for a replicated cache. + self.row_mapping: torch.Tensor | None = None + if sharded: + tp_group = get_tp_group() + self._num_shards = tp_group.world_size + self._shard_idx = tp_group.rank_in_group + num_rows = cdiv(max_num_reqs, self._num_shards) + # Requests never drafted keep a slot-based owner. They have no + # drafts, so any owner works as long as every rank agrees. + slots = np.arange(max_num_reqs, dtype=np.int32) + # [max_num_reqs] The TP rank that caches, and verifies, each request. + self.shard_mapping_np = slots % self._num_shards + self.shard_mapping = torch.from_numpy(self.shard_mapping_np).to(device) + rows_np = np.where( + self.shard_mapping_np == self._shard_idx, + shard_rows(max_num_reqs, self._num_shards), + -1, + ) + self.row_mapping = torch.from_numpy(rows_np.astype(np.int32)).to(device) + # [num_rows] The request last drafted into each row of this rank. + self.inverse_row_mapping = torch.full( + (num_rows,), -1, dtype=torch.int32, device=device + ) + # [num_rows, num_steps, vocab_size] + self.values = torch.full( + (num_rows, num_steps, vocab_size), fill, dtype=dtype, device=device + ) + + def assign(self, idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray) -> None: + """Give a sharded cache's rows out by batch position. Call after + verification and before the draft step writes the draft logits.""" + num_reqs = idx_mapping_np.shape[0] + if self.row_mapping is None or num_reqs == 0: + return + self.shard_mapping_np[idx_mapping_np] = shard_owners(num_reqs, self._num_shards) + block_size = 1024 + _assign_rows_kernel[(triton.cdiv(num_reqs, block_size),)]( + idx_mapping, + self.shard_mapping, + self.row_mapping, + self.inverse_row_mapping, + num_reqs, + NUM_SHARDS=self._num_shards, + SHARD_IDX=self._shard_idx, + BLOCK_SIZE=block_size, + ) + + def owners_of( + self, idx_mapping_np: np.ndarray + ) -> tuple[np.ndarray, torch.Tensor | None]: + """The TP rank that verifies each of a batch's requests: the rank that + cached its draft logits, or for a replicated cache, round-robin over + batch positions. + + Returns: + [num_reqs] each request's owner, and the [max_num_reqs] device owner + table, or None for a replicated cache. + + """ + if self.row_mapping is None: + num_reqs = idx_mapping_np.shape[0] + return shard_owners(num_reqs, get_tp_group().world_size), None + return self.shard_mapping_np[idx_mapping_np], self.shard_mapping + + def stale_drafts(self, input_batch: InputBatch) -> torch.Tensor | None: + """Mark the drafts whose draft logits were overwritten since drafting. + Only valid for requests this rank owns. + + Returns: + [num_logits] True for each stale draft, or None for a replicated + cache, whose rows are never reused. + + """ + if self.row_mapping is None: + return None + expanded_idx_mapping = input_batch.expanded_idx_mapping + num_logits = expanded_idx_mapping.shape[0] + stale = expanded_idx_mapping.new_empty(num_logits, dtype=torch.bool) + if num_logits == 0: + return stale + block_size = 1024 + _stale_drafts_kernel[(triton.cdiv(num_logits, block_size),)]( + stale, + expanded_idx_mapping, + input_batch.expanded_local_pos, + self.row_mapping, + self.inverse_row_mapping, + num_logits, + BLOCK_SIZE=block_size, + ) + return stale diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index d2425ad95a87..a40aa9fab0b6 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -148,7 +148,7 @@ def _sample_logits( sample_pos - 1, apply_temperature=True, is_drafting=True, - logits_cache=self.draft_logits, + logits_cache=self.draft_logits.values, logits_cache_col=self._step_cols[step], use_fp64=self.use_fp64_gumbel, ) @@ -160,11 +160,10 @@ def _sample_logits( sample_pos - 1, apply_temperature=True, is_drafting=True, - logits_cache=self.draft_logits, + logits_cache=self.draft_logits.values, logits_cache_col=self._step_cols[step], use_fp64=self.use_fp64_gumbel, - logits_cache_num_shards=self.draft_logits_num_shards, - logits_cache_shard_idx=self.draft_logits_shard_idx, + logits_cache_row_mapping=self.draft_logits.row_mapping, ) def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index cce3dfbba319..2ee4bf402126 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( rejection_sample, ) @@ -152,7 +153,6 @@ def __init__( device: torch.device, *, watermark_key: int | None = None, - draft_logits_num_shards: int = 1, ): self.sampler = sampler self.watermark_key = watermark_key @@ -165,9 +165,6 @@ def __init__( self.fly_window_size = 0 self.fly_entropy_threshold = spec_config.fly_entropy_threshold self.fly_entropy_top_k = spec_config.fly_entropy_top_k - # Above 1 when each TP rank caches only the draft logits of the - # requests it owns (see DraftModelSpeculator). - self.draft_logits_num_shards = draft_logits_num_shards self.synthetic_conditional_rates: torch.Tensor | None = None if rejection_sample_method == "synthetic": assert spec_config.synthetic_acceptance_rates is not None @@ -254,7 +251,7 @@ def _watermarking_kwargs( def _verify( self, logits: torch.Tensor, - draft_logits: torch.Tensor | None, + draft_logits: DraftLogitsCache | None, draft_sampled: torch.Tensor, pos: torch.Tensor, cu_num_logits: torch.Tensor, @@ -293,7 +290,6 @@ def _verify( fly_window_size=self.fly_window_size, fly_entropy_threshold=self.fly_entropy_threshold, fly_entropy_top_k=self.fly_entropy_top_k, - draft_logits_num_shards=self.draft_logits_num_shards, **self._watermarking_kwargs( draft_sampled, expanded_idx_mapping, expanded_local_pos ), @@ -308,7 +304,7 @@ def _verify_in_chunks( self, logits: torch.Tensor, input_batch: InputBatch, - draft_logits: torch.Tensor | None, + draft_logits: DraftLogitsCache | None, draft_sampled: torch.Tensor, pos: torch.Tensor, max_chunk_logits: int, @@ -426,7 +422,7 @@ def __call__( self, logits: torch.Tensor, input_batch: InputBatch, - draft_logits: torch.Tensor | None = None, + draft_logits: DraftLogitsCache | None = None, invalid_drafts: torch.Tensor | None = None, ) -> SamplerOutput: # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear @@ -446,6 +442,17 @@ def __call__( # accepted length before it reaches a permissive bitmask row. Only # verification sees the copy: `apply_sampling_params` and watermarking # keep the real draft ids. + # A request skipped since its draft step may find its draft logits row + # reused by another request. Reject those drafts. + stale_drafts = ( + draft_logits.stale_drafts(input_batch) if draft_logits is not None else None + ) + if stale_drafts is not None: + invalid_drafts = ( + stale_drafts + if invalid_drafts is None + else invalid_drafts | stale_drafts + ) verify_draft_sampled = draft_sampled if invalid_drafts is not None: verify_draft_sampled = draft_sampled.masked_fill(invalid_drafts, -1) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 0b0139fbba28..af086b0508f0 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -4,13 +4,16 @@ from vllm.triton_utils import tl, tldevice, triton from vllm.v1.spec_decode.fly import compute_fly_entropy -from vllm.v1.worker.gpu.sample.batch_shard import shard_row from vllm.v1.worker.gpu.sample.gumbel import ( gumbel_block_argmax, spec_step_seed, tl_rand32, ) from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax +from vllm.v1.worker.gpu.spec_decode.draft_logits import ( + DraftLogitsCache, + draft_logits_row, +) @triton.jit @@ -142,6 +145,7 @@ def _compute_global_logprobs_and_logsumexp( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping_ptr, # [num_logits, num_blocks] draft_local_max_ptr, draft_local_max_stride, @@ -150,7 +154,6 @@ def _compute_global_logprobs_and_logsumexp( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, - DRAFT_LOGITS_NUM_SHARDS: tl.constexpr, ): target_logit = tl.load( target_logits_ptr + logit_idx * target_logits_stride + token, @@ -172,7 +175,9 @@ def _compute_global_logprobs_and_logsumexp( draft_logit = ( tl.load( draft_logits_ptr - + shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS) + + tl.maximum( + draft_logits_row(req_state_idx, draft_logits_row_mapping_ptr), 0 + ) * draft_logits_stride_0 + draft_step * draft_logits_stride_1 + token, @@ -232,6 +237,7 @@ def _compute_local_logits_stats_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping_ptr, # [num_logits] expanded_idx_mapping_ptr, # [num_logits] @@ -242,7 +248,6 @@ def _compute_local_logits_stats_kernel( num_logits, BLOCK_SIZE: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, - DRAFT_LOGITS_NUM_SHARDS: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) @@ -301,7 +306,9 @@ def _compute_local_logits_stats_kernel( draft_logits = ( tl.load( draft_logits_ptr - + shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS) + + tl.maximum( + draft_logits_row(req_state_idx, draft_logits_row_mapping_ptr), 0 + ) * draft_logits_stride_0 + draft_step_idx * draft_logits_stride_1 + block_offsets, @@ -342,6 +349,7 @@ def _compute_cumulative_log_p_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping_ptr, # [num_logits, num_blocks] draft_local_max_ptr, draft_local_max_stride, @@ -357,7 +365,6 @@ def _compute_cumulative_log_p_kernel( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, - DRAFT_LOGITS_NUM_SHARDS: tl.constexpr, ): req_idx = tl.program_id(0) req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64) @@ -392,6 +399,7 @@ def _compute_cumulative_log_p_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping_ptr, draft_local_max_ptr, draft_local_max_stride, draft_local_sumexp_ptr, @@ -399,7 +407,6 @@ def _compute_cumulative_log_p_kernel( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS, HAS_DRAFT_LOGITS, - DRAFT_LOGITS_NUM_SHARDS, ) ) log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0) @@ -426,6 +433,7 @@ def _compute_local_residual_mass_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping_ptr, # [num_logits, num_blocks] draft_local_max_ptr, draft_local_max_stride, @@ -445,7 +453,6 @@ def _compute_local_residual_mass_kernel( vocab_num_blocks, BLOCK_SIZE: tl.constexpr, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, - DRAFT_LOGITS_NUM_SHARDS: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) # Adaptive verification can end before num_speculative_steps. @@ -487,6 +494,7 @@ def _compute_local_residual_mass_kernel( draft_logits_ptr, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping_ptr, draft_local_max_ptr, draft_local_max_stride, draft_local_sumexp_ptr, @@ -494,7 +502,6 @@ def _compute_local_residual_mass_kernel( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS, True, # HAS_DRAFT_LOGITS - DRAFT_LOGITS_NUM_SHARDS, ) # Compute the residual mass: max(p_i * M_b(x|x_{ tuple[torch.Tensor, torch.Tensor]: assert target_logits.ndim == 2 and target_logits.stride(-1) == 1 + draft_logits = None + draft_logits_row_mapping = None + if draft_logits_cache is not None: + draft_logits = draft_logits_cache.values + draft_logits_row_mapping = draft_logits_cache.row_mapping assert draft_logits is None or ( draft_logits.ndim == 3 and draft_logits.stride(-1) == 1 ) @@ -1225,6 +1236,7 @@ def rejection_sample( draft_logits, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping, expanded_idx_mapping, expanded_local_pos, temperature, @@ -1232,7 +1244,6 @@ def rejection_sample( num_logits, BLOCK_SIZE=VOCAB_BLOCK_SIZE, HAS_DRAFT_LOGITS=has_draft_logits, - DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards, ) # Precompute the running joint ratio and residual mass for block @@ -1258,6 +1269,7 @@ def rejection_sample( draft_logits, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping, draft_local_max, draft_local_max.stride(0), draft_local_sumexp, @@ -1268,7 +1280,6 @@ def rejection_sample( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks, HAS_DRAFT_LOGITS=has_draft_logits, - DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards, num_warps=1, ) @@ -1293,6 +1304,7 @@ def rejection_sample( draft_logits, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping, draft_local_max, draft_local_max.stride(0), draft_local_sumexp, @@ -1306,7 +1318,6 @@ def rejection_sample( vocab_num_blocks, BLOCK_SIZE=VOCAB_BLOCK_SIZE, PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks, - DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards, ) else: local_residual_mass = None @@ -1340,6 +1351,7 @@ def rejection_sample( draft_logits, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping, draft_local_max, draft_local_max.stride(0), draft_local_sumexp, @@ -1358,7 +1370,6 @@ def rejection_sample( vocab_num_blocks, PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks, HAS_DRAFT_LOGITS=has_draft_logits, - DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards, SYNTHETIC_MODE=synthetic_conditional_rates is not None, USE_BLOCK_VERIFICATION=use_block_verification, FLY_WINDOW_SIZE=fly_window_size, @@ -1388,6 +1399,7 @@ def rejection_sample( draft_logits, draft_logits_stride_0, draft_logits_stride_1, + draft_logits_row_mapping, draft_rejected_logsumexp, num_sampled, cu_num_logits, @@ -1406,7 +1418,6 @@ def rejection_sample( vocab_size, BLOCK_SIZE=RESAMPLE_BLOCK_SIZE, HAS_DRAFT_LOGITS=has_draft_logits, - DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards, USE_FP64=use_fp64, USE_BLOCK_VERIFICATION=use_block_verification, CONTEXT_WIDTH=context_width, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index b3d520b9dc19..5d29b207b7fc 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -12,13 +12,12 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.config.compilation import CUDAGraphMode from vllm.distributed.eplb.eplb_state import EplbState -from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.models import supports_multimodal_embeddings -from vllm.utils.math_utils import cdiv from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.watermarking import create_watermarker +from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler from vllm.v1.watermarking.spec_decode import ( DraftWatermarker, create_speculative_draft_watermarker, @@ -39,11 +38,12 @@ from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import ( OnlineAcceptanceEstimator, ) +from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache from vllm.v1.worker.utils import AttentionGroup if TYPE_CHECKING: - from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler from vllm.v1.worker.gpu.pcp_manager import PCPManager + from vllm.v1.worker.gpu.sample.sampler import Sampler logger = init_logger(__name__) @@ -188,12 +188,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # Pre-temperature logits, cached from the previous decode step. # Allocated in load_model. See _init_draft_logits. - self.draft_logits: torch.Tensor | None = None + self.draft_logits: DraftLogitsCache | None = None self.use_draft_logits = ( self.speculative_config.draft_sample_method == "probabilistic" ) - self.draft_logits_num_shards = 1 - self.draft_logits_shard_idx = 0 self.draft_watermarker: DraftWatermarker | None = None watermark_config = getattr(vllm_config, "watermark_config", None) @@ -274,7 +272,7 @@ def load_model(self, target_model: nn.Module) -> None: def supports_batch_sharded_draft_logits(self) -> bool: """Whether every draft logits write on this speculator's path honors - draft_logits_num_shards and draft_logits_shard_idx.""" + DraftLogitsCache.row_mapping.""" return False def _init_draft_logits(self, target_model: nn.Module) -> None: @@ -282,31 +280,25 @@ def _init_draft_logits(self, target_model: nn.Module) -> None: Under batch-sharded sampling, the rank that verifies a request is the only one that reads its draft logits, so each TP rank keeps just the - rows of the requests it owns, at row - req_state_idx // draft_logits_num_shards. + rows of the requests it owns. See DraftLogitsCache. """ parallel_config = self.vllm_config.parallel_config - if ( - batch_sharded_sampling_enabled(self.vllm_config, target_model) - # PCPManager.sync_drafts indexes the draft logits by request state, - # so it must see every request's logits. - and parallel_config.prefill_context_parallel_size == 1 - and self.draft_watermarker is None - and self.supports_batch_sharded_draft_logits() - ): - tp_group = get_tp_group() - self.draft_logits_num_shards = tp_group.world_size - self.draft_logits_shard_idx = tp_group.rank_in_group dtype, fill = self.draft_logits_spec(self.vllm_config) - self.draft_logits = torch.full( - ( - cdiv(self.max_num_reqs, self.draft_logits_num_shards), - self.num_speculative_steps, - self.vocab_size, - ), - fill, + self.draft_logits = DraftLogitsCache( + self.max_num_reqs, + self.num_speculative_steps, + self.vocab_size, dtype=dtype, + fill=fill, device=self.device, + sharded=( + batch_sharded_sampling_enabled(self.vllm_config, target_model) + # PCPManager.sync_drafts indexes the draft logits by request + # state, so it must see every request's logits. + and parallel_config.prefill_context_parallel_size == 1 + and self.draft_watermarker is None + and self.supports_batch_sharded_draft_logits() + ), ) def set_eplb_state(self, eplb_state: EplbState) -> None: @@ -527,7 +519,7 @@ def sample_draft( temperature: torch.Tensor, seeds: torch.Tensor, draft_step: torch.Tensor, - draft_logits: torch.Tensor | None, + draft_logits: DraftLogitsCache | None, spec_step_idx: int = 0, num_steps: int = 1, ) -> torch.Tensor: @@ -544,7 +536,7 @@ def sample_draft( sample_src_positions, apply_temperature=True, is_drafting=True, - logits_cache=draft_logits, + logits_cache=draft_logits.values, logits_cache_col=draft_step, use_fp64=self.use_fp64_gumbel, num_steps=num_steps, @@ -558,11 +550,10 @@ def sample_draft( sample_src_positions, apply_temperature=True, is_drafting=True, - logits_cache=draft_logits, + logits_cache=draft_logits.values, logits_cache_col=draft_step, use_fp64=self.use_fp64_gumbel, - logits_cache_num_shards=self.draft_logits_num_shards, - logits_cache_shard_idx=self.draft_logits_shard_idx, + logits_cache_row_mapping=draft_logits.row_mapping, ) else: sampled = logits.argmax(dim=-1) @@ -597,7 +588,7 @@ def observe_verification( def prepare_watermarking( self, - sampler: "GPUWatermarkSampler", + sampler: GPUWatermarkSampler, idx_mapping: torch.Tensor, ) -> None: if self.draft_watermarker is None: @@ -610,6 +601,24 @@ def prepare_watermarking( sampler.req_states.total_len.gpu, ) + def prepare_propose( + self, + input_batch: InputBatch, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + sampler: "Sampler", + ) -> None: + """Per-step work between verifying the last drafts and the next + `propose`.""" + self.observe_verification(input_batch.idx_mapping, num_sampled, num_rejected) + if isinstance(sampler, GPUWatermarkSampler): + self.prepare_watermarking(sampler, input_batch.idx_mapping) + if self.draft_logits is not None: + # Verification read the previous step's draft logits. + self.draft_logits.assign( + input_batch.idx_mapping, input_batch.idx_mapping_np + ) + def _copy_request_inputs( self, num_reqs: int, From 7ab756940a89a7464ebd6ab26c7c6a1f3b7d08da Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Thu, 29 Oct 2026 12:31:01 +0800 Subject: [PATCH 110/406] [Bugfix][DSv4.1] Budget prefill memory the startup profile misses (#63047) Signed-off-by: Canlin Co-authored-by: Claude Opus 5.5 --- tests/distributed/test_engram_dp_shard.py | 19 +++++-------- .../test_indexer_deepseek_v4_slot_mapping.py | 26 ++++++++++++++++++ .../layers/sparse_attn_indexer.py | 27 ++++++++++++------- vllm/models/deepseek_v41/amd/model.py | 11 ++++---- vllm/models/deepseek_v41/nvidia/model.py | 11 ++++---- vllm/v1/attention/backends/mla/indexer.py | 9 ++++++- 6 files changed, 69 insertions(+), 34 deletions(-) diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 4a175e60098a..aa994bc11dad 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -576,7 +576,8 @@ def unexpected_collective(*args, **kwargs): 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.""" + """A metadata-less replica looks up dead rows, so it joins DP sharding and + the memory profile sees the Engram activations.""" from vllm.models.deepseek_v41.common.engram import NgramHashState if current_platform.is_rocm(): @@ -603,17 +604,11 @@ class Decoder(SimpleNamespace): # **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 - output = hidden.new_zeros( - tokens, len(head_sizes), DIM, dtype=torch.bfloat16 - ) - else: - assert hashes is not None and keep is not None - assert bool(keep.all()) == (dp_rank != 1) - if dp_rank == 1: - assert torch.all(hashes == engram_ops.DEAD_ID) - output = engram.embed(hashes[:, 0]) + assert hashes is not None and keep is not None + assert bool(keep.all()) == (dp_rank != 1) + if dp_rank == 1: + assert torch.all(hashes == engram_ops.DEAD_ID) + output = engram.embed(hashes[:, 0]) # Trailing None is the NVIDIA layer's previous_aux; this stub # captures no aux states. outputs = (output, hidden, hidden, hidden, hidden, None) 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 57164a4380ae..c3e86cecc4d1 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -668,3 +668,29 @@ def test_indexer_prefill_budget_matches_compressed_workspace(compress_ratio): assert len(chunks) > 1 for req_slice, _ in chunks: assert int(compressed_seq_lens[req_slice].sum()) <= workspace_rows + + +def test_indexer_expanded_block_table_sized_by_decode_tokens(): + """Allocated after the memory profile, so it must not scale with the token + budget: only decode tokens (one per request without drafts) are expanded.""" + kv_cache_spec = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + tokens_per_state=4, + ) + vllm_config = create_vllm_config( + max_model_len=1024, max_num_seqs=8, max_num_batched_tokens=8192 + ) + max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1024) + block_table_width = get_block_table_width(max_num_blocks, kv_cache_spec.block_size) + builder = DeepseekV32IndexerMetadataBuilder( + kv_cache_spec=kv_cache_spec, + layer_names=["dummy"], + vllm_config=vllm_config, + device=torch.device("cpu"), + block_table_width=block_table_width, + ) + + assert builder.expanded_block_table_buffer.shape == (8, block_table_width) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index eeedc9edb51f..e4c319abaad8 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -516,22 +516,27 @@ def sparse_attn_indexer( cfg.scheduler_config.max_num_batched_tokens, ) * (1 + cfg.num_speculative_tokens) - # Dummy allocation for the peak logits memory during inference: the - # prefill chunks' budget, or the decode logits, which eager and - # piecewise steps allocate outside the CUDA graph pools. DeepGEMM emits - # them in BF16 on SM100 and FP32 elsewhere, each row holding whole KV - # splits at a 1024-byte stride; the FP32 copy keeps that stride. + # Dummy allocation for the peak logits memory during inference: a + # prefill chunk's, or the decode logits, which eager and piecewise + # steps allocate outside the CUDA graph pools. DeepGEMM emits them in + # BF16 on SM100 and FP32 elsewhere, each row holding whole KV splits at + # a 1024-byte stride; the FP32 copy keeps that stride. Prefill always + # reads the FP32 copy, and the chunk splitter caps it at the budget. if current_platform.is_device_capability_family(100): logits_dtype, split_kv = torch.bfloat16, 384 else: logits_dtype, split_kv = torch.float32, 256 row_align = math.lcm(split_kv, 1024 // logits_dtype.itemsize) - elem_bytes = logits_dtype.itemsize - if logits_dtype != torch.float32 and fp32_decode_logits: - elem_bytes += torch.float32.itemsize + copy_bytes = 0 if logits_dtype == torch.float32 else torch.float32.itemsize + decode_copy_bytes = copy_bytes if fp32_decode_logits else 0 + prefill_budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 max_logits_bytes = max( - envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024, - max_decode_tokens * round_up(max_model_len, row_align) * elem_bytes, + prefill_budget + // torch.float32.itemsize + * (logits_dtype.itemsize + copy_bytes), + max_decode_tokens + * round_up(max_model_len, row_align) + * (logits_dtype.itemsize + decode_copy_bytes), ) _ = torch.empty( max_logits_bytes, dtype=torch.uint8, device=hidden_states.device @@ -830,6 +835,8 @@ def sparse_attn_indexer( cp_kv_cache_interleave_size, row_starts=row_starts, ) + # Free this chunk's logits before the next chunk allocates its own. + del logits if row_shard is not None: # Rows are already fully ranked, so this is a layout gather. diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 68c2f47ec3ef..78a6b9f3c910 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -11,7 +11,6 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import ( - get_engram_dp_size, get_pp_group, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, @@ -706,9 +705,8 @@ def forward( # Engram n-gram hashes for the whole (flattened) batch, computed once # on the full token stream — before any sequence-parallel sharding — - # and consumed by the engram layers (1 and 14) below. Skipped on - # profile runs (KV cache unbound -> no slot space to key the rolling - # compressed-id cache into). + # and consumed by the engram layers (1 and 14) below. Profile runs + # (KV cache unbound) look up dead rows instead. engram_hashes: torch.Tensor | None = None engram_mask: torch.Tensor | None = None if ( @@ -749,9 +747,10 @@ def forward( swa_metadata.slot_mapping, swa_metadata.block_table, ) - elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1: + else: # DP-sharded lookups are collective, so a replica skipping the - # hash still has to reach them. + # hash still has to reach them; the memory profile also has to + # see the Engram activations. engram_hashes, engram_mask = self.engram_hash.dummy_hashes(input_ids) if engram_hashes is not None: # Gather all Engram rows before entering the decoder layers. diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index fb694faf5ea7..100c25170112 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -13,7 +13,6 @@ from vllm.config import VllmConfig from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( - get_engram_dp_size, get_pcp_group, get_pp_group, get_tensor_model_parallel_world_size, @@ -583,6 +582,7 @@ def forward( engram_hashes[:, self.engram.layer_hash_index], engram_mask, ) + del previous_post post_mix, res_mix, x, attn_pre = mhc_pre( residual, self.hc_attn_fn, @@ -1000,8 +1000,8 @@ def forward( # Engram n-gram hashes for the whole (flattened) batch, computed once # on the full token stream — before any sequence-parallel sharding — - # and consumed by the engram layers (1 and 14) below. Skipped on - # profile runs (KV cache unbound). + # and consumed by the engram layers (1 and 14) below. Profile runs + # (KV cache unbound) look up dead rows instead. engram_hashes: torch.Tensor | None = None engram_mask: torch.Tensor | None = None if ( @@ -1042,9 +1042,10 @@ def forward( swa_metadata.slot_mapping, swa_metadata.block_table, ) - elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1: + else: # DP-sharded lookups are collective, so a replica skipping the - # hash still has to reach them. + # hash still has to reach them; the memory profile also has to + # see the Engram activations. engram_hashes, engram_mask = self.engram_hash.dummy_hashes(input_ids) if engram_hashes is not None: # Gather all Engram rows before entering the decoder layers. diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index c353a1f2fbb8..991b91ae2515 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1089,8 +1089,14 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: self.dcp_rank_tensor = torch.tensor( self.dcp_rank, dtype=torch.int32, device=self.device ) + # Only decode tokens are expanded: at most next_n per request. Allocated + # after the memory profile, so it must not scale with the token budget. + max_decode_tokens = min( + scheduler_config.max_num_seqs * next_n, + scheduler_config.max_num_batched_tokens, + ) self.expanded_block_table_buffer = torch.zeros( - (scheduler_config.max_num_batched_tokens, block_table_width), + (max_decode_tokens, block_table_width), dtype=torch.int32, device=self.device, ) @@ -1598,6 +1604,7 @@ def build( decode_metadata = None if num_decodes > 0: + assert num_decode_tokens <= self.expanded_block_table_buffer.shape[0] if not self.supports_varlen: torch.diff( common_attn_metadata.query_start_loc[: num_decodes + 1], From 28a5fba06bfcfee8907f580f4126e81bbd4ff4d4 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Thu, 29 Oct 2026 12:43:17 +0800 Subject: [PATCH 111/406] [Kernel] Update DeepGEMM commit (#64411) Signed-off-by: Jee Jee Li Co-authored-by: Codex --- cmake/external_projects/deepgemm.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 2fbfa157de88..e9d0bbbd2803 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -37,7 +37,7 @@ else() # (vllm-project/DeepGEMM#39), and the MX MegaMoE alignment fix # (vllm-project/DeepGEMM#43). # Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API. - set(_DEEPGEMM_UPSTREAM_TAG "0820030756945b2daec208e8e8255268360f337d") + set(_DEEPGEMM_UPSTREAM_TAG "5cdb145704ec4672ecf4b61441599521ce635df6") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) From 3966924d566340e971a6da313e790d05492aa505 Mon Sep 17 00:00:00 2001 From: Alphi <52458637+HwwwwwwwH@users.noreply.github.com> Date: Thu, 29 Oct 2026 12:46:24 +0800 Subject: [PATCH 112/406] [BugFix]/minicpm ocr downsample mode (#64105) Signed-off-by: zhangxueren Co-authored-by: Cursor Co-authored-by: zhangxueren --- .../multimodal/processing/test_minicpm_ocr.py | 52 ++++++++++++++++--- vllm/model_executor/models/minicpm_ocr.py | 3 ++ 2 files changed, 47 insertions(+), 8 deletions(-) diff --git a/tests/models/multimodal/processing/test_minicpm_ocr.py b/tests/models/multimodal/processing/test_minicpm_ocr.py index 8544112b4473..6f2a1296f38d 100644 --- a/tests/models/multimodal/processing/test_minicpm_ocr.py +++ b/tests/models/multimodal/processing/test_minicpm_ocr.py @@ -50,7 +50,12 @@ class _StubImageProcessor: slice_mode = True def find_best_resize( - self, image_size, scale_resolution, patch_size, allow_upscale=False + self, + image_size, + scale_resolution, + patch_size, + downsample_mode, + allow_upscale=False, ): height, width = image_size if (height * width > scale_resolution * scale_resolution) or allow_upscale: @@ -58,13 +63,19 @@ def find_best_resize( height = int(scale_resolution / math.sqrt(aspect_ratio)) width = int(height * aspect_ratio) # The OCR delta: one 2x2 merge in "4x" mode, two in "16x". - merge_factor = 2 if self.downsample_mode == "4x" else 4 + merge_factor = 2 if downsample_mode == "4x" else 4 best_width = ensure_divide(width, patch_size * merge_factor) best_height = ensure_divide(height, patch_size * merge_factor) return best_height, best_width def get_refine_size( - self, image_size, grid, scale_resolution, patch_size, allow_upscale=False + self, + image_size, + grid, + scale_resolution, + patch_size, + downsample_mode, + allow_upscale=False, ): height, width = image_size grid_y, grid_x = grid @@ -74,6 +85,7 @@ def get_refine_size( (refine_height / grid_y, refine_width / grid_x), scale_resolution, patch_size, + downsample_mode, allow_upscale=allow_upscale, ) return best_height * grid_y, best_width * grid_x @@ -167,10 +179,18 @@ def test_4x_mode_aligns_the_patch_grid_to_two_patches(): # Upscaled to ~448 on the long side, then aligned: "4x" rounds to 28px and # "16x" to 56px, which lands on a different grid for the same image. assert info._pinned_image_processor("4x").find_best_resize( - (200, 300), _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=True + (200, 300), + _SCALE_RESOLUTION, + _PATCH_SIZE, + downsample_mode="4x", + allow_upscale=True, ) == (364, 560) assert info._pinned_image_processor("16x").find_best_resize( - (200, 300), _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=True + (200, 300), + _SCALE_RESOLUTION, + _PATCH_SIZE, + downsample_mode="16x", + allow_upscale=True, ) == (392, 560) # 364*560 / (14*14*4) = 260 against 392*560 / (14*14*16) = 70: the ratio is @@ -227,7 +247,12 @@ def test_sliced_image_token_counts_match_the_transformers_divisor(): assert grids != [0, 0] processor = info._pinned_image_processor("4x") refine = processor.get_refine_size( - (1120, 1680), grids, _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=True + (1120, 1680), + grids, + _SCALE_RESOLUTION, + _PATCH_SIZE, + downsample_mode="4x", + allow_upscale=True, ) expected_patch = ( (refine[0] // grids[0]) @@ -453,12 +478,21 @@ def test_placeholder_matches_the_transformers_formula(width, height, mode): # Rebuild `target_sizes` exactly as `_preprocess` records it. sliced = grids != [0, 0] source_h, source_w = processor.find_best_resize( - (height, width), _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=not sliced + (height, width), + _SCALE_RESOLUTION, + _PATCH_SIZE, + downsample_mode=mode, + allow_upscale=not sliced, ) target_sizes = [[source_h // _PATCH_SIZE, source_w // _PATCH_SIZE]] if sliced: refine_h, refine_w = processor.get_refine_size( - (height, width), grids, _SCALE_RESOLUTION, _PATCH_SIZE, allow_upscale=True + (height, width), + grids, + _SCALE_RESOLUTION, + _PATCH_SIZE, + downsample_mode=mode, + allow_upscale=True, ) patch_h, patch_w = refine_h // grids[0], refine_w // grids[1] target_sizes += [[patch_h // _PATCH_SIZE, patch_w // _PATCH_SIZE]] * ( @@ -500,6 +534,7 @@ def test_pixel_and_target_size_paths_agree_on_the_default_resize(): (height, width), _SCALE_RESOLUTION, _PATCH_SIZE, + downsample_mode="4x", allow_upscale=not sliced, ) target_sizes = [[source_h // _PATCH_SIZE, source_w // _PATCH_SIZE]] @@ -509,6 +544,7 @@ def test_pixel_and_target_size_paths_agree_on_the_default_resize(): grids, _SCALE_RESOLUTION, _PATCH_SIZE, + downsample_mode="4x", allow_upscale=True, ) patch_h, patch_w = refine_h // grids[0], refine_w // grids[1] diff --git a/vllm/model_executor/models/minicpm_ocr.py b/vllm/model_executor/models/minicpm_ocr.py index 75d72896aa43..e737f64218c9 100644 --- a/vllm/model_executor/models/minicpm_ocr.py +++ b/vllm/model_executor/models/minicpm_ocr.py @@ -147,6 +147,7 @@ def _compute_visual_tokens( hf_image_size, scale_res, patch_size, + downsample_mode=image_processor.downsample_mode, allow_upscale=True, ) source_tokens = best_size[0] * best_size[1] // tokens_per_patch @@ -156,6 +157,7 @@ def _compute_visual_tokens( hf_image_size, scale_res, patch_size, + downsample_mode=image_processor.downsample_mode, ) source_tokens = best_resize[0] * best_resize[1] // tokens_per_patch @@ -164,6 +166,7 @@ def _compute_visual_tokens( grids, scale_res, patch_size, + downsample_mode=image_processor.downsample_mode, allow_upscale=True, ) patch_h = refine_size[0] // grids[0] From 1aea05364f4e24bfb249077880c6de6f5a2c21fd Mon Sep 17 00:00:00 2001 From: Sumanth R Hegde <39546518+SumanthRH@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:14:23 -0700 Subject: [PATCH 113/406] [Bugfix] DSA: all-reduce partial hidden_states before PP send (#60477) Signed-off-by: SumanthRH Signed-off-by: aoshen524 Co-authored-by: Claude Opus 5.5 Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: aoshen02 Co-authored-by: aoshen524 --- .../deepseek_v32/test_sequence_parallel.py | 37 +++++++++++++++++++ vllm/models/deepseek_v32/amd/model.py | 22 +++++++---- vllm/models/deepseek_v32/nvidia/model.py | 15 ++++---- vllm/models/dots3_note/nvidia/model.py | 2 +- 4 files changed, 59 insertions(+), 17 deletions(-) diff --git a/tests/models/deepseek_v32/test_sequence_parallel.py b/tests/models/deepseek_v32/test_sequence_parallel.py index c0378dd5e3c2..4ae7f699b681 100644 --- a/tests/models/deepseek_v32/test_sequence_parallel.py +++ b/tests/models/deepseek_v32/test_sequence_parallel.py @@ -9,6 +9,7 @@ from vllm.models.deepseek_v32.nvidia import model as deepseek_v32_model from vllm.models.deepseek_v32.nvidia import mtp as deepseek_v32_mtp +from vllm.sequence import IntermediateTensors class _IdentityNorm(nn.Module): @@ -146,3 +147,39 @@ def test_mtp_projects_sequence_shard_and_restores_full_output(monkeypatch): torch.testing.assert_close(hidden_states, expected) torch.testing.assert_close(recycled_hidden_states, expected) norm.assert_called_once() + + +class _PartialSumLayer(nn.Module): + """Returns a per-TP-rank partial MLP output, like DeepseekV32DecoderLayer.""" + + def forward(self, positions, hidden_states, residual, attn_in=None): + assert residual is None + return torch.ones_like(hidden_states), hidden_states + 1 + + +def test_pp_stage_sends_tp_reduced_residual_stream(monkeypatch): + """PP send/recv needs TP-replicated tensors, so a middle stage must reduce + its partial hidden_states and send the full residual stream only.""" + model = object.__new__(deepseek_v32_model.DeepseekV32Model) + nn.Module.__init__(model) + model.layers = nn.ModuleList([_PartialSumLayer()]) + model.start_layer, model.end_layer = 0, 1 + model.use_sequence_parallel = False + model.aux_hidden_state_layers = () + + pp_group = SimpleNamespace(is_first_rank=False, is_last_rank=False) + monkeypatch.setattr(deepseek_v32_model, "get_pp_group", lambda: pp_group) + monkeypatch.setattr( + deepseek_v32_model, + "tensor_model_parallel_all_reduce", + lambda tensor: tensor * 2, + ) + + received = torch.full((3, 2), 5.0) + output = model( + None, torch.arange(3), IntermediateTensors({"hidden_states": received}) + ) + + assert set(output.tensors) == {"hidden_states"} + # residual (5 + 1) + partial (1) all-reduced over 2 TP ranks. + torch.testing.assert_close(output["hidden_states"], torch.full((3, 2), 8.0)) diff --git a/vllm/models/deepseek_v32/amd/model.py b/vllm/models/deepseek_v32/amd/model.py index 2338b0d030ef..7dc43498fe57 100644 --- a/vllm/models/deepseek_v32/amd/model.py +++ b/vllm/models/deepseek_v32/amd/model.py @@ -9,7 +9,11 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig -from vllm.distributed import get_pp_group, get_tensor_model_parallel_rank +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + tensor_model_parallel_all_reduce, +) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -183,7 +187,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): else: self.norm = PPMissingLayer() self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size + ["hidden_states"], config.hidden_size ) self.aux_hidden_state_layers = tuple[int, ...]() @@ -207,11 +211,11 @@ def forward( else: assert input_ids is not None hidden_states = self.embed_input_ids(input_ids) - residual = None else: assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] + # Later PP stages receive the TP-reduced residual stream, like an embedding. + residual = None aux_hidden_states = [] for idx, layer in enumerate( @@ -219,13 +223,15 @@ def forward( start=self.start_layer, ): if idx in self.aux_hidden_state_layers: - aux_hidden_states.append(hidden_states + residual) + aux_hidden_states.append( + hidden_states if residual is None else hidden_states + residual + ) hidden_states, residual = layer(positions, hidden_states, residual) if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) + # PP send/recv needs TP-replicated tensors; reduce the partial sum. + hidden_states = residual + tensor_model_parallel_all_reduce(hidden_states) + return IntermediateTensors({"hidden_states": hidden_states}) hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm) if len(aux_hidden_states) > 0: diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 7720dfc5d182..fc2236542fcb 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -8,7 +8,7 @@ import vllm.envs as envs from vllm.config import VllmConfig -from vllm.distributed import get_pp_group +from vllm.distributed import get_pp_group, tensor_model_parallel_all_reduce from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_embed_norm import ( fused_embed_norm, @@ -229,7 +229,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): else: self.norm = PPMissingLayer() self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size + ["hidden_states"], config.hidden_size ) self.aux_hidden_state_layers = tuple[int, ...]() @@ -263,11 +263,11 @@ def forward( else: assert input_ids is not None hidden_states = self.embed_input_ids(input_ids) - residual = None else: assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] + # Later PP stages receive the TP-reduced residual stream, like an embedding. + residual = None full_num_tokens = positions.shape[0] if self.use_sequence_parallel: @@ -279,7 +279,6 @@ def forward( hidden_states = sp_shard(hidden_states) if attn_in is not None: attn_in = sp_shard(attn_in) - assert residual is None, "Currently, SP is not supported with PP" aux_hidden_states = [] for idx, layer in enumerate( @@ -297,9 +296,9 @@ def forward( assert not self.use_sequence_parallel, ( "Currently, SP is not supported with PP" ) - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) + # PP send/recv needs TP-replicated tensors; reduce the partial sum. + hidden_states = residual + tensor_model_parallel_all_reduce(hidden_states) + return IntermediateTensors({"hidden_states": hidden_states}) if self.use_sequence_parallel: hidden_states, _ = self.norm(hidden_states, residual) diff --git a/vllm/models/dots3_note/nvidia/model.py b/vllm/models/dots3_note/nvidia/model.py index 15128132befe..828c3bf54ba8 100644 --- a/vllm/models/dots3_note/nvidia/model.py +++ b/vllm/models/dots3_note/nvidia/model.py @@ -612,7 +612,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: else: self.norm = PPMissingLayer() self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size + ["hidden_states"], config.hidden_size ) self.aux_hidden_state_layers = tuple[int, ...]() self.num_redundant_experts = ( From 16777da3475084eda04813c98cfef4365d5dcb58 Mon Sep 17 00:00:00 2001 From: liranschour Date: Thu, 29 Oct 2026 07:23:13 +0200 Subject: [PATCH 114/406] [Bugfix][KV Offload][P2P] Fail a timed-out load at once, defer releasing its slots (#64256) Signed-off-by: Liran Schour --- .../v1/kv_offload/tiering/p2p/test_manager.py | 104 ++++++++++++ .../kv_offload/tiering/p2p/test_sessions.py | 98 ++++++++--- vllm/v1/kv_offload/tiering/p2p/manager.py | 10 +- .../kv_offload/tiering/p2p/session/client.py | 159 ++++++++++-------- .../kv_offload/tiering/p2p/session/session.py | 34 +++- 5 files changed, 310 insertions(+), 95 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 56611497bd4b..6e808a46b242 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -8,6 +8,7 @@ from __future__ import annotations +import threading import time import uuid from concurrent.futures import Future @@ -605,6 +606,79 @@ def capture_finished_jobs(): assert len(observations[("1:p2p", "success")]) == 1 +def test_timed_out_load_holds_chunks_until_abort_ack(monkeypatch): + """A timed-out P2P load fails at once but keeps its primary chunk out of + the free pool: the peer may still be writing it. The AbortAck frees it.""" + from tests.v1.kv_offload.tiering.p2p.test_sessions import ( + FakeConnection, + FakeDataTransport, + _activate, + ) + from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec + from vllm.v1.kv_offload.tiering.p2p.session import P2PSession + from vllm.v1.kv_offload.tiering.p2p.session import client as client_module + from vllm.v1.kv_offload.tiering.p2p.session.client import _PD_LOAD_TIMEOUT_S + from vllm.v1.kv_offload.tiering.p2p.session.protocol import AbortAckMsg + + now = [1000.0] + monkeypatch.setattr(client_module.time, "monotonic", lambda: now[0]) + peer = "10.0.0.1:8000" + p2p_tier = _make_manager() + p2p_tier._control = SimpleNamespace(poll=lambda: []) # type: ignore[assignment] + region = _mock_mmap_region(4) + primary = CPUPrimaryTierOffloadingManager(num_chunks=4, mmap_region=region) + tiering = TieringOffloadingManager( + primary_tier=primary, + secondary_tiers=[p2p_tier], + tier_poll_interval_s=0.0, + ) + conn = FakeConnection(peer_id=peer) + session = P2PSession( + peer_id=peer, + local_id=p2p_tier._local_id, + transport=FakeDataTransport(), + local_block_len=4096, + local_hash_seed="0", + conn=conn, + ) + _activate(session, conn, peer_id=peer) + p2p_tier._sessions[peer] = session + + key = OffloadKey(b"remote-block") + ctx = _req_context(_remote_prefiller_kv_params()) + prepared = primary.prepare_write([key], ctx) + assert prepared is not None + assert isinstance(prepared.store_spec, CPULoadStoreSpec) + job = TransferJob( + job_id=7, + keys=[key], + chunk_ids=prepared.store_spec.chunk_ids, + is_promotion=True, + req_context=ctx, + ) + tiering._register_job(job, tier_idx=0) + p2p_tier.submit_load(job) + free_before = primary._get_num_free_chunks() + + now[0] += _PD_LOAD_TIMEOUT_S + 1.0 + tiering._process_finished_jobs() + assert 7 not in tiering._jobs + assert primary.lookup(key, ctx) is LookupResult.MISS + assert primary.num_deferred_release_chunks == 1 + assert primary._get_num_free_chunks() == free_before + + conn.enqueue( + { + TYPE_KEY: AbortAckMsg.TYPE, + AbortAckMsg.KV_REQUEST_ID: "req-1", + AbortAckMsg.ROUND_SEQ: 0, + } + ) + tiering._process_finished_jobs() + assert primary.num_deferred_release_chunks == 0 + assert primary._get_num_free_chunks() == free_before + 1 + + # --------------------------------------------------------------------------- # Tests for on_request_finished # --------------------------------------------------------------------------- @@ -1660,6 +1734,36 @@ def remove_remote_peer(self_inner, pid): assert peer_dead not in mgr._sessions assert peer_alive in mgr._sessions + def test_failed_load_carries_its_deferred_release(self): + """A session's deferred release reaches the engine's JobResult, so the + TieringManager holds the load's chunks until the event is set.""" + mgr = _make_manager() + peer = "10.0.0.2:9000" + release = threading.Event() + mgr._sessions[peer] = _FakeSession( # type: ignore[assignment] + peer_id=peer, + loads=[ + LoadResult( + job_id=11, + kv_request_id="req-11", + success=False, + deferred_release=release, + ) + ], + ) + + class _Ctrl: + def poll(self_inner): + return [] + + mgr._control = _Ctrl() # type: ignore[assignment] + mgr._poll_once() + + (result,) = mgr._finished_jobs + assert (result.job_id, result.success) == (11, False) + assert result.deferred_release is release + assert "req-11" in mgr._failed_req_ids + def test_new_fetch_id_binds_and_replays_unbound_batches(self): """When session.poll() reports a kv_request_id whose FetchMsg arrived this tick, the manager binds it to that session and diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 368039003ce6..d45edf3954d4 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -326,6 +326,11 @@ def on_request_finished(self, ctx: ReqContext) -> None: self.calls.append(("on_request_finished", ctx.req_id)) +def _released(result: LoadResult) -> bool: + assert result.deferred_release is not None + return result.deferred_release.is_set() + + def _make_session( conn: FakeConnection | None = None, transport: FakeDataTransport | None = None, @@ -684,12 +689,17 @@ def test_finish_preserves_active_load_until_terminal_result(self): assert aborts[0][AbortFetchMsg.ROUND_SEQ] == 0 def test_finish_abort_ack_emits_one_failure(self): + """finish() fails the load at once; the AbortAck only releases it.""" session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=8, kv_request_id="req-1", keys=[OffloadKey(b"k")], block_ids=[0] ) session.finish_request("req-1") + (failure,) = session.poll().loads + assert (failure.job_id, failure.success) == (8, False) + assert not _released(failure) + assert session._client.has_active_loads is True conn.enqueue( { @@ -698,28 +708,30 @@ def test_finish_abort_ack_emits_one_failure(self): AbortAckMsg.KV_REQUEST_ID: "req-1", } ) - loads = session.poll().loads - assert loads == [LoadResult(job_id=8, kv_request_id="req-1", success=False)] + assert session.poll().loads == [] + assert _released(failure) assert session._client.has_active_loads is False assert "req-1" not in session._client._requests - assert session.poll().loads == [] def test_finish_abort_timeout_emits_one_failure(self): + """A missed AbortAck after finish() emits no second result: it ends + the session, and close() releases the held blocks.""" session, conn, _ = _make_session() _activate(session, conn) session.request_blocks( job_id=9, kv_request_id="req-1", keys=[OffloadKey(b"k")], block_ids=[0] ) session.finish_request("req-1") + (failure,) = session.poll().loads _client_load(session, "req-1").aborted_at = ( time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 ) - loads = session.poll().loads - assert loads == [LoadResult(job_id=9, kv_request_id="req-1", success=False)] - assert session._client.has_active_loads is False - assert "req-1" not in session._client._requests assert session.poll().loads == [] + assert not session.alive + assert not _released(failure) + assert session.close().failed_jobs == [] + assert _released(failure) def test_late_transfer_done_after_abort_ack_is_ignored(self): session, conn, _ = _make_session() @@ -736,9 +748,9 @@ def test_late_transfer_done_after_abort_ack_is_ignored(self): AbortAckMsg.KV_REQUEST_ID: "req-1", } ) - assert session.poll().loads == [ - LoadResult(job_id=10, kv_request_id="req-1", success=False) - ] + (failure,) = session.poll().loads + assert (failure.job_id, failure.success) == (10, False) + assert _released(failure) conn.enqueue( { @@ -823,10 +835,9 @@ def test_load_timeout_depends_on_mode(self): assert _client_load(session, "req-pd").aborted_at is None def test_load_abort_ack_timeout_surfaces_failure(self): - """After load timeout sends AbortFetch, if no AbortAck arrives within - _ABORT_ACK_TIMEOUT_S the request is surfaced as failed and removed - from _requests — the engine cannot wait forever on a peer that won't - ack. + """A load timeout fails the load at once, with its blocks held: the + peer may still be writing them. If no AbortAck arrives within + _ABORT_ACK_TIMEOUT_S, no second result is emitted; close() releases. """ session, conn, _ = _make_session() _activate(session, conn) @@ -837,8 +848,9 @@ def test_load_abort_ack_timeout_surfaces_failure(self): _client_load(session, "req-7").submitted_at = ( time.monotonic() - _PD_LOAD_TIMEOUT_S - 1.0 ) - loads = session.poll().loads - assert loads == [] + (failure,) = session.poll().loads + assert (failure.job_id, failure.success) == (7, False) + assert not _released(failure) assert any( m.get(TYPE_KEY) == AbortFetchMsg.TYPE and m[AbortFetchMsg.KV_REQUEST_ID] == "req-7" @@ -851,9 +863,46 @@ def test_load_abort_ack_timeout_surfaces_failure(self): _client_load(session, "req-7").aborted_at = ( time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 ) - loads = session.poll().loads - assert loads == [LoadResult(job_id=7, kv_request_id="req-7", success=False)] - assert "req-7" not in session._client._requests + assert session.poll().loads == [] + assert not _released(failure) + session.close() + assert _released(failure) + + def test_load_abort_ack_timeout_disconnects_peer(self): + """A missed AbortAck faults the peer, so later fetches skip it. + + The peer has stopped confirming that its WRITEs ended; keeping the + session would make every later load wait out both timeouts again. + Loads still in flight on the session fail through close(). + """ + session, conn, _ = _make_session() + _activate(session, conn) + session.request_blocks( + job_id=7, kv_request_id="req-7", keys=[OffloadKey(b"k")], block_ids=[0] + ) + session.request_blocks( + job_id=8, kv_request_id="req-8", keys=[OffloadKey(b"j")], block_ids=[1] + ) + _client_load(session, "req-7").submitted_at = ( + time.monotonic() - _PD_LOAD_TIMEOUT_S - 1.0 + ) + (failure,) = session.poll().loads + assert (failure.job_id, failure.success) == (7, False) + assert session.alive + + _client_load(session, "req-7").aborted_at = ( + time.monotonic() - _ABORT_ACK_TIMEOUT_S - 1.0 + ) + assert session.poll().loads == [] + assert not session.alive + assert session.death_reason == "abort ack timed out" + assert session.data_plane_fault == "abort ack timed out" + assert conn._sent[-1][TYPE_KEY] == DisconnectMsg.TYPE + + result = session.close() + assert result.failed_jobs == [8] + assert sorted(result.failed_req_ids) == ["req-7", "req-8"] + assert _released(failure) def test_load_abort_ack_clears_request(self): """After load timeout sends AbortFetch, an arriving AbortAckMsg from @@ -867,8 +916,9 @@ def test_load_abort_ack_clears_request(self): _client_load(session, "req-8").submitted_at = ( time.monotonic() - _PD_LOAD_TIMEOUT_S - 1.0 ) - # First poll: AbortFetch goes out. - session.poll() + # First poll: AbortFetch goes out and the load is failed at once. + (failure,) = session.poll().loads + assert (failure.job_id, failure.success) == (8, False) assert _client_load(session, "req-8").aborted_at is not None # Peer acks the abort. @@ -879,9 +929,11 @@ def test_load_abort_ack_clears_request(self): AbortAckMsg.KV_REQUEST_ID: "req-8", } ) - loads = session.poll().loads - assert loads == [LoadResult(job_id=8, kv_request_id="req-8", success=False)] + assert session.poll().loads == [] + assert _released(failure) assert "req-8" not in session._client._requests + assert session.alive + assert session.data_plane_fault is None # --------------------------------------------------------------------------- diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 696dde3bf2ee..d6976c9d44b8 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -675,6 +675,10 @@ def drain_jobs(self) -> None: memoryview to be quiescent, not aborted. Results accumulate in ``_finished_jobs`` and are surfaced by the next ``get_finished_jobs()`` call. + + A load already failed with a deferred release stays pending until + its AbortAck, or until a missed AbortAck disconnects the session and + the reap's close() sets the event, so every event is set on return. """ start = time.monotonic() warned = False @@ -1117,7 +1121,11 @@ def _poll_once(self) -> None: result = session.poll() for lr in result.loads: self._finished_jobs.append( - JobResult(job_id=lr.job_id, success=lr.success) + JobResult( + job_id=lr.job_id, + success=lr.success, + deferred_release=lr.deferred_release, + ) ) if not lr.success: self._failed_req_ids.add(lr.kv_request_id) diff --git a/vllm/v1/kv_offload/tiering/p2p/session/client.py b/vllm/v1/kv_offload/tiering/p2p/session/client.py index b545be47aeae..3101bc264821 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/client.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/client.py @@ -11,6 +11,7 @@ from __future__ import annotations +import threading import time from collections.abc import Callable, Sequence from dataclasses import dataclass, field @@ -54,6 +55,12 @@ class _InboundLoadState: # _P2P_LOAD_TIMEOUT_S or _PD_LOAD_TIMEOUT_S, fixed at submit time. timeout_s: float aborted_at: float | None = None + # Set together with aborted_at, when the load is reported failed. The + # peer may still be writing into its blocks, so their release waits on + # this event: set on AbortAck/TransferDone, or when the session closes. + release: threading.Event | None = None + # Its AbortAck timed out; only close() can release it now. + ack_expired: bool = False @dataclass @@ -99,14 +106,18 @@ class LoadResult(NamedTuple): job_id: int kv_request_id: str success: bool + # A failed load whose blocks the peer may still write: release them + # once this is set. + deferred_release: threading.Event | None = None class ClientCloseResult(NamedTuple): """What the manager must fail when a peer session is torn down. - ``failed_jobs`` is the ``job_id`` of every load still in flight; each - becomes a ``JobResult(success=False)``. ``failed_req_ids`` is every - kv_request_id whose lookup() would otherwise defer forever on the dead + ``failed_jobs`` is the ``job_id`` of every load still in flight and not + yet reported; each becomes a ``JobResult(success=False)``. + ``failed_req_ids`` is every kv_request_id whose lookup() would otherwise + defer forever on the dead peer — the in-flight loads plus any request holding an unresolved symmetric-P2P probe. ``failed_jobs`` is the subset of ``failed_req_ids`` that had a load job. @@ -146,6 +157,9 @@ def __init__(self, peer_id: str, send: Callable[[dict], None]) -> None: # Per-request progress deadlines. Only newly resolved probes extend # them; additional keys and repeated scheduler polls do not. self._lookup_deadlines: dict[str, float] = {} + # First load whose AbortAck never arrived. Sticky: the peer has + # stopped confirming that its WRITEs ended, so the session faults it. + self._expired_abort_ack: tuple[str, int] | None = None # ------------------------------------------------------------------ # State helpers @@ -177,6 +191,29 @@ def _on_load_terminal(self, kv_request_id: str, st: _ClientRequestState) -> None self._active_loads.discard(kv_request_id) self._maybe_prune(kv_request_id) + def _abort_load( + self, kv_request_id: str, round_seq: int, load: _InboundLoadState + ) -> None: + """Send AbortFetch and report the load failed now, deferring the + release of its blocks until the peer confirms its WRITEs ended.""" + self._send( + { + TYPE_KEY: AbortFetchMsg.TYPE, + AbortFetchMsg.KV_REQUEST_ID: kv_request_id, + AbortFetchMsg.ROUND_SEQ: round_seq, + } + ) + load.aborted_at = time.monotonic() + load.release = threading.Event() + self._completed_loads.append( + LoadResult( + job_id=load.job_id, + kv_request_id=kv_request_id, + success=False, + deferred_release=load.release, + ) + ) + @property def has_active_loads(self) -> bool: """True if any kv_request_id has a fetch in flight.""" @@ -190,6 +227,11 @@ def get_expired_lookup(self) -> tuple[str, float] | None: return req_id, now - deadline + _LOOKUP_TIMEOUT_S return None + @property + def expired_abort_ack(self) -> tuple[str, int] | None: + """The first ``(kv_request_id, round_seq)`` whose AbortAck timed out.""" + return self._expired_abort_ack + # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ @@ -261,24 +303,16 @@ def finish(self, kv_request_id: str) -> None: while an earlier round's load is still in flight, so both can be owed at once. - Then drop all probe/lookup state. Keep aborted loads until their abort - acknowledgement or timeout produces a terminal result. + Then drop all probe/lookup state. Aborted loads are reported failed + at once and kept until their AbortAck (or the session's close) + releases their blocks. """ st = self._requests.get(kv_request_id) if st is None: return - if st.loads: - for round_seq, load in st.loads.items(): - if load.aborted_at is not None: - continue - self._send( - { - TYPE_KEY: AbortFetchMsg.TYPE, - AbortFetchMsg.KV_REQUEST_ID: kv_request_id, - AbortFetchMsg.ROUND_SEQ: round_seq, - } - ) - load.aborted_at = time.monotonic() + for round_seq, load in st.loads.items(): + if load.aborted_at is None: + self._abort_load(kv_request_id, round_seq, load) if st.peer_lookup_open: st.peer_lookup_open = False self._send( @@ -302,7 +336,11 @@ def on_transfer_done( """Handle a TransferDoneMsg from the peer.""" st = self._requests.get(kv_request_id) load = st.loads.pop(round_seq, None) if st is not None else None - if st is not None and load is not None: + if st is not None and load is not None and load.release is not None: + # Already reported failed; the peer's WRITEs are over. + load.release.set() + self._on_load_terminal(kv_request_id, st) + elif st is not None and load is not None: self._completed_loads.append( LoadResult( job_id=load.job_id, @@ -331,20 +369,10 @@ def on_abort_ack(self, kv_request_id: str, round_seq: int) -> None: st = self._requests.get(kv_request_id) load = st.loads.pop(round_seq, None) if st is not None else None if st is not None and load is not None: - logger.warning( - "P2PSession %s: load request %s (job_id=%d) " - "completed with failure after abort", - self._peer_id, - kv_request_id, - load.job_id, - ) - self._completed_loads.append( - LoadResult( - job_id=load.job_id, - kv_request_id=kv_request_id, - success=False, - ) - ) + # Every abort reported the load failed when it was sent; the + # ack only releases its blocks. + assert load.release is not None + load.release.set() self._on_load_terminal(kv_request_id, st) else: # See on_transfer_done: same ambiguity (duplicate ack @@ -482,54 +510,40 @@ def on_lookup_resp( def collect_results(self) -> list[LoadResult]: """Walk load timeouts and drain completed loads. - Loads past their load timeout (``_P2P_LOAD_TIMEOUT_S`` or - ``_PD_LOAD_TIMEOUT_S``) get an AbortFetchMsg sent and - enter the aborting phase. Aborting loads past - ``_ABORT_ACK_TIMEOUT_S`` are surfaced as failed. + A load past its load timeout (``_P2P_LOAD_TIMEOUT_S`` or + ``_PD_LOAD_TIMEOUT_S``) gets an AbortFetchMsg and is reported failed + at once, with its blocks' release deferred. One whose AbortAck then + misses ``_ABORT_ACK_TIMEOUT_S`` is recorded in ``expired_abort_ack``, + so the session disconnects and close() releases it. """ now = time.monotonic() - to_remove: list[tuple[str, int]] = [] for req_id in self._active_loads: st = self._requests[req_id] assert st.loads for round_seq, load in st.loads.items(): if load.aborted_at is None: if now - load.submitted_at >= load.timeout_s: - load.aborted_at = now logger.warning( "P2PSession %s: %s round=%s timed out, sending abort", self._peer_id, req_id, round_seq, ) - self._send( - { - TYPE_KEY: AbortFetchMsg.TYPE, - AbortFetchMsg.KV_REQUEST_ID: req_id, - AbortFetchMsg.ROUND_SEQ: round_seq, - } - ) - else: - if now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S: - to_remove.append((req_id, round_seq)) - self._completed_loads.append( - LoadResult( - job_id=load.job_id, - kv_request_id=req_id, - success=False, - ) - ) - logger.warning( - "P2PSession %s: abort_ack timed out for " - "kv_request_id=%s round=%s", - self._peer_id, - req_id, - round_seq, - ) - for req_id, round_seq in to_remove: - st = self._requests[req_id] - st.loads.pop(round_seq) - self._on_load_terminal(req_id, st) + self._abort_load(req_id, round_seq, load) + elif ( + not load.ack_expired + and now - load.aborted_at >= _ABORT_ACK_TIMEOUT_S + ): + load.ack_expired = True + if self._expired_abort_ack is None: + self._expired_abort_ack = (req_id, round_seq) + logger.warning( + "P2PSession %s: abort_ack timed out for " + "kv_request_id=%s round=%s", + self._peer_id, + req_id, + round_seq, + ) results = self._completed_loads self._completed_loads = [] @@ -541,10 +555,17 @@ def close(self) -> ClientCloseResult: A request is failed if it has a load in flight (its job fails) or holds an unresolved symmetric-P2P probe (its lookup() would defer forever on an answer that can never arrive). See ``ClientCloseResult``. + The end of the session is the release point for every load already + reported failed: their deferred releases are set here. """ - failed_jobs = [ - load.job_id for st in self._requests.values() for load in st.loads.values() - ] + # Results never collected still owe their job an outcome. + failed_jobs = [r.job_id for r in self._completed_loads] + for st in self._requests.values(): + for load in st.loads.values(): + if load.release is not None: + load.release.set() + else: + failed_jobs.append(load.job_id) failed_req_ids = [ req_id for req_id, st in self._requests.items() diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index d23c306d3b26..63c0a350ca10 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -148,6 +148,7 @@ def __init__( # First cause recorded by _mark_dead; see death_reason. self._death_reason: str | None = None self._peer_fault_at: float | None = None + self._abort_ack_fault: str | None = None now = time.monotonic() self._connected_at: float | None = None @@ -215,8 +216,15 @@ def death_reason(self) -> str: @property def data_plane_fault(self) -> str | None: - """The transport's peer fault, if one was reported (see server role).""" - return self._server.peer_fault + """The data-plane fault that ended this session, if any. + + Either the transport's peer fault (see server role) or an AbortAck + that never arrived, which leaves this side unable to tell whether + the peer's WRITE stopped. + """ + if self._server.peer_fault is not None: + return self._server.peer_fault + return self._abort_ack_fault @property def last_activity_at(self) -> float: @@ -355,6 +363,7 @@ def poll(self) -> SessionPollResult: stores = self._server.collect_results() self._server.drain_pending_aborts() self._escalate_peer_fault() + self._escalate_abort_ack_timeout() # Drain buffered replies and completed transfers before expiring a peer. expired_lookup = self._client.get_expired_lookup() @@ -491,6 +500,27 @@ def _escalate_peer_fault(self) -> None: return self._disconnect(f"data plane: {fault}") + def _escalate_abort_ack_timeout(self) -> None: + """Fault the peer once one of its AbortAcks times out. + + The timed-out load already failed in this poll's results. Without + this, the session stays up and every later fetch to the same stuck + peer waits out the full load and abort-ack timeouts again. + """ + expired = self._client.expired_abort_ack + if expired is None or self._abort_ack_fault is not None: + return + req_id, round_seq = expired + self._abort_ack_fault = "abort ack timed out" + logger.warning( + "P2PSession %s: abort_ack for kv_request_id=%s round=%s timed out; " + "disconnecting peer session", + self.peer_id, + req_id, + round_seq, + ) + self._disconnect(self._abort_ack_fault) + def _disconnect(self, reason: str) -> None: """Best-effort DisconnectMsg, then mark the connection dead.""" if self._conn is not None: From dbd49af2223553d077eb839f610ad1aa84677391 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 28 Oct 2026 22:43:36 -0700 Subject: [PATCH 115/406] [New Model] Add MiniMax-M3.1-Flash Model and Kernels (#64111) Signed-off-by: Yongye Zhu Co-authored-by: Taoyu Zhu Co-authored-by: Hongxia Yang Co-authored-by: qizixi <22851944+zixi-qi@users.noreply.github.com> Co-authored-by: indivats Co-authored-by: andyluo7 <43718156+andyluo7@users.noreply.github.com> Co-authored-by: Thien Tran Co-authored-by: Jee Jee Li Co-authored-by: fxmarty-amd Co-authored-by: Jiejing Zhang <280536569+jiejingzhangamd@users.noreply.github.com> Co-authored-by: Prasun Gera Co-authored-by: Dao Le Co-authored-by: Isotr0py Co-authored-by: Neha Prakriya Co-authored-by: Summer Yang Co-authored-by: Zhewen Li Co-authored-by: Zijing Liu Co-authored-by: xiaozhoupy Co-authored-by: Seungmin Kim <8457324+ehfd@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- .buildkite/test_areas/kernels.yaml | 24 + .buildkite/test_areas/spec_decode.yaml | 3 + benchmarks/kernels/benchmark_mxfp8_tp4.py | 208 ++ cmake/external_projects/fmha_sm100.cmake | 2 +- ...minimax_m3_qknorm_rope_kv_insert_kernel.cu | 1551 ++++++++- csrc/ops.h | 13 +- csrc/torch_bindings.cpp | 14 +- .../test_speculative_draft_hf_overrides.py | 80 +- tests/distributed/test_custom_all_reduce.py | 146 + .../kernels/attention/minimax_m3_nvfp4_ref.py | 148 + tests/kernels/attention/test_minimax_m3.py | 401 +++ .../test_minimax_m3_index_topk_decode.py | 228 ++ ...st_minimax_m3_msa_cutlass_sparse_decode.py | 455 ++- .../test_minimax_m3_norm_rope_nvfp4_insert.py | 210 ++ .../attention/test_minimax_m3_nvfp4_index.py | 429 +++ ...test_minimax_m3_nvfp4_indexer_decode_cp.py | 130 + .../test_minimax_m3_nvfp4_prefill_flydsl.py | 260 ++ .../test_minimax_m3_nvfp4_sparse_decode.py | 243 ++ tests/kernels/mfma_spinner.py | 64 + tests/kernels/moe/test_deep_gemm_mega_moe.py | 2 +- .../quantization/test_rocm_mxfp8_linear.py | 49 +- ..._fused_minimax_m3_qknorm_rope_kv_insert.py | 314 ++ tests/kernels/test_minimax_m3_amd_ops.py | 151 + tests/kernels/test_minimax_m3_dspark_draft.py | 443 +++ tests/models/minimax_m3/__init__.py | 0 .../minimax_m3/test_dense_gemm1_mxfp8.py | 90 + .../minimax_m3/test_dspark_draft_compile.py | 135 + tests/models/minimax_m3/test_mono_host.py | 356 ++ .../test_qkv_index_fused_epilogue.py | 259 ++ .../multimodal/processing/test_minimax_m3.py | 208 +- tests/models/registry.py | 5 + tests/models/test_minimax_m3_mega_moe.py | 295 ++ .../test_qwen3_dflash_mxfp8_context_kv.py | 123 + tests/multimodal/test_video.py | 15 +- tests/quantization/test_modelopt.py | 74 + .../test_minimax_m3_reasoning_parser.py | 37 +- tests/v1/attention/test_flashinfer_ts.py | 229 ++ tests/v1/spec_decode/test_dflash2.py | 8 +- .../test_draft_attention_backend_override.py | 33 +- .../test_eagle3_aux_hidden_states_pp.py | 65 +- .../test_spec_decode_embed_sharing_pp.py | 6 +- vllm/_custom_ops.py | 59 +- vllm/config/attention.py | 21 +- vllm/config/speculative.py | 75 +- vllm/cute_utils/__init__.py | 17 +- vllm/cute_utils/cvt.py | 8 +- .../kernels/linear/mxfp8/flashinfer.py | 4 +- .../kernels/linear/mxfp8/rocm_native.py | 9 +- .../layers/attention/attention.py | 7 +- .../model_executor/layers/fused_moe/config.py | 11 +- vllm/model_executor/layers/fused_moe/layer.py | 5 + .../layers/fused_moe/oracle/nvfp4.py | 34 +- .../layers/fused_moe/prepare_megamoe.py | 156 +- .../compressed_tensors_moe_w4a4_nvfp4.py | 32 +- .../layers/quantization/modelopt.py | 50 + .../layers/quantization/utils/quant_utils.py | 11 + vllm/model_executor/models/config.py | 27 +- vllm/model_executor/models/qwen3_dflash.py | 96 +- vllm/model_executor/models/qwen3_dspark.py | 4 +- vllm/model_executor/models/registry.py | 4 + vllm/model_executor/warmup/kernel_warmup.py | 11 +- .../warmup/minimax_m3_msa_warmup.py | 75 +- .../common/ops/moe_tail_fusion/__init__.py | 5 + .../common/ops/moe_tail_fusion/primitives.py | 671 ++++ .../nvidia/ops/cute_dsl/all_reduce_mhc.py | 57 +- .../nvidia/ops/cute_dsl/primitives.py | 396 --- .../boundary/publishing_gemm.py | 8 +- .../moe/routed_expert_front_cute.py | 3 +- .../moe/routed_expert_gemm.py | 2 +- ...educe_rmsnorm_reduce_scatter_early_exit.py | 126 +- .../fused_add_multicast_gemm.py | 3 +- .../fused_add_multicast_skinny_gemm.py | 6 +- .../cute_dsl/latent_moe_tail/lamport_copy.py | 2 +- .../cute_dsl/latent_moe_tail/primitives.py | 576 ---- .../specialized_decode/cutedsl/common/host.py | 2 +- .../cutedsl/moe/small_batch/front_tp16.py | 3 +- .../specialized_decode/runtime/prefetch.py | 3 +- vllm/models/minimax_m3/amd/indexer_nvfp4.py | 239 ++ vllm/models/minimax_m3/amd/model.py | 234 +- vllm/models/minimax_m3/amd/mono/__init__.py | 4 + vllm/models/minimax_m3/amd/mono/check.py | 68 + vllm/models/minimax_m3/amd/mono/config.py | 108 + vllm/models/minimax_m3/amd/mono/dispatch.py | 216 ++ vllm/models/minimax_m3/amd/mono/envs.py | 20 + .../minimax_m3/amd/mono/kernels/__init__.py | 6 + .../minimax_m3/amd/mono/kernels/common.py | 418 +++ .../amd/mono/kernels/index_score.py | 474 +++ .../minimax_m3/amd/mono/kernels/post_attn.py | 2963 +++++++++++++++++ .../minimax_m3/amd/mono/kernels/pre_attn.py | 902 +++++ vllm/models/minimax_m3/amd/mono/layout.py | 233 ++ .../models/minimax_m3/amd/mono/peer_buffer.py | 52 + vllm/models/minimax_m3/amd/mono/runner.py | 334 ++ vllm/models/minimax_m3/amd/mono/weights.py | 377 +++ .../models/minimax_m3/amd/ops/flydsl_utils.py | 196 ++ .../minimax_m3/amd/ops/index_cp/__init__.py | 2 + .../minimax_m3/amd/ops/index_cp/host.py | 126 + .../amd/ops/index_topk_decode/__init__.py | 2 + .../amd/ops/index_topk_decode/host.py | 113 + .../ops/index_topk_decode/prefill_kernel.py | 387 +++ .../ops/norm_rope_nvfp4_insert/__init__.py | 2 + .../amd/ops/norm_rope_nvfp4_insert/host.py | 91 + .../amd/ops/norm_rope_nvfp4_insert/kernel.py | 581 ++++ .../amd/ops/nvfp4_index_decode/__init__.py | 2 + .../amd/ops/nvfp4_index_decode/host.py | 222 ++ .../amd/ops/nvfp4_index_decode/kernel.py | 564 ++++ .../ops/nvfp4_index_decode/kernel_cp_push.py | 611 ++++ .../amd/ops/nvfp4_index_prefill/__init__.py | 2 + .../amd/ops/nvfp4_index_prefill/host.py | 152 + .../amd/ops/nvfp4_index_prefill/kernel.py | 614 ++++ .../ops/nvfp4_index_prefill/kernel_shared.py | 382 +++ .../amd/ops/nvfp4_prefill/__init__.py | 2 + .../minimax_m3/amd/ops/nvfp4_prefill/host.py | 102 + .../amd/ops/nvfp4_prefill/kernel.py | 507 +++ .../amd/ops/nvfp4_sparse_decode/__init__.py | 2 + .../amd/ops/nvfp4_sparse_decode/host.py | 182 + .../amd/ops/nvfp4_sparse_decode/kernel.py | 576 ++++ .../minimax_m3/amd/sparse_attention_msa.py | 1 - .../minimax_m3/amd/sparse_attention_nvfp4.py | 170 + vllm/models/minimax_m3/common/dspark.py | 413 +++ vllm/models/minimax_m3/common/indexer.py | 60 +- .../models/minimax_m3/common/mm_preprocess.py | 131 +- .../common/ops/dspark_context_kv.py | 170 + .../minimax_m3/common/sparse_attention.py | 34 +- vllm/models/minimax_m3/common/vision_tower.py | 12 +- .../models/minimax_m3/nvidia/indexer_nvfp4.py | 481 +++ .../nvidia/indexer_nvfp4_decode_cp.py | 513 +++ vllm/models/minimax_m3/nvidia/model.py | 709 +++- .../nvidia/msa_cutlass_sparse_decode.py | 81 +- .../nvidia/ops/dense_gemm1_mxfp8.py | 664 ++++ .../nvidia/ops/finalize_allreduce_norm.py | 781 +++++ .../nvidia/ops/gather_nvfp4_index_k.py | 174 + .../nvidia/ops/qkv_index_fused_epilogue.py | 1101 ++++++ .../minimax_m3/nvidia/sparse_attention_msa.py | 413 ++- vllm/reasoning/minimax_m3_reasoning_parser.py | 15 + .../processors/minimax_m3.py | 315 +- vllm/transformers_utils/repo_utils.py | 24 + vllm/utils/flashinfer.py | 8 + vllm/v1/attention/backends/flashinfer_ts.py | 512 +++ vllm/v1/attention/backends/registry.py | 1 + .../gpu/spec_decode/dflash/speculator.py | 4 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 13 +- 141 files changed, 27012 insertions(+), 1941 deletions(-) create mode 100644 benchmarks/kernels/benchmark_mxfp8_tp4.py create mode 100644 tests/kernels/attention/minimax_m3_nvfp4_ref.py create mode 100644 tests/kernels/attention/test_minimax_m3_index_topk_decode.py create mode 100644 tests/kernels/attention/test_minimax_m3_norm_rope_nvfp4_insert.py create mode 100644 tests/kernels/attention/test_minimax_m3_nvfp4_index.py create mode 100644 tests/kernels/attention/test_minimax_m3_nvfp4_indexer_decode_cp.py create mode 100644 tests/kernels/attention/test_minimax_m3_nvfp4_prefill_flydsl.py create mode 100644 tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py create mode 100644 tests/kernels/mfma_spinner.py create mode 100644 tests/kernels/test_minimax_m3_dspark_draft.py create mode 100644 tests/models/minimax_m3/__init__.py create mode 100644 tests/models/minimax_m3/test_dense_gemm1_mxfp8.py create mode 100644 tests/models/minimax_m3/test_dspark_draft_compile.py create mode 100644 tests/models/minimax_m3/test_mono_host.py create mode 100644 tests/models/minimax_m3/test_qkv_index_fused_epilogue.py create mode 100644 tests/models/test_minimax_m3_mega_moe.py create mode 100644 tests/models/test_qwen3_dflash_mxfp8_context_kv.py create mode 100644 tests/v1/attention/test_flashinfer_ts.py create mode 100644 vllm/models/common/ops/moe_tail_fusion/__init__.py create mode 100644 vllm/models/common/ops/moe_tail_fusion/primitives.py delete mode 100644 vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py delete mode 100644 vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py create mode 100644 vllm/models/minimax_m3/amd/indexer_nvfp4.py create mode 100644 vllm/models/minimax_m3/amd/mono/__init__.py create mode 100644 vllm/models/minimax_m3/amd/mono/check.py create mode 100644 vllm/models/minimax_m3/amd/mono/config.py create mode 100644 vllm/models/minimax_m3/amd/mono/dispatch.py create mode 100644 vllm/models/minimax_m3/amd/mono/envs.py create mode 100644 vllm/models/minimax_m3/amd/mono/kernels/__init__.py create mode 100644 vllm/models/minimax_m3/amd/mono/kernels/common.py create mode 100644 vllm/models/minimax_m3/amd/mono/kernels/index_score.py create mode 100644 vllm/models/minimax_m3/amd/mono/kernels/post_attn.py create mode 100644 vllm/models/minimax_m3/amd/mono/kernels/pre_attn.py create mode 100644 vllm/models/minimax_m3/amd/mono/layout.py create mode 100644 vllm/models/minimax_m3/amd/mono/peer_buffer.py create mode 100644 vllm/models/minimax_m3/amd/mono/runner.py create mode 100644 vllm/models/minimax_m3/amd/mono/weights.py create mode 100644 vllm/models/minimax_m3/amd/ops/flydsl_utils.py create mode 100644 vllm/models/minimax_m3/amd/ops/index_cp/__init__.py create mode 100644 vllm/models/minimax_m3/amd/ops/index_cp/host.py create mode 100644 vllm/models/minimax_m3/amd/ops/index_topk_decode/__init__.py create mode 100644 vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py create mode 100644 vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py create mode 100644 vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/__init__.py create mode 100644 vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/host.py create mode 100644 vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/kernel.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/__init__.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/host.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel_cp_push.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/__init__.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/host.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_prefill/__init__.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_prefill/host.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_prefill/kernel.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/__init__.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py create mode 100644 vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/kernel.py create mode 100644 vllm/models/minimax_m3/amd/sparse_attention_nvfp4.py create mode 100644 vllm/models/minimax_m3/common/dspark.py create mode 100644 vllm/models/minimax_m3/common/ops/dspark_context_kv.py create mode 100644 vllm/models/minimax_m3/nvidia/indexer_nvfp4.py create mode 100644 vllm/models/minimax_m3/nvidia/indexer_nvfp4_decode_cp.py create mode 100644 vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py create mode 100644 vllm/models/minimax_m3/nvidia/ops/finalize_allreduce_norm.py create mode 100644 vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py create mode 100644 vllm/models/minimax_m3/nvidia/ops/qkv_index_fused_epilogue.py create mode 100644 vllm/v1/attention/backends/flashinfer_ts.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index a5629f2f8016..789741a10f65 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -262,6 +262,30 @@ steps: - tests/kernels/ - vllm/platforms/rocm.py +- label: ":nvidia: (B200) MiniMax M3 MegaMoE" + key: minimax-m3-mega-moe-test-b200 + timeout_in_minutes: 20 + device: b200-k8s + source_file_dependencies: + - cmake/external_projects/deepgemm.cmake + - vllm/models/minimax_m3/nvidia/ + - vllm/utils/deep_gemm.py + - tests/models/test_minimax_m3_mega_moe.py + commands: + - pytest -v -s models/test_minimax_m3_mega_moe.py + +- label: ":nvidia: (B200) MiniMax-M3 Model Tests" + key: minimax-m3-model-tests-b200 + timeout_in_minutes: 20 + device: b200-k8s + source_file_dependencies: + - vllm/models/minimax_m3/ + - vllm/cute_utils/ + - vllm/utils/flashinfer.py + - tests/models/minimax_m3/ + commands: + - pytest -v -s models/minimax_m3 --ignore=models/minimax_m3/amd_specialized_decode + # Catch-all for test files at the tests/kernels root. This job collects # the whole root so new files are wired by default. # Files with dedicated jobs elsewhere in this file are excluded via --ignore diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index e9c39537246d..1cca6d4160ef 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -352,8 +352,11 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ + - vllm/model_executor/models/qwen3_dflash.py + - tests/models/test_qwen3_dflash_mxfp8_context_kv.py commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ + - pytest -v -s models/test_qwen3_dflash_mxfp8_context_kv.py mirror: amd: label: ":amd: (MI355) Spec Decode Draft Model Nightly" diff --git a/benchmarks/kernels/benchmark_mxfp8_tp4.py b/benchmarks/kernels/benchmark_mxfp8_tp4.py new file mode 100644 index 000000000000..e9ee03c8a431 --- /dev/null +++ b/benchmarks/kernels/benchmark_mxfp8_tp4.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Validate MXFP8 TP4 linears and time HIP graph replay over rotating weights. + +Example (record the baseline and candidate SHAs): + baseline=$(git rev-parse origin/main) + git rev-parse "$baseline" HEAD + git show "$baseline":vllm/model_executor/kernels/linear/mxfp8/rocm_native.py \ + > /tmp/baseline_rocm_native.py + .venv/bin/python benchmarks/kernels/benchmark_mxfp8_tp4.py \ + --baseline /tmp/baseline_rocm_native.py + +Reports cold (>400 MiB rotating weights) and warm timings separately. Each timing +is the median of 30 device-event samples; Python wall time is never ranked. +""" + +import argparse +import importlib.util +import json +import math +import statistics +from functools import partial +from pathlib import Path + +import torch + +from vllm.model_executor.kernels.linear.mxfp8 import rocm_native +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, + mxfp8_e4m3_quantize, +) +from vllm.triton_utils import triton + +SHAPES = [ + (m, n, k) + for n, k in [(2560, 6144), (6144, 2048), (1536, 6144), (6144, 768)] + for m in [1, 8, 16] +] + + +def time_graph(calls): + for call in calls: + call() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(2): + for call in calls: + call() + for _ in range(5): + graph.replay() + events = [ + (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + for _ in range(30) + ] + for start, end in events: + start.record() + for _ in range(5): + graph.replay() + end.record() + # Bound outstanding HIP graph nodes on short-K cases with many slots. + end.synchronize() + torch.accelerator.synchronize() + return statistics.median( + start.elapsed_time(end) / (10 * len(calls)) for start, end in events + ) + + +def launch(module, xq, xs, w, ws, out, config): + bm, bn, bk, nw, ns = config + m, k = xq.shape + n = w.shape[0] + return module._mxfp8_linear_kernel[(triton.cdiv(m, bm), triton.cdiv(n, bn))]( + xq, + xs, + w, + ws, + out, + m, + n, + k, + *xq.stride(), + *xs.stride(), + *w.stride(), + *ws.stride(), + *out.stride(), + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + num_warps=nw, + num_stages=ns, + ) + + +def emit(**data): + print(json.dumps(data), flush=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline", type=Path) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--shape", type=int, choices=range(len(SHAPES))) + parser.add_argument("--tune", action="store_true") + parser.add_argument("--profile", action="store_true") + parser.add_argument("--asm-dir", type=Path) + args = parser.parse_args() + torch.accelerator.set_device_index(args.device) + torch.manual_seed(42) + variants = {"candidate": rocm_native} + if args.baseline: + spec = importlib.util.spec_from_file_location( + "vllm.model_executor.kernels.linear.mxfp8.baseline_rocm_native", + args.baseline, + ) + baseline = importlib.util.module_from_spec(spec) + spec.loader.exec_module(baseline) + variants = {"baseline": baseline, **variants} + emit( + torch=torch.__version__, + triton=triton.__version__, + device=str(torch.cuda.get_device_properties(args.device)), + ) + for index, (m, n, k) in enumerate(SHAPES): + if args.shape is not None and index != args.shape: + continue + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5 + w, ws = mxfp8_e4m3_quantize( + torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.05 + ) + xq, xs = mxfp8_e4m3_quantize(x) + reference = ( + dequant_mxfp8_to_bf16(xq, xs).float() + @ dequant_mxfp8_to_bf16(w, ws).float().T + ) + out = torch.empty(m, n, device="cuda", dtype=x.dtype) + slots = math.ceil(400 * 1024**2 / (n * k)) + weights = [(w.clone(), ws.clone()) for _ in range(slots)] + calls = {} + for name, module in variants.items(): + configs = [module._select_cfg(m, n, k)] + if args.tune: + configs += [ + (16, 16, 1024, 2, 2), + (16, 16, 1024, 2, 3), + (16, 16, 1024, 2, 4), + (16, 16, 512, 2, 2), + (16, 32, 256, 4, 3), + (16, 16, 256, 2, 3), + (128, 128, 128, 8, 2), + ] + for config in dict.fromkeys(configs): + if k % config[2]: + continue + fn = partial(launch, module, xq, xs, config=config, out=out) + obj = fn(w=w, ws=ws) + rel = ((out.float() - reference).norm() / reference.norm()).item() + snr = -20 * math.log10(rel) if rel else float("inf") + assert rel < 0.005 and snr >= 30, (name, config, rel, snr) + emit( + shape=[m, n, k], + variant=name, + config=config, + rel=rel, + snr_db=snr, + registers=obj.n_regs, + spills=obj.n_spills, + lds=obj.metadata.shared, + ) + if args.asm_dir: + args.asm_dir.mkdir(parents=True, exist_ok=True) + (args.asm_dir / f"{name}-{m}-{n}-{k}-{config}.amdgcn").write_text( + obj.asm["amdgcn"] + ) + if args.profile: + for _ in range(10): + fn(w=w, ws=ws) + torch.accelerator.synchronize() + continue + calls[f"{name}-kernel-{config}"] = [ + partial(fn, w=wi, ws=si) for wi, si in weights + ] + if not args.profile and not args.tune: + wrapper = partial(module._mxfp8_dot_scaled_linear, x) + actual = wrapper(w, ws) + assert ((actual.float() - reference).norm() / reference.norm()) < 0.005 + calls[f"{name}-wrapper"] = [ + partial(wrapper, wi, si) for wi, si in weights + ] + for repeat in range(2): + order = list(calls.items()) + if repeat: + order.reverse() + for name, functions in order: + cold = time_graph(functions) + warm = time_graph([functions[0]] * 32) + emit( + shape=[m, n, k], + variant=name, + repeat=repeat, + cold_ms=cold, + warm_ms=warm, + slots=slots, + ) + + +if __name__ == "__main__": + main() diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index 6b84b35c9049..2336d66d23dc 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 968061ab16cb14991a5c7d4c0e8eabf531be0dcf + GIT_TAG 25c144700f727b638067e98066a60507081b63ad GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index 09484daf3315..11e76e6d9873 100644 --- a/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -48,9 +48,11 @@ * ``qkv`` tensor. Caches (bf16) are scatter-written by slot. */ +#include #include #include #include +#include #include "torch_utils.h" @@ -271,11 +273,43 @@ __device__ __forceinline__ void storeElemsFp8( // Match scaled_fp8_quant(q_out): materialize q in scalar_t before applying the // inverse dequantization scale and converting it to E4M3. +#ifndef USE_ROCM +// storeScaledQElemsFp8 with inv_scale = 1, from the packed model-dtype values: +// E4M3 of a scalar_t equals E4M3 of its float, the max takes NaN and values +// below -448 to -448 as storeElemsFp8's fmaxf does, and the cvt saturates the +// top at 448. +template +__device__ __forceinline__ void storeUnitScaleQElemsFp8( + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { + using Converter = vllm::_typeConvert; + __nv_fp8x2_storage_t out2[kElemsPerLane / 2]; + #pragma unroll + for (int i = 0; i < kElemsPerLane / 2; i++) { + auto const v = + Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1])); + if constexpr (std::is_same_v) { + out2[i] = __nv_cvt_bfloat16raw2_to_fp8x2( + __hmax2(v, __float2bfloat162_rn(-448.0f)), __NV_SATFINITE, __NV_E4M3); + } else { + out2[i] = __nv_cvt_halfraw2_to_fp8x2( + __hmax2(v, __float2half2_rn(-448.0f)), __NV_SATFINITE, __NV_E4M3); + } + } + *reinterpret_cast(dst) = *reinterpret_cast(out2); +} +#endif + template __device__ __forceinline__ void storeScaledQElemsFp8( uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane], float const inv_scale) { using Converter = vllm::_typeConvert; +#ifndef USE_ROCM + if (inv_scale == 1.0f) { + storeUnitScaleQElemsFp8(dst, elems); + return; + } +#endif float scaled[kElemsPerLane]; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -424,27 +458,14 @@ __device__ __forceinline__ uint16_t clearNegZeroNibbles(uint16_t packed) { return packed & (((packed & 0x7777u) + 0x7777u) | 0x7777u); } -// Quantize one K or V head row of one token into its NVFP4 page, where -// ``scale`` is the dequantization (global) scale. Every lane of the warp must -// call this (the scale-group reductions shuffle). +// The E2M1 codes (low nibble first) of four values x, already rounded to the +// model dtype, of a 16-value group with absolute maximum group_max, and the +// group's E4M3 scale in ``sf_byte``, where ``scale`` is the dequantization +// (global) scale. template -__device__ __forceinline__ void storeNvfp4CacheElems( - uint8_t* __restrict__ page, int const kv, int const head, int const token, - int const block_size, int const laneId, float const (&elems)[kElemsPerLane], - float const scale) { - // Quantize the model-dtype value that the unfused path reads back from qkv. - using Converter = vllm::_typeConvert; - float x[kElemsPerLane]; - float group_max = 0.0f; - #pragma unroll - for (int i = 0; i < kElemsPerLane; i++) { - x[i] = static_cast(Converter::convert(elems[i])); - group_max = fmaxf(group_max, fabsf(x[i])); - } - group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 1, 32)); - group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 2, 32)); - - uint8_t sf_byte; +__device__ __forceinline__ uint16_t +nvfp4EncodeElems(float const (&x)[kElemsPerLane], float const group_max, + float const scale, uint8_t& sf_byte) { float sf_rcp; float q[kElemsPerLane]; // (x / scale) / sf_value, E2M1-encoded below if (scale == 1.0f) { @@ -516,8 +537,49 @@ __device__ __forceinline__ void storeNvfp4CacheElems( for (int i = 0; i < kElemsPerLane; i++) codes |= e2m1Code(q[i]) << (4 * i); packed = static_cast(codes); #endif - packed = clearNegZeroNibbles(packed); + return clearNegZeroNibbles(packed); +} + +// Offset, within a K (kv 0) or V (kv 1) head's scale block, of the E4M3 scale +// of a token's 16-value group: linear for K, the SM100 trtllm-gen 4x4 +// token-quad swizzle for V. +__device__ __forceinline__ int nvfp4ScaleOffset(int const kv, int const token, + int const group) { + if (kv == 0) return token * kNvfp4ScaleBytes + group; + constexpr int kGroupsPerQuad = kNvfp4ScaleBytes / 4; + int const swizzled_t = (token / 4) * 4 + group / kGroupsPerQuad; + int const swizzled_s = (group % kGroupsPerQuad) * 4 + token % 4; + return swizzled_t * kNvfp4ScaleBytes + swizzled_s; +} + +// The E2M1 codes of this lane's four values of a K or V head row, rounded to +// the model dtype, and in ``sf_byte`` its 16-value group's E4M3 scale, where +// ``scale`` is the dequantization (global) scale. Every lane of the warp must +// call this (the scale-group reductions shuffle). +template +__device__ __forceinline__ uint16_t nvfp4QuantizeElems( + float const (&elems)[kElemsPerLane], float const scale, uint8_t& sf_byte) { + // Quantize the model-dtype value that the unfused path reads back from qkv. + using Converter = vllm::_typeConvert; + float x[kElemsPerLane]; + float group_max = 0.0f; + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + x[i] = static_cast(Converter::convert(elems[i])); + group_max = fmaxf(group_max, fabsf(x[i])); + } + group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 1, 32)); + group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 2, 32)); + + return nvfp4EncodeElems(x, group_max, scale, sf_byte); +} +// Store the codes and scale from nvfp4QuantizeElems of lane laneId's four +// values of a token's K (kv 0) or V (kv 1) head row. +__device__ __forceinline__ void nvfp4StoreElems( + uint8_t* __restrict__ page, int const kv, int const head, int const token, + int const block_size, int const laneId, uint16_t const packed, + uint8_t const sf_byte) { uint8_t* const slot = page + static_cast(2 * head + kv) * block_size * (kNvfp4DataBytes + kNvfp4ScaleBytes); @@ -526,19 +588,57 @@ __device__ __forceinline__ void storeNvfp4CacheElems( *reinterpret_cast(data) = packed; if (laneId % kNvfp4LanesPerGroup == 0) { - int const group = laneId / kNvfp4LanesPerGroup; - int offset; - if (kv == 0) { - offset = token * kNvfp4ScaleBytes + group; - } else { - constexpr int kGroupsPerQuad = kNvfp4ScaleBytes / 4; - int const swizzled_t = (token / 4) * 4 + group / kGroupsPerQuad; - int const swizzled_s = (group % kGroupsPerQuad) * 4 + token % 4; - offset = swizzled_t * kNvfp4ScaleBytes + swizzled_s; - } - slot[block_size * kNvfp4DataBytes + offset] = sf_byte; + slot[block_size * kNvfp4DataBytes + + nvfp4ScaleOffset(kv, token, laneId / kNvfp4LanesPerGroup)] = sf_byte; } } + +// nvfp4QuantizeElems for one lane holding a whole 16-value group x (already +// rounded to the model dtype): the group max is order-independent, so this is +// the same codes and scale without the shuffles. Low nibble first. +template +__device__ __forceinline__ uint2 nvfp4QuantizeGroup(float const (&x)[16], + float const scale, + uint8_t& sf_byte) { + float group_max = 0.0f; + #pragma unroll + for (int i = 0; i < 16; i++) group_max = fmaxf(group_max, fabsf(x[i])); + uint32_t codes[4]; + #pragma unroll + for (int k = 0; k < 4; k++) { + float const xk[kElemsPerLane] = {x[4 * k], x[4 * k + 1], x[4 * k + 2], + x[4 * k + 3]}; + codes[k] = nvfp4EncodeElems(xk, group_max, scale, sf_byte); + } + return make_uint2(codes[0] | codes[1] << 16, codes[2] | codes[3] << 16); +} + +// Store group ``group`` (its 8 code bytes and its scale) of a token's K +// (kv 0) or V (kv 1) head row into its NVFP4 page. +__device__ __forceinline__ void nvfp4StoreGroup( + uint8_t* __restrict__ page, int const kv, int const head, int const token, + int const block_size, int const group, uint2 const codes, + uint8_t const sf_byte) { + uint8_t* const slot = page + static_cast(2 * head + kv) * + block_size * + (kNvfp4DataBytes + kNvfp4ScaleBytes); + *reinterpret_cast(slot + token * kNvfp4DataBytes + 8 * group) = codes; + slot[block_size * kNvfp4DataBytes + nvfp4ScaleOffset(kv, token, group)] = + sf_byte; +} + +// Quantize one K or V head row of one token into its NVFP4 page, where +// ``scale`` is the dequantization (global) scale. Every lane of the warp must +// call this (the scale-group reductions shuffle). +template +__device__ __forceinline__ void storeNvfp4CacheElems( + uint8_t* __restrict__ page, int const kv, int const head, int const token, + int const block_size, int const laneId, float const (&elems)[kElemsPerLane], + float const scale) { + uint8_t sf_byte; + uint16_t const packed = nvfp4QuantizeElems(elems, scale, sf_byte); + nvfp4StoreElems(page, kv, head, token, block_size, laneId, packed, sf_byte); +} #endif // USE_ROCM // ──────────────────────────────────────────────────────────────────────────── @@ -553,18 +653,21 @@ __device__ __forceinline__ void storeNvfp4CacheElems( // Slots per token: // Q : nq (always — norm+RoPE) // K : nkv (always — norm+RoPE; +K-cache insert) -// V : nkv only if kInsertKV (V-cache insert; no warps in dense) +// V : nkv only if kInsertKV (V-cache insert) // IQ: niq only if kProcessIndex (norm+RoPE) // IK: 1 only if kProcessIndex (norm+RoPE; +index-cache insert) // cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx: // indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte). +// kNvfp4Idx (implies kFp8Idx): the index-K cache is instead packed NVFP4 pages +// [nb, bs, 72] (bs*64 B E2M1, then bs*8 B linear E4M3 scales, no global scale), +// the layout of the MSA Q8KV4 indexer; index-Q stays e4m3. // kHasIndex means the qkv row is laid out as sparse [q|k|v|index_q|index_k]. // kProcessIndex controls whether this launch actually norms/ropes the index // branch and writes index_q/index_k outputs. Skip-index-topk reuse layers keep // kHasIndex=true but set kProcessIndex=false. template + bool kProcessIndex, bool kFp8Idx, bool kNvfp4Idx> __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse) scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr @@ -588,6 +691,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( // NHD/HND layout choice is captured by the head/token strides. int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, int64_t const kv_s_dim, + // NVFP4 index cache only: page stride in bytes. + int64_t const idx_s_block, // NVFP4 main cache only: per-layer K/V dequantization (global) scales. float const* __restrict__ kv_k_scale, float const* __restrict__ kv_v_scale) { @@ -604,8 +709,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( static_assert(!kProcessIndex || kHasIndex, "index processing requires sparse row layout"); + static_assert(!kNvfp4Idx || (kFp8Idx && kInsertKV), + "NVFP4 index cache requires fp8 index-Q and cache insert"); - // Slot layout (compile-time gated: dense has neither V nor index slots). + // Slot layout (compile-time gated: V slots only when inserting, index + // slots only on the sparse layer). int const v_slots = kInsertKV ? nkv : 0; int const idx_slots = kProcessIndex ? niq + 1 : 0; int const slots_per_token = nq + nkv + v_slots + idx_slots; @@ -743,7 +851,7 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( } } - // ── Cache inserts (sparse serving only). ─────────────────────────────── + // ── Cache inserts. ───────────────────────────────────────────────────── if constexpr (kInsertKV) { // Guard (not early-return) so every thread reaches the PDL trigger below. int64_t sm = -1; @@ -754,7 +862,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( } if (sm >= 0) { // skip padded / unscheduled tokens if (isIK) { - if constexpr (kFp8Idx) { + if constexpr (kNvfp4Idx) { +#ifndef USE_ROCM + // A single-head K side: data [bs, 64] then linear scales [bs, 8]. + storeNvfp4CacheElems( + reinterpret_cast(index_cache) + + (sm / block_size) * idx_s_block, + /*kv=*/0, /*head=*/0, static_cast(sm % block_size), + block_size, laneId, elems, /*scale=*/1.0f); +#endif + } else if constexpr (kFp8Idx) { storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems); } else { storeElems(index_cache + sm * kHeadDim + dim_base, elems); @@ -779,6 +896,15 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( int64_t const off = b * kv_s_block + head * kv_s_head + t * kv_s_token + (kv * kHeadDim + dim_base) * kv_s_dim; + if constexpr (!kHasIndex) { + // Dense layers stand in for reshape_and_cache_flash, which + // quantizes the k already materialized in the model dtype. + using Converter = vllm::_typeConvert; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + elems[i] = Converter::convert(Converter::convert(elems[i])); + } + } storeCacheElems(kv_cache + off, elems); } } @@ -794,6 +920,754 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( #endif } +#ifndef USE_ROCM +// ──────────────────────────────────────────────────────────────────────────── +// Persistent row-group kernel +// ──────────────────────────────────────────────────────────────────────────── +// Writes the same bytes as the kernel above, for every instantiation of it, +// when rotary_dim == 64. A token's work is a group of warps with fixed roles: +// the q heads kQHeads (4 or 8) per warp, then K, V and index_q four heads per +// warp, then index_k. A branch of fewer than four heads, and index_k, takes +// one warp per head in the layout of the kernel above (lane l holds dims +// 4l .. 4l + 3); the others run four-row steps, where lane (g, v) holds dims +// 32c + 4v .. 32c + 4v + 3 (c = 0..3) of row g, the dims lane 8c + v holds +// above, so the stores call the same helpers with that lane id. A CTA runs +// some or all warps of the warp groups of tokens x, x + gridDim.x, ..., and +// each warp copies the qkv rows and cos/sin row of its next tokens into a +// kStages-deep shared-memory ring with cp.async. +namespace rg { + +constexpr int kRows = 4; +constexpr int kChunks = 4; +constexpr int kRotaryDim = 64; +constexpr int kStages = 3; // ring depth: two tokens in flight per warp +// A ring row: 256 B padded to 320 B, so that rows 1 apart start 16 banks apart +// and a 4-row step's 8-byte loads are conflict-free. +constexpr int kRowPitch = kHeadDim * 2 + 64; +constexpr int kCosSinBytes = kRotaryDim * 2; +enum Role : int { kQ = 0, kK, kV, kIQ, kIK }; + +// SM100's packed FP32 instructions where they exist; the same RN results. +__device__ __forceinline__ float2 mul2(float2 const a, float2 const b) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + return __fmul2_rn(a, b); + #else + return make_float2(__fmul_rn(a.x, b.x), __fmul_rn(a.y, b.y)); + #endif +} + +__device__ __forceinline__ float2 fma2(float2 const a, float2 const b, + float2 const c) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + return __ffma2_rn(a, b, c); + #else + return make_float2(__fmaf_rn(a.x, b.x, c.x), __fmaf_rn(a.y, b.y, c.y)); + #endif +} + +__device__ __forceinline__ float2 add2(float2 const a, float2 const b) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + return __fadd2_rn(a, b); + #else + return make_float2(__fadd_rn(a.x, b.x), __fadd_rn(a.y, b.y)); + #endif +} + +// A head row as one lane holds it, paired for packed FP32 math: +// x[k][i] = element i of chunks 2k and 2k + 1. +using Row = float2[2][4]; + +template +__device__ __forceinline__ void loadRow(uint8_t const* row, int const v, + Row& x) { + float e[kChunks][kElemsPerLane]; + #pragma unroll + for (int c = 0; c < kChunks; c++) { + loadElems(reinterpret_cast(row + 64 * c + 8 * v), + e[c]); + } + #pragma unroll + for (int k = 0; k < 2; k++) { + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + x[k][i] = make_float2(e[2 * k][i], e[2 * k + 1][i]); + } + } +} + +__device__ __forceinline__ void chunkElems(Row const& x, int const c, + float (&e)[kElemsPerLane]) { + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + e[i] = c % 2 ? x[c / 2][i].y : x[c / 2][i].x; + } +} + +// normAndRope (Gemma RMSNorm, then NeoX RoPE on dims [0, 64)) with its +// rounding as compiled: an FFMA sum of squares per lane of the one-warp-per-row +// layout (one chunk here), the xor-16/8/4/2/1 tree (16 pairs chunks c and +// c ^ 2, 8 pairs c and c ^ 1), FFMA(sum, 1/128, eps), rsqrtf, x * rms * w, and +// x * cos -/+ RN(partner * sin) as one FFMA. +__device__ __forceinline__ void normRope(Row& x, Row const& w1, float const eps, + float const (&cos4)[4], + float const (&sin4)[4]) { + float2 sq[2]; + #pragma unroll + for (int k = 0; k < 2; k++) { + sq[k] = mul2(x[k][0], x[k][0]); // = FFMA(x, x, 0) + #pragma unroll + for (int i = 1; i < 4; i++) sq[k] = fma2(x[k][i], x[k][i], sq[k]); + } + float2 const xor16 = add2(sq[0], sq[1]); + float sum = __fadd_rn(xor16.x, xor16.y); + sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 4, 32)); + sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 2, 32)); + sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 1, 32)); + float const rms_rcp = + rsqrtf(__fmaf_rn(sum, 1.0f / static_cast(kHeadDim), eps)); + float2 const r2 = make_float2(rms_rcp, rms_rcp); + #pragma unroll + for (int k = 0; k < 2; k++) { + #pragma unroll + for (int i = 0; i < 4; i++) x[k][i] = mul2(mul2(x[k][i], r2), w1[k][i]); + } + // x[0][i] = (x1, x2): dims 4v + i and 32 + 4v + i. + #pragma unroll + for (int i = 0; i < 4; i++) { + float2 const x_sin = mul2(x[0][i], make_float2(sin4[i], sin4[i])); + x[0][i] = make_float2(__fmaf_rn(x[0][i].x, cos4[i], -x_sin.y), + __fmaf_rn(x[0][i].y, cos4[i], x_sin.x)); + } +} + +__device__ __forceinline__ void cpAsync16(uint32_t const dst, void const* src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(dst), + "l"(src) + : "memory"); +} + +__device__ __forceinline__ void cpAsyncCommit() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cpAsyncWait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(kPending) : "memory"); +} + +// Warp roles of a token's warp group: the q heads kQHeads at a time, then K, +// V and index_q four heads at a time, then index_k. +// A branch of fewer than kRows heads (and index_k) takes one warp per head in +// the layout of the kernel above instead of a part-empty four-row step. +struct Roles { + int nq, nkv, niq; + bool kv_one, iq_one; // one warp per head + int q_warps, kv_warps, v_warps, iq_warps, ik_warps; + int warps; // one token's warp group + + __host__ __device__ Roles(int const q_heads, int const nq_, int const nkv_, + int const niq_, bool const insert_kv, + bool const process_index) + : nq(nq_), nkv(nkv_), niq(niq_) { + kv_one = nkv < kRows; + iq_one = niq < kRows; + q_warps = (nq + q_heads - 1) / q_heads; + kv_warps = kv_one ? nkv : (nkv + kRows - 1) / kRows; + v_warps = insert_kv ? kv_warps : 0; + iq_warps = process_index ? (iq_one ? niq : (niq + kRows - 1) / kRows) : 0; + ik_warps = process_index ? 1 : 0; + warps = q_warps + kv_warps + v_warps + iq_warps + ik_warps; + } + + // Shared-memory offset of warp w's ring: the q warps' rings, then the + // others'. + __host__ __device__ int ringOffset(int const w, int const q_ring, + int const side_ring) const { + return min(w, q_warps) * q_ring + max(w - q_warps, 0) * side_ring; + } +}; + +// CTA size limit, for 80 registers per thread. +constexpr int kMaxWarps = 24; + +template +constexpr int stageBytes() { + return kQHeads * kRowPitch + kCosSinBytes; +} + +// What a warp writes besides the caches, fixed per launch so that the stores +// of the other modes are not compiled into the loop: q (kQFp8Unit: only +// q_fp8_out, unit scale; kQFp8: only q_fp8_out; kQAny: decided per store), +// index_q (kIQOut: index_q_out; kIQInPlace: back into qkv; kNone), and the +// in-place k / index_k rewrite (kInPlace or kNone). +enum Out : int { + kNone = 0, + kInPlace, + kQFp8Unit, + kQFp8, + kQAny, + kIQOut, + kIQInPlace +}; + +// One warp's loop over its CTA's tokens: rows [head0, head0 + nrows) of one +// branch, kSteps four-row steps per token. +template +__device__ __forceinline__ void runWarp( + int const head0, int const nrows, int const qkv_head0, uint8_t* const ring, + int const stage_bytes, scalar_t* __restrict__ qkv, + scalar_t* __restrict__ q_out, uint8_t* __restrict__ q_fp8_out, + out_idx_t* __restrict__ index_q_out, scalar_t const* __restrict__ norm_w, + scalar_t const* __restrict__ cos_sin_cache, + int64_t const* __restrict__ positions, + int64_t const* __restrict__ slot_mapping, cache_t* __restrict__ kv_cache, + out_idx_t* __restrict__ index_cache, float const eps, + float const q_fp8_inv_scale, int const num_tokens, int const nq, + int const niq, int const qkv_row, int const block_size, + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, + int64_t const kv_s_dim, int64_t const idx_s_block, float const kv_scale, + bool const write_qkv) { + constexpr bool kNorm = kRole != kV; + constexpr bool kSlots = + kInsertKV && (kRole == kK || kRole == kV || kRole == kIK); + // kOneRow: a single row in the layout of the kernel above (lane l holds + // dims 4l .. 4l + 3), for branches of one head. + constexpr int kStepRows = kOneRow ? 1 : kRows; + constexpr int kLaneChunks = kOneRow ? 1 : kChunks; + int const lane = threadIdx.x % 32; + int const g = kOneRow ? 0 : lane / 8; + int const v = kOneRow ? lane : lane % 8; + constexpr int kCopies = (kSteps * kStepRows * 16 + 31) / 32; + uint32_t const ring_smem = + static_cast(__cvta_generic_to_shared(ring)); + + // Not produced by the predecessor: loaded before the PDL wait, used after. + uint2 w_raw[kChunks]; + if constexpr (kNorm && !kOneRow) { + #pragma unroll + for (int c = 0; c < kChunks; c++) { + w_raw[c] = __ldg(reinterpret_cast(norm_w + 32 * c + 4 * v)); + } + } + + // PDL: qkv comes from the predecessor (the qkv GEMM). + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); + #endif + + Row w1; + if constexpr (kNorm && !kOneRow) { + float e[kChunks][kElemsPerLane]; + #pragma unroll + for (int c = 0; c < kChunks; c++) { + loadElems(reinterpret_cast(&w_raw[c]), e[c]); + } + #pragma unroll + for (int k = 0; k < 2; k++) { + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + w1[k][i] = make_float2(1.0f + e[2 * k][i], 1.0f + e[2 * k + 1][i]); + } + } + } + + int const stride = gridDim.x; + int const t0 = blockIdx.x; + int const n = (num_tokens - t0 + stride - 1) / stride; + // Lane j holds the position (slot) of the token of iteration base + j. + auto fetch = [&](int64_t const* src, int const base, int64_t const fill) { + int64_t const t = t0 + static_cast(base + lane) * stride; + return t < num_tokens ? static_cast(__ldg( + reinterpret_cast(src) + t)) + : fill; + }; + int pos_base = 0, slot_base = 0; + int64_t pos_l = 0, slot_l = -1; + if constexpr (kNorm) pos_l = fetch(positions, 0, 0); + if constexpr (kSlots) slot_l = fetch(slot_mapping, 0, -1); + + scalar_t const* const src0 = + qkv + static_cast(t0) * qkv_row + qkv_head0 * kHeadDim; + int64_t const src_step = static_cast(stride) * qkv_row; + auto issue_rows = [&](int const it) { + if (it >= n) return; + scalar_t const* const src = src0 + it * src_step; + uint32_t const stage = ring_smem + (it % kStages) * stage_bytes; + #pragma unroll + for (int k = 0; k < kCopies; k++) { + int const r = lane / 16 + 2 * k; + if (r < nrows && (!kOneRow || lane < 16)) { + cpAsync16(stage + r * kRowPitch + lane % 16 * 16, + src + r * kHeadDim + lane % 16 * 8); + } + } + }; + // The cos/sin row of iteration it, then the commit that closes its group. + auto issue_cos_sin = [&](int const it) { + if constexpr (kNorm) { + if (it < n) { + if (it - pos_base == 32) { + pos_base = it; + pos_l = fetch(positions, pos_base, 0); + } + int64_t const pos = __shfl_sync( + FINAL_MASK, static_cast(pos_l), it - pos_base, 32); + if (lane < kCosSinBytes / 16) { + cpAsync16(ring_smem + (it % kStages) * stage_bytes + + kSteps * kStepRows * kRowPitch + lane * 16, + cos_sin_cache + pos * kRotaryDim + lane * 8); + } + } + } + cpAsyncCommit(); + }; + + // The prologue's rows do not wait for the positions. + #pragma unroll + for (int it = 0; it < kStages - 1; it++) issue_rows(it); + #pragma unroll + for (int it = 0; it < kStages - 1; it++) issue_cos_sin(it); + + // This lane's dims 4v .. 4v + 3 (+ 32c) of row g of the first step, at the + // first token; each iteration advances them by stride tokens. Dims + // 32c + 4v of a row are lane 8c + v's dims in the kernel above. + scalar_t* in_place = qkv + static_cast(t0) * qkv_row + + (qkv_head0 + g) * kHeadDim + 4 * v; + int const head_rows = kRole == kQ ? nq : kRole == kIQ ? niq : 0; + int64_t const out_lane = + (static_cast(t0) * head_rows + head0 + g) * kHeadDim + 4 * v; + int64_t const out_step = static_cast(stride) * head_rows * kHeadDim; + int64_t out = out_lane; + + for (int it = 0; it < n; it++) { + issue_rows(it + kStages - 1); + issue_cos_sin(it + kStages - 1); + cpAsyncWait(); + __syncwarp(); + uint8_t const* const stage = ring + (it % kStages) * stage_bytes; + float cos4[4], sin4[4]; + if constexpr (kNorm) { + scalar_t const* const cs = reinterpret_cast( + stage + kSteps * kStepRows * kRowPitch); + if constexpr (!kOneRow) { + loadElems(cs + 4 * v, cos4); + loadElems(cs + kRotaryDim / 2 + 4 * v, sin4); + } + } + // Slot of the token: block and offset with one 32-bit division (the host + // checks that every slot fits). + int sm = -1; + uint32_t blk = 0, tok = 0; + if constexpr (kSlots) { + if (it - slot_base == 32) { + slot_base = it; + slot_l = fetch(slot_mapping, slot_base, -1); + } + sm = static_cast(__shfl_sync( + FINAL_MASK, static_cast(slot_l), it - slot_base, 32)); + blk = static_cast(sm) / static_cast(block_size); + tok = static_cast(sm) - blk * block_size; + } + + #pragma unroll + for (int step = 0; step < kSteps; step++) { + int const r = step * kStepRows + g; // row of the warp + int const head = head0 + r; + bool const row_ok = r < nrows; + int const so = step * kStepRows * kHeadDim; // step offset, elements + float e[kChunks][kElemsPerLane]; + if constexpr (kOneRow) { + loadElems( + reinterpret_cast(stage + r * kRowPitch) + 4 * v, + e[0]); + if constexpr (kNorm) { + normAndRope( + e[0], lane, eps, norm_w, /*do_rope=*/true, kRotaryDim, + reinterpret_cast(stage + kSteps * kRowPitch), + /*apply_norm=*/true); + } + } else { + Row x; + loadRow(stage + r * kRowPitch, v, x); + if constexpr (kNorm) normRope(x, w1, eps, cos4, sin4); + #pragma unroll + for (int c = 0; c < kChunks; c++) chunkElems(x, c, e[c]); + } + if constexpr (kRole == kQ) { + if (row_ok) { + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + if constexpr (kOut == kQFp8Unit) { + storeUnitScaleQElemsFp8(q_fp8_out + out + so + 32 * c, + e[c]); + } else if constexpr (kOut == kQFp8) { + storeScaledQElemsFp8(q_fp8_out + out + so + 32 * c, + e[c], q_fp8_inv_scale); + } else { + if (q_out != nullptr) { + storeElems(q_out + out + so + 32 * c, e[c]); + } else if (q_fp8_out == nullptr && write_qkv) { + storeElems(in_place + so + 32 * c, e[c]); + } + if (q_fp8_out != nullptr) { + storeScaledQElemsFp8(q_fp8_out + out + so + 32 * c, + e[c], q_fp8_inv_scale); + } + } + } + } + } else if constexpr (kRole == kIQ) { + if (row_ok) { + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + if constexpr (kOut == kIQOut) { + if constexpr (kFp8Idx) { + storeIndexQElemsFp8( + reinterpret_cast(index_q_out + out + so + 32 * c), + e[c]); + } else { + storeElems(index_q_out + out + so + 32 * c, e[c]); + } + } else if constexpr (kOut == kIQInPlace) { + storeElems(in_place + so + 32 * c, e[c]); + } + } + } + } else if constexpr (kRole == kK || kRole == kIK) { + if constexpr (kOut == kInPlace) { + if (row_ok) { + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + storeElems(in_place + so + 32 * c, e[c]); + } + } + } + } + + if constexpr (kInsertKV && kRole == kIK) { + if constexpr (kNvfp4Idx) { + uint8_t* const page = + reinterpret_cast(index_cache) + blk * idx_s_block; + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + uint8_t sf; + uint16_t const packed = + nvfp4QuantizeElems(e[c], 1.0f, sf); + if (sm >= 0 && row_ok) { + nvfp4StoreElems(page, /*kv=*/0, /*head=*/0, tok, block_size, + kOneRow ? lane : 8 * c + v, packed, sf); + } + } + } else if (sm >= 0 && row_ok) { + int64_t const off = static_cast(sm) * kHeadDim + 4 * v; + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + if constexpr (kFp8Idx) { + storeElemsFp8( + reinterpret_cast(index_cache) + off + 32 * c, e[c]); + } else { + storeElems(index_cache + off + 32 * c, e[c]); + } + } + } + } else if constexpr (kInsertKV && (kRole == kK || kRole == kV)) { + constexpr int kv = kRole == kK ? 0 : 1; + if constexpr (kNvfp4 && !kOneRow) { + // One 16-value scale group per lane, regrouped through the ring + // row: lane (g, v) takes dims 16v .. 16v + 15 of row g. + uint8_t* const row = + ring + (it % kStages) * stage_bytes + r * kRowPitch; + if constexpr (kRole == kK) { + // The k in the model dtype, as the helper would quantize it (V + // rows already hold their model-dtype values). + #pragma unroll + for (int c = 0; c < kChunks; c++) { + storeElems( + reinterpret_cast(row + 64 * c + 8 * v), e[c]); + } + } + __syncwarp(); + float x[16]; + #pragma unroll + for (int k = 0; k < 4; k++) { + float q[kElemsPerLane]; + loadElems( + reinterpret_cast(row + 32 * v + 8 * k), q); + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) x[4 * k + i] = q[i]; + } + uint8_t sf; + uint2 const codes = nvfp4QuantizeGroup(x, kv_scale, sf); + if (sm >= 0 && row_ok) { + nvfp4StoreGroup( + reinterpret_cast(kv_cache) + blk * kv_s_block, kv, + head, tok, block_size, v, codes, sf); + } + } else if constexpr (kNvfp4) { + uint8_t* const page = + reinterpret_cast(kv_cache) + blk * kv_s_block; + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + uint8_t sf; + uint16_t const packed = + nvfp4QuantizeElems(e[c], kv_scale, sf); + if (sm >= 0 && row_ok) { + nvfp4StoreElems(page, kv, head, tok, block_size, + kOneRow ? lane : 8 * c + v, packed, sf); + } + } + } else if (sm >= 0 && row_ok) { + cache_t* const dst = kv_cache + blk * kv_s_block + head * kv_s_head + + tok * kv_s_token + + (kv * kHeadDim + 4 * v) * kv_s_dim; + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + if constexpr (!kHasIndex) { + using Converter = vllm::_typeConvert; + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + e[c][i] = Converter::convert(Converter::convert(e[c][i])); + } + } + storeCacheElems(dst + 32 * c * kv_s_dim, + e[c]); + } + } + } + } + in_place += src_step; + out += out_step; + __syncwarp(); + } +} + +template +__global__ void __launch_bounds__(kMaxWarps * 32) fusedMiniMaxM3RowGroupKernel( + scalar_t* __restrict__ qkv, scalar_t* __restrict__ q_out, + uint8_t* __restrict__ q_fp8_out, out_idx_t* __restrict__ index_q_out, + scalar_t const* __restrict__ q_norm_w, + scalar_t const* __restrict__ k_norm_w, + scalar_t const* __restrict__ iq_norm_w, + scalar_t const* __restrict__ ik_norm_w, + scalar_t const* __restrict__ cos_sin_cache, + int64_t const* __restrict__ positions, + int64_t const* __restrict__ slot_mapping, + int64_t const* __restrict__ index_slot_mapping, + cache_t* __restrict__ kv_cache, out_idx_t* __restrict__ index_cache, + float const eps, float const q_fp8_inv_scale, int const rotary_dim, + int const num_tokens, int const nq, int const nkv, int const niq, + int const block_size, int64_t const kv_s_block, int64_t const kv_s_head, + int64_t const kv_s_token, int64_t const kv_s_dim, int64_t const idx_s_block, + float const* __restrict__ kv_k_scale, float const* __restrict__ kv_v_scale, + bool const write_qkv) { + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + extern __shared__ __align__(16) uint8_t smem[]; + Roles const roles(kQHeads, nq, nkv, niq, kInsertKV, kProcessIndex); + int const qkv_row = (nq + 2 * nkv + (kHasIndex ? niq + 1 : 0)) * kHeadDim; + // CTA (x, y) runs warps [y * cta_warps, (y + 1) * cta_warps) of the warp + // groups of tokens x, x + gridDim.x, ... + int const cta_warps = blockDim.x / 32; + int w = blockIdx.y * cta_warps + threadIdx.x / 32; + if (w >= roles.warps) return; + constexpr int kQStage = stageBytes(); + constexpr int kSideStage = stageBytes(); + int const q_ring = kStages * kQStage; + int const side_ring = kStages * kSideStage; + uint8_t* const ring_w = + smem + roles.ringOffset(w, q_ring, side_ring) - + roles.ringOffset(blockIdx.y * cta_warps, q_ring, side_ring); + #define RG_ARGS(NORM_W, SLOTS, CACHE_SCALE) \ + qkv, q_out, q_fp8_out, index_q_out, NORM_W, cos_sin_cache, positions, \ + SLOTS, kv_cache, index_cache, eps, q_fp8_inv_scale, num_tokens, nq, \ + niq, qkv_row, block_size, kv_s_block, kv_s_head, kv_s_token, \ + kv_s_dim, idx_s_block, CACHE_SCALE, write_qkv + #define RG_RUN(ROLE, STEPS, OUT) \ + runWarp + #define RG_RUN1(ROLE, OUT) \ + runWarp + if (w < roles.q_warps) { + int const h0 = w * kQHeads; + constexpr int kSteps = kQHeads / kRows; + uint8_t* const ring = ring_w; + int const rows = min(kQHeads, nq - h0); + if (q_out == nullptr && q_fp8_out != nullptr) { + if (q_fp8_inv_scale == 1.0f) { + RG_RUN(kQ, kSteps, kQFp8Unit) + (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f)); + } else { + RG_RUN(kQ, kSteps, kQFp8) + (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f)); + } + } else { + RG_RUN(kQ, kSteps, kQAny) + (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f)); + } + } else { + uint8_t* const ring = ring_w; + w -= roles.q_warps; + if (w < roles.kv_warps) { + int const h0 = roles.kv_one ? w : w * kRows; + float const scale = kNvfp4 ? *kv_k_scale : 1.0f; + int const rows = roles.kv_one ? 1 : min(kRows, nkv - h0); + if (rows == 1) { + if (write_qkv) { + RG_RUN1(kK, kInPlace) + (h0, 1, nq + h0, ring, kSideStage, + RG_ARGS(k_norm_w, slot_mapping, scale)); + } else { + RG_RUN1(kK, kNone) + (h0, 1, nq + h0, ring, kSideStage, + RG_ARGS(k_norm_w, slot_mapping, scale)); + } + } else if (write_qkv) { + RG_RUN(kK, 1, kInPlace) + (h0, rows, nq + h0, ring, kSideStage, + RG_ARGS(k_norm_w, slot_mapping, scale)); + } else { + RG_RUN(kK, 1, kNone) + (h0, rows, nq + h0, ring, kSideStage, + RG_ARGS(k_norm_w, slot_mapping, scale)); + } + } else if ((w -= roles.kv_warps) < roles.v_warps) { + if constexpr (kInsertKV) { + int const h0 = roles.kv_one ? w : w * kRows; + int const rows = roles.kv_one ? 1 : min(kRows, nkv - h0); + float const scale = kNvfp4 ? *kv_v_scale : 1.0f; + if (rows == 1) { + RG_RUN1(kV, kNone) + (h0, 1, nq + nkv + h0, ring, kSideStage, + RG_ARGS(nullptr, slot_mapping, scale)); + } else { + RG_RUN(kV, 1, kNone) + (h0, rows, nq + nkv + h0, ring, kSideStage, + RG_ARGS(nullptr, slot_mapping, scale)); + } + } + } else if ((w -= roles.v_warps) < roles.iq_warps) { + if constexpr (kProcessIndex) { + int const h0 = roles.iq_one ? w : w * kRows; + int const rows = roles.iq_one ? 1 : min(kRows, niq - h0); + int const qh0 = nq + 2 * nkv + h0; + if (rows == 1) { + if (index_q_out != nullptr) { + RG_RUN1(kIQ, kIQOut) + (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } else if (write_qkv) { + RG_RUN1(kIQ, kIQInPlace) + (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } else { + RG_RUN1(kIQ, kNone) + (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } + } else if (index_q_out != nullptr) { + RG_RUN(kIQ, 1, kIQOut) + (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } else if (write_qkv) { + RG_RUN(kIQ, 1, kIQInPlace) + (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } else { + RG_RUN(kIQ, 1, kNone) + (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } + } + } else if constexpr (kProcessIndex) { + if (write_qkv) { + RG_RUN1(kIK, kInPlace) + (0, 1, nq + 2 * nkv + niq, ring, kSideStage, + RG_ARGS(ik_norm_w, index_slot_mapping, 1.0f)); + } else { + RG_RUN1(kIK, kNone) + (0, 1, nq + 2 * nkv + niq, ring, kSideStage, + RG_ARGS(ik_norm_w, index_slot_mapping, 1.0f)); + } + } + } + #undef RG_RUN + #undef RG_RUN1 + #undef RG_ARGS + #if __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); + #endif + #endif +} + +// Shared memory of the largest CTA of cta_warps warps. +template +__host__ inline int smemBytes(Roles const& r, int const cta_warps) { + int const q_ring = kStages * stageBytes(); + int const side_ring = kStages * stageBytes(); + int bytes = 0; + for (int w0 = 0; w0 < r.warps; w0 += cta_warps) { + int const w1 = std::min(w0 + cta_warps, r.warps); + bytes = std::max(bytes, r.ringOffset(w1, q_ring, side_ring) - + r.ringOffset(w0, q_ring, side_ring)); + } + return bytes; +} + +// Launch with the grid this instantiation's occupancy allows: the CTAs walk +// the tokens with stride gridDim.x. +template +void launch(cudaLaunchConfig_t config, int const num_tokens, Args... args) { + auto* const kernel = + fusedMiniMaxM3RowGroupKernel; + // Per CTA shape (warps, shared memory), computed once: the resident CTAs + // per SM. + static bool const opted_in = + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + get_device_prop()->sharedMemPerBlockOptin) == + cudaSuccess; + static int smem_of[kMaxWarps + 1] = {}; + static int resident_of[kMaxWarps + 1] = {}; + int const warps = config.blockDim.x / 32; + int const smem = static_cast(config.dynamicSmemBytes); + if (smem_of[warps] != smem) { + (void)opted_in; + int n = 0; + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&n, kernel, warps * 32, smem); + resident_of[warps] = n > 0 ? n : 1; + smem_of[warps] = smem; + } + int const ctas_per_token = config.gridDim.y; + config.gridDim.x = std::max( + 1, std::min(num_tokens, get_device_prop()->multiProcessorCount * + resident_of[warps] / ctas_per_token)); + cudaLaunchKernelEx(&config, kernel, args...); +} + +} // namespace rg +#endif // USE_ROCM + +#ifndef USE_ROCM +// q heads per q warp: four (one step, shorter warps) for a launch of at most +// 128 q rows, or while a token has at most 32 q heads and there are few q rows +// per SM; eight otherwise (fewer, longer warps keep more CTAs resident). +inline int rgQHeads(int const num_tokens, int const nq) { + static int const num_sms = get_device_prop()->multiProcessorCount; + int64_t const q_rows = static_cast(num_tokens) * nq; + return q_rows <= 128 || (nq <= 32 && q_rows < 64 * num_sms) ? 4 : 8; +} +// Warps per CTA: a token's whole warp group once there are enough tokens to +// give every SM CTAs; before that, eight, which spreads a token over SMs. +inline int rgCtaWarps(int const num_tokens) { + static int const num_sms = get_device_prop()->multiProcessorCount; + return 2 * num_tokens >= 3 * num_sms ? rg::kMaxWarps : 8; +} +#endif + // ──────────────────────────────────────────────────────────────────────────── // Launch wrapper // ──────────────────────────────────────────────────────────────────────────── @@ -809,9 +1683,10 @@ void launchFusedMiniMaxM3( float const q_fp8_inv_scale, int const rotary_dim, int const num_tokens, int const nq, int const nkv, int const niq, int const block_size, int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, - int64_t const kv_s_dim, bool const has_index, bool const insert_kv, - bool const process_index, bool const fp8_idx, float const* kv_k_scale, - float const* kv_v_scale, cudaStream_t stream) { + int64_t const kv_s_dim, int64_t const idx_s_block, bool const has_index, + bool const insert_kv, bool const process_index, bool const fp8_idx, + bool const nvfp4_idx, float const* kv_k_scale, float const* kv_v_scale, + bool const write_qkv, bool const slots_fit_int32, cudaStream_t stream) { // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the // void* pointers per instantiation in the LAUNCH macro. // Slot count must match the kernel's compile-time gating. @@ -843,73 +1718,437 @@ void launchFusedMiniMaxM3( config.attrs = attrs; config.numAttrs = (sm_version >= 90) ? 1 : 0; - #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \ - cudaLaunchKernelEx( \ - &config, \ - fusedMiniMaxM3QNormRopeKVInsertKernel, \ - qkv, q_out, q_fp8_out, reinterpret_cast(index_q_out), \ - q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ + // The row-group kernel: same bytes, for rotary_dim == 64 on SM80+ (cp.async). + int const q_heads = rgQHeads(num_tokens, nq); + rg::Roles const roles(q_heads, nq, nkv, niq, insert_kv, process_index); + auto const rg_smem = [&](int const warps) { + return q_heads == 8 ? rg::smemBytes<8>(roles, warps) + : rg::smemBytes<4>(roles, warps); + }; + int cta_warps = + std::min({roles.warps, rgCtaWarps(num_tokens), rg::kMaxWarps}); + // Eight warps take at most 63 KB; a whole group can exceed the 99 KB per + // block of SM 8.6 / 8.9 / 12.x. + static int const max_block_smem = get_device_prop()->sharedMemPerBlockOptin; + if (rg_smem(cta_warps) > max_block_smem) cta_warps = std::min(roles.warps, 8); + auto const aligned = [](void const* p, uintptr_t const n) { + return reinterpret_cast(p) % n == 0; + }; + bool const row_groups = + sm_version >= 80 && rotary_dim == rg::kRotaryDim && slots_fit_int32 && + aligned(qkv, 16) && aligned(cos_sin_cache, 16) && aligned(q_norm_w, 8) && + aligned(k_norm_w, 8) && aligned(iq_norm_w, 8) && aligned(ik_norm_w, 8); + cudaLaunchConfig_t rg_config = config; + rg_config.blockDim = dim3(cta_warps * 32); + rg_config.gridDim = dim3(1, (roles.warps + cta_warps - 1) / cta_warps); + rg_config.dynamicSmemBytes = rg_smem(cta_warps); + + #define LAUNCH_ARGS(OUT_T) \ + qkv, q_out, q_fp8_out, reinterpret_cast(index_q_out), q_norm_w, \ + k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ slot_mapping, index_slot_mapping, kv_cache, \ reinterpret_cast(index_cache), eps, q_fp8_inv_scale, \ rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ - kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale) + kv_s_head, kv_s_token, kv_s_dim, idx_s_block, kv_k_scale, kv_v_scale + #define LAUNCH_RG(QH, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, \ + OUT_T) \ + rg::launch(rg_config, num_tokens, \ + LAUNCH_ARGS(OUT_T), write_qkv) + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T) \ + do { \ + if (!row_groups) { \ + cudaLaunchKernelEx( \ + &config, \ + fusedMiniMaxM3QNormRopeKVInsertKernel< \ + scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT, \ + PROCESS_INDEX, FP8, NVFP4_IDX>, \ + LAUNCH_ARGS(OUT_T)); \ + } else if (q_heads == 8) { \ + LAUNCH_RG(8, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T); \ + } else { \ + LAUNCH_RG(4, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T); \ + } \ + } while (0) #else // ROCm: standard kernel launch syntax (no PDL/stream serialization). // clang-format off - #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \ + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T) \ fusedMiniMaxM3QNormRopeKVInsertKernel< \ scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT, \ - PROCESS_INDEX, FP8><<>>( \ + PROCESS_INDEX, FP8, NVFP4_IDX><<>>( \ qkv, q_out, q_fp8_out, reinterpret_cast(index_q_out), \ q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ slot_mapping, index_slot_mapping, kv_cache, \ reinterpret_cast(index_cache), eps, q_fp8_inv_scale, \ rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ - kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale) + kv_s_head, kv_s_token, kv_s_dim, idx_s_block, kv_k_scale, \ + kv_v_scale) // clang-format on #endif if constexpr (kNvfp4) { // NVFP4 main cache: sparse serving only (the host checks insert_kv). if (!process_index) { - LAUNCH(true, true, false, false, scalar_t); + LAUNCH(true, true, false, false, false, scalar_t); + } else if (nvfp4_idx) { + LAUNCH(true, true, true, true, true, uint8_t); } else if (fp8_idx) { - LAUNCH(true, true, true, true, uint8_t); + LAUNCH(true, true, true, true, false, uint8_t); } else { - LAUNCH(true, true, true, false, scalar_t); + LAUNCH(true, true, true, false, false, scalar_t); } } else if (has_index) { if (!process_index) { if (insert_kv) { - LAUNCH(true, true, false, false, scalar_t); + LAUNCH(true, true, false, false, false, scalar_t); } else { - LAUNCH(true, false, false, false, scalar_t); + LAUNCH(true, false, false, false, false, scalar_t); } } else if (insert_kv) { - if (fp8_idx) { - LAUNCH(true, true, true, true, + if (nvfp4_idx) { + LAUNCH(true, true, true, true, true, + uint8_t); // sparse serving, NVFP4 index-K, fp8 index_q + } else if (fp8_idx) { + LAUNCH(true, true, true, true, false, uint8_t); // sparse serving, fp8 index outputs } else { - LAUNCH(true, true, true, false, scalar_t); // sparse serving, bf16 + LAUNCH(true, true, true, false, false, + scalar_t); // sparse serving, bf16 } } else { if (fp8_idx) { - LAUNCH(true, false, true, true, + LAUNCH(true, false, true, true, false, uint8_t); // sparse profiling, fp8 index_q } else { - LAUNCH(true, false, true, false, scalar_t); // sparse profiling, bf16 + LAUNCH(true, false, true, false, false, + scalar_t); // sparse profiling, bf16 } } + } else if (insert_kv) { + LAUNCH(false, true, false, false, false, scalar_t); // dense, KV insert } else { - // Dense layer: never has an index branch and never inserts here (the - // generic Attention layer owns the KV insert). - LAUNCH(false, false, false, false, scalar_t); + LAUNCH(false, false, false, false, false, scalar_t); // dense } #undef LAUNCH +#undef LAUNCH_RG +#undef LAUNCH_ARGS +} + +#ifndef USE_ROCM +// ──────────────────────────────────────────────────────────────────────────── +// DSpark draft context K/V: every draft layer in one launch (bf16 in). +// ──────────────────────────────────────────────────────────────────────────── +// The draft projects the target's context states to the K/V of all its layers +// with one GEMM, ``all_kv [N, L, 2, nkv, 128]``. blockIdx.z picks the layer, +// blockIdx.y the head, blockIdx.x a run of tokens. Eight lanes own one token; +// lane l holds two 8-dim chunks of K and of V, A = dims [8l, 8l + 8) and +// B = dims [64 + 8l, 64 + 8l + 8), so with rotary_dim == 64 every lane rotates +// its A chunk (partner lane l ^ 4) and no lane idles through the RoPE. K gets +// the layer's Gemma k-norm (rounded to bf16 like the unfused norm) and partial +// NeoX RoPE (fp32 math rounded once, like rotary_embedding); K and V go into +// the token's cache row like reshape_and_cache_flash (fp8: e4m3(bf16(x) * +// (1 / scale)); a unit scale converts the bf16 pairs directly). +// Layer caches are separate allocations of one layout, logical shape +// [nb, nkv, bs, 2*128] with K at [..., :128] and V at [..., 128:]; they are +// addressed as layer 0's cache plus a per-layer element offset (likewise the +// fp8 k/v scales from layer 0's k scale), so the launch is CUDA-graph safe. +constexpr int kCtxLanesPerRow = 8; +constexpr int kCtxElems = 8; // bf16 per 16-byte chunk +constexpr int kCtxHalfVecs = kHeadDim / 2 / kCtxElems; // 8: chunk B offset +constexpr int kCtxBlockSize = 256; +constexpr int kCtxRowsPerBlock = kCtxBlockSize / kCtxLanesPerRow; // 32 + +__device__ __forceinline__ float bf16lo(uint32_t v) { + return __uint_as_float(v << 16); +} +__device__ __forceinline__ float bf16hi(uint32_t v) { + return __uint_as_float(v & 0xffff0000u); +} +__device__ __forceinline__ uint32_t packBf16(float lo, float hi) { + __nv_bfloat162 const r = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&r); +} +// Loads of what the PDL predecessor may still be writing: coherent (an .nc +// load may be hoisted above griddepcontrol.wait). +__device__ __forceinline__ uint4 ldgV4(uint4 const* p) { + uint4 r; + asm volatile("ld.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} +// Low 32 bits of an int64 (positions and slots fit). +__device__ __forceinline__ int ldgLo32(int64_t const* p) { + int r; + asm volatile("ld.global.b32 %0, [%1];" : "=r"(r) : "l"(p)); + return r; +} +// Same approximation as flashinfer's gemma_rmsnorm. +__device__ __forceinline__ float rsqrtApproxFtz(float x) { + float r; + asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} +// e4m3 of four floats, low byte first. +__device__ __forceinline__ uint32_t packE4m3(float a, float b, float c, + float d) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890 + uint32_t r; + asm("{\n" + " .reg .b16 lo, hi;\n" + " cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" + " cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" + " mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(r) + : "f"(a), "f"(b), "f"(c), "f"(d)); + return r; + #else + uint32_t const lo = + __nv_cvt_float2_to_fp8x2(make_float2(a, b), __NV_SATFINITE, __NV_E4M3); + uint32_t const hi = + __nv_cvt_float2_to_fp8x2(make_float2(c, d), __NV_SATFINITE, __NV_E4M3); + return lo | (hi << 16); + #endif +} +// Eight bf16 (4 packed pairs) * inv_scale -> e4m3, low byte first. +__device__ __forceinline__ uint2 bf16x8ToE4m3(uint4 const p, + float const inv_scale) { + uint32_t const w[4] = {p.x, p.y, p.z, p.w}; + uint32_t r[2]; + #pragma unroll + for (int i = 0; i < 2; i++) { + r[i] = packE4m3(bf16lo(w[2 * i]) * inv_scale, bf16hi(w[2 * i]) * inv_scale, + bf16lo(w[2 * i + 1]) * inv_scale, + bf16hi(w[2 * i + 1]) * inv_scale); + } + return make_uint2(r[0], r[1]); +} + #if ((defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && \ + __CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000) || \ + (defined(__CUDA_ARCH_SPECIFIC__) && __CUDA_ARCH_SPECIFIC__ >= 1000)) && \ + defined(CUDART_VERSION) && CUDART_VERSION >= 13010 + #define VLLM_CTX_KV_BF16X2_TO_E4M3X2 + #endif +// Eight bf16 -> e4m3 with a unit scale (exact); one convert per pair where +// the arch and toolkit (PTX ISA 9.1+) convert bf16x2 -> e4m3x2. +__device__ __forceinline__ uint2 bf16x8ToE4m3(uint4 const p) { + #ifdef VLLM_CTX_KV_BF16X2_TO_E4M3X2 + uint32_t r[2]; + uint32_t const w[4] = {p.x, p.y, p.z, p.w}; + #pragma unroll + for (int i = 0; i < 2; i++) { + asm("{\n" + " .reg .b16 lo, hi;\n" + " cvt.rn.satfinite.e4m3x2.bf16x2 lo, %1;\n" + " cvt.rn.satfinite.e4m3x2.bf16x2 hi, %2;\n" + " mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(r[i]) + : "r"(w[2 * i]), "r"(w[2 * i + 1])); + } + return make_uint2(r[0], r[1]); + #else + return bf16x8ToE4m3(p, 1.0f); + #endif } +template +__global__ void __launch_bounds__(kCtxBlockSize) + fusedMiniMaxM3DSparkContextKVKernel( + uint4 const* __restrict__ all_kv, // bf16 [N, L, 2, nkv, 128] + uint4 const* __restrict__ k_norm_w, // bf16 [L, 128] + uint4 const* __restrict__ cos_sin_cache, // bf16 [max_pos, rot] + int64_t const* __restrict__ positions, // [N] + int64_t const* __restrict__ slot_mapping, // [N], < 0: no write + uint8_t* __restrict__ kv_cache, // layer 0's cache, bytes + int64_t const* __restrict__ cache_offsets, // [L] elements from it + float const* __restrict__ scales, // layer 0's k scale + int64_t const* __restrict__ scale_offsets, // [2, L] k, v; floats + float const eps, int const rotary_dim, int const num_tokens, + int const num_layers, int const block_shift, int const s_token, + int const s_layer, int const s_side, int const s_head, + int const kv_s_block, int const kv_s_head, int const kv_s_token) { + // Strides are in 16-byte (8 x bf16) units for all_kv / cos_sin, elements + // for the cache. + constexpr int kElemBytes = kFp8 ? 1 : 2; + int const row0 = blockIdx.x * kCtxRowsPerBlock; + // Warps with no token exit; the others run their tail rows too (the + // shuffles need every lane) and skip the store. + if (row0 + threadIdx.x / 32 * (32 / kCtxLanesPerRow) >= num_tokens) return; + int const head = blockIdx.y; + int const layer = blockIdx.z; + int const lane = threadIdx.x % kCtxLanesPerRow; + int const dim0 = lane * kCtxElems; // chunk A; chunk B is 64 dims later + int const t = row0 + threadIdx.x / kCtxLanesPerRow; + + int const half = rotary_dim / 2; + bool const in_rope = dim0 < rotary_dim; + bool const first_half = dim0 < half; + int const cs_vec = + in_rope ? (first_half ? dim0 : dim0 - half) / kCtxElems : 0; + int const lane_xor = half / kCtxElems; + int const rot_vecs = rotary_dim / kCtxElems; + // x * cos - partner * sin in the first half, + in the second. + uint32_t const sin_sign = first_half ? 0x80008000u : 0u; + + uint8_t* const cache = + kv_cache + (cache_offsets[layer] + head * kv_s_head + dim0) * kElemBytes; + float k_scale = 1.0f, v_scale = 1.0f; + if constexpr (kFp8) { + int64_t const k_off = scale_offsets[layer]; + int64_t const v_off = scale_offsets[num_layers + layer]; + k_scale = scales[k_off]; + v_scale = scales[v_off]; + } + float g[2][kCtxElems]; // 1 + w, chunks A and B + #pragma unroll + for (int c = 0; c < 2; c++) { + uint4 const w = + k_norm_w[layer * (kHeadDim / kCtxElems) + c * kCtxHalfVecs + lane]; + uint32_t const wp[4] = {w.x, w.y, w.z, w.w}; + #pragma unroll + for (int i = 0; i < 4; i++) { + g[c][2 * i] = 1.0f + bf16lo(wp[i]); + g[c][2 * i + 1] = 1.0f + bf16hi(wp[i]); + } + } + + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); + #endif + + bool const valid = t < num_tokens; + int const ts = valid ? t : 0; + int const pos = ldgLo32(positions + ts); + int const slot = valid ? ldgLo32(slot_mapping + ts) : -1; + uint4 const* const src = + all_kv + ts * s_token + layer * s_layer + head * s_head + lane; + uint4 kr[2], vr[2]; + #pragma unroll + for (int c = 0; c < 2; c++) { + kr[c] = ldgV4(src + c * kCtxHalfVecs); + vr[c] = ldgV4(src + s_side + c * kCtxHalfVecs); + } + // Lanes past the rotary dims load (and later discard) a valid row too. + uint4 const* const cs = cos_sin_cache + pos * rot_vecs + cs_vec; + uint4 const cr = cs[0]; + uint4 const sr = cs[lane_xor]; + + bool const store = slot >= 0; + // Power-of-two block size: shift and mask, not division. + uint8_t* const dst = + cache + (static_cast(slot >> block_shift) * kv_s_block + + (slot & ((1 << block_shift) - 1)) * kv_s_token) * + kElemBytes; + constexpr int kChunkB = kHeadDim / 2 * kElemBytes; + constexpr int kV = kHeadDim * kElemBytes; + float const k_inv = __fdividef(1.0f, k_scale); + float const v_inv = __fdividef(1.0f, v_scale); + #ifdef VLLM_CTX_KV_BF16X2_TO_E4M3X2 + bool const unit_scale = k_scale == 1.0f && v_scale == 1.0f; + #else + constexpr bool unit_scale = false; + #endif + // Eight bf16 -> the cache row at byte offset `off`. + auto put = [&](int const off, uint4 const x, float const inv_scale) { + if constexpr (kFp8) { + uint2 const y = unit_scale ? bf16x8ToE4m3(x) : bf16x8ToE4m3(x, inv_scale); + if (store) *reinterpret_cast(dst + off) = y; + } else { + if (store) *reinterpret_cast(dst + off) = x; + } + }; + + // V first, so its registers free up before the K math. + #pragma unroll + for (int c = 0; c < 2; c++) put(kV + c * kChunkB, vr[c], v_inv); + + // K: Gemma RMSNorm over the 8 lanes of the row. + float x[2][kCtxElems]; + float sumsq = 0.0f; + #pragma unroll + for (int c = 0; c < 2; c++) { + uint32_t const kp[4] = {kr[c].x, kr[c].y, kr[c].z, kr[c].w}; + #pragma unroll + for (int j = 0; j < 4; j++) { + x[c][2 * j] = bf16lo(kp[j]); + x[c][2 * j + 1] = bf16hi(kp[j]); + sumsq = fmaf(x[c][2 * j], x[c][2 * j], sumsq); + sumsq = fmaf(x[c][2 * j + 1], x[c][2 * j + 1], sumsq); + } + } + #pragma unroll + for (int mask = kCtxLanesPerRow / 2; mask > 0; mask >>= 1) { + sumsq += __shfl_xor_sync(FINAL_MASK, sumsq, mask, 32); + } + float const rr = rsqrtApproxFtz(sumsq * (1.0f / kHeadDim) + eps); + // bf16-rounded norm output, packed. + uint32_t kn[2][4]; + #pragma unroll + for (int c = 0; c < 2; c++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + kn[c][j] = packBf16(x[c][2 * j] * rr * g[c][2 * j], + x[c][2 * j + 1] * rr * g[c][2 * j + 1]); + } + } + // RoPE on chunk A; the partner chunk is half / 8 lanes away. + uint32_t const cp[4] = {cr.x, cr.y, cr.z, cr.w}; + uint32_t const sp[4] = {sr.x ^ sin_sign, sr.y ^ sin_sign, sr.z ^ sin_sign, + sr.w ^ sin_sign}; + #pragma unroll + for (int j = 0; j < 4; j++) { + uint32_t const pn = __shfl_xor_sync(FINAL_MASK, kn[0][j], lane_xor, 32); + uint32_t const rot = + packBf16(__fmaf_rn(bf16lo(kn[0][j]), bf16lo(cp[j]), + __fmul_rn(bf16lo(pn), bf16lo(sp[j]))), + __fmaf_rn(bf16hi(kn[0][j]), bf16hi(cp[j]), + __fmul_rn(bf16hi(pn), bf16hi(sp[j])))); + kn[0][j] = in_rope ? rot : kn[0][j]; + } + #pragma unroll + for (int c = 0; c < 2; c++) { + put(c * kChunkB, make_uint4(kn[c][0], kn[c][1], kn[c][2], kn[c][3]), k_inv); + } + + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); + #endif +} + +template +void launchFusedMiniMaxM3DSparkContextKV( + uint4 const* all_kv, uint4 const* k_norm_w, uint4 const* cos_sin_cache, + int64_t const* positions, int64_t const* slot_mapping, uint8_t* kv_cache, + int64_t const* cache_offsets, float const* scales, + int64_t const* scale_offsets, float const eps, int const rotary_dim, + int const num_tokens, int const num_layers, int const nkv, + int const block_shift, int const s_token, int const s_layer, + int const s_side, int const s_head, int const kv_s_block, + int const kv_s_head, int const kv_s_token, cudaStream_t stream) { + if (num_tokens == 0) return; + static int const sm_version = getSMVersion(); + cudaLaunchConfig_t config; + config.gridDim = dim3((num_tokens + kCtxRowsPerBlock - 1) / kCtxRowsPerBlock, + nkv, num_layers); + config.blockDim = dim3(kCtxBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = (sm_version >= 90) ? 1 : 0; + cudaLaunchKernelEx(&config, fusedMiniMaxM3DSparkContextKVKernel, all_kv, + k_norm_w, cos_sin_cache, positions, slot_mapping, kv_cache, + cache_offsets, scales, scale_offsets, eps, rotary_dim, + num_tokens, num_layers, block_shift, s_token, s_layer, + s_side, s_head, kv_s_block, kv_s_head, kv_s_token); +} +#endif // !USE_ROCM + } // namespace minimax_m3_fused_ops } // namespace vllm @@ -948,8 +2187,9 @@ void launchFusedMiniMaxM3( static_cast(eps), 1.0f / static_cast(q_fp8_scale), \ static_cast(rotary_dim), num_tokens, nq, nkv, niq, \ static_cast(block_size), kv_s_block, kv_s_head, kv_s_token, \ - kv_s_dim, has_index, insert_kv, process_index, fp8_idx, kv_k_scale_ptr, \ - kv_v_scale_ptr, stream) + kv_s_dim, idx_s_block, has_index, insert_kv, process_index, fp8_idx, \ + nvfp4_idx, kv_k_scale_ptr, kv_v_scale_ptr, write_qkv, slots_fit_int32, \ + stream) #define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \ CALL_FUSED_MINIMAX_M3_IMPL(CACHE_T, KV_DTYPE, false) // clang-format on @@ -980,7 +2220,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert( q_fp8_out, // [N, nq*128] contiguous E4M3 double q_fp8_scale, std::optional kv_k_scale, // [] f32, NVFP4 only - std::optional kv_v_scale) { + std::optional kv_v_scale, bool write_qkv) { STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(), "qkv must be contiguous CUDA"); STD_TORCH_CHECK( @@ -1032,11 +2272,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert( " + num_index_heads + 1) * 128 for sparse, " "(num_heads + 2*num_kv_heads) * 128 for dense"); - // Only the sparse layer inserts here (dense lets the generic Attention layer - // own the KV write); there is no dense+insert kernel instantiation. - STD_TORCH_CHECK( - !insert_kv || has_index, - "insert mode (kv_cache) requires the index branch (sparse layer)"); + STD_TORCH_CHECK(!nvfp4_kv || has_index, + "nvfp4 kv_cache requires the index branch (sparse layer)"); STD_TORCH_CHECK(has_index || !skip_index_branch, "skip_index_branch requires sparse qkv rows"); if (process_index) { @@ -1084,14 +2321,18 @@ void fused_minimax_m3_qknorm_rope_kv_insert( kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte, "fp8 kv_cache must use uint8 storage"); } - // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3. + // Indexer index-K cache: independent dtype -- qkv dtype, fp8 e4m3, or + // uint8 packed NVFP4 pages. if (process_index) { STD_TORCH_CHECK( index_cache.has_value() && (index_cache->scalar_type() == qkv.scalar_type() || index_cache->scalar_type() == - torch::headeronly::ScalarType::Float8_e4m3fn), - "insert mode requires index_cache matching qkv dtype or fp8 e4m3"); + torch::headeronly::ScalarType::Float8_e4m3fn || + index_cache->scalar_type() == + torch::headeronly::ScalarType::Byte), + "insert mode requires index_cache matching qkv dtype, fp8 e4m3, " + "or uint8 (nvfp4)"); } STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1, "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous " @@ -1150,6 +2391,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert( STD_TORCH_CHECK(std::isfinite(q_fp8_scale) && q_fp8_scale > 0.0, "q_fp8_scale must be finite and positive"); } + STD_TORCH_CHECK(write_qkv || (!q_out.has_value() && q_fp8_out.has_value()), + "write_qkv=False writes q only to q_fp8_out: pass q_fp8_out " + "and no q_out"); if (index_q_out.has_value()) { STD_TORCH_CHECK(process_index, "index_q_out requires index branch processing"); @@ -1166,16 +2410,36 @@ void fused_minimax_m3_qknorm_rope_kv_insert( } // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while - // q/k/v + q_out stay qkv dtype. Both index outputs must agree. + // q/k/v + q_out stay qkv dtype. Both index outputs must agree. The NVFP4 + // index path keeps e4m3 index-Q and stores index-K as packed NVFP4 pages. auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn; + bool const nvfp4_idx = + process_index && insert_kv && + index_cache->scalar_type() == torch::headeronly::ScalarType::Byte; bool const fp8_idx = process_index && - ((index_cache.has_value() && index_cache->scalar_type() == kFp8) || + (nvfp4_idx || + (index_cache.has_value() && index_cache->scalar_type() == kFp8) || (index_q_out.has_value() && index_q_out->scalar_type() == kFp8)); - if (fp8_idx) { + int64_t idx_s_block = 0; + if (nvfp4_idx) { +#ifdef USE_ROCM + STD_TORCH_CHECK(false, "nvfp4 index_cache is not supported on ROCm"); +#endif + int64_t const page_width = kHeadDim / 2 + kHeadDim / 16; STD_TORCH_CHECK( - !index_cache.has_value() || index_cache->scalar_type() == kFp8, - "fp8 index path: index_cache must be fp8 e4m3"); + index_cache->dim() == 3 && index_cache->size(1) == block_size && + index_cache->size(2) == page_width && index_cache->stride(2) == 1 && + index_cache->stride(1) == page_width && + index_cache->stride(0) >= block_size * page_width, + "nvfp4 index_cache must be [nb, block_size, 72] uint8 with " + "contiguous pages"); + idx_s_block = index_cache->stride(0); + } + if (fp8_idx) { + STD_TORCH_CHECK(!index_cache.has_value() || nvfp4_idx || + index_cache->scalar_type() == kFp8, + "fp8 index path: index_cache must be fp8 e4m3"); STD_TORCH_CHECK( !index_q_out.has_value() || index_q_out->scalar_type() == kFp8, "fp8 index path: index_q_out must be fp8 e4m3"); @@ -1188,6 +2452,12 @@ void fused_minimax_m3_qknorm_rope_kv_insert( float const* kv_v_scale_ptr = nvfp4_kv ? kv_v_scale->const_data_ptr() : nullptr; + // The row-group kernel indexes cache slots in 32 bits. + bool const slots_fit_int32 = + !insert_kv || + (kv_cache->size(0) * block_size <= INT32_MAX && + (!process_index || index_cache->size(0) * block_size <= INT32_MAX)); + const torch::stable::accelerator::DeviceGuard device_guard( qkv.get_device_index()); auto stream = get_current_cuda_stream(qkv.get_device_index()); @@ -1216,3 +2486,122 @@ void fused_minimax_m3_qknorm_rope_kv_insert( #undef CALL_FUSED_MINIMAX_M3 #undef CALL_FUSED_MINIMAX_M3_IMPL + +// ──────────────────────────────────────────────────────────────────────────── +// DSpark context K/V torch op +// ──────────────────────────────────────────────────────────────────────────── +void fused_minimax_m3_dspark_context_kv( + torch::stable::Tensor const& all_kv, // bf16 [N, L, 2, nkv, 128] + torch::stable::Tensor const& k_norm_weight, // bf16 [L, 128] + torch::stable::Tensor const& cos_sin_cache, // bf16 [max_pos, rotary_dim] + torch::stable::Tensor const& positions, // [N] i64 + torch::stable::Tensor const& slot_mapping, // [N] i64 + torch::stable::Tensor& kv_cache, // layer 0's [nb, nkv, bs, 2*128] + torch::stable::Tensor const& cache_offsets, // [L] i64, elements + torch::stable::Tensor const& k_scale, // layer 0's [] f32 k scale + torch::stable::Tensor const& scale_offsets, // [2, L] i64, floats + int64_t rotary_dim, double eps, const std::string& kv_cache_dtype) { +#ifdef USE_ROCM + STD_TORCH_CHECK(false, "fused_minimax_m3_dspark_context_kv is CUDA-only"); +#else + using torch::headeronly::ScalarType; + int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim; + STD_TORCH_CHECK(all_kv.is_cuda() && all_kv.dim() == 5 && + all_kv.size(2) == 2 && all_kv.size(4) == kHeadDim && + all_kv.stride(4) == 1 && + all_kv.scalar_type() == ScalarType::BFloat16, + "all_kv must be bf16 CUDA [N, L, 2, nkv, 128], last dim " + "contiguous"); + for (int d = 0; d < 4; d++) { + STD_TORCH_CHECK(all_kv.stride(d) % 8 == 0, + "all_kv strides must be multiples of 8 elements"); + } + STD_TORCH_CHECK(reinterpret_cast(all_kv.data_ptr()) % 16 == 0, + "all_kv must be 16-byte aligned"); + int const num_tokens = static_cast(all_kv.size(0)); + int const num_layers = static_cast(all_kv.size(1)); + int const nkv = static_cast(all_kv.size(3)); + STD_TORCH_CHECK(k_norm_weight.is_contiguous() && + k_norm_weight.scalar_type() == ScalarType::BFloat16 && + k_norm_weight.numel() == num_layers * kHeadDim, + "k_norm_weight must be contiguous bf16 [L, 128]"); + STD_TORCH_CHECK(cos_sin_cache.is_contiguous() && + cos_sin_cache.scalar_type() == ScalarType::BFloat16 && + cos_sin_cache.dim() == 2 && + cos_sin_cache.size(1) == rotary_dim, + "cos_sin_cache must be contiguous bf16 [max_pos, " + "rotary_dim]"); + STD_TORCH_CHECK(rotary_dim == 16 || rotary_dim == 32 || rotary_dim == 64, + "rotary_dim must be 16, 32 or 64"); + for (auto const* t : {&positions, &slot_mapping}) { + STD_TORCH_CHECK(t->is_cuda() && t->is_contiguous() && + t->scalar_type() == ScalarType::Long && + t->numel() == num_tokens, + "positions and slot_mapping must be contiguous int64 [N]"); + } + STD_TORCH_CHECK(cache_offsets.is_cuda() && cache_offsets.is_contiguous() && + cache_offsets.scalar_type() == ScalarType::Long && + cache_offsets.numel() == num_layers, + "cache_offsets must be int64 [L]"); + STD_TORCH_CHECK(scale_offsets.is_cuda() && scale_offsets.is_contiguous() && + scale_offsets.scalar_type() == ScalarType::Long && + scale_offsets.numel() == 2 * num_layers, + "scale_offsets must be int64 [2, L]"); + STD_TORCH_CHECK( + k_scale.is_cuda() && k_scale.scalar_type() == ScalarType::Float, + "k_scale must be a float32 CUDA tensor"); + bool const fp8 = kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3"; + STD_TORCH_CHECK(fp8 || kv_cache_dtype == "auto", + "kv_cache_dtype must be auto, fp8 or fp8_e4m3"); + STD_TORCH_CHECK(fp8 ? kv_cache.scalar_type() == ScalarType::Byte + : kv_cache.scalar_type() == ScalarType::BFloat16, + "kv_cache must be uint8 (fp8) or bf16 (auto)"); + STD_TORCH_CHECK( + kv_cache.dim() == 4 && kv_cache.size(1) == nkv && + kv_cache.size(3) == 2 * kHeadDim && kv_cache.stride(3) == 1 && + kv_cache.stride(0) % 16 == 0 && kv_cache.stride(1) % 16 == 0 && + kv_cache.stride(2) % 16 == 0, + "kv_cache must be [nb, nkv, bs, 256] with a contiguous last dim"); + int64_t const block_size = kv_cache.size(2); + STD_TORCH_CHECK(block_size > 0 && (block_size & (block_size - 1)) == 0, + "kv_cache block size must be a power of two"); + int block_shift = 0; + while ((int64_t{1} << block_shift) < block_size) block_shift++; + int64_t constexpr kIntMax = std::numeric_limits::max(); + STD_TORCH_CHECK(all_kv.numel() / 8 <= kIntMax && + kv_cache.stride(0) <= kIntMax && + kv_cache.size(0) * block_size <= kIntMax, + "all_kv / kv_cache too large for 32-bit indexing"); + + const torch::stable::accelerator::DeviceGuard device_guard( + all_kv.get_device_index()); + auto stream = get_current_cuda_stream(all_kv.get_device_index()); + // all_kv strides in 16-byte vectors. + auto launch = [&](auto fp8_tag) { + constexpr bool kFp8 = decltype(fp8_tag)::value; + vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3DSparkContextKV( + reinterpret_cast(all_kv.data_ptr()), + reinterpret_cast(k_norm_weight.data_ptr()), + reinterpret_cast(cos_sin_cache.data_ptr()), + reinterpret_cast(positions.data_ptr()), + reinterpret_cast(slot_mapping.data_ptr()), + reinterpret_cast(kv_cache.data_ptr()), + reinterpret_cast(cache_offsets.data_ptr()), + reinterpret_cast(k_scale.data_ptr()), + reinterpret_cast(scale_offsets.data_ptr()), + static_cast(eps), static_cast(rotary_dim), num_tokens, + num_layers, nkv, block_shift, static_cast(all_kv.stride(0) / 8), + static_cast(all_kv.stride(1) / 8), + static_cast(all_kv.stride(2) / 8), + static_cast(all_kv.stride(3) / 8), + static_cast(kv_cache.stride(0)), + static_cast(kv_cache.stride(1)), + static_cast(kv_cache.stride(2)), stream); + }; + if (fp8) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } +#endif +} diff --git a/csrc/ops.h b/csrc/ops.h index 401f6b4996ab..993d834e038b 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -413,7 +413,18 @@ void fused_minimax_m3_qknorm_rope_kv_insert( const std::string& kv_cache_dtype, bool skip_index_branch, std::optional q_fp8_out, double q_fp8_scale, std::optional kv_k_scale, - std::optional kv_v_scale); + std::optional kv_v_scale, bool write_qkv); + +void fused_minimax_m3_dspark_context_kv( + torch::stable::Tensor const& all_kv, + torch::stable::Tensor const& k_norm_weight, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& positions, + torch::stable::Tensor const& slot_mapping, torch::stable::Tensor& kv_cache, + torch::stable::Tensor const& cache_offsets, + torch::stable::Tensor const& k_scale, + torch::stable::Tensor const& scale_offsets, int64_t rotary_dim, double eps, + const std::string& kv_cache_dtype); #ifdef VLLM_ENABLE_FUSED_KDA_DECODE void fused_kda_decode( diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index fef3076056ee..99a742f856de 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -538,7 +538,17 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int block_size, Tensor!? q_out, Tensor!? index_q_out, " "str kv_cache_dtype, bool skip_index_branch=False, " "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, " - "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()"); + "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None, " + "bool write_qkv=True) -> ()"); + + // MiniMax-M3 DSpark draft: Gemma k-norm + partial NeoX RoPE + K/V cache + // insert for every draft layer's context K/V in one launch. + ops.def( + "fused_minimax_m3_dspark_context_kv(" + "Tensor all_kv, Tensor k_norm_weight, Tensor cos_sin_cache, " + "Tensor positions, Tensor slot_mapping, Tensor! kv_cache, " + "Tensor cache_offsets, Tensor k_scale, Tensor scale_offsets, " + "int rotary_dim, float eps, str kv_cache_dtype) -> ()"); #ifdef VLLM_ENABLE_FUSED_KDA_DECODE ops.def( @@ -879,6 +889,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { TORCH_BOX(&fused_muse_spark_qknorm_rope_kv_insert)); ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); + ops.impl("fused_minimax_m3_dspark_context_kv", + TORCH_BOX(&fused_minimax_m3_dspark_context_kv)); #ifdef VLLM_ENABLE_FUSED_KDA_DECODE ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); #endif diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index 01ea1193953b..4334f8a4bb5b 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -52,20 +52,25 @@ def test_none_overrides_fall_back_to_arch_mapping(): def _make_speculative_config( - hf_kwargs: dict, architecture: str, **speculative_kwargs + hf_kwargs: dict, + architecture: str, + target_model: str = "target", + **speculative_kwargs, ) -> SpeculativeConfig: hf_config = SpeculativeConfig.hf_config_override(_make_hf_config(**hf_kwargs)) draft = MagicMock( hf_config=hf_config, architectures=hf_config.architectures, max_model_len=128 ) draft.registry.inspect_model_cls.return_value = (None, architecture) - target = MagicMock(max_model_len=128, quantization=None, hf_overrides={}) + target = MagicMock( + model=target_model, max_model_len=128, quantization=None, hf_overrides={} + ) target.hf_config.model_type = hf_kwargs["model_type"] target.using_transformers_backend.return_value = False + speculative_kwargs.setdefault("model", "draft") with patch("vllm.config.speculative.ModelConfig", return_value=draft): return SpeculativeConfig( - model="draft", **speculative_kwargs, target_model_config=target, target_parallel_config=ParallelConfig(), @@ -118,6 +123,75 @@ def test_dspark_width_is_independent_of_mtp_stages( assert config.draft_model_config.hf_config.architectures == [architecture] +def _make_minimax_dspark_hf_kwargs() -> dict: + return dict( + model_type="minimax_m3_vl", + architectures=["DSparkMiniMaxDraftModel"], + quantization_config={ + "quant_method": "mxfp8", + "ignored_layers": ["language_model.model.dspark.fc"], + }, + text_config=_make_hf_config( + architectures=["DSparkMiniMaxDraftModel"], + model_type="", + num_hidden_layers=5, + dspark_block_size=7, + dspark_markov_rank=256, + dflash_config={"block_size": 7}, + ), + ) + + +@pytest.mark.cpu_test +def test_minimax_dspark_draft_promotes_text_config(): + """The MiniMax DSpark draft ships a VL wrapper config. It must keep its own + draft class (not MiniMax MTP or DeepSeek DSpark), with the DSpark keys and + the wrapper's quantization config on the promoted text config.""" + hf_kwargs = _make_minimax_dspark_hf_kwargs() + config = _make_speculative_config( + hf_kwargs, "DSparkMiniMaxDraftModel", method="dspark" + ) + + hf_config = config.draft_model_config.hf_config + assert hf_config is hf_kwargs["text_config"] + assert hf_config.architectures == ["DSparkMiniMaxDraftModel"] + assert config.num_speculative_tokens == 7 + assert hf_config.quantization_config == hf_kwargs["quantization_config"] + assert hf_config.markov_rank == 256 + assert hf_config.dflash_config["causal"] is False + + +@pytest.mark.cpu_test +def test_minimax_dspark_defaults_to_dspark_subfolder(tmp_path): + """With only method="dspark", MiniMax-M3 loads the draft checkpoint shipped + in the target's dspark/ subfolder, with its own quantization config.""" + (tmp_path / "dspark").mkdir() + (tmp_path / "dspark" / "config.json").write_text("{}") + config = _make_speculative_config( + _make_minimax_dspark_hf_kwargs(), + "DSparkMiniMaxDraftModel", + target_model=str(tmp_path), + model=None, + method="dspark", + ) + + assert config.model == str(tmp_path / "dspark") + assert config.quantization is None + assert config.num_speculative_tokens == 7 + + +@pytest.mark.cpu_test +def test_minimax_dspark_without_subfolder_raises(tmp_path): + with pytest.raises(ValueError, match="no draft checkpoint in the 'dspark/'"): + _make_speculative_config( + _make_minimax_dspark_hf_kwargs(), + "DSparkMiniMaxDraftModel", + target_model=str(tmp_path), + model=None, + method="dspark", + ) + + @pytest.mark.cpu_test def test_mtp_stages_are_independent_of_dspark_width(): hf_kwargs = dict( diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index dd462c380d19..57d22a47ce2f 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/tests/distributed/test_custom_all_reduce.py @@ -182,6 +182,152 @@ def test_all_reduce_mhc_matches_unfused_path(monkeypatch): multi_process_parallel(monkeypatch, 4, 1, _all_reduce_mhc) +def _swizzled_scale_rows(scales: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 scales read out of the 128x4 swizzled layout.""" + row = torch.arange(rows, device=scales.device)[:, None] + col = torch.arange(cols, device=scales.device)[None, :] + offset = ( + col % 4 + + col // 4 * 512 + + row % 32 * 16 + + row % 128 // 32 * 4 + + row // 128 * 128 * cols + ) + return scales[offset] + + +@ray.remote(num_gpus=1, max_calls=1) +def _finalize_all_reduce_gemma_norm_mxfp8( + monkeypatch, tp_size, pp_size, rank, distributed_init_port +): + from vllm.model_executor.layers.fused_moe.moe_output import ( + MoEOutput, + UnfinalizedMoEOutput, + ) + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + from vllm.models.minimax_m3.nvidia.ops.finalize_allreduce_norm import ( + FinalizeAllReduceGemmaRMSNormMxfp8, + ) + + hidden, top_k, eps = 6144, 4, 1e-6 + with monkeypatch.context() as m: + m.delenv("CUDA_VISIBLE_DEVICES", raising=False) + device = torch.device(f"cuda:{rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port) + ensure_model_parallel_initialized(tp_size, pp_size) + op = FinalizeAllReduceGemmaRMSNormMxfp8( + hidden_size=hidden, + max_num_tokens=32, + top_k=top_k, + fp32_weights=False, + device=device, + ) + + def run(n): + torch.manual_seed(7 + rank) + # A padded permuted GEMM2 buffer, like the MoE's. + rows = n * top_k + 5 + gemm2 = torch.randn(rows, hidden, device=device, dtype=torch.bfloat16) + permuted = torch.randperm(rows, device=device)[: n * top_k] + permuted = permuted.view(n, top_k).int() + # A route to an expert this rank does not hold. + permuted[0, -1] = -1 + weights = torch.rand(n, top_k, device=device).bfloat16() + shared = torch.randn(n, hidden, device=device, dtype=torch.bfloat16) + torch.manual_seed(123) + residual = torch.randn(n, hidden, device=device, dtype=torch.bfloat16) + weight = torch.randn(hidden, device=device, dtype=torch.bfloat16) * 0.1 + moe_output = MoEOutput( + routed=UnfinalizedMoEOutput(gemm2, weights, permuted), + shared_output=shared, + ) + + def fused(): + return op(moe_output, residual, weight, eps) + + def check(outputs): + quantized, residual_out = outputs + # One FP32 FMA per route in route order, the shared add, one + # rounding: the partial each rank publishes. + acc = torch.zeros(n, hidden, device=device) + for k in range(top_k): + valid = (permuted[:, k] >= 0).unsqueeze(-1) + rows_k = gemm2[permuted[:, k].clamp_min(0).long()].float() + torch.addcmul( + acc, + torch.where(valid, rows_k, 0.0), + torch.where(valid, weights[:, k : k + 1].float(), 0.0), + out=acc, + ) + partial = (acc + shared.float()).bfloat16() + peers = get_tp_group().all_gather(partial, dim=0).view(tp_size, n, -1) + reduced = peers[0].float() + for peer in peers[1:]: + reduced = reduced + peer.float() + prenorm = (reduced + residual.float()).bfloat16() + assert torch.equal(residual_out, prenorm) + + inv_rms = torch.rsqrt( + prenorm.float().square().mean(-1, keepdim=True) + eps + ) + normed = (prenorm.float() * inv_rms * (weight.float() + 1)).bfloat16() + ref_values, ref_scales = mxfp8_e4m3_quantize( + normed, is_sf_swizzled_layout=True + ) + # The norm sums the squares in another order with an approximate + # rsqrt, so a value can move by one BF16 ULP, and with it rarely + # a block's scale. Padding rows' scales are zero in both. + assert quantized.scale.shape == ref_scales.shape + assert (quantized.scale == ref_scales).float().mean() >= 0.99 + codes = quantized.data.view(torch.uint8) + assert (codes == ref_values.view(torch.uint8)).float().mean() >= 0.99 + cols = hidden // 32 + scales = _swizzled_scale_rows( + quantized.scale, quantized.scale.numel() // cols, cols + ) + assert not scales[n:].any() + scales = scales[:n] + step = torch.exp2(scales.float() - 127).repeat_interleave(32, dim=1) + dequantized = quantized.data.float() * step + expected = normed.float() + # E4M3 rounding (3 mantissa bits) plus one BF16 ULP, and the + # smallest E4M3 subnormal step of the block. + tolerance = 0.07 * expected.abs() + step * 2.0**-9 + assert torch.all((dequantized - expected).abs() <= tolerance) + + check(fused()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(5): + fused() + captured = fused() + for _ in range(20): + graph.replay() + check(captured) + # Replays must pick up in-place input changes. + gemm2.mul_(0.5) + residual.mul_(2) + graph.replay() + check(captured) + + # Changing shapes cover shrinking and growing batches in one mailbox. + for n in (1, 6, 32, 3, 17, 2, 1): + run(n) + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), reason="Requires SM100" +) +def test_finalize_all_reduce_gemma_norm_mxfp8_matches_unfused_path(monkeypatch): + """MiniMax-M3's fused MoE finalize + all-reduce + Gemma RMSNorm + MXFP8.""" + if torch.accelerator.device_count() < 4: + pytest.skip("Requires four GPUs with NVLink multicast") + multi_process_parallel(monkeypatch, 4, 1, _finalize_all_reduce_gemma_norm_mxfp8) + + @pytest.mark.parametrize( ("dtype", "expected"), [ diff --git a/tests/kernels/attention/minimax_m3_nvfp4_ref.py b/tests/kernels/attention/minimax_m3_nvfp4_ref.py new file mode 100644 index 000000000000..05f264eeb882 --- /dev/null +++ b/tests/kernels/attention/minimax_m3_nvfp4_ref.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Torch reference of MiniMax-M3's NVFP4 K/V and index-K pages on ROCm. + +A 128-token page is ``[K side | V side]``; a side holds every KV head's e2m1 +data (8192 B per head), then every head's e4m3 block scales (1024 B per head). +Within a head: K data ``[token][64 B]``, K scales in the FlyDSL prefill's lane +order, V data ``[token octet][dim octet][8 dims][4 B]`` (a byte holds two +consecutive tokens of one dim), V scales ``[group][token]``. An index page is +``[128 x 64 B data][128 x 8 B scales]``. +""" + +import torch + +BLOCK = 128 +D = 128 +E2M1_GRID = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def quantize_nvfp4(x: torch.Tensor, global_scale: float): + """E2M1 codes ``[T, 128]`` and e4m3 scale bytes ``[T, 8]`` of fp32 rows. + + The CUDA encoder's scale math and a nearest-grid E2M1 rounding with ties to + the even code, written independently of the kernels' bit arithmetic. + """ + groups = x.view(-1, 8, 16) + one_sixth = torch.tensor(1 / 6, dtype=torch.float32) + # fp32 throughout, like the kernels: 1 / 0.002 is not 500 in fp32. + gs = torch.tensor(global_scale, dtype=torch.float32, device=x.device) + scale = (groups.abs().amax(-1) * one_sixth * (1.0 / gs)).clamp(max=448) + scale_e4m3 = scale.to(torch.float8_e4m3fn) + scale_value = scale_e4m3.float() + quant = torch.where(scale_value != 0, 1.0 / (scale_value * gs), 0.0) + y = groups * quant[..., None] + grid = torch.tensor(E2M1_GRID, device=x.device) + dist = (y.abs()[..., None] - grid).abs() + is_min = dist == dist.amin(-1, keepdim=True) + ids = torch.arange(8, device=x.device) + # Neighbouring codes alternate parity, so every tie has an even candidate. + even_min = torch.where(is_min & (ids % 2 == 0), ids, 8).amin(-1) + mag = torch.where(even_min < 8, even_min, dist.argmin(-1)) + codes = mag | (torch.signbit(y).long() << 3) + return codes.view(-1, D), scale_e4m3.view(torch.uint8).view(-1, 8) + + +def dequant(data: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Token-major NVFP4 rows -> fp32 e2m1 x block scale, ``[..., 128]``.""" + codes = torch.stack((data & 15, data >> 4), -1).flatten(-2).long() + grid = torch.tensor(E2M1_GRID, device=data.device) + mag = grid[codes & 7] * torch.where(codes >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() + return (mag.unflatten(-1, (8, 16)) * s[..., None]).flatten(-2) + + +def head_offsets(side: int, device): + """Byte offsets inside one head's regions of side 0 (K) or 1 (V). + + Returns the data byte and nibble shift of each (token, dim) ``[128, 128]`` + and the scale byte of each (token, group) ``[128, 8]``. + """ + t = torch.arange(BLOCK, device=device)[:, None] + d = torch.arange(D, device=device)[None, :] + g = torch.arange(8, device=device)[None, :] + if side == 0: + r = t % 32 + n = 4 * (r // 8) + r % 4 + data = t * 64 + d // 2 + shift = (d % 2) * 4 + 0 * t + scale = ((((t // 32) * 4 + g // 2) * 16 + n) * 2 + (r % 8) // 4) * 2 + g % 2 + else: + data = ((t // 8 * 16 + d // 8) * 8 + d % 8) * 4 + (t % 8) // 2 + shift = (t % 2) * 4 + 0 * d + scale = g * BLOCK + t + return data, shift, scale + + +def _region_bases(side: int, heads: int, device): + """``[H, 1, 1]`` page offsets of each head's data and scale regions.""" + h = torch.arange(heads, device=device)[:, None, None] + side0 = side * heads * BLOCK * 72 + return side0 + h * BLOCK * 64, side0 + heads * BLOCK * 64 + h * BLOCK * 8 + + +def _flat(cache: torch.Tensor) -> torch.Tensor: + """Writable 1-D view of a cache whose block stride may exceed one page.""" + n = (cache.shape[0] - 1) * cache.stride(0) + cache[0].numel() + return cache.as_strided((n,), (1,)) + + +def _insert_side(cache, rows, slots, side, global_scale): + tokens, heads, _ = rows.shape + codes, scales = quantize_nvfp4(rows.reshape(-1, D).float(), global_scale) + slots = slots.long() + keep = slots >= 0 + codes = codes.view(tokens, heads, D)[keep].to(torch.uint8) + scales = scales.view(tokens, heads, 8)[keep] + blk, pos = slots[keep] // BLOCK, slots[keep] % BLOCK + data_off, shift, scale_off = head_offsets(side, rows.device) + data_base, scale_base = _region_bases(side, heads, rows.device) + page = (blk * cache.stride(0))[:, None, None] + byte = page + data_base[:, 0] + data_off[pos][:, None] # [n, H, 128] + sh = shift[pos][:, None].expand_as(byte) + flat = _flat(cache) + for s in (0, 4): # each byte takes one nibble per pass + m = sh == s + flat[byte[m]] = (flat[byte[m]] & (0xF0 >> s)) | (codes[m] << s) + flat[page + scale_base[:, 0] + scale_off[pos][:, None]] = scales + + +def insert_nvfp4_kv(k, v, kv_cache, slots, k_scale, v_scale): + """Quantize K / V ``[T, Hkv, 128]`` into their pages; slots < 0 skip.""" + _insert_side(kv_cache, k, slots, 0, float(k_scale)) + _insert_side(kv_cache, v, slots, 1, float(v_scale)) + + +def insert_nvfp4_index(index_cache, keys, slots): + """Quantize index keys ``[T, 128]`` into their pages (no global scale).""" + codes, scales = quantize_nvfp4(keys.float(), 1.0) + slots = slots.long() + keep = slots >= 0 + codes, scales = codes[keep].to(torch.uint8), scales[keep] + blk, pos = slots[keep] // BLOCK, slots[keep] % BLOCK + base = (blk * index_cache.stride(0))[:, None] + flat = _flat(index_cache) + flat[base + pos[:, None] * 64 + torch.arange(64, device=keys.device)] = codes[ + :, 0::2 + ] | (codes[:, 1::2] << 4) + flat[base + BLOCK * 64 + pos[:, None] * 8 + torch.arange(8, device=keys.device)] = ( + scales + ) + + +def nvfp4_kv_cache_token_major(kv_cache: torch.Tensor): + """``(k_data, k_scale, v_data, v_scale)`` of a main cache, token-major. + + Data ``[blocks, Hkv, 128, 64]`` (the low nibble the even dim), scales + ``[blocks, Hkv, 128, 8]``, all uint8. + """ + blocks, heads = kv_cache.shape[0], kv_cache.shape[1] // 2 + pages = kv_cache.reshape(blocks, -1) + out = [] + for side in (0, 1): + data_off, shift, scale_off = head_offsets(side, kv_cache.device) + data_base, scale_base = _region_bases(side, heads, kv_cache.device) + codes = (pages[:, data_base + data_off] >> shift.to(torch.uint8)) & 0xF + out.append(codes[..., 0::2] | (codes[..., 1::2] << 4)) + out.append(pages[:, scale_base + scale_off]) + return tuple(out) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index d46224bd5fd8..29c0e3d24aa3 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -1179,6 +1179,407 @@ def test_aiter_indexer_impl_matches_reference( assert torch.equal(sparse_ctx.cpu(), expected_ctx) +# NVFP4 index cache (SM100/SM103): MSA Q8KV4 decode on the packed pages, FP8 +# gather + Q8KV8 prefill. Both rank pages from e4m3_rn_satfinite(e2m1 * scale). +_E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_NVFP4_PAGE_WIDTH = HEAD_DIM // 2 + HEAD_DIM // 16 +_requires_q8_indexer = pytest.mark.skipif( + not ( + current_platform.is_device_capability(100) + or current_platform.is_device_capability(103) + or current_platform.is_device_capability(107) + ), + reason="The MSA Q8KV4/Q8KV8 indexers require SM100, SM103 or SM107.", +) + + +def _random_nvfp4_index_pages( + num_pages: int, scale_max: float = 0.5, page_pad: int = 0 +) -> torch.Tensor: + """``[num_pages, 128, 72]`` uint8 pages: random E2M1 codes, then E4M3 scales.""" + data_bytes = BLOCK_SIZE * HEAD_DIM // 2 + page_bytes = BLOCK_SIZE * _NVFP4_PAGE_WIDTH + raw = torch.zeros( + num_pages, page_bytes + page_pad, dtype=torch.uint8, device="cuda" + ) + raw[:, :data_bytes] = torch.randint( + 0, 256, (num_pages, data_bytes), dtype=torch.uint8, device="cuda" + ) + scale = torch.rand(num_pages, page_bytes - data_bytes, device="cuda") * scale_max + raw[:, data_bytes:page_bytes] = scale.to(torch.float8_e4m3fn).view(torch.uint8) + return raw[:, :page_bytes].view(num_pages, BLOCK_SIZE, _NVFP4_PAGE_WIDTH) + + +def _dequant_nvfp4_index_pages(pages: torch.Tensor) -> torch.Tensor: + """Reference ``[P, 128, 128]`` e4m3: e2m1 * scale, clamped, rounded once.""" + num_pages = pages.shape[0] + flat = pages.reshape(num_pages, -1) + data = flat[:, : BLOCK_SIZE * HEAD_DIM // 2].view(num_pages, BLOCK_SIZE, -1) + scale = flat[:, BLOCK_SIZE * HEAD_DIM // 2 :].view(torch.float8_e4m3fn) + codes = torch.stack((data & 0xF, data >> 4), dim=-1).reshape( + num_pages, BLOCK_SIZE, HEAD_DIM + ) + magnitude = torch.tensor(_E2M1_VALUES, device="cuda")[(codes & 7).long()] + value = torch.where(codes >= 8, -magnitude, magnitude) + value = value * scale.float().view(num_pages, BLOCK_SIZE, -1).repeat_interleave( + 16, dim=-1 + ) + return value.clamp_(-448.0, 448.0).to(torch.float8_e4m3fn) + + +@_requires_q8_indexer +@pytest.mark.parametrize("use_qmul4", [False, True]) +@pytest.mark.parametrize("scale_max", [0.5, 448.0]) +def test_gather_nvfp4_index_k_fp8_bitwise( + scale_max: float, use_qmul4: bool, monkeypatch +): + """Each request's leading pages land at its compact offset, bit-identical + to the reference dequant (incl. saturation), and nothing else is written, + on both the QMUL4 and the fp32 dequant paths.""" + import vllm.models.minimax_m3.nvidia.ops.gather_nvfp4_index_k as gather_mod + from vllm.models.minimax_m3.nvidia.ops.gather_nvfp4_index_k import ( + minimax_m3_gather_nvfp4_index_k_fp8, + ) + + if use_qmul4 and not gather_mod._qmul4_supported(0): + pytest.skip("Triton's Blackwell ptxas emits PTX older than 9.4 (no QMUL4).") + monkeypatch.setattr(gather_mod, "_qmul4_supported", lambda _: use_qmul4) + + torch.manual_seed(0) + num_physical = 40 + kv = _random_nvfp4_index_pages(num_physical, scale_max, page_pad=32) + block_table = torch.randperm(num_physical, device="cuda", dtype=torch.int32)[ + :36 + ].view(3, 12) + num_pages = torch.tensor([5, 0, 12], dtype=torch.int32, device="cuda") + page_offsets = torch.cumsum(num_pages, 0, dtype=torch.int32) - num_pages + total = int(num_pages.sum()) + out = torch.full((total + 1, BLOCK_SIZE, HEAD_DIM), 1.0, device="cuda").to( + torch.float8_e4m3fn + ) + + minimax_m3_gather_nvfp4_index_k_fp8( + kv, block_table, num_pages, page_offsets, 12, out + ) + + expected = _dequant_nvfp4_index_pages(kv) + for req, (count, offset) in enumerate( + zip(num_pages.tolist(), page_offsets.tolist()) + ): + pages = block_table[req, :count].long() + assert torch.equal( + out[offset : offset + count].view(torch.uint8), + expected[pages].view(torch.uint8), + ) + assert torch.all(out[total].float() == 1.0) + + +def _assert_forced_tail_topk( + actual: torch.Tensor, scores: torch.Tensor, lengths: torch.Tensor +) -> None: + """Check the Q8 top-k contract on ``[rows, 16]`` output. + + ``scores[row, :lengths[row] - 1]`` are the reference history-page scores. + Rows with <= 16 candidates are ``0 .. length - 1`` then ``-1``. Longer rows + hold 15 distinct history pages, then the local page. The ranked pages must + be the top 15 up to the kernel's 16-bit score-key quantization. + """ + actual = actual.long() + rows, max_pages = scores.shape + cols = torch.arange(TOPK, device=actual.device) + short = lengths <= TOPK + expected_short = torch.where(cols < lengths[:, None], cols, -1) + assert torch.equal(actual[short], expected_short[short]) + + long_rows = (~short).nonzero().flatten() + if long_rows.numel() == 0: + return + actual, scores, lengths = actual[long_rows], scores[long_rows], lengths[long_rows] + assert torch.equal(actual[:, -1], lengths - 1) + ranked = actual[:, :-1] + assert torch.all((ranked >= 0) & (ranked < (lengths - 1)[:, None])) + assert torch.all(ranked.sort(dim=1).values.diff(dim=1) > 0) + valid = torch.arange(max_pages, device=scores.device) < (lengths - 1)[:, None] + history = scores.masked_fill(~valid, float("-inf")) + high = history.amax(dim=1) + low = scores.masked_fill(~valid, float("inf")).amin(dim=1) + tol = 1.5 * (high - low) / 65534.0 + 1e-5 * high.abs() + threshold = history.topk(TOPK - 1, dim=1).values[:, -1] + assert torch.all(history.gather(1, ranked) >= (threshold - tol)[:, None]) + must_pick = history > (threshold + tol)[:, None] + picked = torch.zeros_like(must_pick).scatter_(1, ranked, True) + assert not torch.any(must_pick & ~picked) + + +@_requires_q8_indexer +@pytest.mark.parametrize("num_idx_heads", [1, 2, 4]) +@pytest.mark.parametrize("decode_query_len", [1, 3, 8]) +def test_nvfp4_indexer_impl_matches_reference( + num_idx_heads, decode_query_len, monkeypatch, workspace_init +): + """Drive MiniMaxM3IndexerNvfp4Impl through its metadata builder on a mixed + batch (short and full 8-query decodes, fresh and chunked prefills, a request + without history pages, several gather chunks) and check every row against + a reference computed from the dequantized NVFP4 keys.""" + import vllm.models.minimax_m3.common.indexer as indexer_mod + from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, + ) + from vllm.config import set_current_vllm_config + from vllm.forward_context import set_forward_context + from vllm.models.minimax_m3.nvidia.indexer_nvfp4 import ( + MiniMaxM3IndexerNvfp4Impl, + MiniMaxM3IndexerNvfp4MetadataBuilder, + ) + + torch.manual_seed(0) + device = torch.device("cuda") + monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1) + vllm_config = create_vllm_config( + block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192 + ) + vllm_config.model_config.hf_config.sparse_attention_config = { + "sparse_num_index_heads": num_idx_heads + } + + decode_seq_lens = [2305, 2561, 9] + prefill = [(2624, 64), (2720, 96), (100, 100), (1500, 1500)] + batch = BatchSpec( + seq_lens=decode_seq_lens + [s for s, _ in prefill], + query_lens=[decode_query_len] * len(decode_seq_lens) + [q for _, q in prefill], + ) + common = create_common_attn_metadata(batch, BLOCK_SIZE, device) + num_blocks = common.block_table_tensor.shape[1] + num_physical = batch.batch_size * num_blocks + block_table = torch.randperm(num_physical, device=device, dtype=torch.int32) + common.block_table_tensor = block_table.view(batch.batch_size, num_blocks) + num_tokens = batch.compute_num_tokens() + + index_cache = _random_nvfp4_index_pages(num_physical) + index_q = ( + torch.randn(num_tokens, num_idx_heads * HEAD_DIM, device=device) * 0.5 + ).to(torch.float8_e4m3fn) + + impl_kwargs = dict( + num_kv_heads=num_idx_heads, + scale=HEAD_DIM**-0.5, + topk_blocks=TOPK, + sparse_block_size=BLOCK_SIZE, + num_index_heads=num_idx_heads, + index_head_dim=HEAD_DIM, + init_blocks=0, + local_blocks=1, + indexer_kv_dtype="nvfp4", + ) + spec = MLAAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=_NVFP4_PAGE_WIDTH, + dtype=torch.uint8, + ) + with set_current_vllm_config(vllm_config): + impl = MiniMaxM3IndexerNvfp4Impl(prefix="idx_nvfp4", **impl_kwargs) + builder = MiniMaxM3IndexerNvfp4MetadataBuilder( + spec, [impl.index_cache.prefix], vllm_config, device + ) + index_spec = impl.index_cache.get_kv_cache_spec(vllm_config) + assert isinstance(index_spec, MLAAttentionSpec) + assert index_spec.page_size_bytes == BLOCK_SIZE * _NVFP4_PAGE_WIDTH + assert index_spec.dtype == torch.uint8 + # Spec decode would raise the threshold to 1 + num_speculative_tokens. + builder.reorder_batch_threshold = decode_query_len + # Fit the longest request's 21 history pages but not two such requests. + builder.max_gather_pages = 30 + impl.index_cache.kv_cache = index_cache + impl.topk_indices_buffer = torch.full( + (num_tokens, num_idx_heads, TOPK), -2, dtype=torch.int32, device=device + ) + + # Planning reads only host-side metadata: any device sync is an error. + torch.cuda.set_sync_debug_mode("error") + try: + md = builder.build(0, common) + finally: + torch.cuda.set_sync_debug_mode("default") + assert md.num_decode_tokens == decode_query_len * len(decode_seq_lens) + assert len(md.prefill_chunks) == 3 + with set_forward_context({impl.index_cache.prefix: md}, vllm_config): + assert impl(index_q) == (None, None) + + # Reference: max over each history page's 128 tokens of q . k. + keys = _dequant_nvfp4_index_pages(index_cache).float() + q = index_q.float().view(num_tokens, num_idx_heads, HEAD_DIM) + max_pages = num_blocks + scores = torch.full( + (num_tokens, num_idx_heads, max_pages), float("-inf"), device=device + ) + lengths = torch.empty(num_tokens, dtype=torch.long, device=device) + token = 0 + for req, (seq_len, q_len) in enumerate(zip(batch.seq_lens, batch.query_lens)): + pages = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE + k = keys[common.block_table_tensor[req, :pages].long()] + rows = slice(token, token + q_len) + scores[rows, :, :pages] = torch.einsum("nhd,pkd->nhpk", q[rows], k).amax(dim=-1) + positions = torch.arange(seq_len - q_len, seq_len, device=device) + lengths[rows] = positions // BLOCK_SIZE + 1 + token += q_len + _assert_forced_tail_topk( + impl.topk_indices_buffer.view(-1, TOPK), + scores.view(-1, max_pages), + lengths.repeat_interleave(num_idx_heads), + ) + + +@_requires_q8_indexer +@pytest.mark.parametrize("decode_query_len", [1, 3, 8]) +@pytest.mark.parametrize("num_idx_heads", [1, 2, 4]) +def test_nvfp4_indexer_decode_cuda_graph_replay( + decode_query_len, num_idx_heads, monkeypatch, workspace_init +): + """A full CUDA graph captures the decode indexer once per padded batch + size; re-planning a later step (new lengths, pages and keys, one + zero-length padding request) outside capture and replaying must give that + step's top-k, which fails if any captured tensor was re-allocated. The + decode reads index_q and writes the top-k in place, copying nothing.""" + import vllm.models.minimax_m3.common.indexer as indexer_mod + from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, + ) + from vllm.config import set_current_vllm_config + from vllm.forward_context import set_forward_context + from vllm.models.minimax_m3.nvidia.indexer_nvfp4 import ( + MiniMaxM3IndexerNvfp4Impl, + MiniMaxM3IndexerNvfp4MetadataBuilder, + ) + + torch.manual_seed(0) + device = torch.device("cuda") + monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1) + vllm_config = create_vllm_config( + block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192 + ) + vllm_config.model_config.hf_config.sparse_attention_config = { + "sparse_num_index_heads": num_idx_heads + } + num_reqs, max_blocks = 4, 24 + num_tokens = num_reqs * decode_query_len + num_physical = num_reqs * max_blocks + index_cache = _random_nvfp4_index_pages(num_physical) + + with set_current_vllm_config(vllm_config): + impl = MiniMaxM3IndexerNvfp4Impl( + prefix="idx_nvfp4_graph", + num_kv_heads=num_idx_heads, + scale=HEAD_DIM**-0.5, + topk_blocks=TOPK, + sparse_block_size=BLOCK_SIZE, + num_index_heads=num_idx_heads, + index_head_dim=HEAD_DIM, + init_blocks=0, + local_blocks=1, + indexer_kv_dtype="nvfp4", + ) + builder = MiniMaxM3IndexerNvfp4MetadataBuilder( + MLAAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=_NVFP4_PAGE_WIDTH, + dtype=torch.uint8, + ), + [impl.index_cache.prefix], + vllm_config, + device, + ) + builder.reorder_batch_threshold = decode_query_len + impl.index_cache.kv_cache = index_cache + impl.topk_indices_buffer = torch.full( + (num_tokens, num_idx_heads, TOPK), -2, dtype=torch.int32, device=device + ) + index_q = torch.empty(num_tokens, num_idx_heads * HEAD_DIM, device=device).to( + torch.float8_e4m3fn + ) + + def step(seq_lens, query_lens): + common = create_common_attn_metadata( + BatchSpec(seq_lens=seq_lens, query_lens=query_lens), BLOCK_SIZE, device + ) + common.block_table_tensor = torch.randperm( + num_physical, device=device, dtype=torch.int32 + )[: num_reqs * max_blocks].view(num_reqs, max_blocks) + index_q.copy_( + (torch.randn(num_tokens, num_idx_heads * HEAD_DIM, device=device) * 0.5).to( + torch.float8_e4m3fn + ) + ) + return common, builder.build(0, common) + + # Warm up (compiles), then capture the forward of a full batch. + common, md = step([2305, 2561, 2944, 700], [decode_query_len] * num_reqs) + with set_forward_context({impl.index_cache.prefix: md}, vllm_config): + impl(index_q) + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as prof: + impl(index_q) + torch.accelerator.synchronize() + copies = [ + e.name for e in prof.events() if "Memcpy" in e.name or "copy_" in e.name + ] + assert not copies, copies + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + impl(index_q) + + # A later step: the last request pads the graph batch with no tokens. + seq_lens = [2800, 1500, 3000, 0] + query_lens = [decode_query_len] * 3 + [0] + common, md = step(seq_lens, query_lens) + assert md.decode_q8kv4.num_reqs == num_reqs + # An eager step of another batch size runs before the replay; every decode + # wrapper shares one page-score scratch, which this one overwrites. + other = create_common_attn_metadata( + BatchSpec(seq_lens=[1800, 900], query_lens=[decode_query_len] * 2), + BLOCK_SIZE, + device, + ) + other.block_table_tensor = common.block_table_tensor[:2].clone() + other_md = builder.build(0, other) + with set_forward_context({impl.index_cache.prefix: other_md}, vllm_config): + impl(index_q[: 2 * decode_query_len]) + assert len(builder.decode_wrappers) == 2 + impl.topk_indices_buffer.fill_(-2) + graph.replay() + + num_real = 3 * decode_query_len + keys = _dequant_nvfp4_index_pages(index_cache).float() + q = index_q[:num_real].float().view(num_real, num_idx_heads, HEAD_DIM) + scores = torch.full( + (num_real, num_idx_heads, max_blocks), float("-inf"), device=device + ) + lengths = torch.empty(num_real, dtype=torch.long, device=device) + for req in range(3): + pages = (seq_lens[req] + BLOCK_SIZE - 1) // BLOCK_SIZE + k = keys[common.block_table_tensor[req, :pages].long()] + rows = slice(req * decode_query_len, (req + 1) * decode_query_len) + scores[rows, :, :pages] = torch.einsum("nhd,pkd->nhpk", q[rows], k).amax(dim=-1) + positions = torch.arange( + seq_lens[req] - decode_query_len, seq_lens[req], device=device + ) + lengths[rows] = positions // BLOCK_SIZE + 1 + _assert_forced_tail_topk( + impl.topk_indices_buffer[:num_real].view(-1, TOPK), + scores.view(-1, max_blocks), + lengths.repeat_interleave(num_idx_heads), + ) + + @pytest.mark.parametrize( ("decode_query_len", "max_decode_query_len"), [ diff --git a/tests/kernels/attention/test_minimax_m3_index_topk_decode.py b/tests/kernels/attention/test_minimax_m3_index_topk_decode.py new file mode 100644 index 000000000000..afbd90c009f3 --- /dev/null +++ b/tests/kernels/attention/test_minimax_m3_index_topk_decode.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3 index top-k (FlyDSL) against AITER's ``pa_sparse_block_topk`` +block ids, bit for bit, on realistic and adversarial decode and prefill rows.""" + +import math + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm (gfx950) FlyDSL kernels", allow_module_level=True) + +from vllm.platforms.rocm import on_gfx950 # noqa: E402 + +if not on_gfx950(): + pytest.skip("gfx950 instructions", allow_module_level=True) + +from vllm.models.minimax_m3.amd.ops.index_topk_decode.host import ( # noqa: E402 + minimax_m3_index_topk_decode_indices, + minimax_m3_index_topk_prefill_indices, + supports_index_topk, +) + +BLOCK = 128 + + +def _width(max_seq_len: int) -> int: + n = math.ceil(math.ceil(max_seq_len / BLOCK) / 64) + return (1 << (n - 1).bit_length()) * 64 + + +def _scores(kind, heads, rows, width, nblk, gen): + s = torch.randn(heads, rows, width, generator=gen) + if kind == "real": # a few outliers, the init block 1e30, the last two 1e29 + s = torch.where(torch.rand(s.shape, generator=gen) < 4.0 / width, s + 6, s) + for r in range(rows): + b = min(int(nblk[r]), width) + if b > 0: + s[:, r, max(b - 2, 0) : b] = 1e29 + s[:, r, 0] = 1e30 + elif kind == "ties": + s = torch.randint(-3, 4, s.shape, generator=gen).float() + elif kind == "equal": + s = torch.full_like(s, 0.5) + elif kind == "neginf": + s = torch.full_like(s, float("-inf")) + elif kind == "nan": + s[torch.rand(s.shape, generator=gen) < 0.3] = float("nan") + neg_nan = torch.tensor([-4194303], dtype=torch.int32).view(torch.float32) + s[torch.rand(s.shape, generator=gen) < 0.1] = neg_nan.item() + elif kind == "zeros": # +0 / -0 / denormals + vals = torch.tensor([0.0, -0.0, 1e-40, -1e-40, 1e-45, -1e-45]) + s = vals[torch.randint(0, 6, s.shape, generator=gen)] + elif kind == "sentinels": + s = torch.where(torch.rand(s.shape, generator=gen) < 0.2, 1e29, s) + s = torch.where(torch.rand(s.shape, generator=gen) < 0.05, 1e30, s) + elif kind == "inf": + r = torch.rand(s.shape, generator=gen) + s = torch.where(r < 0.1, float("inf"), s) + s = torch.where((r > 0.1) & (r < 0.2), float("-inf"), s) + elif kind == "rising": # increasing with position + s = torch.arange(width, dtype=torch.float32).expand(heads, rows, width).clone() + # blocks past a row's length hold garbage the kernel must not read as scores + tail = torch.randn(s.shape, generator=gen) * 1e6 + tail[torch.rand(s.shape, generator=gen) < 0.1] = float("nan") + live = torch.arange(width)[None, None, :] < nblk[None, :, None] + return torch.where(live, s, tail) + + +def _topk_out(heads, rows, pad): + """``topk_idx`` and the tensor compared: with ``pad`` a view (non-zero row + start, larger head stride) of a larger persistent buffer, compared whole + (nothing outside the view is written).""" + if not pad: + t = torch.full((heads, rows, 16), -7, dtype=torch.int32, device="cuda") + return t, t + full = torch.full((heads, rows + 37, 16), -7, dtype=torch.int32, device="cuda") + return full[:, 5 : 5 + rows], full + + +def _aiter_topk(score, bt, seq, msl, nkv, pad, **kw): + from aiter.ops.msa_attention import pa_sparse_block_topk + + heads, rows, _ = score.shape + topk, compared = _topk_out(heads, rows, pad) + sbt = torch.empty(rows * nkv, 128, dtype=torch.int32, device="cuda") + sctx = torch.empty(rows * nkv, dtype=torch.int32, device="cuda") + pa_sparse_block_topk( + score, topk, bt, seq, sbt, sctx, max_seq_len=msl, block_size=BLOCK, + num_kv_heads=nkv, pages_per_block=8, **kw, + ) # fmt: skip + return compared + + +SHAPES = [ + # reqs, query_len, max_seq_len, (lo, hi) seq_len, num_kv_heads, index heads + (40, 4, 1 << 20, (30000, 210000), 1, 1), # the QPS decode batch + (8, 4, 262144, (1, 262144), 1, 1), + (5, 4, 1 << 20, (1 << 19, 1 << 20), 1, 1), # 4 windows of 2048 blocks + (16, 1, 1 << 20, (2040 * BLOCK, 2100 * BLOCK), 1, 1), # the window edge + (12, 1, 4000, (1, 4000), 1, 1), + (4, 2, 300000, (1, 300000), 2, 2), + (7, 4, 20000, (20000, 40000), 1, 1), # rows past max_seq_len + (256, 4, 1 << 20, (1, 3000), 1, 1), # 1024 short rows + (32, 6, 204800, (190000, 204800), 4, 4), # 6 rows, 4 heads + # rows past max_seq_len inside the padded width (2048 blocks): the window + # sized by max_seq_len must go on to the next one, not cut them + (8, 1, 200000, (200000, 262000), 4, 1), + (4, 2, 0, (0, 0), 1, 1), # a zero max_seq_len (dummy batch) + (1, 1, 1, (0, 1), 1, 1), +] +KINDS = [ + "real", + "ties", + "equal", + "neginf", + "nan", + "zeros", + "sentinels", + "inf", + "rising", +] + + +def _check_decode(kind, shape, pad=False): + reqs, ql, msl, (lo, hi), nkv, heads = shape + gen = torch.Generator().manual_seed(100 * KINDS.index(kind) + SHAPES.index(shape)) + seq = torch.randint(lo, hi + 1, (reqs,), generator=gen) + if reqs > 1: + seq[-1] = 0 # a cudagraph padding request + tok = torch.arange(reqs * ql) % ql + kv = (seq.repeat_interleave(ql) - ql + tok + 1).clamp(min=0) + score = _scores(kind, heads, reqs * ql, _width(msl), (kv + BLOCK - 1) // BLOCK, gen) + # the table covers every row's blocks (rows may run past max_seq_len) + cols = max(1, math.ceil(max(msl, int(seq.max())) / BLOCK)) + bt = torch.randint(0, 1 << 22, (reqs, cols), generator=gen).int().cuda() + score, seq = score.cuda(), seq.int().cuda() + ref = _aiter_topk(score, bt, seq, msl, nkv, pad, query_len=ql) + topk, got = _topk_out(heads, reqs * ql, pad) + minimax_m3_index_topk_decode_indices( + score, topk, seq, max_seq_len=msl, query_len=ql + ) + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("shape", SHAPES) +def test_index_topk_decode_matches_aiter(kind, shape): + _check_decode(kind, shape) + + +@pytest.mark.parametrize("kind", ["real", "nan", "ties"]) +@pytest.mark.parametrize("shape", [SHAPES[0], SHAPES[5]]) +def test_index_topk_decode_padded_output(kind, shape): + _check_decode(kind, shape, pad=True) + + +def test_index_topk_decode_all_padding(): + """A captured dummy batch: every request padding (seq_len 0).""" + reqs, ql, msl = 48, 4, 1 << 20 + score = torch.randn(1, reqs * ql, _width(msl), device="cuda") + bt = torch.zeros(reqs, msl // BLOCK, dtype=torch.int32, device="cuda") + seq = torch.zeros(reqs, dtype=torch.int32, device="cuda") + ref = _aiter_topk(score, bt, seq, msl, 1, False, query_len=ql) + topk, _ = _topk_out(1, reqs * ql, False) + minimax_m3_index_topk_decode_indices( + score, topk, seq, max_seq_len=msl, query_len=ql + ) + torch.testing.assert_close(topk, ref, rtol=0, atol=0) + + +PREFILL_SHAPES = [ + # (q_len, seq_len) per request, max_seq_len, num_kv_heads, index heads + ([(1000, 1000)], 1000, 1, 1), # a fresh prompt: every row diagonal + ([(700, 30000), (1500, 1500)], 30000, 4, 4), # cached prefix + fresh + ([(2000, 9000), (513, 513), (3000, 20000)], 20000, 4, 1), # ragged + ([(300, 262144), (64, 200000)], 262144, 4, 4), # one window, 2048 blocks + ([(129, 2049 * BLOCK), (40, 1 << 20)], 1 << 20, 2, 2), # multi-window rows + ([(1, 1), (5, 5), (128, 128), (129, 129)], 4096, 1, 1), # tiny rows + # 3K new on a 200K prefix beside a fresh 3K prompt + ([(3072, 203072), (3072, 3072), (3072, 207872)], 207872, 4, 1), + ([(64, 262000), (3072, 203072)], 200000, 4, 1), # rows past max_seq_len +] + + +def _check_prefill(kind, shape, pad=False): + reqs, msl, nkv, heads = shape + gen = torch.Generator().manual_seed( + 100 * KINDS.index(kind) + PREFILL_SHAPES.index(shape) + ) + # every request's chunk ends its sequence: row positions seq - q .. seq - 1 + pos = torch.cat([torch.arange(s - q, s) for q, s in reqs]) + rid = torch.cat([torch.full((q,), r) for r, (q, _) in enumerate(reqs)]) + nvp = pos // BLOCK + 1 + seq = torch.tensor([s for _, s in reqs]) + score = _scores(kind, heads, pos.numel(), _width(msl), nvp, gen).cuda() + cols = max(1, math.ceil(max(msl, int(seq.max())) / BLOCK)) + bt = torch.randint(0, 1 << 22, (len(reqs), cols), generator=gen).int().cuda() + nvp, rid, kvl = (x.int().cuda() for x in (nvp, rid, pos + 1)) + ref = _aiter_topk( + score, bt, seq.int().cuda(), msl, nkv, pad, + num_valid_pages=nvp, row_req_id=rid, kv_lens=kvl, + ) # fmt: skip + topk, got = _topk_out(heads, pos.numel(), pad) + minimax_m3_index_topk_prefill_indices( + score, topk, max_seq_len=msl, num_valid_pages=nvp + ) + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("shape", PREFILL_SHAPES) +def test_index_topk_prefill_matches_aiter(kind, shape): + _check_prefill(kind, shape) + + +@pytest.mark.parametrize("kind", ["real", "nan", "ties"]) +@pytest.mark.parametrize("shape", [PREFILL_SHAPES[2], PREFILL_SHAPES[4]]) +def test_index_topk_prefill_padded_output(kind, shape): + _check_prefill(kind, shape, pad=True) + + +def test_supports_index_topk(): + assert supports_index_topk(16, BLOCK) + assert not supports_index_topk(8, BLOCK) and not supports_index_topk(16, 64) diff --git a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py index c7a4916737c0..47ae7854b227 100644 --- a/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py +++ b/tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py @@ -4,6 +4,7 @@ NVFP4 KV cache).""" import math +from dataclasses import fields, replace from types import SimpleNamespace import pytest @@ -29,11 +30,13 @@ ) from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import ( + NVFP4_BLOCK_SCALE_SHIFT, MSACutlassDecodePlanCache, msa_cutlass_sparse_decode, nvfp4_kv_cache_views, prepare_decode_metadata, should_prepare_decode_metadata, + uses_q8kv4_decode, ) from vllm.models.minimax_m3.nvidia.sparse_attention_msa import ( MiniMaxM3SparseMSABackend, @@ -41,6 +44,8 @@ MiniMaxM3SparseMSAImpl, MiniMaxM3SparseMSAMetadataBuilder, MiniMaxM3SparseMSANvfp4Backend, + MiniMaxM3SparseMSANvfp4Impl, + MiniMaxM3SparseMSAPrefillMetadata, ) from vllm.platforms import current_platform from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -179,6 +184,26 @@ def test_msa_cutlass_decode_static_dispatch_requires_sm100( ) +def test_msa_cutlass_decode_static_dispatch_query_head_limits() -> None: + """FP8 plans fit the kv_mode 3 work list (65536 query-head rows); NVFP4 + plans run on the Q8KV4 kernel, which takes 65535 query tokens.""" + + def planned(batch_size: int, kv_cache_dtype: str) -> bool: + return should_prepare_decode_metadata( + batch_size, + 8, + decode_backend="cutlass", + num_q_heads=16, + num_kv_heads=1, + kv_cache_dtype=kv_cache_dtype, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + ) + + assert planned(512, "fp8") and not planned(513, "fp8") + assert planned(8191, "nvfp4") and not planned(8192, "nvfp4") + + @pytest.mark.parametrize("decode_backend", ["cutlass", "triton"]) def test_msa_cutlass_decode_static_dispatch_nvfp4_has_no_fallback( decode_backend: str, @@ -205,7 +230,7 @@ def test_nvfp4_selects_msa_nvfp4_backend_and_requires_sm100( num_kv_heads=1, ) assert backend_cls is MiniMaxM3SparseMSANvfp4Backend - assert impl_cls is MiniMaxM3SparseMSAImpl + assert impl_cls is MiniMaxM3SparseMSANvfp4Impl with pytest.raises(ValueError, match="nvfp4_4over6"): select_main_backend_and_impl_cls( topk_blocks=TOPK, @@ -229,12 +254,13 @@ def test_nvfp4_selects_msa_nvfp4_backend_and_requires_sm100( @pytest.mark.parametrize( ("backend", "expected"), [ + (None, "cutlass"), (AttentionBackendEnum.CUTLASS_MSA, "cutlass"), (AttentionBackendEnum.TRITON_MSA, "triton"), ], ) def test_msa_attention_backend_alias( - backend: AttentionBackendEnum, + backend: AttentionBackendEnum | None, expected: str, ) -> None: config = AttentionConfig(backend=backend) @@ -275,6 +301,7 @@ def test_msa_metadata_builder_prepares_cutlass_for_regular_decode( seq_lens = torch.full((batch,), 257, dtype=torch.int32, device="cuda") base_metadata = SimpleNamespace( num_decodes=batch, + prefill=None, decode=SimpleNamespace( block_table=block_table, seq_lens=seq_lens, @@ -300,6 +327,7 @@ def test_msa_metadata_builder_prepares_cutlass_for_regular_decode( builder.kv_cache_spec = SimpleNamespace(num_kv_heads=4) builder.kv_cache_dtype = "fp8_e4m3" builder.decode_backend = "cutlass" + builder.q8kv4_decode = False builder.msa_cutlass_plan_cache = object() builder.msa_cutlass_min_batch_size = None @@ -315,6 +343,39 @@ def test_msa_metadata_builder_prepares_cutlass_for_regular_decode( assert metadata.decode.msa_cutlass is expected_metadata +@pytest.mark.parametrize("is_sm107", [True, False]) +def test_msa_metadata_builder_reserves_prefill_workspace_only_on_sm107( + monkeypatch: pytest.MonkeyPatch, + is_sm107: bool, +) -> None: + monkeypatch.setattr( + MiniMaxM3SparseMetadataBuilder, "__init__", lambda *args, **kwargs: None + ) + monkeypatch.setattr( + current_platform, + "is_device_capability", + lambda capability: is_sm107 and capability == 107, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace( + num_attention_heads=16, + sparse_attention_config={"sparse_topk_blocks": TOPK}, + ) + ), + parallel_config=SimpleNamespace(tensor_parallel_size=1), + cache_config=SimpleNamespace(cache_dtype="nvfp4"), + attention_config=SimpleNamespace(minimax_m3_msa_decode_backend="cutlass"), + scheduler_config=SimpleNamespace(max_num_batched_tokens=256), + ) + + builder = MiniMaxM3SparseMSAMetadataBuilder( + SimpleNamespace(num_kv_heads=1), [], vllm_config, torch.device("cuda") + ) + + assert (builder.prefill_workspace is not None) == is_sm107 + + def test_msa_cutlass_plan_cache_keys_query_len( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -382,7 +443,7 @@ def fake_build_plan(**kwargs): assert built_query_lens == [1, 2] -def test_query_fp8_stays_valid_when_cutlass_plan_appears_on_replay( +def test_query_stays_valid_when_cutlass_plan_appears_on_replay( monkeypatch: pytest.MonkeyPatch, ) -> None: import vllm.envs as envs @@ -430,20 +491,21 @@ def test_query_fp8_stays_valid_when_cutlass_plan_appears_on_replay( observed_ptrs: list[int] = [] @eager_break_during_capture - def run_attention(query_fp8: torch.Tensor | None) -> None: + def run_attention(query: torch.Tensor) -> None: if impl.should_use_msa_decode(layer_name): - assert query_fp8 is not None - observed_ptrs.append(query_fp8.data_ptr()) + assert query.dtype == torch.float8_e4m3fn + observed_ptrs.append(query.data_ptr()) stream = torch.cuda.Stream() with torch.cuda.stream(stream), override_forward_context(forward_context): capture = BreakableCUDAGraphCapture() with capture: - query_fp8 = MiniMaxM3SparseAttention._allocate_query_fp8(attention, qkv) - if query_fp8 is not None: - query_fp8.zero_() - run_attention(query_fp8) + query = MiniMaxM3SparseAttention._allocate_query( + attention, num_tokens, qkv.dtype, qkv.device + ) + query.zero_() + run_attention(query) qkv.zero_() assert capture.num_graphs == 2 @@ -797,18 +859,18 @@ def _make_nvfp4_kv_cache( return kv_cache, torch.cat(dequantized, dim=-1).to(torch.bfloat16) -def _e4m3_staged_kv_reference( +def _q8kv4_kv_reference( kv_reference: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor ) -> torch.Tensor: - """The dequantized NVFP4 cache as an E4M3-query kernel reads it: each - code * block_scale rounded to E4M3, then the global scale.""" + """The dequantized cache as the Q8KV4 kernels read it: block scales staged + by 2**NVFP4_BLOCK_SCALE_SHIFT, and each code * block_scale rounded to E4M3 + for the FP8 tensor core.""" sides = [] for side, global_scale in zip( kv_reference.float().split(HEAD_DIM, dim=-1), (k_scale, v_scale) ): - sides.append( - (side / global_scale).to(torch.float8_e4m3fn).float() * global_scale - ) + unit = global_scale * 2.0**NVFP4_BLOCK_SCALE_SHIFT + sides.append((side / unit).to(torch.float8_e4m3fn).float() * unit) return torch.cat(sides, dim=-1).to(torch.bfloat16) @@ -862,6 +924,8 @@ def test_nvfp4_kv_cache_views_address_head_slots() -> None: pytest.param(64, 4, 1, 1, True, id="tp1-small-batch-query-len-1"), pytest.param(64, 4, 8, 8, False, id="tp1-query-len-8"), pytest.param(16, 1, 16, 4, True, id="tp4-query-len-4"), + # 1024 x 8 x 16 = 131072 query-head rows, past the kv_mode 3 work list. + pytest.param(16, 1, 512, 8, True, id="tp4-query-len-8-1024-requests"), ], ) def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache( @@ -871,8 +935,9 @@ def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache( query_len: int, capture: bool, ) -> None: - """CUTLASS decode reads vLLM's packed NVFP4 pages, including swizzled V - scales and device global scales, as their exact dequantized values.""" + """Q8KV4 decode reads vLLM's packed NVFP4 pages, including swizzled V + scales and device global scales, as their exact dequantized values, and a + reused plan follows new lengths and zero-length padding requests.""" torch.manual_seed(0) seq_lens_list = [257, 513] * num_request_pairs seq_lens_cpu = torch.tensor(seq_lens_list, dtype=torch.int32) @@ -890,18 +955,27 @@ def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache( ) query_fp8 = query.to(torch.float8_e4m3fn) topk = _make_topk(seq_lens_list, num_kv_heads, query_len) - expected = torch.empty_like(query) - minimax_m3_sparse_attn_decode( - query_fp8.to(torch.bfloat16), - kv_reference, - topk.transpose(0, 1), - block_table, - seq_lens, - num_kv_heads, - SM_SCALE, - expected, - query_len, - ) + kv_reference = _q8kv4_kv_reference(kv_reference, k_scale, v_scale) + + def reference(lens: list[int]) -> torch.Tensor: + # Padding requests (length 0) trail the batch and are not compared. + num_reqs = sum(1 for n in lens if n > 0) + rows = num_reqs * query_len + expected = torch.empty_like(query[:rows]) + minimax_m3_sparse_attn_decode( + query_fp8[:rows].to(torch.bfloat16), + kv_reference, + topk[:rows].transpose(0, 1), + block_table[:num_reqs], + seq_lens[:num_reqs], + num_kv_heads, + SM_SCALE, + expected, + query_len, + ) + return expected + + expected = reference(seq_lens_list) assert should_prepare_decode_metadata( len(seq_lens_list), @@ -913,16 +987,24 @@ def test_msa_cutlass_decode_nvfp4_matches_triton_on_dequantized_cache( page_size=BLOCK_SIZE, topk_blocks=TOPK, ) - metadata = prepare_decode_metadata( - block_table, - seq_lens, - seq_lens_cpu, - query_len, - num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, - page_size=BLOCK_SIZE, - topk_blocks=TOPK, - ) + plan_cache = MSACutlassDecodePlanCache() + + def prepare(lens_cpu: torch.Tensor): + return prepare_decode_metadata( + block_table, + seq_lens, + lens_cpu, + query_len, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + page_size=BLOCK_SIZE, + topk_blocks=TOPK, + q8kv4=True, + plan_cache=plan_cache, + ) + + metadata = prepare(seq_lens_cpu) + assert "q8kv4" in metadata.plan[3] actual = torch.zeros_like(query) def run() -> None: @@ -940,6 +1022,7 @@ def run() -> None: v_scale=v_scale, ) + graph = None if capture: run() graph = torch.cuda.CUDAGraph() @@ -952,6 +1035,25 @@ def run() -> None: current_platform.synchronize() torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.05) + # Next step on the same plan: shorter requests and, as a padded full CUDA + # graph batch has, a trailing zero-length request. + updated_seq_lens_list = [129, 385] * num_request_pairs + updated_seq_lens_list[-1] = 0 + seq_lens.copy_(torch.tensor(updated_seq_lens_list, dtype=torch.int32)) + topk.copy_(_make_topk(updated_seq_lens_list, num_kv_heads, query_len)) + updated_metadata = prepare(torch.tensor(updated_seq_lens_list, dtype=torch.int32)) + assert updated_metadata.plan is metadata.plan + expected = reference(updated_seq_lens_list) + actual.zero_() + if graph is not None: + graph.replay() + else: + run() + current_platform.synchronize() + torch.testing.assert_close( + actual[: expected.shape[0]], expected, atol=0.05, rtol=0.05 + ) + def _make_prefill_metadata( q_lens: list[int], @@ -1003,18 +1105,27 @@ def _make_prefill_metadata( return metadata, topk +def _skip_without_q8kv4_prefill(num_q_heads: int, num_kv_heads: int) -> None: + blocker = sparse_attention_msa_module._q8kv4_prefill_blocker( + num_q_heads, num_kv_heads, BLOCK_SIZE, TOPK + ) + if blocker is not None: + pytest.skip(blocker) + + @pytest.mark.parametrize( ("num_q_heads", "num_kv_heads"), [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], ) def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( - monkeypatch: pytest.MonkeyPatch, num_q_heads: int, num_kv_heads: int, ) -> None: - """MSA prefill reads NVFP4 pages as their dequantized values, staged to E4M3 - for the E4M3 q, including chunked prefill over cached context; it attends - with the fp8 q (the only one the fused insert writes on this path).""" + """MSA prefill reads NVFP4 pages as their exact dequantized values, + including chunked prefill over cached context, with the fp8 q the fused + insert writes (the only q it writes on this path), on the CUTLASS Q8KV4 + kernel.""" + _skip_without_q8kv4_prefill(num_q_heads, num_kv_heads) torch.manual_seed(0) q_lens = [300, 128, 77] seq_lens_list = [300, 640, 1100] @@ -1025,10 +1136,10 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( num_pages, num_kv_heads, "nvfp4", k_scale, v_scale ) + total_q = sum(q_lens) metadata, topk = _make_prefill_metadata( q_lens, seq_lens_list, block_table, num_kv_heads ) - total_q = sum(q_lens) layer_name = "model.layers.0.self_attn.attn" q_scale = 0.5 @@ -1046,7 +1157,7 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( / q_scale ).to(torch.float8_e4m3fn) query = query_fp8.to(torch.bfloat16) * q_scale - kv_reference = _e4m3_staged_kv_reference(kv_reference, k_scale, v_scale) + kv_reference = _q8kv4_kv_reference(kv_reference, k_scale, v_scale) common = dict( num_heads=num_q_heads, head_size=HEAD_DIM, @@ -1055,7 +1166,6 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( topk_blocks=TOPK, sparse_block_size=BLOCK_SIZE, ) - query_dtypes = _record_query_dtypes(monkeypatch, "sparse_atten_nvfp4_kv_func") forward_context = ForwardContext( no_compile_layers={}, attn_metadata={layer_name: metadata}, @@ -1065,100 +1175,98 @@ def test_msa_prefill_nvfp4_matches_triton_on_dequantized_cache( expected = MiniMaxM3SparseTritonImpl(**common).forward( layer, query, kv_reference, torch.empty_like(query) ) - actual = MiniMaxM3SparseMSAImpl(**common, kv_cache_dtype="nvfp4").forward( - layer, - torch.full_like(query, float("nan")), - kv_cache, - torch.zeros_like(query), - query_fp8=query_fp8, + # As the model calls it for NVFP4: the fp8 q is the only query. + actual = MiniMaxM3SparseMSANvfp4Impl(**common).forward( + layer, query_fp8, kv_cache, torch.zeros_like(query) ) current_platform.synchronize() - assert query_dtypes == [torch.float8_e4m3fn] _assert_close_fp8_prefill(actual, expected) -def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None: - """With CUTLASS decode on but not planned for the batch, the Triton decode - fallback attends with the dequantized fp8 q, the only q written.""" +@pytest.mark.parametrize( + ("num_q_heads", "num_kv_heads"), + [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], +) +def test_msa_prefill_nvfp4_uses_the_builders_workspace( + num_q_heads: int, + num_kv_heads: int, +) -> None: + """The NVFP4 prefill hands the builder's split workspace to ``run_prefill``: + a workspace reserved for the step's token budget is neither replaced nor + grown by the call, and the output equals the per-call-allocation path's.""" + from vllm.third_party.fmha_sm100.prefill_q8kv4 import PrefillWorkspace + + _skip_without_q8kv4_prefill(num_q_heads, num_kv_heads) torch.manual_seed(0) - num_q_heads, num_kv_heads, q_scale = 16, 1, 0.5 - seq_lens_list = [257, 513, 129] - seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda") + q_lens = [300, 128, 77] + seq_lens_list = [300, 640, 1100] block_table, num_pages = _make_paged_layout(seq_lens_list) - kv_cache = ( - torch.randn( - num_pages, - num_kv_heads, - BLOCK_SIZE, - 2 * HEAD_DIM, - dtype=torch.bfloat16, - device="cuda", - ) - * 0.25 - ).to(torch.float8_e4m3fn) - num_tokens = len(seq_lens_list) - metadata = MiniMaxM3SparseMetadata( - seq_lens=seq_lens, - max_seq_len=max(seq_lens_list), - slot_mapping=torch.empty(num_tokens, dtype=torch.int64, device="cuda"), - num_actual_tokens=num_tokens, - num_decodes=num_tokens, - num_decode_tokens=num_tokens, - num_prefills=0, - num_prefill_tokens=0, - decode=MiniMaxM3SparseMSADecodeMetadata( - seq_lens=seq_lens, - block_table=block_table, - decode_query_len=1, - msa_cutlass=None, + k_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") + v_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") + kv_cache, _ = _make_nvfp4_kv_cache( + num_pages, num_kv_heads, "nvfp4", k_scale, v_scale + ) + metadata, topk = _make_prefill_metadata( + q_lens, seq_lens_list, block_table, num_kv_heads + ) + total_q = sum(q_lens) + max_num_batched_tokens = 2048 + workspace = PrefillWorkspace(torch.device("cuda"), block_scale_shift=0) + workspace.split_buffers(TOPK, max_num_batched_tokens, num_q_heads) + reserved = workspace.split_buffers(TOPK, total_q, num_q_heads)[0].data_ptr() + assert isinstance(metadata.prefill, MiniMaxM3SparsePrefillMetadata) + with_workspace = replace( + metadata, + prefill=MiniMaxM3SparseMSAPrefillMetadata( + **{ + f.name: getattr(metadata.prefill, f.name) + for f in fields(metadata.prefill) + }, + workspace=workspace, ), ) + layer_name = "model.layers.0.self_attn.attn" layer = SimpleNamespace( layer_name=layer_name, - topk_indices_buffer=_make_topk(seq_lens_list, num_kv_heads, 1), - _q_scale_float=q_scale, + topk_indices_buffer=topk, + _q_scale_float=1.0, + _k_scale=k_scale, + _v_scale=v_scale, ) query_fp8 = torch.randn( - num_tokens, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda" + total_q, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda" ).to(torch.float8_e4m3fn) - query = query_fp8.to(torch.bfloat16) * q_scale - common = dict( + impl = MiniMaxM3SparseMSANvfp4Impl( num_heads=num_q_heads, head_size=HEAD_DIM, scale=SM_SCALE, num_kv_heads=num_kv_heads, - kv_cache_dtype="fp8", topk_blocks=TOPK, sparse_block_size=BLOCK_SIZE, ) - impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass") - assert impl.use_cutlass_decode - forward_context = ForwardContext( - no_compile_layers={}, - attn_metadata={layer_name: metadata}, - slot_mapping={}, - ) - with override_forward_context(forward_context): - expected = MiniMaxM3SparseTritonImpl(**common).forward( - layer, query, kv_cache, torch.empty_like(query) + outputs = [] + for md in (metadata, with_workspace): + forward_context = ForwardContext( + no_compile_layers={}, + attn_metadata={layer_name: md}, + slot_mapping={}, ) - actual = impl.forward( - layer, - torch.full_like(query, float("nan")), - kv_cache, - torch.zeros_like(query), - query_fp8=query_fp8, + output = torch.zeros( + total_q, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda" ) + with override_forward_context(forward_context): + outputs.append(impl.forward(layer, query_fp8, kv_cache, output)) current_platform.synchronize() - torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) + assert torch.equal(outputs[0], outputs[1]) + assert workspace.split_buffers(TOPK, total_q, num_q_heads)[0].data_ptr() == reserved @pytest.mark.parametrize( ("num_q_heads", "num_kv_heads"), [pytest.param(64, 4, id="tp1"), pytest.param(16, 1, id="tp4")], ) -def test_msa_prefill_reads_query_fp8_over_fp8_cache( +def test_msa_prefill_reads_the_e4m3_query_over_an_fp8_cache( monkeypatch: pytest.MonkeyPatch, num_q_heads: int, num_kv_heads: int, @@ -1219,13 +1327,126 @@ def test_msa_prefill_reads_query_fp8_over_fp8_cache( expected = MiniMaxM3SparseTritonImpl(**common).forward( layer, query, kv_cache, torch.empty_like(query) ) - actual = impl.forward( - layer, - torch.full_like(query, float("nan")), - kv_cache, - torch.zeros_like(query), - query_fp8=query_fp8, - ) + # As the model calls it: the fp8 q is the only query the insert wrote. + actual = impl.forward(layer, query_fp8, kv_cache, torch.zeros_like(query)) current_platform.synchronize() assert query_dtypes == [torch.float8_e4m3fn] _assert_close_fp8_prefill(actual, expected) + + +def test_nvfp4_kv_cache_never_takes_a_bf16_query() -> None: + """NVFP4 attention always reads an E4M3 q on the CUTLASS Q8KV4 kernels, and + a head layout either kernel cannot serve fails at init.""" + assert uses_q8kv4_decode(kv_cache_dtype="nvfp4", num_q_heads=16, num_kv_heads=1) + assert not uses_q8kv4_decode( + kv_cache_dtype="fp8_e4m3", num_q_heads=16, num_kv_heads=1 + ) + with pytest.raises(ValueError, match="Q8KV4"): + uses_q8kv4_decode(kv_cache_dtype="nvfp4", num_q_heads=4, num_kv_heads=1) + # 8 Q heads per KV head (TP8) decodes on Q8KV4, but Q8KV4 cannot prefill it. + with pytest.raises(ValueError, match="CUTLASS Q8KV4 prefill"): + MiniMaxM3SparseMSANvfp4Impl( + num_heads=8, + head_size=HEAD_DIM, + scale=SM_SCALE, + num_kv_heads=1, + topk_blocks=TOPK, + sparse_block_size=BLOCK_SIZE, + ) + + +@pytest.mark.parametrize( + "architecture", + ["MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration"], +) +def test_nvfp4_kv_cache_with_tp8_fails_at_config_time(architecture: str) -> None: + """An NVFP4 KV cache needs 16 query heads per KV head on each rank: the config + check rejects TP8 (8 per KV head) before the workers start; TP1/2/4 and other + KV cache dtypes pass.""" + from vllm.model_executor.models.config import MODELS_CONFIG_MAP + + def vllm_config(tensor_parallel_size: int, cache_dtype: str) -> SimpleNamespace: + text_config = SimpleNamespace(num_attention_heads=64, num_key_value_heads=4) + return SimpleNamespace( + model_config=SimpleNamespace(hf_text_config=text_config), + cache_config=SimpleNamespace(cache_dtype=cache_dtype), + parallel_config=SimpleNamespace(tensor_parallel_size=tensor_parallel_size), + ) + + verify = MODELS_CONFIG_MAP[architecture].verify_and_update_config + for tensor_parallel_size in (1, 2, 4): + verify(vllm_config(tensor_parallel_size, "nvfp4")) + verify(vllm_config(8, "fp8")) + with pytest.raises(ValueError, match="NVFP4 KV cache with tensor_parallel_size=8"): + verify(vllm_config(8, "nvfp4")) + + +def test_msa_triton_decode_fallback_reads_dequantized_query_fp8() -> None: + """With CUTLASS decode on but not planned for the batch, the Triton decode + fallback attends with the dequantized E4M3 query.""" + torch.manual_seed(0) + num_q_heads, num_kv_heads, q_scale = 16, 1, 0.5 + seq_lens_list = [257, 513, 129] + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda") + block_table, num_pages = _make_paged_layout(seq_lens_list) + kv_cache = ( + torch.randn( + num_pages, + num_kv_heads, + BLOCK_SIZE, + 2 * HEAD_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.25 + ).to(torch.float8_e4m3fn) + num_tokens = len(seq_lens_list) + metadata = MiniMaxM3SparseMetadata( + seq_lens=seq_lens, + max_seq_len=max(seq_lens_list), + slot_mapping=torch.empty(num_tokens, dtype=torch.int64, device="cuda"), + num_actual_tokens=num_tokens, + num_decodes=num_tokens, + num_decode_tokens=num_tokens, + num_prefills=0, + num_prefill_tokens=0, + decode=MiniMaxM3SparseMSADecodeMetadata( + seq_lens=seq_lens, + block_table=block_table, + decode_query_len=1, + msa_cutlass=None, + ), + ) + layer_name = "model.layers.0.self_attn.attn" + layer = SimpleNamespace( + layer_name=layer_name, + topk_indices_buffer=_make_topk(seq_lens_list, num_kv_heads, 1), + _q_scale_float=q_scale, + ) + query_fp8 = torch.randn( + num_tokens, num_q_heads * HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ).to(torch.float8_e4m3fn) + query = query_fp8.to(torch.bfloat16) * q_scale + common = dict( + num_heads=num_q_heads, + head_size=HEAD_DIM, + scale=SM_SCALE, + num_kv_heads=num_kv_heads, + kv_cache_dtype="fp8", + topk_blocks=TOPK, + sparse_block_size=BLOCK_SIZE, + ) + impl = MiniMaxM3SparseMSAImpl(**common, msa_decode_backend="cutlass") + assert impl.use_cutlass_decode + forward_context = ForwardContext( + no_compile_layers={}, + attn_metadata={layer_name: metadata}, + slot_mapping={}, + ) + with override_forward_context(forward_context): + expected = MiniMaxM3SparseTritonImpl(**common).forward( + layer, query, kv_cache, torch.empty_like(query) + ) + actual = impl.forward(layer, query_fp8, kv_cache, torch.zeros_like(query)) + current_platform.synchronize() + torch.testing.assert_close(actual, expected, atol=0.02, rtol=0.02) diff --git a/tests/kernels/attention/test_minimax_m3_norm_rope_nvfp4_insert.py b/tests/kernels/attention/test_minimax_m3_norm_rope_nvfp4_insert.py new file mode 100644 index 000000000000..dd60ed7c9129 --- /dev/null +++ b/tests/kernels/attention/test_minimax_m3_norm_rope_nvfp4_insert.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The fused FlyDSL norm + RoPE + NVFP4 insert against the HIP norm + RoPE +kernel followed by the reference NVFP4 inserts: every output byte, stale page +bytes included, and qkv left untouched.""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm (gfx950) FlyDSL kernel", allow_module_level=True) + +from vllm.platforms.rocm import on_gfx950 # noqa: E402 + +if not on_gfx950(): + pytest.skip("gfx950 instructions", allow_module_level=True) + +from tests.kernels.attention.minimax_m3_nvfp4_ref import ( # noqa: E402 + insert_nvfp4_index, + insert_nvfp4_kv, +) +from vllm import _custom_ops as ops # noqa: E402 +from vllm.models.minimax_m3.amd.ops.norm_rope_nvfp4_insert.host import ( # noqa: E402 + minimax_m3_norm_rope_nvfp4_insert, +) + +D = 128 +BLK = 128 +EPS = 1e-6 +ROT = 64 + + +def _request_slots(starts, lens, num_pages, gen): + """VLLM's slot mapping: request r's tokens at positions starts[r].. of its + own random pages.""" + perm = torch.randperm(num_pages, generator=gen) + out, p = [], 0 + for st, n in zip(starts, lens): + npg = (st + n + BLK - 1) // BLK + pos = torch.arange(st, st + n) + out.append(perm[p : p + npg][pos // BLK] * BLK + pos % BLK) + p += npg + return torch.cat(out) if out else torch.empty(0, dtype=torch.int64) + + +class _Batch: + def __init__(self, starts, lens, pad, heads, scales, seed=0): + g = torch.Generator().manual_seed(seed) + self.nq, self.nkv, self.niq = heads + n = self.n = sum(lens) + pad + row = (self.nq + 2 * self.nkv + self.niq + 1) * D + # magnitudes 1e-3 .. 3e3 per 16-value group, some all-zero groups + x = torch.randn(n, row // 16, 16, generator=g) + mags = torch.tensor([1e-3, 0.05, 1.0, 1.0, 30.0, 3e3]) + x *= mags[torch.randint(0, len(mags), (n, row // 16), generator=g)][..., None] + x[torch.rand(n, row // 16, generator=g) < 0.02] = 0.0 + self.qkv = x.view(n, row).to(torch.bfloat16).cuda() + self.w = [ + (torch.randn(D, generator=g) * 0.5).to(torch.bfloat16).cuda() + for _ in range(4) + ] + self.cs = (torch.rand(4096, ROT, generator=g) * 2 - 1).to(torch.bfloat16).cuda() + self.pos = torch.randint(0, 4096, (n,), generator=g).cuda() + pages = sum((s + m + BLK - 1) // BLK for s, m in zip(starts, lens)) + 2 + pad_slots = torch.full((pad,), -1, dtype=torch.int64) + self.slots = torch.cat( + [_request_slots(starts, lens, pages, g), pad_slots] + ).cuda() + self.islots = torch.cat( + [_request_slots(starts, lens, pages, g), pad_slots] + ).cuda() + # stale bytes; the layer views have a block stride past one page + main = 2 * self.nkv * BLK * 72 + self.kv_full = torch.randint(0, 256, (pages, main + 256), generator=g) + self.kv_full = self.kv_full.to(torch.uint8).cuda() + self.idx_full = torch.randint(0, 256, (pages, BLK * 72 + 64), generator=g) + self.idx_full = self.idx_full.to(torch.uint8).cuda() + self.ks = torch.tensor(scales[0], dtype=torch.float32, device="cuda") + self.vs = torch.tensor(scales[1], dtype=torch.float32, device="cuda") + + def kv(self, full): + return full[:, : 2 * self.nkv * BLK * 72].view(-1, 2 * self.nkv, BLK, 72) + + def idx(self, full): + return full[:, : BLK * 72].view(-1, BLK, 72) + + def reference(self): + qkv = self.qkv.clone() + q = torch.empty(self.n, self.nq * D, dtype=torch.bfloat16, device="cuda") + iq = torch.empty(self.n, self.niq * D, dtype=torch.float8_e4m3fn, device="cuda") + kv_full, idx_full = self.kv_full.clone(), self.idx_full.clone() + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, self.w[0], self.w[1], self.cs, self.pos, self.nq, self.nkv, ROT, + EPS, self.w[2], self.w[3], self.niq, q_out=q, index_q_out=iq, + ) # fmt: skip + k0, v0 = self.nq * D, (self.nq + self.nkv) * D + ik0 = (self.nq + 2 * self.nkv + self.niq) * D + insert_nvfp4_kv( + qkv[:, k0:v0].view(self.n, self.nkv, D), + qkv[:, v0 : v0 + self.nkv * D].view(self.n, self.nkv, D), + self.kv(kv_full), + self.slots, + self.ks, + self.vs, + ) + insert_nvfp4_index(self.idx(idx_full), qkv[:, ik0 : ik0 + D], self.islots) + return q, iq, kv_full, idx_full + + def run(self, qkv, pos, slots, islots, kv_full, idx_full, q, iq): + minimax_m3_norm_rope_nvfp4_insert( + qkv, self.w[0], self.w[1], self.w[2], self.w[3], self.cs, pos, + slots, islots, self.kv(kv_full), self.idx(idx_full), self.ks, + self.vs, q, iq, num_heads=self.nq, num_kv_heads=self.nkv, + num_index_heads=self.niq, eps=EPS, + ) # fmt: skip + + def fused(self, iq): + qkv = self.qkv.clone() + q = torch.empty(self.n, self.nq * D, dtype=torch.bfloat16, device="cuda") + kv_full, idx_full = self.kv_full.clone(), self.idx_full.clone() + self.run(qkv, self.pos, self.slots, self.islots, kv_full, idx_full, q, iq) + assert torch.equal(qkv.view(torch.int16), self.qkv.view(torch.int16)) + return q, kv_full, idx_full + + +def _check(batch, iq_mode="contiguous"): + q_ref, iq_ref, kv_ref, idx_ref = batch.reference() + n, w = batch.n, batch.niq * D + if iq_mode == "rank": # the context-parallel indexer's [4, n, 128] buffer + buf = torch.full((4, n, w), 0x7F, dtype=torch.uint8, device="cuda") + iq = buf[2] + elif iq_mode == "cols": + buf = torch.full((n, 3 * w), 0x7F, dtype=torch.uint8, device="cuda") + iq = buf[:, w : 2 * w] + else: + buf = iq = torch.full((n, w), 0x7F, dtype=torch.uint8, device="cuda") + q, kv, idx = batch.fused(iq.view(torch.float8_e4m3fn)) + assert torch.equal(q.view(torch.int16), q_ref.view(torch.int16)) + assert torch.equal(iq, iq_ref.view(torch.uint8)) + assert torch.equal(kv, kv_ref) + assert torch.equal(idx, idx_ref) + iq.fill_(0x7F) + assert bool((buf == 0x7F).all()), "bytes outside index_q_out written" + + +TP4, TP1 = (16, 1, 1), (64, 4, 4) +DSPARK = ([3 + 9 * i for i in range(43)], [4] * 43) # 172 decode rows + + +@pytest.mark.parametrize("heads", [TP4, TP1]) +@pytest.mark.parametrize( + ("starts", "lens", "pad"), + [ + ([5], [1], 0), # one decode token + ([0, 13], [7, 4], 1), # a run of 7, runs across an octet edge + ([125, 254, 4091], [6, 6, 6], 12), # Q6 crossing page boundaries + (*DSPARK, 20), # DSpark decode rows, graph padding + ([0, 3, 250, 128], [131, 61, 9, 99], 0), # unaligned prefill chunks + ([0, 100], [4096, 4000], 3), # the prefill tier + ], +) +@pytest.mark.parametrize("scales", [(1.0, 1.0), (0.002, 3.0)]) +def test_matches_two_kernel_path(heads, starts, lens, pad, scales): + _check(_Batch(starts, lens, pad, heads, scales)) + + +@pytest.mark.parametrize("iq_mode", ["rank", "cols"]) +@pytest.mark.parametrize(("starts", "lens"), [DSPARK, ([7], [600])]) +def test_index_q_out_view(iq_mode, starts, lens): + _check(_Batch(starts, lens, 5, TP4, (1.0, 1.0)), iq_mode) + + +def test_64k_tokens(): + _check(_Batch([0, 777], [32768, 32768], 0, TP4, (1.0, 1.0))) + + +def test_all_padding(): + """A graph-capture dummy batch: every slot -1, the pages stay as they were.""" + batch = _Batch([], [], 64, TP4, (1.0, 1.0)) + _check(batch) + iq = torch.empty(batch.n, D, dtype=torch.float8_e4m3fn, device="cuda") + _, kv, idx = batch.fused(iq) + assert torch.equal(kv, batch.kv_full) and torch.equal(idx, batch.idx_full) + + +def test_cuda_graph_replay(): + """Captured once, replayed on new qkv / positions / slots in place.""" + a = _Batch(*DSPARK, 4, TP4, (1.0, 1.0), seed=1) + b = _Batch(*DSPARK, 4, TP4, (1.0, 1.0), seed=2) + b.w, b.cs = a.w, a.cs # captured with the layer's weights and RoPE cache + qkv, pos = a.qkv.clone(), a.pos.clone() + slots, islots = a.slots.clone(), a.islots.clone() + kv_full, idx_full = b.kv_full.clone(), b.idx_full.clone() + q = torch.empty(a.n, a.nq * D, dtype=torch.bfloat16, device="cuda") + iq = torch.empty(a.n, D, dtype=torch.float8_e4m3fn, device="cuda") + args = (qkv, pos, slots, islots, kv_full, idx_full, q, iq) + a.run(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + a.run(*args) + new = (b.qkv, b.pos, b.slots, b.islots, b.kv_full, b.idx_full) + for dst, src in zip(args[:6], new): + dst.copy_(src) + graph.replay() + q_ref, iq_ref, kv_ref, idx_ref = b.reference() + assert torch.equal(q.view(torch.int16), q_ref.view(torch.int16)) + assert torch.equal(iq.view(torch.uint8), iq_ref.view(torch.uint8)) + assert torch.equal(kv_full, kv_ref) and torch.equal(idx_full, idx_ref) diff --git a/tests/kernels/attention/test_minimax_m3_nvfp4_index.py b/tests/kernels/attention/test_minimax_m3_nvfp4_index.py new file mode 100644 index 000000000000..86c0632a79af --- /dev/null +++ b/tests/kernels/attention/test_minimax_m3_nvfp4_index.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3 NVFP4 index scorers (FlyDSL) against AITER's FP8 scorers and an +fp32 reference over the same e4m3 keys; the context-parallel push scorer +against the single-head scorer, bit for bit, on one GPU.""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm (gfx950) FlyDSL kernels", allow_module_level=True) + +from vllm.platforms.rocm import on_gfx950 # noqa: E402 + +if not on_gfx950(): + pytest.skip("gfx950 instructions", allow_module_level=True) + +from tests.kernels.attention.minimax_m3_nvfp4_ref import dequant # noqa: E402 +from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode import ( # noqa: E402 + host as decode_host, +) +from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode import ( # noqa: E402 + kernel_cp_push as cp, +) +from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill import ( # noqa: E402 + host as prefill_host, +) + +minimax_m3_nvfp4_index_score_decode = decode_host.minimax_m3_nvfp4_index_score_decode +minimax_m3_nvfp4_index_score_prefill = prefill_host.minimax_m3_nvfp4_index_score_prefill + +BLOCK = 128 +D = 128 +SPECIAL_LENS = [0, 1, 2, 3, 127, 128, 129, 1024] # padding, short and edge rows + + +def _index_cache(num_blocks: int, page_step: int = 1) -> torch.Tensor: + """Random NVFP4 index pages ``page_step`` pages apart (a view of a larger + cache).""" + n = num_blocks * page_step + cache = torch.randint(0, 256, (n, BLOCK * 72), device="cuda").to(torch.uint8) + s = torch.rand(n, BLOCK * 8, device="cuda") * 1.75 + 0.25 + cache[:, BLOCK * 64 :] = s.to(torch.float8_e4m3fn).view(torch.uint8) + return cache.view(n, BLOCK, 72)[::page_step] + + +def _fp8_keys(cache: torch.Tensor) -> torch.Tensor: + """Every page's keys as e4m3_rn_satfinite(e2m1 x scale), [nb, 128, 128].""" + flat = cache.reshape(cache.shape[0], -1) + keys = dequant( + flat[:, : BLOCK * 64].view(-1, BLOCK, 64), + flat[:, BLOCK * 64 :].view(-1, BLOCK, 8), + ) + return keys.clamp(-448, 448).to(torch.float8_e4m3fn) + + +def _block_table(seq_lens, share=False): + """Shuffled pages per request; ``share``: some requests start with another + request's pages (a shared prefix).""" + widths = [(s + BLOCK - 1) // BLOCK for s in seq_lens] + num_blocks = sum(widths) + 2 + perm = torch.randperm(num_blocks, device="cuda", dtype=torch.int32) + table = torch.zeros(len(seq_lens), max(widths), dtype=torch.int32, device="cuda") + start = 0 + for r, w in enumerate(widths): + table[r, :w] = perm[start : start + w] + start += w + for r in range(len(seq_lens) if share else 0): + src = 0 if r % 2 == 0 else 3 * (r // 6) + n = min(widths[r], widths[src], 10 if r % 2 == 0 else 6) + if r != src: + table[r, :n] = table[src, :n] + return table, num_blocks + + +def _lens(num_reqs: int, hi: int) -> list[int]: + lens = torch.randint(1, hi, (num_reqs,)).tolist() + for i, v in enumerate(SPECIAL_LENS[: num_reqs // 2]): + lens[2 * i + 1] = v + return lens + + +def _ref_block_scores(q_rows, keys, kv_lens, init_blocks, local_blocks): + """[H, R, nb] per-block max of q . k over each row's visible tokens, -inf + where none is visible, AITER's init (1e30) / local (1e29) sentinels.""" + nb = keys.shape[0] // BLOCK + dots = torch.einsum("rhd,td->hrt", q_rows.float(), keys) + tok = torch.arange(nb * BLOCK, device=keys.device) + dots = dots.masked_fill(tok[None, None, :] >= kv_lens[None, :, None], -torch.inf) + mx = dots.view(dots.shape[0], dots.shape[1], nb, BLOCK).amax(-1) + rb = (kv_lens.clamp(min=0) + BLOCK - 1) // BLOCK # visible blocks per row + blk = torch.arange(nb, device=keys.device)[None, :] + vis = blk < rb[:, None] + init = vis & (blk < init_blocks) + local = vis & (blk >= (rb - local_blocks).clamp(min=0)[:, None]) + mx = torch.where(init[None], torch.full_like(mx, 1e30), mx) + return torch.where(local[None], torch.full_like(mx, 1e29), mx) + + +def _assert_scores(got, want, msg=None): + """Sentinels and -inf identical; finite scores close.""" + exact = ~torch.isfinite(want) | (want >= 1e28) + assert torch.equal(got[exact], want[exact]), msg + torch.testing.assert_close(got[~exact], want[~exact], rtol=2e-4, atol=1e-3, msg=msg) + + +def _check_decode(actual, q, keys, table, lens, query_len, init_blocks, local_blocks): + """Each request's rows against the reference; nothing past its blocks + written.""" + for r, s in enumerate(lens): + nb = (s + BLOCK - 1) // BLOCK + sl = slice(r * query_len, (r + 1) * query_len) + assert actual[:, sl, nb:].isnan().all(), (r, s) + if nb == 0: + continue + kv_lens = (s - query_len + 1 + torch.arange(query_len, device="cuda")).clamp( + min=0 + ) + page_keys = keys[table[r, :nb].long()].reshape(-1, D) + want = _ref_block_scores(q[sl], page_keys, kv_lens, init_blocks, local_blocks) + _assert_scores(actual[:, sl, :nb], want, (r, s)) + + +@pytest.mark.parametrize( + ("num_heads", "query_len"), + [(1, 1), (1, 4), (1, 16), (2, 4), (4, 1), (4, 4), (4, 5), (4, 8)], +) +@pytest.mark.parametrize("page_step", [1, 2]) +def test_index_decode_matches_aiter_fp8(num_heads, query_len, page_step): + """Every block a row can see against AITER's FP8 decode scorer.""" + from aiter.ops.msa_attention import pa_sparse_block_score_decode + + torch.manual_seed(0) + lens = [5, 300, 2100, 9000] + table, num_blocks = _block_table(lens) + seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda") + cache = _index_cache(num_blocks, page_step) + rows = len(lens) * query_len + q = (torch.randn(rows, num_heads, D, device="cuda") * 0.5).to(torch.float8_e4m3fn) + width = (max(lens) + BLOCK - 1) // BLOCK + args = (table, seq_lens, 1, 2, query_len, max(lens)) + actual = torch.full((num_heads, rows, width), float("nan"), device="cuda") + minimax_m3_nvfp4_index_score_decode(q, cache, actual, *args) + expected = torch.full_like(actual, float("nan")) + pages = _fp8_keys(cache) + for h0 in range(0, num_heads, 2): # AITER builds one- and two-head scorers + pa_sparse_block_score_decode( + q[:, h0 : h0 + 2].contiguous(), pages, expected[h0 : h0 + 2], *args + ) + torch.accelerator.synchronize() + for r, s in enumerate(lens): + for t in range(query_len): + row = r * query_len + t + seen = (max(s - query_len + t + 1, 0) + BLOCK - 1) // BLOCK + got, want = actual[:, row, :seen], expected[:, row, :seen] + finite = want < 1e28 + assert torch.equal(got[~finite], want[~finite]) + torch.testing.assert_close(got[finite], want[finite], rtol=1e-5, atol=1e-4) + + +@pytest.mark.parametrize( + ("num_heads", "query_len"), [(1, 1), (1, 4), (2, 3), (4, 1), (4, 4), (4, 6), (4, 8)] +) +@pytest.mark.parametrize("num_reqs", [1, 40, 130]) +@pytest.mark.parametrize(("init_blocks", "local_blocks"), [(1, 2), (0, 1)]) +@pytest.mark.parametrize("buffered", [False, True]) +def test_index_decode_matches_reference( + num_heads, query_len, num_reqs, init_blocks, local_blocks, buffered, monkeypatch +): + """Request chunks, shared prefixes, padding requests, sequences shorter + than the query; ``buffered`` forces the LDS-buffered variant.""" + if buffered: + monkeypatch.setattr(decode_host, "BUFFER_MIN_BYTES", 0) + monkeypatch.setattr(decode_host, "UNITS_PER_SLOT", 1) + torch.manual_seed(num_reqs) + lens = _lens(num_reqs, 6000) + table, num_blocks = _block_table(lens, share=True) + seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda") + cache = _index_cache(num_blocks, 2 if num_reqs == 40 else 1) + rows = num_reqs * query_len + q = (torch.randn(rows, num_heads, D, device="cuda") * 0.5).to(torch.float8_e4m3fn) + width = (max(lens) + BLOCK - 1) // BLOCK + 2 + actual = torch.full((num_heads, rows, width), float("nan"), device="cuda") + minimax_m3_nvfp4_index_score_decode( + q, + cache, + actual, + table, + seq_lens, + init_blocks, + local_blocks, + query_len, + max(lens), + ) + keys = _fp8_keys(cache).float() + torch.accelerator.synchronize() + _check_decode(actual, q, keys, table, lens, query_len, init_blocks, local_blocks) + + +@pytest.mark.parametrize(("num_heads", "query_len"), [(1, 4), (4, 4)]) +def test_index_decode_graph_replays_new_lengths(num_heads, query_len): + """Captured once at a large max_seq_len, replayed with new lengths and + queries.""" + torch.manual_seed(5) + sets = [ + [140000, 6, 70001, 129, 7, 5000, 0, 128], + [6, 127, 2048, 140000, 40000, 300, 129, 0], + ] + table, num_blocks = _block_table([max(x) for x in zip(*sets)]) + cache = _index_cache(num_blocks) + seq_lens = torch.tensor(sets[0], dtype=torch.int32, device="cuda") + rows = len(sets[0]) * query_len + q = torch.zeros(rows, num_heads, D, device="cuda", dtype=torch.float8_e4m3fn) + max_seq_len = 1 << 18 + out = torch.empty(num_heads, rows, max_seq_len // BLOCK, device="cuda") + args = (q, cache, out, table, seq_lens, 1, 2, query_len, max_seq_len) + minimax_m3_nvfp4_index_score_decode(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + minimax_m3_nvfp4_index_score_decode(*args) + keys = _fp8_keys(cache).float() + for lens in sets + sets[:1]: + seq_lens.copy_(torch.tensor(lens, dtype=torch.int32)) + q.copy_((torch.randn(q.shape, device="cuda") * 0.5).to(q.dtype)) + out.fill_(float("nan")) + graph.replay() + torch.accelerator.synchronize() + _check_decode(out, q, keys, table, lens, query_len, 1, 2) + + +def _check_cp_push(lens, page_step, init_blocks, local_blocks, grid): + """The context-parallel scorer on one GPU, the other ranks' index-q rows + and flags pre-set locally: per rank, the single-head scorer's bits on its + blocks, -inf past the length inside its last run, nothing elsewhere; the + four ranks cover every block once.""" + num_reqs = len(lens) + table, num_blocks = _block_table(lens, share=True) + seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda") + cache = _index_cache(num_blocks, page_step) + query_len, heads = cp.QL, cp.HEADS + rows = num_reqs * query_len + q = (torch.randn(rows, heads, D, device="cuda") * 0.5).to(torch.float8_e4m3fn) + max_seq_len = max(lens) + width = max(64, ((max_seq_len + BLOCK - 1) // BLOCK + 63) // 64 * 64) + expected = torch.full((heads, rows, width), float("nan"), device="cuda") + for h in range(heads): + decode_host.minimax_m3_nvfp4_index_score_decode( + q[:, h : h + 1], cache, expected[h : h + 1], table, seq_lens, + init_blocks, local_blocks, query_len, max_seq_len, + ) # fmt: skip + blocks = torch.arange(width, device="cuda") + nblk = ((seq_lens + BLOCK - 1) // BLOCK).repeat_interleave(query_len)[:, None] + seen = torch.zeros(heads, rows, width, dtype=torch.int32, device="cuda") + for rank in range(cp.WORLD): + counts = torch.zeros(cp.CNT_WORDS, dtype=torch.int32, device="cuda") + flags = torch.zeros(cp.FLAG_BYTES // 4, dtype=torch.int32, device="cuda") + flags[cp.SC_FLAGS : cp.SC_FLAGS + cp.WORLD] = 1 << 30 + flags[cp.SC_FLAGS + rank] = 0 + iq_buf = torch.full((heads, rows, D), 0x7F, dtype=torch.uint8, device="cuda") + iq_buf = iq_buf.view(torch.float8_e4m3fn) + for p in range(heads): + if p != rank: + iq_buf[p] = q[:, p] + w = cp.SC_IQ + p * cp.SC_MAX_REQS + flags[w : w + cp.SC_MAX_REQS] = 1 << 30 + outs = [] + for _ in range(2): + actual = torch.full_like(expected, float("nan")) + dests = [actual[h].data_ptr() for h in range(heads)] + decode_host.minimax_m3_nvfp4_index_score_push_cp( + q[:, rank].contiguous(), cache, dests, width, rows * width * 4, + table, seq_lens, init_blocks, local_blocks, rank, counts, + [flags.data_ptr()] * 4, [iq_buf.data_ptr()] * 4, rows, grid=grid, + ) # fmt: skip + outs.append(actual) + torch.accelerator.synchronize() + actual = outs[0] + assert torch.equal(actual.view(torch.int32), outs[1].view(torch.int32)) + mine = (blocks // cp.CP_CH) % cp.WORLD == rank + in_runs = mine & (blocks // cp.CP_CH * cp.CP_CH < nblk) + inside = in_runs & (blocks < nblk) + assert actual[:, ~in_runs].isnan().all(), rank + assert (actual[:, in_runs & ~inside] == float("-inf")).all(), rank + assert torch.equal( + actual[:, inside].view(torch.int32), expected[:, inside].view(torch.int32) + ), rank + seen += inside.int() + assert torch.equal(seen, (~expected.isnan()).int()) + + +@pytest.mark.parametrize("num_reqs", [1, 5, 17, 45]) +@pytest.mark.parametrize("page_step", [1, 2]) +@pytest.mark.parametrize(("init_blocks", "local_blocks"), [(1, 2), (0, 1)]) +@pytest.mark.parametrize("grid", [3, 768]) +def test_index_decode_cp_push_matches_single_head( + num_reqs, page_step, init_blocks, local_blocks, grid +): + """Grid 3: many runs and requests per workgroup.""" + torch.manual_seed(num_reqs) + _check_cp_push(_lens(num_reqs, 9000), page_step, init_blocks, local_blocks, grid) + + +@pytest.mark.parametrize("grid", [3, 768]) +def test_index_decode_cp_push_all_padding(grid): + """Padding requests only (graph capture): nothing written.""" + _check_cp_push([0] * 5, 1, 1, 2, grid) + + +def _prefill_case(q_lens, seq_lens_list, num_heads, page_step=1): + torch.manual_seed(0) + table, num_blocks = _block_table(seq_lens_list) + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda") + cu = torch.zeros(len(q_lens) + 1, dtype=torch.int32, device="cuda") + cu[1:] = torch.tensor(q_lens, dtype=torch.int32, device="cuda").cumsum(0) + cache = _index_cache(num_blocks, page_step) + q = torch.randn(sum(q_lens), num_heads, D, device="cuda") * 0.5 + return table, seq_lens, cu, cache, q.to(torch.float8_e4m3fn) + + +def _causal_blocks(q_len, s, num_heads): + """[H, q_len, nb] blocks each row of a chunk ending its sequence sees.""" + nb = (s + BLOCK - 1) // BLOCK + row_pos = s - q_len + torch.arange(q_len, device="cuda") + seen = torch.arange(nb, device="cuda")[None, :] <= (row_pos // BLOCK)[:, None] + return row_pos, seen[None].expand(num_heads, -1, -1) + + +@pytest.mark.parametrize("num_heads", [1, 2, 4]) +@pytest.mark.parametrize( + ("q_lens", "seq_lens_list"), + [ + ([129, 300, 5], [129, 2500, 2049]), # 128-row tiles, ragged + ([1000, 700], [4300, 700]), # 256-row tiles, unaligned prefix + ], +) +@pytest.mark.parametrize("page_step", [1, 2]) +def test_index_prefill_matches_aiter_fp8(num_heads, q_lens, seq_lens_list, page_step): + """Every block a row can see against AITER's FP8 prefill scorer.""" + from aiter.ops.msa_attention import pa_sparse_block_score_prefill + + table, seq_lens, cu, cache, q = _prefill_case( + q_lens, seq_lens_list, num_heads, page_step + ) + msl = max(seq_lens_list) + args = (table, cu, seq_lens, 1, 2, max(q_lens), msl) + width = (msl + BLOCK - 1) // BLOCK + actual = torch.full((num_heads, sum(q_lens), width), float("nan"), device="cuda") + minimax_m3_nvfp4_index_score_prefill(q, cache, actual, *args) + expected = torch.full_like(actual, float("nan")) + pages = _fp8_keys(cache) + for h in range(num_heads): + pa_sparse_block_score_prefill( + q[:, h : h + 1].contiguous(), pages, expected[h : h + 1], *args + ) + torch.accelerator.synchronize() + q_start = 0 + for q_len, s in zip(q_lens, seq_lens_list): + _, seen = _causal_blocks(q_len, s, num_heads) + rows = slice(q_start, q_start + q_len) + got = actual[:, rows, : seen.shape[2]][seen] + want = expected[:, rows, : seen.shape[2]][seen] + finite = want < 1e28 + assert torch.equal(got[~finite], want[~finite]) + torch.testing.assert_close(got[finite], want[finite], rtol=1e-5, atol=1e-4) + q_start += q_len + + +@pytest.mark.parametrize( + ("q_lens", "seq_lens_list"), + [ + ([4096], [4096]), # fresh prompt: every tile diagonal + ([2000, 513, 3000], [9000, 513, 20000]), # ragged, unaligned prefixes + ([700, 1500], [30000, 1500]), # long cached prefix + fresh + ], +) +@pytest.mark.parametrize(("init_blocks", "local_blocks"), [(0, 1), (2, 2)]) +@pytest.mark.parametrize("num_heads", [1, 2, 4]) +def test_index_prefill_matches_reference( + q_lens, seq_lens_list, init_blocks, local_blocks, num_heads, monkeypatch +): + """Every kernel the host can pick (``kernel_shared.py`` with 128- and + 256-row tiles, ``kernel.py``) bit for bit alike, and against an fp32 + reference.""" + table, seq_lens, cu, cache, q = _prefill_case( + q_lens, seq_lens_list, num_heads, page_step=2 if num_heads == 2 else 1 + ) + width = (max(seq_lens_list) + BLOCK - 1) // BLOCK + 3 + args = ( + table, + cu, + seq_lens, + init_blocks, + local_blocks, + max(q_lens), + max(seq_lens_list), + ) + + def run(shared, mq): + monkeypatch.setattr(prefill_host, "PIN_SHARED", shared) + monkeypatch.setattr(prefill_host, "PIN_MQ", mq) + out = torch.full((num_heads, sum(q_lens), width), float("nan"), device="cuda") + minimax_m3_nvfp4_index_score_prefill(q, cache, out, *args) + return out + + variants = [(True, 1), (True, 2), (False, None)] + outs = [run(*v) for v in variants] + keys = _fp8_keys(cache).float() + torch.accelerator.synchronize() + q_start = 0 + for req, (q_len, s) in enumerate(zip(q_lens, seq_lens_list)): + row_pos, seen = _causal_blocks(q_len, s, num_heads) + nb = seen.shape[2] + rows = slice(q_start, q_start + q_len) + for v, out in zip(variants[1:], outs[1:]): + got, ref = out[:, rows, :nb][seen], outs[0][:, rows, :nb][seen] + assert torch.equal(got, ref), (v, q_len, s) + page_keys = keys[table[req, :nb].long()].reshape(-1, D) + for r0 in range(0, q_len, 512): # [H, rows, tokens] dots in slices + r1 = min(q_len, r0 + 512) + want = _ref_block_scores( + q[q_start + r0 : q_start + r1], page_keys, row_pos[r0:r1] + 1, + init_blocks, local_blocks, + ) # fmt: skip + got = outs[0][:, q_start + r0 : q_start + r1, :nb] + sl = seen[:, r0:r1] + _assert_scores(got[sl], want[sl], (q_len, s)) + q_start += q_len diff --git a/tests/kernels/attention/test_minimax_m3_nvfp4_indexer_decode_cp.py b/tests/kernels/attention/test_minimax_m3_nvfp4_indexer_decode_cp.py new file mode 100644 index 000000000000..1e658fa72407 --- /dev/null +++ b/tests/kernels/attention/test_minimax_m3_nvfp4_indexer_decode_cp.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 indexer decode CP at TP4: each rank's top-k must equal an exact fp32 +top-15 over the full-context scores, across repeated runs (sentinel reset).""" + +import os + +import pytest +import torch +import torch.multiprocessing as mp + +from vllm.platforms import current_platform + +if not current_platform.is_cuda(): + pytest.skip("MiniMax M3 NVFP4 indexer requires CUDA.", allow_module_level=True) + +PAGE = 128 +SLOTS = 8 +H = 4 + + +def _cache(num_pages: int, gen: torch.Generator) -> torch.Tensor: + kv = torch.randint( + 0, 256, (num_pages, PAGE, 72), dtype=torch.uint8, device="cuda", generator=gen + ) + # Positive normal E4M3 block scales keep every score finite. + s = torch.rand(num_pages, PAGE * 8, device="cuda", generator=gen) * 1.75 + 0.25 + kv.view(num_pages, -1)[:, 8192:] = s.to(torch.float8_e4m3fn).view(torch.uint8) + return kv + + +def _reference(scores, seq_lens, query_len, head) -> torch.Tensor: + num_reqs = seq_lens.shape[0] + out = torch.full((num_reqs * query_len, 16), -1, dtype=torch.int32) + for b in range(num_reqs): + for i in range(query_len): + t = b * query_len + i + own = (int(seq_lens[b]) - query_len + i) // PAGE + if own + 1 <= 16: + out[t, : own + 1] = torch.arange(own + 1) + continue + row = scores[b, (SLOTS - query_len + i) * H + head, :own] + out[t, :15] = torch.sort(-row, stable=True).indices[:15].to(torch.int32) + out[t, 15] = own + return out + + +def _run(rank: int, port: int, query_len: int, seq_lens_list: list[int]) -> None: + from vllm.distributed import ( + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.models.minimax_m3.nvidia.indexer_nvfp4_decode_cp import ( + Nvfp4IndexerDecodeCP, + ) + + torch.accelerator.set_device_index(rank) + init_distributed_environment( + world_size=4, + rank=rank, + local_rank=rank, + distributed_init_method=f"tcp://127.0.0.1:{port}", + ) + initialize_model_parallel(tensor_model_parallel_size=4) + + gen = torch.Generator(device="cuda").manual_seed(0) # same on every rank + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cuda") + num_reqs = seq_lens.shape[0] + pages = (int(seq_lens.max()) + PAGE - 1) // PAGE + kv = _cache(num_reqs * pages + 8, gen) + # Sized to the longest context, as vLLM's block table is: the longest + # request then owns the last page of its rank's compact table. + block_table = ( + torch.randperm(num_reqs * pages, device="cuda", generator=gen) + .to(torch.int32) + .view(num_reqs, pages) + ) + rows = num_reqs * query_len + + cp = Nvfp4IndexerDecodeCP(max_reqs=num_reqs, device=torch.device("cuda", rank)) + layers = [cp.new_layer() for _ in range(2)] + plan = cp.plan(block_table, seq_lens) + + from vllm.third_party.fmha_sm100.sparse import BatchDecodeIndexerQ8KV4Wrapper + + full = BatchDecodeIndexerQ8KV4Wrapper(num_heads=H) + full.plan(block_table, seq_lens) + for step in range(3): + for layer in layers: + q = (torch.randn(rows, H, 128, device="cuda", generator=gen) * 0.5).to( + torch.float8_e4m3fn + ) + out = torch.full((rows, 1, 16), -7, dtype=torch.int32, device="cuda") + cp.forward( + layer, plan, q[:, rank : rank + 1].contiguous(), kv, out, query_len + ) + q_tile = torch.zeros( + num_reqs, SLOTS, H, 128, dtype=torch.float8_e4m3fn, device="cuda" + ) + q_tile[:, SLOTS - query_len :] = q.view(num_reqs, query_len, H, 128) + # NaN in the columns the kernel leaves alone, so a selection + # that reads one cannot match the reference. + full._scores.fill_(float("nan")) + scores = full._run_scores(q_tile.view(-1, H, 128), kv) + expected = _reference(scores.cpu(), seq_lens.cpu(), query_len, rank) + torch.testing.assert_close(out[:, 0].cpu(), expected, rtol=0, atol=0) + + +def _worker(rank: int, port: int, query_len: int, seq_lens_list: list[int]) -> None: + from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config + + config = VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=4)) + with set_current_vllm_config(config): + _run(rank, port, query_len, seq_lens_list) + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 4 + or not ( + current_platform.is_device_capability(100) + or current_platform.is_device_capability(103) + ), + reason="Needs 4 SM100/SM103 GPUs.", +) +@pytest.mark.parametrize("query_len", [1, 4, 8]) +def test_nvfp4_indexer_decode_cp_matches_exact_topk(query_len: int): + # Trivial rows (<= 16 pages), page boundaries, and long contexts. + seq_lens = [9, 2049, 2048 + query_len, 20000, 131072, 300001] + port = 29500 + (os.getpid() % 1000) + mp.spawn(_worker, args=(port, query_len, seq_lens), nprocs=4, join=True) diff --git a/tests/kernels/attention/test_minimax_m3_nvfp4_prefill_flydsl.py b/tests/kernels/attention/test_minimax_m3_nvfp4_prefill_flydsl.py new file mode 100644 index 000000000000..8e3b368dba38 --- /dev/null +++ b/tests/kernels/attention/test_minimax_m3_nvfp4_prefill_flydsl.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3 NVFP4 sparse prefill (FlyDSL) against fp64 attention over the +exactly dequantized K / V, and bitwise stable across runs and call splits.""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm (gfx950) FlyDSL kernel", allow_module_level=True) + +from vllm.platforms.rocm import on_gfx950 # noqa: E402 + +if not on_gfx950(): + pytest.skip("gfx950 instructions", allow_module_level=True) + +from tests.kernels.attention.minimax_m3_nvfp4_ref import ( # noqa: E402 + dequant, + insert_nvfp4_kv, + nvfp4_kv_cache_token_major, +) +from tests.kernels.mfma_spinner import mfma_spinner # noqa: E402 +from vllm.models.minimax_m3.amd.ops.nvfp4_prefill.host import ( # noqa: E402 + minimax_m3_nvfp4_sparse_prefill, + supports_nvfp4_sparse_prefill, +) + +BLOCK = 128 +D = 128 +GROUP = 16 # query heads per KV head +TOPK = 16 +SM_SCALE = D**-0.5 + + +def _case(q_lens, ctx_lens, hkv, k_scale, v_scale, seed=0): + """The last q_lens[i] of ctx_lens[i] tokens are queries; shuffled pages; + each token picks its own block and 15 random earlier ones.""" + gen = torch.Generator().manual_seed(seed) + nbs = [(c + BLOCK - 1) // BLOCK for c in ctx_lens] + n_pages = sum(nbs) + 2 + cache = torch.zeros(n_pages, 2 * hkv, BLOCK, 72, dtype=torch.uint8, device="cuda") + slots = torch.arange(n_pages * BLOCK, device="cuda") + k = (torch.randn(slots.numel(), hkv, D, generator=gen) * 2).bfloat16().cuda() + v = (torch.randn(slots.numel(), hkv, D, generator=gen) * 2).bfloat16().cuda() + insert_nvfp4_kv(k, v, cache, slots, k_scale, v_scale) + perm = torch.randperm(n_pages, generator=gen).int() + table = torch.zeros(len(q_lens), max(nbs), dtype=torch.int32) + off = 0 + for r, nb in enumerate(nbs): + table[r, :nb] = perm[off : off + nb] + off += nb + pos = torch.cat([torch.arange(c - q, c) for q, c in zip(q_lens, ctx_lens)]).int() + req = torch.cat([torch.full((q,), r) for r, q in enumerate(q_lens)]).int() + rows = [] + for _ in range(hkv): + for p in pos.tolist(): + cur = p // BLOCK + sel = torch.cat( + [torch.tensor([cur]), torch.randperm(cur, generator=gen)[: TOPK - 1]] + ) + rows.append(torch.cat([sel, torch.full((TOPK - sel.numel(),), -1)])) + topk = torch.stack(rows).int().view(hkv, pos.numel(), TOPK) + q = (torch.randn(pos.numel(), hkv * GROUP, D, generator=gen) * 2).bfloat16() + return cache, table.cuda(), topk.cuda(), q.cuda(), pos.cuda(), req.cuda() + + +def _prefill(q, cache, topk, table, req, pos, hkv, k_scale=None, v_scale=None): + out = torch.empty_like(q) + minimax_m3_nvfp4_sparse_prefill( + q, cache, topk, table, req, pos, hkv, SM_SCALE, out, k_scale=k_scale, + v_scale=v_scale, + ) # fmt: skip + return out + + +def _fp64(q, cache, topk, table, req, pos, hkv, k_scale, v_scale, rows): + """fp64 attention of the token rows ``rows`` over the dequantized K / V.""" + k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_token_major(cache) + kv = [ + (dequant(data, sf) * s).to(torch.bfloat16).double() # [pages, hkv, 128, 128] + for data, sf, s in ((k_data, k_sf, k_scale), (v_data, v_sf, v_scale)) + ] + ref = torch.empty(len(rows), hkv * GROUP, D, dtype=torch.float64, device="cuda") + for i, r in enumerate(rows): + for h in range(hkv): + blocks = [b for b in topk[h, r].tolist() if b >= 0] + pages = [int(table[int(req[r]), b]) for b in blocks] + kk = torch.cat([kv[0][p, h] for p in pages]) + vv = torch.cat([kv[1][p, h] for p in pages]) + key_pos = torch.cat( + [torch.arange(b * BLOCK, (b + 1) * BLOCK) for b in blocks] + ).cuda() + s = (q[r, h * GROUP : (h + 1) * GROUP].double() @ kk.T) * SM_SCALE + s[:, key_pos > int(pos[r])] = float("-inf") + ref[i, h * GROUP : (h + 1) * GROUP] = torch.softmax(s, -1) @ vv + return ref + + +def _rel(out, ref, rows): + o = out[rows].double() + return float((o - ref).norm() / ref.norm()) + + +MAX_REL = 3e-3 # fp16 P and the bf16 output: ~2e-3 + + +@pytest.mark.parametrize("scales", [(1.0, 1.0), (0.5, 2.0)]) +def test_matches_fp64(scales): + k_scale, v_scale = (torch.tensor(s, device="cuda") for s in scales) + cache, table, topk, q, pos, req = _case([2048], [2048], 1, k_scale, v_scale) + out = _prefill(q, cache, topk, table, req, pos, 1, k_scale, v_scale) + assert torch.isfinite(out).all() + rows = list(range(0, 2048, 37)) + list(range(2040, 2048)) + ref = _fp64(q, cache, topk, table, req, pos, 1, k_scale, v_scale, rows) + assert _rel(out, ref, rows) < MAX_REL + + +@pytest.mark.parametrize("hkv", [1, 2, 4]) +def test_batch_with_prefix_and_tiles_across_requests(hkv): + one = torch.tensor(1.0, device="cuda") + q_lens, ctx_lens = [1027, 514, 3, 262], [3000, 514, 1100, 900] + cache, table, topk, q, pos, req = _case(q_lens, ctx_lens, hkv, one, one, seed=1) + out = _prefill(q, cache, topk, table, req, pos, hkv, one, one) + assert torch.isfinite(out).all() + rows = [0, 1, 2, 3, 1026, 1027, 1028, 1540, 1541, 1542, 1543, 1544, 1805] + ref = _fp64(q, cache, topk, table, req, pos, hkv, one, one, rows) + assert _rel(out, ref, rows) < MAX_REL + + +@pytest.mark.parametrize("n_tok", [1, 3, 5]) +def test_few_tokens(n_tok): + one = torch.tensor(1.0, device="cuda") + cache, table, topk, q, pos, req = _case([n_tok], [700], 1, one, one, seed=2) + out = _prefill(q, cache, topk, table, req, pos, 1) # unit scales by default + rows = list(range(n_tok)) + ref = _fp64(q, cache, topk, table, req, pos, 1, one, one, rows) + assert _rel(out, ref, rows) < MAX_REL + + +@pytest.mark.parametrize("hkv", [1, 4]) +def test_block_outer_pages_bitwise(hkv): + """A page inside a wider block (stride(0) past one page): the dense + cache's output.""" + one = torch.tensor(1.0, device="cuda") + q_lens, ctx_lens = [1027, 3, 262], [3000, 1100, 900] + cache, table, topk, q, pos, req = _case(q_lens, ctx_lens, hkv, one, one, seed=3) + page = cache[0].numel() + wide = torch.full( + (cache.shape[0], 3 * page), 0x7F, dtype=torch.uint8, device="cuda" + ) + wide = wide[:, page : 2 * page].view(cache.shape) + wide.copy_(cache) + out = _prefill(q, cache, topk, table, req, pos, hkv) + assert torch.isfinite(out).all() + assert torch.equal(out, _prefill(q, wide, topk, table, req, pos, hkv)) + + +def _shared_picks(pos, req, hkv, seed): + """Indexer-like top-k: the own block, then a per-request base of 15 + earlier blocks, each swapped for a random one with probability 0.1.""" + gen = torch.Generator().manual_seed(seed) + first_blk: dict[int, int] = {} + for p, r in zip(pos.tolist(), req.tolist()): + first_blk.setdefault(r, p // BLOCK) + base = { + r: torch.randperm(max(b, TOPK), generator=gen)[: TOPK - 1].tolist() + for r, b in first_blk.items() + } + rows = [] + for _ in range(hkv): + for p, r in zip(pos.tolist(), req.tolist()): + cur = p // BLOCK + sel = [cur] + for b in base[r]: + if torch.rand(1, generator=gen).item() < 0.1: + b = int(torch.randint(0, max(cur, 1), (1,), generator=gen)) + if b < cur and b not in sel: + sel.append(b) + rows.append(sel + [-1] * (TOPK - len(sel))) + return torch.tensor(rows, dtype=torch.int32).view(hkv, pos.numel(), TOPK).cuda() + + +# Grids above / below the CU count: the two- and one-round combine. +SHARED_CASES = [ + ( + [1501, 6, 2050], + [9000, 2100, 6200], + [0, 1, 2, 3, 1499, 1500, 1501, 1502, 1506, 3556], + ), # fmt: skip + ( + [201, 6, 150], + [5000, 2100, 3000], + [0, 1, 2, 3, 199, 200, 201, 202, 206, 207, 356], + ), # fmt: skip +] + + +@pytest.mark.parametrize("hkv", [1, 2]) +@pytest.mark.parametrize("case", SHARED_CASES) +def test_shared_picks_deterministic(hkv, case): + """Pages shared within a tile, tiles across requests: matches fp64, + bitwise repeatable.""" + one = torch.tensor(1.0, device="cuda") + q_lens, ctx_lens, rows = case + cache, table, _, q, pos, req = _case(q_lens, ctx_lens, hkv, one, one, seed=3) + topk = _shared_picks(pos, req, hkv, seed=4) + out = _prefill(q, cache, topk, table, req, pos, hkv) + ref = _fp64(q, cache, topk, table, req, pos, hkv, one, one, rows) + assert _rel(out, ref, rows) < MAX_REL + assert torch.equal(out, _prefill(q, cache, topk, table, req, pos, hkv)) + + +def test_fresh_prompt_calls_bitwise(): + """A fresh 4096-token prompt in one call equals the same rows in calls of + 1024 / 2048 / 1024 tokens, bitwise (DBO microbatches).""" + one = torch.tensor(1.0, device="cuda") + n = 4096 + cache, table, _, q, pos, req = _case([n], [n], 1, one, one, seed=5) + topk = _shared_picks(pos, req, 1, seed=6) + assert (topk[0, : 15 * BLOCK] == -1).any(1).all() + + def run(a, b): + return _prefill(q[a:b], cache, topk[:, a:b], table, req[a:b], pos[a:b], 1) + + full = run(0, n) + assert torch.equal(full, torch.cat([run(0, 1024), run(1024, 3072), run(3072, n)])) + + +@pytest.mark.parametrize("n_tok, ctx", [(1024, 31024), (4096, 64096)]) +def test_bitwise_next_to_mfma_waves(n_tok, ctx): + """Another kernel's MFMA waves on every CU leave the output unchanged + (v_pk_mul_f32 gives wrong products there on MI355X).""" + one = torch.tensor(1.0, device="cuda") + cache, table, topk, q, pos, req = _case([n_tok], [ctx], 1, one, one, seed=7) + + def run(stream): + with torch.cuda.stream(stream): + return _prefill(q, cache, topk, table, req, pos, 1) + + # A high-priority queue, so the spinner never waits behind the attention. + main, side = torch.cuda.Stream(), torch.cuda.Stream(priority=-1) + ref = run(main) + torch.accelerator.synchronize() + spin = mfma_spinner(torch.cuda.get_device_properties(0).multi_processor_count) + for _ in range(8): + spin(200_000, side) # 2 ms + out = run(main) + torch.accelerator.synchronize() + assert torch.equal(out, ref) + + +def test_supported_shapes(): + """16 query heads per KV head, 16 selected 128-token blocks only.""" + assert supports_nvfp4_sparse_prefill(16, 1, D, TOPK, BLOCK) + assert supports_nvfp4_sparse_prefill(32, 2, D, TOPK, BLOCK) + assert not supports_nvfp4_sparse_prefill(8, 1, D, TOPK, BLOCK) + assert not supports_nvfp4_sparse_prefill(16, 1, D, 8, BLOCK) + assert not supports_nvfp4_sparse_prefill(16, 1, 64, TOPK, BLOCK) diff --git a/tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py b/tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py new file mode 100644 index 000000000000..8b368f36f3f8 --- /dev/null +++ b/tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3 NVFP4 sparse decode (FlyDSL) against fp64 attention over +K / V = e4m3_rn_satfinite(e2m1 x block scale) x global scale.""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip("ROCm (gfx950) FlyDSL kernel", allow_module_level=True) + +from vllm.platforms.rocm import on_gfx950 # noqa: E402 + +if not on_gfx950(): + pytest.skip("gfx950 instructions", allow_module_level=True) + +from tests.kernels.attention.minimax_m3_nvfp4_ref import ( # noqa: E402 + dequant, + head_offsets, + insert_nvfp4_kv, + nvfp4_kv_cache_token_major, +) +from vllm.models.minimax_m3.amd.ops.nvfp4_sparse_decode.host import ( # noqa: E402 + minimax_m3_nvfp4_sparse_decode_fly, +) + +BLOCK = 128 +D = 128 +GROUP = 16 # query heads per KV head +TOPK = 16 +SM_SCALE = D**-0.5 + + +def _nvfp4_cache(num_blocks, scales, num_kv_heads=1): + """Every slot of the cache filled with random K / V.""" + k_scale, v_scale = scales + cache = torch.zeros( + num_blocks, 2 * num_kv_heads, BLOCK, 72, dtype=torch.uint8, device="cuda" + ) + slots = torch.arange(num_blocks * BLOCK, device="cuda") + shape = (slots.numel(), num_kv_heads, D) + k = (torch.randn(shape, device="cuda") * 2).bfloat16() + v = (torch.randn(shape, device="cuda") * 2).bfloat16() + insert_nvfp4_kv(k, v, cache, slots, k_scale, v_scale) + return cache + + +def _block_table(seq_lens): + widths = [(s + BLOCK - 1) // BLOCK for s in seq_lens] + num_blocks = sum(widths) + 3 + perm = torch.randperm(num_blocks, device="cuda", dtype=torch.int32) + table = torch.zeros(len(seq_lens), max(widths), dtype=torch.int32, device="cuda") + start = 0 + for r, w in enumerate(widths): + table[r, :w] = perm[start : start + w] + start += w + return table, num_blocks + + +def _scales(): + return ( + torch.tensor(0.5, dtype=torch.float32, device="cuda"), + torch.tensor(2.0, dtype=torch.float32, device="cuda"), + ) + + +def _selections(seq_lens, qlen, num_kv_heads=1, seed=0): + """[Hkv, rows, 16] picks, mostly shared within a request, with -1 padding + and an out-of-range entry; head h's picks rotated by h.""" + g = torch.Generator(device="cuda").manual_seed(seed) + topk = torch.full( + (1, len(seq_lens) * qlen, TOPK), -1, dtype=torch.int32, device="cuda" + ) + for r, s in enumerate(seq_lens): + nb = (s + BLOCK - 1) // BLOCK + base = torch.randperm(nb, device="cuda", generator=g).to(torch.int32) + for q in range(qlen): + cur = max(s - qlen + q, 0) // BLOCK + older = base[(base != cur) & (base <= cur)] + pick = torch.cat([older[:10], older[10 + 3 * q : 13 + 3 * q]]) + sel = torch.cat([torch.tensor([cur], device="cuda"), pick]) + sel = torch.cat([sel, torch.tensor([nb + 5], device="cuda")])[:TOPK] + topk[0, r * qlen + q, : sel.numel()] = sel.to(torch.int32) + return torch.cat([topk.roll(h, -1) for h in range(num_kv_heads)]) + + +def _reference(q, cache, topk, table, seq_lens, scales, qlen): + """fp64 attention over the dequantized K / V; rows that see nothing give 0.""" + k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_token_major(cache) + k_all, v_all = ( + dequant(data, sf).clamp(-448, 448).to(torch.float8_e4m3fn).double() * s.item() + for data, sf, s in ((k_data, k_sf, scales[0]), (v_data, v_sf, scales[1])) + ) + hkv, rows, _ = topk.shape + req = torch.arange(rows, device="cuda") // qlen + pos = seq_lens.long()[req] - qlen + torch.arange(rows, device="cuda") % qlen + offs = torch.arange(BLOCK, device="cuda") + out = torch.zeros(q.shape, dtype=torch.float64, device="cuda") + for h in range(hkv): + blk = topk[h, :rows].long() + ok = (pos[:, None] >= 0) & (blk >= 0) & (blk <= (pos // BLOCK)[:, None]) + blk = torch.where(ok, blk, 0) + tok = blk[..., None] * BLOCK + offs # [rows, TOPK, BLOCK] + ok = (ok[..., None] & (tok <= pos[:, None, None])).flatten(1) + page = table[req[:, None], blk].long()[..., None].expand_as(tok) + k = k_all[page, h, tok % BLOCK].flatten(1, 2) + v = v_all[page, h, tok % BLOCK].flatten(1, 2).masked_fill(~ok[..., None], 0) + heads = slice(h * GROUP, (h + 1) * GROUP) + s = torch.einsum("nhd,ntd->nht", q[:, heads].double(), k) * SM_SCALE + s = s.masked_fill(~ok[:, None, :], float("-inf")) + p = torch.softmax(s, -1).nan_to_num(0.0) + out[:, heads] = torch.einsum("nht,ntd->nhd", p, v) + return out + + +def _rel(a, b): + a, b = a.double(), b.double() + return float((a - b).norm() / b.norm()) + + +def _decode(q, cache, topk, table, seq_lens, qlen, out=None, **kw): + k_scale, v_scale = _scales() + out = torch.full_like(q, float("nan")) if out is None else out + minimax_m3_nvfp4_sparse_decode_fly( + q, cache, topk, table, seq_lens, SM_SCALE, out, k_scale=k_scale, + v_scale=v_scale, decode_query_len=qlen, **kw, + ) # fmt: skip + return out + + +@pytest.mark.parametrize("decode_query_len", [1, 4, 8]) +@pytest.mark.parametrize("splits", [None, 1, 2, 4, 8]) +@pytest.mark.parametrize("random_bytes", [False, True]) +@pytest.mark.parametrize("xcd_count", [0, 8]) +def test_decode_matches_reference(decode_query_len, splits, random_bytes, xcd_count): + """Realistic and random-byte caches: close to fp64, bitwise repeatable.""" + torch.manual_seed(0) + lens = [9, 300, 2100, 20000, 131, 4096] + table, num_blocks = _block_table(lens) + seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda") + topk = _selections(lens, decode_query_len) + rows = topk.shape[1] + if random_bytes: + cache = torch.randint( + 0, 256, (num_blocks, 2, BLOCK, 72), dtype=torch.uint8, device="cuda" + ) + for side in range(2): + sf = cache.view(num_blocks, -1)[:, side * 9216 + 8192 : (side + 1) * 9216] + sf.masked_fill_((sf & 0x7F) == 0x7F, 0x38) # no NaN scales + q = (torch.randn(rows, GROUP, D, device="cuda") * 0.01).bfloat16() + else: + cache = _nvfp4_cache(num_blocks, _scales()) + q = torch.randn(rows, GROUP, D, device="cuda").bfloat16() + kw = dict(splits=splits, xcd_count=xcd_count) + out = _decode(q, cache, topk, table, seq_lens, decode_query_len, **kw) + again = _decode(q, cache, topk, table, seq_lens, decode_query_len, **kw) + ref = _reference(q, cache, topk, table, seq_lens, _scales(), decode_query_len) + assert torch.isfinite(out).all() + assert torch.equal(out, again) + assert _rel(out, ref) < 5e-3 + + +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +@pytest.mark.parametrize("decode_query_len", [1, 4, 6]) +@pytest.mark.parametrize("splits", [1, 4]) +def test_decode_padding_rows_and_kv_heads(num_kv_heads, decode_query_len, splits): + """Padding requests (seq_len 0) give zero rows; nothing past the rows is + written.""" + torch.manual_seed(1) + lens = [700, 0, 3000, 0, 0, 129] + table, num_blocks = _block_table([max(s, 1) for s in lens]) + seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda") + topk = _selections([max(s, 1) for s in lens], decode_query_len, num_kv_heads) + cache = _nvfp4_cache(num_blocks, _scales(), num_kv_heads) + rows = topk.shape[1] + q = torch.randn(rows, GROUP * num_kv_heads, D, device="cuda").bfloat16() + out = torch.full((rows + 3, *q.shape[1:]), float("nan"), device="cuda").bfloat16() + _decode( + q, cache, topk, table, seq_lens, decode_query_len, out[:rows], splits=splits + ) + ref = _reference(q, cache, topk, table, seq_lens, _scales(), decode_query_len) + real = (seq_lens.repeat_interleave(decode_query_len) > 0).cpu() + assert _rel(out[:rows][real], ref[real]) < 5e-3 + assert (out[:rows][~real] == 0).all() + assert out[rows:].isnan().all() + zeros = torch.zeros_like(seq_lens) + pad = _decode(q, cache, topk, table, zeros, decode_query_len, splits=splits) + assert (pad == 0).all() + + +@pytest.mark.parametrize( + "num_kv_heads,qlen,splits,xcd_count", + [(1, 1, 1, 0), (1, 4, 4, 8), (4, 6, 1, 8), (4, 8, 4, 0)], +) +def test_garbage_past_the_length_is_ignored(num_kv_heads, qlen, splits, xcd_count): + """NaN scale bytes past each request's length leave the output unchanged.""" + torch.manual_seed(42) + lens = [131, 4097, 2100, 777, 65] + table, blocks = _block_table(lens) + seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda") + topk = _selections(lens, qlen, num_kv_heads) + cache = _nvfp4_cache(blocks, _scales(), num_kv_heads) + q = torch.randn(len(lens) * qlen, num_kv_heads * GROUP, D, device="cuda").bfloat16() + kw = dict(splits=splits, xcd_count=xcd_count) + out = _decode(q, cache, topk, table, seq_lens, qlen, **kw) + dirty = cache.clone().view(blocks, -1) + for side in range(2): + start = side * num_kv_heads * 9216 + num_kv_heads * 8192 + region = dirty[:, start : start + num_kv_heads * 1024] + region = region.view(blocks, num_kv_heads, 1024) + _, _, offsets = head_offsets(side, cache.device) + for req, length in enumerate(lens): + page = table[req, length // BLOCK] + region[page, :, offsets[length % BLOCK :].reshape(-1)] = 0x7F + dirty_out = _decode(q, dirty.view_as(cache), topk, table, seq_lens, qlen, **kw) + assert torch.isfinite(out).all() + assert torch.equal(out, dirty_out) + + +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +@pytest.mark.parametrize("decode_query_len", [1, 4]) +def test_decode_block_outer_pages_bitwise(num_kv_heads, decode_query_len): + """A page inside a wider block (stride(0) past one page): the dense + cache's output.""" + torch.manual_seed(2) + lens = [300, 2100, 20000, 131] + table, num_blocks = _block_table(lens) + seq_lens = torch.tensor(lens, dtype=torch.int32, device="cuda") + topk = _selections(lens, decode_query_len, num_kv_heads) + cache = _nvfp4_cache(num_blocks, _scales(), num_kv_heads) + page = cache[0].numel() + wide = torch.full((num_blocks, 3 * page), 0x7F, dtype=torch.uint8, device="cuda") + wide_cache = wide[:, page : 2 * page].view(cache.shape) + wide_cache.copy_(cache) + q = torch.randn(topk.shape[1], GROUP * num_kv_heads, D, device="cuda").bfloat16() + outs = [ + _decode(q, c, topk, table, seq_lens, decode_query_len) + for c in (cache, wide_cache) + ] + assert torch.isfinite(outs[0]).all() + assert torch.equal(outs[0], outs[1]) diff --git a/tests/kernels/mfma_spinner.py b/tests/kernels/mfma_spinner.py new file mode 100644 index 000000000000..09b7e07b20ac --- /dev/null +++ b/tests/kernels/mfma_spinner.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""An MFMA co-runner for gfx950 kernel tests: on MI355X, v_pk_mul_f32 can +return wrong products while another kernel's MFMA waves share the CU.""" + +import torch + + +def mfma_spinner(n_wg): + """``n_wg`` workgroups looping v_mfma_f32_16x16x32_bf16 for ``ticks`` + (s_memrealtime, 100 per us). + + Launch it on a high-priority stream: on a normal one it can share a hardware + queue with the kernel under test and never overlap it. + """ + import flydsl.compiler as flyc + import flydsl.expr as fx + from flydsl._mlir.dialects import llvm + from flydsl.expr import range_constexpr + from flydsl.expr.typing import T + from flydsl.expr.typing import Vector as Vec + + @fx.struct + class Shared: + a: fx.Array[fx.Int32, 4096, 16] + + def now(): + return fx.Int64( + llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], []) + ) + + @flyc.kernel(name="test_mfma_spinner", known_block_size=[256, 1, 1]) + def spinner(arg_out: fx.Int64, i64_ticks: fx.Int64): + sm = fx.SharedAllocator().allocate(Shared).peek() + tx = fx.thread_idx.x + sm.a.ptr[tx] = tx + a4 = Vec.from_elements([tx, tx + 1, tx + 2, tx + 3], fx.Int32) + acc = Vec.filled(4, 0.0, fx.Float32) + t0 = now() + t = now() + while t - t0 < i64_ticks: + for _ in range_constexpr(16): + ab = Vec(a4).bitcast(fx.BFloat16) + acc = Vec( + fx.rocdl.mfma_f32_16x16x32_bf16( + Vec.make_type(4, fx.Float32), [ab, ab, acc, 0, 0, 0] + ) + ) + t = now() + f32p = fx.PointerType.get( + T.f32, address_space=fx.AddressSpace.Global, alignment=4 + ) + out = fx.inttoptr(f32p, arg_out) + out[fx.block_idx.x * 256 + tx] = acc[0] + fx.Float32(sm.a.ptr[255 - tx]) + + @flyc.jit + def launch(arg_out: fx.Int64, i64_ticks: fx.Int64, stream: fx.Stream): + spinner(arg_out, i64_ticks).launch( + grid=(n_wg, 1, 1), block=(256, 1, 1), stream=stream + ) + + buf = torch.zeros(n_wg * 256, device="cuda") + compiled = flyc.compile(launch, buf.data_ptr(), 1, torch.cuda.current_stream()) + return lambda ticks, stream: compiled(buf.data_ptr(), ticks, stream) diff --git a/tests/kernels/moe/test_deep_gemm_mega_moe.py b/tests/kernels/moe/test_deep_gemm_mega_moe.py index 0f707e16ea9f..dee9478615e2 100644 --- a/tests/kernels/moe/test_deep_gemm_mega_moe.py +++ b/tests/kernels/moe/test_deep_gemm_mega_moe.py @@ -940,7 +940,7 @@ def test_prepare_megamoe_inputs_masks_padding(nonfinite_padding, num_tokens, qua @requires_sm100 @pytest.mark.parametrize("shared", [None, "mxfp8", "bf16"]) -@pytest.mark.parametrize("num_tokens", [0, 7, 65, 135]) +@pytest.mark.parametrize("num_tokens", [0, 7, 65, 135, 600]) def test_prepare_megamoe_nvfp4_inputs_is_bitwise_exact(num_tokens, shared): """Per-token NVFP4 dispatch, shared-expert inputs, unrouted padding and the L2 alphas match DeepGEMM's reference quantizers bit for bit.""" diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py index 8c9f4a6d0055..03cbc96428a1 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -36,6 +36,13 @@ ("wo_b", 5120, 8192 // 4), ] +M31_TP4_SHAPES = [ + ("fused_qkv", 2560, 6144), + ("o_proj", 6144, 2048), + ("shared_gate_up", 1536, 6144), + ("shared_down", 6144, 768), +] + def _make_layer(n: int, k: int, device: str) -> torch.nn.Module: weight_bf16 = torch.randn(n, k, device=device, dtype=torch.bfloat16) * 0.05 @@ -46,8 +53,46 @@ def _make_layer(n: int, k: int, device: str) -> torch.nn.Module: return layer -@pytest.mark.parametrize("name,n,k", V41_ATTN_SHAPES, ids=lambda v: str(v)) -@pytest.mark.parametrize("num_tokens", [1, 2, 7, 16, 31, 32, 128, 1024]) +@pytest.mark.parametrize( + "n,k,selected,fallback", + [ + (2560, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 2)), + (1536, 6144, (16, 16, 1024, 2, 3), (16, 16, 1024, 2, 2)), + (6144, 768, (16, 16, 256, 2, 3), (16, 32, 256, 4, 3)), + ], +) +def test_rocm_mxfp8_linear_selects_m31_tp4_tiles(n, k, selected, fallback): + """The gfx950 small-M tuning is shape-scoped and ends after M=16.""" + from vllm.model_executor.kernels.linear.mxfp8 import rocm_native + from vllm.platforms.rocm import on_gfx950 + + if not on_gfx950(): + pytest.skip("MiniMax-M3.1 small-M tuning requires gfx950") + for num_tokens in (1, 8, 16): + assert rocm_native._select_cfg(num_tokens, n, k) == selected + assert rocm_native._select_cfg(17, n, k) == fallback + + +@pytest.mark.parametrize( + "n,k,fallback", + [ + (2560, 6144, (16, 16, 1024, 2, 2)), + (1536, 6144, (16, 16, 1024, 2, 2)), + (6144, 768, (16, 32, 256, 4, 3)), + ], +) +def test_rocm_mxfp8_linear_preserves_non_gfx950_tiles(monkeypatch, n, k, fallback): + from vllm.model_executor.kernels.linear.mxfp8 import rocm_native + + monkeypatch.setattr(rocm_native, "_ON_GFX950", False) + for num_tokens in (1, 8, 16, 17): + assert rocm_native._select_cfg(num_tokens, n, k) == fallback + + +@pytest.mark.parametrize( + "name,n,k", V41_ATTN_SHAPES + M31_TP4_SHAPES, ids=lambda v: str(v) +) +@pytest.mark.parametrize("num_tokens", [1, 2, 7, 8, 16, 31, 32, 128, 1024]) @torch.inference_mode() def test_rocm_mxfp8_linear_matches_dequant_reference(name, n, k, num_tokens): """Output is finite and matches an exact MXFP8 GEMM of the same inputs. diff --git a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py index d6a96c6c9960..711e4d2bda9e 100644 --- a/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py @@ -156,6 +156,74 @@ def test_dense_norm_rope(num_tokens, num_heads, num_kv_heads): torch.testing.assert_close(v_out, v_in, rtol=0, atol=0) +@pytest.mark.parametrize("num_tokens", [1, 7, 513]) +@pytest.mark.parametrize("block_size", [16, 128]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_dense_insert_matches_reshape_and_cache_flash( + num_tokens, block_size, kv_cache_dtype +): + """A dense layer's insert stores exactly what reshape_and_cache_flash + writes from the normed/roped k (unit scales), and leaves q/k/v as the + insert-free call does.""" + torch.manual_seed(0) + device, dtype, eps = "cuda", torch.bfloat16, 1e-6 + num_heads, num_kv_heads, max_pos = 8, 2, 4096 + q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 + k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, 5_000_000.0, dtype, device) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device + ) + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM + qkv = torch.randn(num_tokens, qsz + 2 * kvsz, dtype=dtype, device=device) + qkv_ref = qkv.clone() + num_blocks = (num_tokens + block_size - 1) // block_size + 2 + slot_mapping = torch.randperm(num_blocks * block_size, device=device)[ + :num_tokens + ].to(torch.int64) + slot_mapping[::4] = -1 # padding: no write + cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype + kv_cache = torch.zeros( + num_blocks, + num_kv_heads, + block_size, + 2 * HEAD_DIM, + dtype=cache_dtype, + device=device, + ) + ref_cache = torch.zeros_like(kv_cache) + + args = (qkv, q_w, k_w, cos_sin, positions, num_heads, num_kv_heads) + ops.fused_minimax_m3_qknorm_rope_kv_insert( + *args, + ROTARY_DIM, + eps, + slot_mapping=slot_mapping, + kv_cache=kv_cache, + block_size=block_size, + kv_cache_dtype=kv_cache_dtype, + ) + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv_ref, *args[1:], ROTARY_DIM, eps, kv_cache_dtype="auto" + ) + assert torch.equal(qkv, qkv_ref) + + _, k, v = qkv_ref.split([qsz, kvsz, kvsz], dim=-1) + key_cache, value_cache = ref_cache.transpose(1, 2).split(HEAD_DIM, dim=-1) + one = torch.ones((), dtype=torch.float32, device=device) + ops.reshape_and_cache_flash( + k.reshape(num_tokens, num_kv_heads, HEAD_DIM), + v.reshape(num_tokens, num_kv_heads, HEAD_DIM), + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + one, + one, + ) + assert torch.equal(kv_cache, ref_cache) + + # ── Test 2: sparse mode (full: index branch + cache inserts) ───────────────── @@ -756,3 +824,249 @@ def test_sparse_full_nvfp4(num_tokens, num_kv_heads, block_size, kv_scales): torch.testing.assert_close( index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0 ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="The NVFP4 reshape_and_cache_flash reference needs SM100", +) +@pytest.mark.parametrize("num_tokens", [1, 7, 513]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "nvfp4"]) +@pytest.mark.parametrize("page_pad", [0, 48]) +def test_sparse_full_nvfp4_index(num_tokens, kv_cache_dtype, page_pad): + """The NVFP4 index-K insert writes MiniMax-M3.1 KV4 bytes as a single-head + NVFP4 K side with no global scale (E2M1 page, then linear E4M3 scales: the + MSA Q8KV4 page), honours a padded page stride and skips padded slots, while + index-Q and everything else match the fp8 index path exactly.""" + torch.manual_seed(6) + device, dtype, eps = "cuda", torch.bfloat16, 1e-6 + base, max_pos = 5_000_000.0, 4096 + num_heads, num_kv_heads, num_idx_heads, block_size = 16, 4, 4, 128 + page_width = HEAD_DIM // 2 + HEAD_DIM // 16 + + weights = [ + torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 for _ in range(4) + ] + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device + ) + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM + iqsz = num_idx_heads * HEAD_DIM + qkv0 = torch.randn( + num_tokens, qsz + 2 * kvsz + iqsz + HEAD_DIM, dtype=dtype, device=device + ) + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.randperm( + num_blocks * block_size, dtype=torch.int64, device=device + )[:num_tokens] + index_slot_mapping = torch.randperm( + num_blocks * block_size, dtype=torch.int64, device=device + )[:num_tokens] + index_slot_mapping[::5] = -1 # padded / unscheduled tokens are not written + nvfp4_kwargs = {} + if kv_cache_dtype == "nvfp4": + nvfp4_kwargs = dict( + kv_k_scale=torch.tensor(0.5, dtype=torch.float32, device=device), + kv_v_scale=torch.tensor(2.0, dtype=torch.float32, device=device), + ) + + def run(index_cache): + qkv = qkv0.clone() + if kv_cache_dtype == "nvfp4": + kv_cache = torch.zeros( + num_blocks, + 2 * num_kv_heads, + block_size, + page_width, + dtype=torch.uint8, + device=device, + ) + else: + kv_cache = torch.zeros( + num_blocks, + num_kv_heads, + block_size, + 2 * HEAD_DIM, + dtype=dtype, + device=device, + ) + index_q = torch.empty( + num_tokens, iqsz, dtype=torch.float8_e4m3fn, device=device + ) + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + weights[0], + weights[1], + cos_sin, + positions, + num_heads, + num_kv_heads, + ROTARY_DIM, + eps, + weights[2], + weights[3], + num_idx_heads, + slot_mapping, + index_slot_mapping, + kv_cache, + index_cache, + block_size, + torch.empty(num_tokens, qsz, dtype=dtype, device=device), + index_q, + kv_cache_dtype, + **nvfp4_kwargs, + ) + return qkv, kv_cache, index_q + + fp8_index_cache = torch.zeros( + num_blocks, block_size, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device + ) + qkv_fp8, kv_cache_fp8, index_q_fp8 = run(fp8_index_cache) + raw = torch.zeros( + num_blocks, + block_size * page_width + page_pad, + dtype=torch.uint8, + device=device, + ) + index_cache = raw[:, : block_size * page_width].view( + num_blocks, block_size, page_width + ) + qkv, kv_cache, index_q = run(index_cache) + + assert torch.equal(qkv, qkv_fp8) + assert torch.equal(kv_cache, kv_cache_fp8) + assert torch.equal(index_q.view(torch.uint8), index_q_fp8.view(torch.uint8)) + + # A one-head NVFP4 page's K side is [bs, 64] E2M1 then [bs, 8] linear + # scales, which is exactly the index page. + index_k = qkv.split([qsz, kvsz, kvsz, iqsz, HEAD_DIM], dim=-1)[-1] + expected = nvfp4_pages( + [(index_k.reshape(num_tokens, 1, HEAD_DIM), 1.0)], + index_slot_mapping, + num_blocks, + block_size, + ) + assert torch.equal(index_cache, expected[:, 0]), ( + f"{int((index_cache != expected[:, 0]).sum())} NVFP4 index bytes differ" + ) + assert not raw[:, block_size * page_width :].any() + + +# ── Test 7: write_qkv=False, for callers that read only the FP8 q ──────────── + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="NVFP4 KV caches need SM100", +) +@pytest.mark.parametrize("num_tokens", [1, 7, 513]) +@pytest.mark.parametrize("num_heads,num_kv_heads", [(64, 4), (32, 2), (16, 1)]) +@pytest.mark.parametrize("index_cache_dtype", ["nvfp4", "fp8"]) +@pytest.mark.parametrize("scales", [(1.0, 1.0, 1.0), (0.5, 2.0, 0.5)]) +def test_write_qkv_false( + num_tokens, num_heads, num_kv_heads, index_cache_dtype, scales +): + """write_qkv=False without q_out writes the same q_fp8_out, index_q_out and + cache bytes as a default call, and leaves qkv either untouched or as the + default call rewrites it (k and index_k need not be rewritten).""" + torch.manual_seed(7) + device, dtype, eps = "cuda", torch.bfloat16, 1e-6 + base, max_pos = 5_000_000.0, 4096 + num_idx_heads, block_size = num_kv_heads, 128 + page_width = HEAD_DIM // 2 + HEAD_DIM // 16 + k_scale, v_scale, q_scale = scales + + weights = [ + torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1 for _ in range(4) + ] + cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device + ) + qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM + iqsz = num_idx_heads * HEAD_DIM + qkv0 = torch.randn( + num_tokens, qsz + 2 * kvsz + iqsz + HEAD_DIM, dtype=dtype, device=device + ) + qkv0[:, qsz + kvsz : qsz + 2 * kvsz] = nvfp4_edge_values( + num_tokens * num_kv_heads, v_scale + ).view(num_tokens, kvsz) + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.randperm( + num_blocks * block_size, dtype=torch.int64, device=device + )[:num_tokens] + slot_mapping[::5] = -1 + index_slot_mapping = torch.randperm( + num_blocks * block_size, dtype=torch.int64, device=device + )[:num_tokens] + kv_cache0 = torch.randint( + 0, + 256, + (num_blocks, 2 * num_kv_heads, block_size, page_width), + dtype=torch.uint8, + device=device, + ) + if index_cache_dtype == "nvfp4": + index_cache0 = torch.randint( + 0, + 256, + (num_blocks, block_size, page_width), + dtype=torch.uint8, + device=device, + ) + else: + index_cache0 = torch.randn(num_blocks, block_size, HEAD_DIM, device=device).to( + torch.float8_e4m3fn + ) + + def run(write_qkv): + qkv = qkv0.clone() + kv_cache, index_cache = kv_cache0.clone(), index_cache0.clone() + q_fp8 = torch.empty(num_tokens, qsz, dtype=torch.float8_e4m3fn, device=device) + index_q = torch.empty( + num_tokens, iqsz, dtype=torch.float8_e4m3fn, device=device + ) + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + weights[0], + weights[1], + cos_sin, + positions, + num_heads, + num_kv_heads, + ROTARY_DIM, + eps, + weights[2], + weights[3], + num_idx_heads, + slot_mapping, + index_slot_mapping, + kv_cache, + index_cache, + block_size, + torch.empty(num_tokens, qsz, dtype=dtype, device=device) + if write_qkv + else None, + index_q, + "nvfp4", + q_fp8_out=q_fp8, + q_fp8_scale=q_scale, + kv_k_scale=torch.tensor(k_scale, dtype=torch.float32, device=device), + kv_v_scale=torch.tensor(v_scale, dtype=torch.float32, device=device), + write_qkv=write_qkv, + ) + return qkv, q_fp8, index_q, kv_cache, index_cache + + expected = run(True) + qkv, *got = run(False) + assert torch.equal(qkv.view(torch.int16), qkv0.view(torch.int16)) or ( + torch.equal(qkv.view(torch.int16), expected[0].view(torch.int16)) + ) + for name, g, e in zip( + ("q_fp8", "index_q", "kv_cache", "index_cache"), got, expected[1:] + ): + assert torch.equal(g.view(torch.uint8), e.view(torch.uint8)), ( + f"{name}: {int((g.view(torch.uint8) != e.view(torch.uint8)).sum())} " + "bytes differ" + ) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 4e277470a2d8..d5d8ecf00e90 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -31,6 +31,7 @@ import pytest import torch +from transformers import PreTrainedConfig from vllm.platforms import current_platform @@ -46,6 +47,9 @@ mxfp8_e4m3_quantize, ) from vllm.models.minimax_m3.amd.model import ( # noqa: E402 + MiniMAXGemmaRMSNorm, + MiniMaxM3SparseAttention, + _build_rotary_emb, _can_fuse_input_norm_mxfp8, _kv_insert_operand, ) @@ -57,6 +61,7 @@ swiglu_oai_split, ) from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps # noqa: E402 +from vllm.platforms.rocm import on_gfx942 # noqa: E402 DEVICE = "cuda" EPS = 1e-6 @@ -843,3 +848,149 @@ def test_kv_insert_operand_copies_unsupported_layouts(): flat = k.reshape(32, -1) assert _kv_insert_operand(flat).is_contiguous() + + +@pytest.mark.skipif(not on_gfx942(), reason="Requires gfx942.") +@torch.inference_mode() +def test_gfx942_native_qknorm_rope_writes_main_and_index_caches(): + """The gfx942 fallback writes both caches and ignores padded slots.""" + from types import SimpleNamespace + + torch.manual_seed(0) + num_tokens = 3 + num_heads = 2 + num_kv_heads = num_idx_heads = 1 + head_dim = idx_head_dim = 16 + block_size = 4 + + attention = object.__new__(MiniMaxM3SparseAttention) + torch.nn.Module.__init__(attention) + attention.num_heads = num_heads + attention.num_kv_heads = num_kv_heads + attention.num_idx_heads = num_idx_heads + attention.head_dim = head_dim + attention.idx_head_dim = idx_head_dim + attention.q_size = num_heads * head_dim + attention.kv_size = num_kv_heads * head_dim + attention.index_q_size = num_idx_heads * idx_head_dim + attention.layer_name = "model.layers.0.self_attn" + attention.skip_index_topk = False + attention.kv_cache_dtype = "fp8" + + attention.q_norm = MiniMAXGemmaRMSNorm(head_dim).to( + device=DEVICE, dtype=torch.bfloat16 + ) + attention.k_norm = MiniMAXGemmaRMSNorm(head_dim).to( + device=DEVICE, dtype=torch.bfloat16 + ) + attention.index_q_norm = MiniMAXGemmaRMSNorm(idx_head_dim).to( + device=DEVICE, dtype=torch.bfloat16 + ) + attention.index_k_norm = MiniMAXGemmaRMSNorm(idx_head_dim).to( + device=DEVICE, dtype=torch.bfloat16 + ) + config = PreTrainedConfig( + rope_theta=10000.0, + partial_rotary_factor=0.5, + max_position_embeddings=32, + ) + from vllm.config import VllmConfig, set_current_vllm_config + + with set_current_vllm_config(VllmConfig()): + attention.rotary_emb = _build_rotary_emb(config, head_dim).to(DEVICE) + attention.index_rotary_emb = attention.rotary_emb + + fp8_dtype = current_platform.fp8_dtype() + attention.kv_cache = torch.zeros( + 2, + num_kv_heads, + block_size, + 2 * head_dim, + device=DEVICE, + dtype=fp8_dtype, + ) + attention._k_scale = torch.ones((), device=DEVICE, dtype=torch.float32) + attention._v_scale = torch.ones((), device=DEVICE, dtype=torch.float32) + index_cache = torch.zeros( + 2, block_size, idx_head_dim, device=DEVICE, dtype=torch.bfloat16 + ) + index_prefix = "model.layers.0.self_attn.indexer" + attention.indexer = SimpleNamespace( + index_cache=SimpleNamespace( + dtype=torch.bfloat16, + kv_cache=index_cache, + prefix=index_prefix, + ) + ) + + width = ( + attention.q_size + + 2 * attention.kv_size + + attention.index_q_size + + attention.idx_head_dim + ) + qkv = torch.randn(num_tokens, width, device=DEVICE, dtype=torch.bfloat16).mul_(0.25) + positions = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.long) + main_slots = torch.tensor([0, -1, 5], device=DEVICE, dtype=torch.long) + index_slots = torch.tensor([2, -1, 6], device=DEVICE, dtype=torch.long) + + q_raw, k_raw, v_raw, index_q_raw, index_k_raw = qkv.clone().split( + [ + attention.q_size, + attention.kv_size, + attention.kv_size, + attention.index_q_size, + attention.idx_head_dim, + ], + dim=-1, + ) + expected_q = attention.q_norm(q_raw.view(num_tokens, num_heads, head_dim)) + expected_k = attention.k_norm(k_raw.view(num_tokens, num_kv_heads, head_dim)) + expected_q, expected_k = attention.rotary_emb.forward_native( + positions, expected_q, expected_k + ) + expected_index_q = attention.index_q_norm( + index_q_raw.view(num_tokens, num_idx_heads, idx_head_dim) + ) + expected_index_k = attention.index_k_norm( + index_k_raw.view(num_tokens, idx_head_dim) + ) + expected_index_q, expected_index_k = attention.index_rotary_emb.forward_native( + positions, expected_index_q, expected_index_k + ) + assert expected_k is not None + assert expected_index_k is not None + + q, index_q = attention._native_gfx942_qknorm_rope_cache( + qkv, + positions, + { + attention.layer_name: main_slots, + index_prefix: index_slots, + }, + ) + + torch.testing.assert_close(q, expected_q.reshape(num_tokens, -1)) + assert index_q is not None + torch.testing.assert_close(index_q, expected_index_q.reshape(num_tokens, -1)) + + expected_v = v_raw.view(num_tokens, num_kv_heads, head_dim) + for token, slot in ((0, 0), (2, 5)): + block, offset = divmod(slot, block_size) + assert torch.equal( + attention.kv_cache[block, 0, offset, :head_dim], + expected_k[token, 0].to(fp8_dtype), + ) + assert torch.equal( + attention.kv_cache[block, 0, offset, head_dim:], + expected_v[token, 0].to(fp8_dtype), + ) + assert torch.count_nonzero(attention.kv_cache[0, 0, 1]) == 0 + + for token, slot in ((0, 2), (2, 6)): + block, offset = divmod(slot, block_size) + assert torch.equal( + index_cache[block, offset], + expected_index_k[token].to(index_cache.dtype), + ) + assert torch.count_nonzero(index_cache[0, 1]) == 0 diff --git a/tests/kernels/test_minimax_m3_dspark_draft.py b/tests/kernels/test_minimax_m3_dspark_draft.py new file mode 100644 index 000000000000..207a72ff8378 --- /dev/null +++ b/tests/kernels/test_minimax_m3_dspark_draft.py @@ -0,0 +1,443 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3.1 DSpark draft fusions against the unfused ops they replace. + +- The context-KV CUDA kernel + (vllm/models/minimax_m3/common/ops/dspark_context_kv.py) against flashinfer + gemma_rmsnorm per layer, vllm._custom_ops.rotary_embedding (partial NeoX) and + vllm._custom_ops.reshape_and_cache_flash into FlashAttention-layout caches + [num_blocks, nkv, block_size, 2 * hd] viewed as + kv_cache.transpose(1, 2).split(hd, dim=-1). +- The draft attention's fused q/k Gemma norm + RoPE against DFlash's q-norm, + k-norm and RoPE. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm._custom_ops as ops +from vllm.models.minimax_m3.common.ops.dspark_context_kv import ( + DSparkContextKVWriter, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_default_torch_dtype + +flashinfer_norm = pytest.importorskip("flashinfer.norm") + +HEAD_DIM = 128 +ROTARY_DIM = 64 +BLOCK_SIZE = 16 +EPS = 1e-6 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA not available" +) + + +def make_cos_sin_cache(max_pos, rotary_dim, dtype, device): + inv_freq = 1.0 / ( + 5_000_000.0 + ** ( + torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=device) + / rotary_dim + ) + ) + t = torch.arange(max_pos, dtype=torch.float32, device=device) + freqs = torch.einsum("i,j->ij", t, inv_freq) + return torch.cat((freqs.cos(), freqs.sin()), dim=-1).to(dtype) + + +def make_slots(num_ctx, num_blocks, device): + """Unique valid slots with ~1/4 padding (-1), at least one valid.""" + num_slots = num_blocks * BLOCK_SIZE + slots = torch.randperm(num_slots, device=device)[:num_ctx].to(torch.int64) + num_pad = max(1, num_ctx // 4) + slots[:num_pad] = -1 + return slots + + +def build_layers(kv_cache_dtype, L, nkv, num_blocks, device, seed, unit_scale=False): + g = torch.Generator(device="cpu").manual_seed(seed) + if kv_cache_dtype == "fp8" and not unit_scale: + dtype = torch.uint8 + scale = lambda: ( # noqa: E731 + torch.rand((), generator=g).to(device) * 1.5 + 0.25 + ) + else: + dtype = torch.uint8 if kv_cache_dtype == "fp8" else torch.bfloat16 + scale = lambda: torch.ones((), device=device) # noqa: E731 + layers = [] + for _ in range(L): + layers.append( + SimpleNamespace( + kv_cache=torch.zeros( + num_blocks, + nkv, + BLOCK_SIZE, + 2 * HEAD_DIM, + dtype=dtype, + device=device, + ), + _k_scale=scale(), + _v_scale=scale(), + kv_cache_dtype=kv_cache_dtype, + ) + ) + return layers + + +def reference_update( + all_kv, k_norm_w, positions, cos_sin, slot_mapping, layers, kv_cache_dtype +): + """The unfused path, run on `layers`' caches.""" + num_ctx, L, _, nkv, hd = all_kv.shape + all_k = torch.stack( + [ + flashinfer_norm.gemma_rmsnorm(all_kv[:, layer, 0], k_norm_w[layer], EPS) + for layer in range(L) + ] + ) # [L, num_ctx, nkv, hd] + all_v = all_kv[:, :, 1] # [num_ctx, L, nkv, hd] + k_flat = all_k.view(L * num_ctx, nkv * hd) + ops.rotary_embedding( + positions.repeat(L), + k_flat, + None, + hd, + cos_sin.to(k_flat.dtype), + True, + ) + all_k = k_flat.view(L, num_ctx, nkv, hd) + for layer, a in enumerate(layers): + kc, vc = a.kv_cache.transpose(1, 2).split(hd, dim=-1) + ops.reshape_and_cache_flash( + all_k[layer], + all_v[:, layer].contiguous(), + kc, + vc, + slot_mapping, + kv_cache_dtype, + a._k_scale, + a._v_scale, + ) + + +def assert_bf16_close(got, exp): + diff = (got.view(torch.int16).int() - exp.view(torch.int16).int()).abs() + assert int(diff.max()) <= 1, ( + f"bf16 cache differs by {int(diff.max())} ulp " + f"in {int((diff > 0).sum())} elements" + ) + + +def assert_fp8_close(got, exp): + equal = got == exp + frac = equal.float().mean().item() + assert frac >= 0.999, f"fp8 cache matches only {frac:.4%} of bytes" + bad = ~equal + if bad.any(): + byte_diff = (got[bad].int() - exp[bad].int()).abs() + g = got[bad].view(torch.float8_e4m3fn).float() + e = exp[bad].view(torch.float8_e4m3fn).float() + ok = (byte_diff <= 1) | ((g == 0) & (e == 0)) + assert bool(ok.all()), ( + f"fp8 cache differs by more than 1 step in {int((~ok).sum())} bytes" + ) + + +def run_case( + kv_cache_dtype, num_ctx, nkv, L, seed=0, rotary_dim=ROTARY_DIM, unit_scale=False +): + torch.manual_seed(seed) + device = "cuda" + num_blocks = (num_ctx * 2 + BLOCK_SIZE - 1) // BLOCK_SIZE + 1 + + layers, ref_layers = ( + build_layers(kv_cache_dtype, L, nkv, num_blocks, device, seed, unit_scale) + for _ in range(2) + ) + + all_kv = torch.randn( + num_ctx, L, 2, nkv, HEAD_DIM, dtype=torch.bfloat16, device=device + ) + k_norm_w = torch.randn(L, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.2 + positions = torch.randint(0, 2048, (num_ctx,), dtype=torch.int64, device=device) + cos_sin = make_cos_sin_cache(2048, rotary_dim, torch.bfloat16, device) + slot_mapping = make_slots(num_ctx, num_blocks, device) + + writer = DSparkContextKVWriter(layers) + writer(all_kv, k_norm_w, EPS, positions, cos_sin, rotary_dim, slot_mapping) + reference_update( + all_kv, k_norm_w, positions, cos_sin, slot_mapping, ref_layers, kv_cache_dtype + ) + + for a, b in zip(layers, ref_layers): + if kv_cache_dtype == "fp8": + assert_fp8_close(a.kv_cache, b.kv_cache) + else: + assert_bf16_close(a.kv_cache, b.kv_cache) + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("num_ctx", [1, 7, 64, 300]) +@pytest.mark.parametrize("nkv", [1, 2]) +@pytest.mark.parametrize("L", [1, 5]) +def test_dspark_context_kv(kv_cache_dtype, num_ctx, nkv, L): + run_case(kv_cache_dtype, num_ctx, nkv, L) + + +@pytest.mark.parametrize("num_ctx", [7, 300]) +def test_dspark_context_kv_fp8_unit_scale(num_ctx): + """Unit k/v scales take the kernel's direct bf16 -> e4m3 conversion.""" + run_case("fp8", num_ctx, nkv=2, L=5, unit_scale=True) + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("rotary_dim", [16, 32]) +def test_dspark_context_kv_short_rope(kv_cache_dtype, rotary_dim): + """Only some lanes rotate when rotary_dim < 64.""" + run_case(kv_cache_dtype, 64, nkv=2, L=2, rotary_dim=rotary_dim) + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_dspark_context_kv_cuda_graph(kv_cache_dtype): + torch.manual_seed(0) + device = "cuda" + num_ctx, nkv, L = 16, 1, 5 + num_blocks = 8 + + layers = build_layers(kv_cache_dtype, L, nkv, num_blocks, device, seed=0) + all_kv = torch.randn( + num_ctx, L, 2, nkv, HEAD_DIM, dtype=torch.bfloat16, device=device + ) + k_norm_w = torch.randn(L, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.2 + positions = torch.randint(0, 2048, (num_ctx,), dtype=torch.int64, device=device) + cos_sin = make_cos_sin_cache(2048, ROTARY_DIM, torch.bfloat16, device) + slot_mapping = make_slots(num_ctx, num_blocks, device) + + writer = DSparkContextKVWriter(layers) + writer._ensure_tables() # build pointer tables outside the graph + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + writer(all_kv, k_norm_w, EPS, positions, cos_sin, ROTARY_DIM, slot_mapping) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + writer(all_kv, k_norm_w, EPS, positions, cos_sin, ROTARY_DIM, slot_mapping) + + # Change every input in place, replay, and compare against the unfused + # reference run on the new inputs. + all_kv.copy_(torch.randn_like(all_kv)) + positions.copy_( + torch.randint(0, 2048, (num_ctx,), dtype=torch.int64, device=device) + ) + slot_mapping.copy_(make_slots(num_ctx, num_blocks, device)) + + sentinel = 0x5A if kv_cache_dtype == "fp8" else 0x3F80 + for a in layers: + a.kv_cache.fill_(sentinel) + graph.replay() + + ref_layers = build_layers(kv_cache_dtype, L, nkv, num_blocks, device, seed=0) + for a in ref_layers: + a.kv_cache.fill_(sentinel) + reference_update( + all_kv, k_norm_w, positions, cos_sin, slot_mapping, ref_layers, kv_cache_dtype + ) + + torch.accelerator.synchronize() + for a, b in zip(layers, ref_layers): + if kv_cache_dtype == "fp8": + assert_fp8_close(a.kv_cache, b.kv_cache) + else: + assert_bf16_close(a.kv_cache, b.kv_cache) + + +def test_dspark_context_kv_supported_by_cache_layout(): + """The writer engages on the [nb, nkv, bs, 2 * 128] page that FLASH_ATTN + and FLASHINFER_TS share, whatever the backend, and not on a 5-D cache.""" + layers = build_layers("fp8", 5, 1, 4, "cuda", seed=0, unit_scale=True) + assert DSparkContextKVWriter.supported(layers) + layers[0].kv_cache = torch.zeros( + 4, 2, BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.uint8, device="cuda" + ) + assert not DSparkContextKVWriter.supported(layers) + + +class _CaptureAttention(torch.nn.Module): + """Stands in for the draft's Attention: records the q/k/v it is given.""" + + _q_scale_float = 1.0 + + def __init__(self, layer_name): + super().__init__() + self.layer_name = layer_name + + def forward(self, q, k, v, output_dtype=None): + self.qkv = tuple(None if t is None else t.clone() for t in (q, k, v)) + return torch.zeros_like(q, dtype=output_dtype) + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +def test_draft_attention_fused_qk_norm_rope( + dist_init, default_vllm_config, kv_cache_dtype +): + """One fused launch gives the q of DFlash's q-norm and RoPE and stores its + k/v in a bf16 or fp8 cache as the layer's reshape_and_cache_flash would; + the attention layer gets q alone.""" + from vllm.forward_context import set_forward_context + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Attention + from vllm.models.minimax_m3.common.dspark import MiniMaxM3DSparkAttention + from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig + + torch.manual_seed(0) + config = MiniMaxM3TextConfig( + hidden_size=512, + num_attention_heads=8, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=8192, + ) + with torch.device("cuda"), set_default_torch_dtype(torch.bfloat16): + attn = MiniMaxM3DSparkAttention( + config, + sliding_window=None, + causal=False, + cache_config=None, + quant_config=None, + prefix="model.layers.0.self_attn", + ) + for param in (attn.qkv_proj.weight, attn.q_norm.weight, attn.k_norm.weight): + param.data.normal_(std=0.05) + layer = attn.attn + cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else torch.bfloat16 + layer.kv_cache = torch.zeros( + 4, 2, BLOCK_SIZE, 2 * HEAD_DIM, dtype=cache_dtype, device="cuda" + ) + attn.attn = _CaptureAttention(layer.layer_name) + + hidden = torch.randn(7, config.hidden_size, device="cuda", dtype=torch.bfloat16) + positions = torch.tensor([0, 1, 127, 128, 4000, 8190, 9], device="cuda") + slots = torch.tensor([5, -1, 17, 18, 40, 63, 0], device="cuda") + DFlashQwen3Attention.forward(attn, positions, hidden) + q_ref, k_ref, v_ref = attn.attn.qkv + assert k_ref is not None and v_ref is not None + with set_forward_context( + None, default_vllm_config, slot_mapping={layer.layer_name: slots} + ): + attn(positions, hidden) + q, k, v = attn.attn.qkv + torch.testing.assert_close(q, q_ref, atol=2e-2, rtol=2e-2) + assert k is None and v is None + expected = torch.zeros_like(layer.kv_cache) + key_cache, value_cache = expected.transpose(1, 2).split(HEAD_DIM, dim=-1) + one = torch.ones((), dtype=torch.float32, device="cuda") + ops.reshape_and_cache_flash( + k_ref.reshape(7, 2, HEAD_DIM), + v_ref.reshape(7, 2, HEAD_DIM), + key_cache, + value_cache, + slots, + kv_cache_dtype, + one, + one, + ) + got, rtol = layer.kv_cache, 2e-2 + if kv_cache_dtype == "fp8": + # k/v_ref come from DFlash's unfused kernels: allow one e4m3 step. + got = got.view(torch.float8_e4m3fn).float() + expected = expected.view(torch.float8_e4m3fn).float() + rtol = 0.13 + torch.testing.assert_close(got, expected, atol=2e-2, rtol=rtol) + + +class _FixedOutputAttention(torch.nn.Module): + """Stands in for the draft's Attention: returns a fixed output in the + requested dtype and records the dtypes asked for.""" + + def __init__(self, layer_name, out): + super().__init__() + self.layer_name, self.out, self.output_dtypes = layer_name, out, [] + + def forward(self, q, k, v, output_dtype): + self.output_dtypes.append(output_dtype) + return self.out.to(output_dtype) + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="Native MXFP8 kernels require SM100-family CUDA", +) +@torch.inference_mode() +def test_draft_fp16_attention_output_feeds_mxfp8_o_proj(dist_init, default_vllm_config): + """A backend that writes an FP8 query's output as FP16 hands it to the + MXFP8 o_proj without a conversion pass; o_proj quantizes it and still + returns bf16, bit-identical to the bf16 path for values both dtypes hold.""" + from vllm.forward_context import set_forward_context + from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + _mxfp8_e4m3_quantize_torch, + ) + from vllm.models.minimax_m3.common.dspark import MiniMaxM3DSparkAttention + from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig + + torch.manual_seed(0) + config = MiniMaxM3TextConfig( + hidden_size=512, + num_attention_heads=8, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=8192, + ) + quant_config = ModelOptMxFp8Config.from_config( + { + "quant_method": "mxfp8", + "activation_scheme": "dynamic", + "weight_block_size": [1, 32], + "ignored_layers": [], + } + ) + with torch.device("cuda"), set_default_torch_dtype(torch.bfloat16): + attn = MiniMaxM3DSparkAttention( + config, + sliding_window=None, + causal=False, + cache_config=None, + quant_config=quant_config, + prefix="model.layers.0.self_attn", + query_pre_quantized=True, + ) + for proj in (attn.qkv_proj, attn.o_proj): + weight, scale = _mxfp8_e4m3_quantize_torch( + torch.randn(proj.weight.shape, device="cuda") / 16 + ) + proj.weight.data.copy_(weight) + proj.weight_scale.data.copy_(scale) + proj.quant_method.process_weights_after_loading(proj) + layer = attn.attn + layer.kv_cache = torch.zeros( + 4, 2, BLOCK_SIZE, 2 * HEAD_DIM, dtype=torch.uint8, device="cuda" + ) + num_tokens = 7 + out = torch.randint(-256, 256, (num_tokens, attn.q_size), device="cuda") / 128 + attn.attn = _FixedOutputAttention(layer.layer_name, out.to(torch.bfloat16)) + + hidden = torch.randn( + num_tokens, config.hidden_size, device="cuda", dtype=torch.bfloat16 + ) + positions = torch.arange(num_tokens, device="cuda") + outputs = [] + for attn_output_dtype in (torch.float16, None): + attn._attn_output_dtype = attn_output_dtype + with set_forward_context( + None, default_vllm_config, slot_mapping={layer.layer_name: positions} + ): + outputs.append(attn(positions, hidden)) + assert attn.attn.output_dtypes == [torch.float16, torch.bfloat16] + assert outputs[0].dtype == torch.bfloat16 + torch.testing.assert_close(outputs[0], outputs[1], atol=0, rtol=0) diff --git a/tests/models/minimax_m3/__init__.py b/tests/models/minimax_m3/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/minimax_m3/test_dense_gemm1_mxfp8.py b/tests/models/minimax_m3/test_dense_gemm1_mxfp8.py new file mode 100644 index 000000000000..cd7516a82559 --- /dev/null +++ b/tests/models/minimax_m3/test_dense_gemm1_mxfp8.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness test for the MiniMax-M3.1 fused dense GEMM1 kernel +(vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py). +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not (current_platform.is_cuda() and current_platform.has_device_capability(100)), + reason="dense GEMM1 MXFP8 kernel is SM100+ CUDA only", +) + + +def _mxfp8_quantize_f32(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8-quantize fp32 [M, N] with the kernel's integer-exact scale math: + scale byte = E - 8 + (amax > 1.75 * 2^E) with E the biased exponent of the + 32-elem group amax (0 for amax=0), clamped to [0, 254]. Returns q and the + unswizzled [M, N//32] scale bytes. + """ + M, N = x.shape + g = x.float().view(M, N // 32, 32) + amax = g.abs().amax(-1) + e_bias = amax.view(torch.int32) >> 23 + thr = torch.ldexp(torch.full_like(amax, 1.75), e_bias - 127) + sb = (e_bias - 8 + (amax > thr).to(torch.int32)).clamp(0, 254) + q = g * torch.ldexp(torch.ones_like(amax), 127 - sb)[..., None] + q = q.clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + return q.view(M, N), sb.to(torch.uint8) + + +# Per-rank I: TP4 shared expert, TP4 dense / TP1 shared, TP1 dense. 12288 is +# the only one whose weight exceeds half of L2 (grouped rasterization). +@pytest.mark.parametrize("num_tokens", [1, 37, 128, 1000, 4096]) +@pytest.mark.parametrize("intermediate", [768, 3072, 12288]) +def test_dense_gemm1_mxfp8_correctness(num_tokens: int, intermediate: int): + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + swizzle_mxfp8_scale, + ) + from vllm.models.minimax_m3.nvidia.ops.dense_gemm1_mxfp8 import run + + limit, alpha, beta = 7.0, 1.702, 1.0 + hidden = 6144 + torch.manual_seed(num_tokens) + x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16) * 0.5 + w1 = ( + torch.randn(intermediate * 2, hidden, device="cuda", dtype=torch.bfloat16) + * 0.05 + ) + + # What Mxfp8OnlineLinearMethod.process_weights_after_loading produces. + w1q, wsf_2d = mxfp8_e4m3_quantize(w1, is_sf_swizzled_layout=False) + w1sf = swizzle_mxfp8_scale(wsf_2d, M=intermediate * 2, K=hidden) + + M, N = num_tokens, intermediate + e8 = torch.float8_e8m0fnu + q = torch.empty((M, N), dtype=torch.float8_e4m3fn, device="cuda") + sf = torch.zeros( + ((M + 127) // 128) * 128 * (N // 32), dtype=torch.uint8, device="cuda" + ) + xq, xsf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True) + run(xq, w1q, xsf.view(-1).view(e8), w1sf.view(e8), q, sf, limit, alpha, beta) + + # fp32 mirror chain: same MXFP8 GEMM (fp32 out) -> fp32 clamped SwiGLU + # -> MXFP8 quantize with the kernel's scale math. + gemm = torch._scaled_mm( + xq, + w1q.t(), + scale_a=xsf.view(e8), + scale_b=w1sf.view(e8), + out_dtype=torch.float32, + ) + gate = gemm[:, :N].clamp(max=limit) + up = gemm[:, N:].clamp(-limit, limit) + act = gate * torch.sigmoid(gate * alpha) * (up + beta) + q_ref, sf_ref_2d = _mxfp8_quantize_f32(act) + sf_ref = swizzle_mxfp8_scale(sf_ref_2d, M=M, K=N) + + # Measured envelope: sf bytes identical; q bytes differ on <=0.007% of + # elements, always by 1 fp8 ulp. Bounds below carry >10x margin. + sf_mism = (sf != sf_ref).sum().item() + assert sf_mism <= sf.numel() * 1e-4, f"sf byte mismatches: {sf_mism}" + q_mism = q.view(torch.uint8) != q_ref.view(torch.uint8) + assert q_mism.sum().item() <= q.numel() * 1e-3, ( + f"q byte mismatches: {q_mism.sum().item()}" + ) diff --git a/tests/models/minimax_m3/test_dspark_draft_compile.py b/tests/models/minimax_m3/test_dspark_draft_compile.py new file mode 100644 index 000000000000..dc24bfb18da0 --- /dev/null +++ b/tests/models/minimax_m3/test_dspark_draft_compile.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The MiniMax-M3.1 DSpark draft attention under vLLM's torch.compile. + +vLLM traces the draft model once, in the memory-profiling run, whose forward +context has no slot mapping, and reuses that graph afterwards. The draft +tokens' k/v insert must still happen on the later calls. +""" + +import pytest +import torch +from torch import nn + +from vllm.utils.torch_utils import set_default_torch_dtype + +HEAD_DIM = 128 +BLOCK_SIZE = 16 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA not available" +) + + +class _QueryAttention(nn.Module): + """Stands in for the draft's Attention: returns its query, so the layer + output shows the q the attention got.""" + + def __init__(self, layer_name): + super().__init__() + self.layer_name = layer_name + + def forward(self, q, k, v, output_dtype=None): + return q.to(output_dtype) + + +@pytest.mark.parametrize( + "kv_cache_dtype,query_pre_quantized", + [("auto", False), ("fp8", False), ("fp8", True)], +) +def test_draft_attention_stores_kv_after_compile( + dist_init, monkeypatch, kv_cache_dtype, query_pre_quantized +): + from vllm.compilation.decorators import support_torch_compile + from vllm.config import ( + CompilationConfig, + CompilationMode, + CUDAGraphMode, + VllmConfig, + set_current_vllm_config, + ) + from vllm.forward_context import set_forward_context + from vllm.models.minimax_m3.common.dspark import MiniMaxM3DSparkAttention + from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig + + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1") + config = MiniMaxM3TextConfig( + hidden_size=512, + num_attention_heads=8, + num_key_value_heads=2, + head_dim=HEAD_DIM, + max_position_embeddings=8192, + ) + + @support_torch_compile + class DraftAttention(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs): + super().__init__() + self.self_attn = MiniMaxM3DSparkAttention( + config, + sliding_window=None, + causal=False, + cache_config=None, + quant_config=None, + prefix=f"{prefix}.self_attn", + query_pre_quantized=query_pre_quantized, + ) + + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + return self.self_attn(positions, hidden_states) + + vllm_config = VllmConfig( + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, cudagraph_mode=CUDAGraphMode.NONE + ) + ) + torch.manual_seed(0) + with ( + set_current_vllm_config(vllm_config), + torch.device("cuda"), + set_default_torch_dtype(torch.bfloat16), + ): + model = DraftAttention(vllm_config=vllm_config, prefix="model.layers.0") + attn = model.self_attn + for param in ( + attn.qkv_proj.weight, + attn.q_norm.weight, + attn.k_norm.weight, + attn.o_proj.weight, + ): + param.data.normal_(std=0.05) + layer = attn.attn + cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else torch.bfloat16 + layer.kv_cache = torch.zeros( + 4, 2, BLOCK_SIZE, 2 * HEAD_DIM, dtype=cache_dtype, device="cuda" + ) + attn.attn = _QueryAttention(layer.layer_name) + + hidden = torch.randn(7, config.hidden_size, device="cuda", dtype=torch.bfloat16) + positions = torch.tensor([0, 1, 127, 128, 4000, 8190, 9], device="cuda") + slots = {layer.layer_name: torch.tensor([5, -1, 17, 18, 40, 63, 0], device="cuda")} + + with torch.inference_mode(): + with set_forward_context(None, vllm_config, slot_mapping=slots): + expected_out = attn(positions, hidden) + expected_cache = layer.kv_cache.clone() + assert expected_cache.any() + layer.kv_cache.zero_() + + # The memory-profiling run: no slot mapping, nothing is stored. + with set_forward_context(None, vllm_config): + model(positions, hidden) + assert not layer.kv_cache.any() + + with set_forward_context(None, vllm_config, slot_mapping=slots): + out = model(positions, hidden) + assert layer.kv_cache.any(), "the compiled draft attention stored no k/v" + got, expected, rtol = layer.kv_cache, expected_cache, 2e-2 + if kv_cache_dtype == "fp8": + got = got.view(torch.float8_e4m3fn).float() + expected = expected.view(torch.float8_e4m3fn).float() + rtol = 0.13 + torch.testing.assert_close(got, expected, atol=2e-2, rtol=rtol) + torch.testing.assert_close(out, expected_out, atol=2e-2, rtol=2e-2) diff --git a/tests/models/minimax_m3/test_mono_host.py b/tests/models/minimax_m3/test_mono_host.py new file mode 100644 index 000000000000..df845b274d17 --- /dev/null +++ b/tests/models/minimax_m3/test_mono_host.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host-side logic of the MiniMax-M3 mono decode path that needs no GPU. + +The kernel package imports flydsl at module import, so these tests run where it +is installed (the ROCm images); elsewhere they are skipped. +""" + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("flydsl") + +from vllm.models.minimax_m3.amd.mono import config as mono_config # noqa: E402 +from vllm.models.minimax_m3.amd.mono import dispatch as mono_dispatch # noqa: E402 +from vllm.models.minimax_m3.amd.mono import envs as mono_envs # noqa: E402 +from vllm.models.minimax_m3.amd.mono.config import ( # noqa: E402 + MAX_CONTEXT, + MAX_TOKENS, + SPARSE_BLOCK, + TP, + MonoUnsupported, +) +from vllm.models.minimax_m3.amd.mono.dispatch import ( # noqa: E402 + _config_refusal, + _step_refusal, +) +from vllm.models.minimax_m3.amd.mono.runner import ( # noqa: E402 + dense_prefix, + token_rows, +) + + +def _layer(moe: bool, sparse: bool): + attn = ( + SimpleNamespace( + layer_name="attn", + indexer=SimpleNamespace(index_cache=SimpleNamespace(prefix="idx")), + ) + if sparse + else SimpleNamespace() + ) + return SimpleNamespace(is_moe_layer=moe, self_attn=attn) + + +def test_token_rows_single_query_token_passes_tables_through(): + bt = torch.arange(12, dtype=torch.int32).view(3, 4) + seq = torch.tensor([10, 20, 30], dtype=torch.int32) + rows_bt, rows_seq = token_rows(bt, seq, q=1, n=3) + assert rows_bt is bt and rows_seq is seq + + +def test_token_rows_expands_speculative_verify_rows(): + bt = torch.arange(8, dtype=torch.int32).view(2, 4) + seq = torch.tensor([10, 20], dtype=torch.int32) + rows_bt, rows_seq = token_rows(bt, seq, q=3, n=6) + # token j of a request sees seq_len - (q - 1 - j) keys + assert rows_seq.tolist() == [8, 9, 10, 18, 19, 20] + assert rows_bt.shape == (6, 4) and torch.equal(rows_bt[:3], bt[0].expand(3, 4)) + assert rows_bt.stride(0) == rows_bt.shape[1] + + +def test_dense_prefix_counts_leading_dense_ffn_layers(): + layers = [_layer(False, True)] * 3 + [_layer(True, True)] * 5 + assert dense_prefix(layers) == 3 + assert dense_prefix([_layer(True, True)] * 2) == 0 + + +@pytest.mark.parametrize( + "layers", + [ + [_layer(False, True)] * 3, # no MoE layer at all + [ + _layer(True, True), + _layer(False, True), + _layer(True, True), + ], # dense after MoE + [ + _layer(False, True), + _layer(True, False), + ], # MoE layer without sparse attention + ], +) +def test_dense_prefix_refuses_other_layer_mixes(layers): + with pytest.raises(MonoUnsupported): + dense_prefix(layers) + + +def test_configure_accepts_only_the_two_kv_page_strides(): + saved = (mono_config.KV_PAGE_STRIDE, mono_config.STATIC_KV_SCALE) + try: + mono_config.configure(kv_page_stride=16, static_kv_scale=True) + assert (mono_config.KV_PAGE_STRIDE, mono_config.STATIC_KV_SCALE) == (16, True) + with pytest.raises(AssertionError): + mono_config.configure(kv_page_stride=4, static_kv_scale=True) + finally: + mono_config.configure(kv_page_stride=saved[0], static_kv_scale=saved[1]) + + +def _vllm_config(**overrides): + cg = SimpleNamespace(requires_piecewise_compilation=lambda: False) + cfg = SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=TP, + data_parallel_size=1, + pipeline_parallel_size=1, + enable_expert_parallel=False, + ), + cache_config=SimpleNamespace(cache_dtype="fp8", block_size=SPARSE_BLOCK), + compilation_config=SimpleNamespace(cudagraph_mode=cg), + model_config=SimpleNamespace( + max_model_len=32768, hf_text_config=SimpleNamespace() + ), + attention_config=SimpleNamespace( + resolve_indexer_kv_dtype=lambda default: "fp8" + ), + ) + for path, value in overrides.items(): + obj, _, attr = path.rpartition(".") + setattr(eval(f"cfg.{obj}") if obj else cfg, attr, value) + return cfg + + +def _model(fse: bool = True, shared_mlp: bool = True, aux=()): + moe = _layer(True, True) + moe.block_sparse_moe = SimpleNamespace( + is_fused_shared_expert_enabled=fse, + shared_experts=object() if shared_mlp else None, + ) + return SimpleNamespace( + is_fused_shared_expert_enabled=fse, + aux_hidden_state_layers=aux, + layers=[_layer(False, True)] * 3 + [moe] * 2, + ) + + +@pytest.mark.parametrize("fse", [True, False]) +@pytest.mark.parametrize("aux", [(), (2, 30, 57)]) +def test_config_refusal_accepts_the_supported_deployment(fse, aux): + assert _config_refusal(_vllm_config(), _model(fse=fse, aux=aux)) is None + + +@pytest.mark.parametrize( + "cfg, model, why", + [ + (_vllm_config(**{"parallel_config.tensor_parallel_size": 8}), _model(), "TP"), + (_vllm_config(**{"cache_config.cache_dtype": "auto"}), _model(), "kv cache"), + (_vllm_config(**{"cache_config.block_size": 64}), _model(), "block size"), + ( + _vllm_config(**{"model_config.max_model_len": MAX_CONTEXT + 1}), + _model(), + "max_model_len", + ), + (_vllm_config(), _model(fse=False, shared_mlp=False), "neither fused nor"), + ], +) +def test_config_refusal_names_the_unmet_condition(cfg, model, why): + reason = _config_refusal(cfg, model) + assert reason is not None and why in reason + + +def _fwd(n: int, q: int = 1, reqs: int | None = None, prefills: int = 0): + reqs = n // q if reqs is None else reqs + decode = SimpleNamespace( + decode_query_len=q, + block_table=torch.zeros(reqs, 8, dtype=torch.int32), + seq_lens=torch.ones(reqs, dtype=torch.int32), + ) + md = SimpleNamespace(decode=decode, num_prefills=prefills) + return SimpleNamespace( + ubatch_slices=None, + attn_metadata={"attn": md}, + slot_mapping={"attn": torch.zeros(n, dtype=torch.int64)}, + ) + + +def test_step_refusal_accepts_a_pure_decode(): + pos = torch.zeros(4, dtype=torch.int64) + assert _step_refusal(_fwd(4), "attn", 4, pos) is None + assert _step_refusal(_fwd(6, q=3), "attn", 6, pos) is None + + +def _step_reason(fwd, n: int, positions: torch.Tensor) -> str: + reason = _step_refusal(fwd, "attn", n, positions) + assert reason is not None + return reason + + +def test_step_refusal_rejects_what_the_kernels_cannot_serve(): + pos = torch.zeros(4, dtype=torch.int64) + assert "tokens" in _step_reason(_fwd(MAX_TOKENS + 1), MAX_TOKENS + 1, pos) + assert "prefill" in _step_reason(_fwd(4, prefills=1), 4, pos) + assert "per request" in _step_reason(_fwd(4, q=3, reqs=2), 4, pos) + assert "shorter" in _step_reason(_fwd(4, reqs=2), 4, pos) + assert "positions" in _step_reason(_fwd(4), 4, pos.to(torch.int32)) + fwd = _fwd(4) + fwd.attn_metadata["attn"].decode.block_table = fwd.attn_metadata[ + "attn" + ].decode.block_table.t() + assert "contiguous" in _step_reason(fwd, 4, pos) + + +# ---- MonoDecode host contract: enable, runner life cycle, slot check, routing ---- +# Everything the kernels are not: these pin the dispatcher's behaviour so a refactor +# of the host side (or of the switch) must keep passing them. + + +class _FakeRunner: + built = 0 + + def __init__(self, model, config, *, stale=False, fp8=False): + _FakeRunner.built += 1 + self._stale = stale + self.shared_fp8 = fp8 + self.sparse = [] + self.rank = 0 + self.closed = False + self.calls = [] + + def stale(self): + return self._stale + + def close(self): + self.closed = True + + def forward(self, input_ids, positions, inputs_embeds): + self.calls.append(positions.numel()) + return "mono-out" + + +@pytest.fixture +def mono(monkeypatch): + """An enabled dispatcher on a supported deployment, outside graph capture, + with the runner build stubbed.""" + monkeypatch.setattr(mono_dispatch.envs, "MONO_ENABLE", True) + monkeypatch.setattr(mono_dispatch.envs, "MONO_CHECK", False) + monkeypatch.setattr(mono_dispatch, "MonoDecodeRunner", _FakeRunner) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + _FakeRunner.built = 0 + return mono_dispatch.MonoDecode(_model(), _vllm_config()) + + +def _fwd_shared(n: int, shared: bool = True): + fwd = _fwd(n) + slots = fwd.slot_mapping["attn"] + fwd.slot_mapping["idx"] = slots.clone() if shared else slots + 1 + return fwd + + +def test_env_flag_parsing(monkeypatch): + for raw, want in (("1", True), ("true", True), ("True", True), ("0", False)): + monkeypatch.setenv("VLLM_ROCM_MINIMAX_M3_MONO_TESTFLAG", raw) + assert mono_envs._flag("VLLM_ROCM_MINIMAX_M3_MONO_TESTFLAG") is want + monkeypatch.delenv("VLLM_ROCM_MINIMAX_M3_MONO_TESTFLAG") + assert mono_envs._flag("VLLM_ROCM_MINIMAX_M3_MONO_TESTFLAG") is False + + +def test_disabled_without_the_switch(monkeypatch): + monkeypatch.setattr(mono_dispatch.envs, "MONO_ENABLE", False) + mono = mono_dispatch.MonoDecode(_model(), _vllm_config()) + assert not mono._enabled + assert not mono.supports(None, torch.zeros(1, dtype=torch.int64), None, None) + + +def test_unsupported_deployment_falls_back(monkeypatch): + monkeypatch.setattr(mono_dispatch.envs, "MONO_ENABLE", True) + bad = _vllm_config(**{"parallel_config.tensor_parallel_size": 8}) + assert not mono_dispatch.MonoDecode(_model(), bad)._enabled + + +def test_enabled_dispatcher_reads_the_first_sparse_layer(mono): + assert mono._enabled + assert mono._layer_name == "attn" and mono._index_prefix == "idx" + assert mono._runner is None # built lazily, on the first eligible step + + +def test_runner_is_built_once_outside_capture(mono): + assert mono._ensure_runner() and _FakeRunner.built == 1 + assert mono._ensure_runner() and _FakeRunner.built == 1 + + +def test_runner_is_not_built_inside_capture(mono, monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert not mono._ensure_runner() and _FakeRunner.built == 0 + assert mono._enabled # not a refusal: the next uncaptured step builds it + + +def test_stale_runner_is_rebuilt_after_a_kv_rebind(mono, monkeypatch): + first = _FakeRunner(None, None, stale=True) + mono._runner, mono._slots_verified = first, True + assert mono._ensure_runner() + assert first.closed and mono._runner is not first + assert not mono._slots_verified # the new cache's slots are checked again + + +def test_stale_runner_is_kept_while_capturing(mono, monkeypatch): + first = _FakeRunner(None, None, stale=True) + mono._runner = first + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert not mono._ensure_runner() and not first.closed + + +@pytest.mark.parametrize("exc", [MonoUnsupported("no CUs"), RuntimeError("boom")]) +def test_runner_build_failure_disables_mono(mono, monkeypatch, exc): + def boom(model, config): + raise exc + + monkeypatch.setattr(mono_dispatch, "MonoDecodeRunner", boom) + assert not mono._ensure_runner() + assert not mono._enabled and mono._runner is None + + +def test_slot_check_runs_once_and_caches(mono): + assert mono._slots_shared(_fwd_shared(4), 4) and mono._slots_verified + assert mono._slots_shared(_fwd_shared(4, shared=False), 4) # cached verdict + + +@pytest.mark.parametrize("missing", [True, False]) +def test_unshared_index_slots_disable_mono(mono, missing): + fwd = _fwd_shared(4, shared=False) # mismatched ids ... + if missing: + del fwd.slot_mapping["idx"] # ... or no index slots at all + assert not mono._slots_shared(fwd, 4) and not mono._enabled + + +def test_supports_routes_only_pure_decodes(mono, monkeypatch): + pos = torch.zeros(4, dtype=torch.int64) + monkeypatch.setattr(mono_dispatch, "get_forward_context", lambda: _fwd_shared(4)) + assert mono.supports(None, pos, None, None) and _FakeRunner.built == 1 + assert not mono.supports(None, pos, object(), None) # pipeline-parallel input + monkeypatch.setattr( + mono_dispatch, "get_forward_context", lambda: _fwd_shared(4, shared=True) + ) + fwd = _fwd(4, prefills=1) + monkeypatch.setattr(mono_dispatch, "get_forward_context", lambda: fwd) + assert not mono.supports(None, pos, None, None) and mono._enabled # just declined + + +def test_refused_step_does_not_build_the_runner(mono, monkeypatch): + fwd = _fwd(MAX_TOKENS + 1) + monkeypatch.setattr(mono_dispatch, "get_forward_context", lambda: fwd) + pos = torch.zeros(MAX_TOKENS + 1, dtype=torch.int64) + assert not mono.supports(None, pos, None, None) and _FakeRunner.built == 0 + + +def test_forward_runs_the_runner_or_the_self_check(mono, monkeypatch): + mono._ensure_runner() + pos = torch.zeros(2, dtype=torch.int64) + assert mono.forward(None, pos, None) == "mono-out" and mono._runner.calls == [2] + monkeypatch.setattr(mono_dispatch.envs, "MONO_CHECK", True) + monkeypatch.setattr( + mono_dispatch, "forward_checked", lambda runner, *a: ("checked", runner) + ) + assert mono.forward(None, pos, None) == ("checked", mono._runner) diff --git a/tests/models/minimax_m3/test_qkv_index_fused_epilogue.py b/tests/models/minimax_m3/test_qkv_index_fused_epilogue.py new file mode 100644 index 000000000000..4f4689eb2896 --- /dev/null +++ b/tests/models/minimax_m3/test_qkv_index_fused_epilogue.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3 QKV+indexer MXFP8 GEMM with fused epilogue +(ops/qkv_index_fused_epilogue.py) vs the production chain: mxfp8_e4m3_quantize +-> flashinfer mm_mxfp8 -> fused_minimax_m3_qknorm_rope_kv_insert. + +Each case runs both tile configs (1-SM and 2-SM), which must match bitwise, +and compares them to the chain: bf16 q within tolerances, the quantized +outputs (fp8 q / index_q, NVFP4 caches) code by code. The chain rounds the +GEMM output to bf16 first, so values near a rounding boundary land one code +apart. +""" + +import pytest +import torch + +import vllm._custom_ops as ops +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not (current_platform.is_cuda() and current_platform.has_device_capability(100)), + reason="QKV+indexer MXFP8 GEMM is SM100+ CUDA only", +) + +HIDDEN = 6144 +HEAD_DIM = 128 +ROTARY_DIM = 64 +EPS = 1e-6 +MAX_POS = 8192 +Q_FP8_SCALE = 2.0 +# per-rank (q, kv, index_q) heads; N = (nq + 2 * nkv + niq + 1) * HEAD_DIM +TP_HEADS = {1: (64, 4, 4), 2: (32, 2, 2), 4: (16, 1, 1)} +# Caches start as 0xFF: never a scale byte (E4M3 scales saturate at 0x7E), so +# a stray or missing write shows up in the code comparison. +CACHE_FILL = 0xFF + +BF16_RTOL = 2**-6 +BF16_ATOL = 0.05 +# fp8 values near zero (RoPE cancellation, where an E4M3 step is 2^-7..2^-9) +# can land a few codes apart but stay within 2^-6. +FP8_MAX_MISMATCH_FRAC = 0.05 +FP8_SMALL_ABS = 2**-6 +NVFP4_MAX_MISMATCH_FRAC = 0.02 +NVFP4_SF_MAX_MISMATCH_FRAC = 0.05 + + +def _make_inputs(nq, nkv, niq, M, nb, bs, slots): + """MXFP8 activations / weight, norm weights, positions, the RoPE cache and + the KV / index-cache slots: random unique with ~3% -1 (no insert), the + last M of the cache, or all -1.""" + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + swizzle_mxfp8_scale, + ) + + # CPU RNG: the CUDA one's output depends on the GPU's SM count + gen = torch.Generator().manual_seed(M) + + def randn(*shape): + return torch.randn(*shape, generator=gen).cuda() + + n = (nq + 2 * nkv + niq + 1) * HEAD_DIM + # activations with a few outlier channels, like real hidden states + x = randn(M, HIDDEN) * 0.5 + x[:, torch.randperm(HIDDEN, generator=gen)[:16].cuda()] *= 20 + x = x.bfloat16() + w = (randn(n, HIDDEN) * 0.05).bfloat16() + xq, xsf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True) + wq, wsf = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=False) + wsf = swizzle_mxfp8_scale(wsf, M=n, K=HIDDEN) + # Gemma norm weights (scale 1 + w) as in the checkpoint: q / k near 0, + # index_q down to -1 (scale near 0), index_k up to ~2.5 + offset_std = torch.tensor( + [[0.1, 0.1], [0.2, 0.1], [-0.4, 0.25], [0.9, 0.5]], device="cuda" + ) + norm_weights = offset_std[:, :1] + offset_std[:, 1:] * randn(4, HEAD_DIM) + norm_weights = norm_weights.bfloat16() + positions = torch.randint(0, MAX_POS, (M,), generator=gen).cuda() + inv_freq = 1.0 / ( + 5e6 ** (torch.arange(0, ROTARY_DIM, 2, device="cuda").float() / ROTARY_DIM) + ) + freqs = torch.outer(torch.arange(MAX_POS, device="cuda").float(), inv_freq) + cos_sin = torch.cat((freqs.cos(), freqs.sin()), dim=-1).bfloat16() + + if slots == "none": + sm = ism = torch.full((M,), -1, device="cuda") + elif slots == "last_page": + sm = ism = torch.arange(nb * bs - M, nb * bs, device="cuda") + else: + sm, ism = (torch.randperm(nb * bs, generator=gen)[:M] for _ in range(2)) + sm[torch.rand(M, generator=gen) < 0.03] = -1 + ism[torch.rand(M, generator=gen) < 0.03] = -1 + sm, ism = sm.cuda(), ism.cuda() + return xq, xsf, wq, wsf, norm_weights, positions, cos_sin, sm, ism + + +def _ordered(codes, sign_bit): + """Sign-magnitude codes -> integers ordered like their values (+0 == -0), + so adjacent representable values differ by 1.""" + codes = codes.to(torch.int16) + mag = codes & (sign_bit - 1) + return torch.where(codes & sign_bit != 0, -mag, mag) + + +def _check_fp8(name, got, want): + """E4M3 codes: NaN masks exact, at most FP8_MAX_MISMATCH_FRAC differ, each + within one code or FP8_SMALL_ABS of the chain's value.""" + a, b = got.float(), want.float() + assert (a.isnan() == b.isnan()).all(), f"{name}: NaN masks differ" + fin = ~a.isnan() + d = _ordered(got.view(torch.uint8), 0x80) - _ordered(want.view(torch.uint8), 0x80) + d = d.abs()[fin] + mism = (d > 0).float().mean() + assert mism <= FP8_MAX_MISMATCH_FRAC, f"{name}: {mism:.2%} E4M3 codes differ" + far = (d > 1) & ((a - b).abs()[fin] > FP8_SMALL_ABS) + assert not far.any(), f"{name}: {far.sum()} E4M3 codes off by > 1 code" + + +def _check_cache(name, got, want, slots, bs): + """NVFP4 cache [nb, (nslots,) bs, 72] vs the chain's: pages without an + inserted token are byte-identical; in the others, E2M1 codes and E4M3 + scale bytes are each within one code, and few differ. A scale one code + apart re-codes its whole group, so E2M1 mismatches only count in groups + with equal scales.""" + pages = torch.zeros(got.shape[0], dtype=torch.bool, device=got.device) + pages[slots[slots >= 0] // bs] = True + changed = (got != want).flatten(1).any(1) + assert not (changed & ~pages).any(), f"{name}: writes to pages with no token" + + def split(c): + """Each (page, slot) holds bs tokens' E2M1 data (8 groups of 8B), then + their E4M3 group scales (8B; V's swizzled in 4-token quads) -> codes + [rows, bs, 8, 16] and scales [rows, bs, 8].""" + c = c[pages].reshape(-1, bs * 72) + data = c[:, : bs * 64].reshape(-1, bs, 8, 8) + codes = _ordered(torch.stack((data & 0xF, data >> 4), -1), 0x8) + sf = c[:, bs * 64 :].reshape(-1, bs, 8).to(torch.int16) + if got.dim() == 4: # K/V cache: odd slots are V + # byte (t // 4 * 4 + g // 2) * 8 + g % 2 * 4 + t % 4 -> [t, g] + k, v = sf.reshape(-1, 2, bs, 8).unbind(1) + v = v.reshape(-1, bs // 4, 4, 2, 4).permute(0, 1, 4, 2, 3) + sf = torch.stack((k, v.reshape(-1, bs, 8)), 1).reshape(-1, bs, 8) + return codes.flatten(-2), sf + + (g, g_sf), (w, w_sf) = split(got), split(want) + d, d_sf = (g - w).abs(), (g_sf - w_sf).abs() + assert (d <= 1).all() and (d_sf <= 1).all(), f"{name}: codes off by > 1 code" + written = (g_sf != CACHE_FILL) | (w_sf != CACHE_FILL) + for what, mism, frac in ( + ("E2M1 codes", d[written & (d_sf == 0)] > 0, NVFP4_MAX_MISMATCH_FRAC), + ("scale bytes", d_sf[written] > 0, NVFP4_SF_MAX_MISMATCH_FRAC), + ): + if mism.numel() == 0: + continue + # 5 sigma of binomial slack: small M writes only a few groups + rate = mism.float().mean().item() + bound = frac + 5 * (frac / mism.numel()) ** 0.5 + assert rate <= bound, f"{name}: {rate:.2%} of written {what} differ" + + +@pytest.mark.parametrize( + "case", + [ + pytest.param(dict(tp=1, M=1), id="tp1-m1"), + pytest.param(dict(tp=1, M=37), id="tp1-m37"), + pytest.param(dict(tp=1, M=1000), id="tp1-m1000"), + pytest.param(dict(tp=1, M=4096, q="fp8"), id="tp1-m4096-fp8q"), + pytest.param(dict(tp=1, M=32768, q="fp8"), id="tp1-m32768-fp8q"), + pytest.param(dict(tp=2, M=1024), id="tp2-m1024"), + pytest.param(dict(tp=4, M=1024, q="bf16"), id="tp4-m1024-bf16q"), + pytest.param(dict(tp=4, M=4100, q="fp8"), id="tp4-m4100-fp8q"), + pytest.param(dict(tp=4, M=1024, kv_scales=(0.5, 1.7)), id="tp4-kv-scales"), + pytest.param(dict(tp=4, M=1024, slots="none"), id="tp4-no-insert"), + pytest.param(dict(tp=1, M=1024, layers=3), id="tp1-cross-layer"), + pytest.param(dict(tp=4, M=1024, block_size=64), id="tp4-bs64"), + # > 2**31 bytes: writers in the last page exercise int64 byte offsets + pytest.param( + dict(tp=1, M=128, nb=30000, slots="last_page"), id="tp1-big-cache" + ), + ], +) +def test_qkv_index_fused_epilogue(case): + import flashinfer + + from vllm.models.minimax_m3.nvidia.ops.qkv_index_fused_epilogue import ( + minimax_m3_qkv_index_fused_epilogue, + ) + + tp, M = case["tp"], case["M"] + q_mode = case.get("q", "both") + kv_scales = case.get("kv_scales", (1.0, 1.0)) + bs, layers = case.get("block_size", 128), case.get("layers", 1) + nb = case.get("nb", (M + M // 2 + bs - 1) // bs + 1) + nq, nkv, niq = TP_HEADS[tp] + + xq, xsf, wq, wsf, norm_weights, positions, cos_sin, sm, ism = _make_inputs( + nq, nkv, niq, M, nb, bs, case.get("slots", "random") + ) + + def caches(): + # the last layer of [nb, layers, *page] buffers: page-strided if layers > 1 + kv = torch.full((nb, layers, 2 * nkv, bs, 72), CACHE_FILL, dtype=torch.uint8) + ic = torch.full((nb, layers, bs, 72), CACHE_FILL, dtype=torch.uint8) + return kv.cuda()[:, -1], ic.cuda()[:, -1] + + # the production chain + kv_ref, ic_ref = caches() + q_ref = torch.empty(M, nq * HEAD_DIM, dtype=torch.bfloat16, device="cuda") + q8_ref = torch.empty(M, nq * HEAD_DIM, dtype=torch.float8_e4m3fn, device="cuda") + iq_ref = torch.empty(M, niq * HEAD_DIM, dtype=torch.float8_e4m3fn, device="cuda") + qkv = flashinfer.mm_mxfp8( + xq, wq.t(), xsf, wsf, out_dtype=torch.bfloat16, backend="cute-dsl" + ) + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, norm_weights[0], norm_weights[1], cos_sin, positions, nq, nkv, + ROTARY_DIM, EPS, norm_weights[2], norm_weights[3], niq, sm, ism, + kv_ref, ic_ref, bs, q_ref, iq_ref, "nvfp4", + q_fp8_out=q8_ref, q_fp8_scale=Q_FP8_SCALE, + kv_k_scale=torch.tensor([kv_scales[0]], device="cuda"), + kv_v_scale=torch.tensor([kv_scales[1]], device="cuda"), + ) # fmt: skip + + outs = [] + for cta_group in (1, 2): + kv, ic = caches() + q = torch.empty_like(q_ref) if q_mode != "fp8" else None + q8 = torch.empty_like(q8_ref) if q_mode != "bf16" else None + iq = torch.empty_like(iq_ref) + minimax_m3_qkv_index_fused_epilogue( + xq, xsf, wq, wsf, positions, cos_sin, norm_weights[0], + norm_weights[1], norm_weights[2], norm_weights[3], EPS, + slot_mapping=sm, index_slot_mapping=ism, kv_cache=kv, index_cache=ic, + q_out=q, index_q_out=iq, q_fp8_out=q8, q_fp8_scale=Q_FP8_SCALE, + k_scale=kv_scales[0], v_scale=kv_scales[1], cta_group=cta_group, + ) # fmt: skip + outs.append((q, q8, iq, kv, ic)) + + # both tile configs write identical bytes + for a, b in zip(*outs): + if a is not None: + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + + q, q8, iq, kv, ic = outs[0] + if q is not None: + torch.testing.assert_close( + q, q_ref, rtol=BF16_RTOL, atol=BF16_ATOL, equal_nan=True + ) + if q8 is not None: + _check_fp8("q_fp8", q8, q8_ref) + _check_fp8("index_q", iq, iq_ref) + _check_cache("kv_cache", kv, kv_ref, sm, bs) + _check_cache("index_cache", ic, ic_ref, ism, bs) + # the other layers of a cross-layer buffer stay untouched + for c in (kv, ic): + storage = torch.empty(0, dtype=torch.uint8, device="cuda") + storage.set_(c.untyped_storage()) + outside = torch.ones_like(storage, dtype=torch.bool) + outside.as_strided(c.size(), c.stride(), c.storage_offset()).fill_(False) + assert (storage[outside] == CACHE_FILL).all(), "writes outside the view" diff --git a/tests/models/multimodal/processing/test_minimax_m3.py b/tests/models/multimodal/processing/test_minimax_m3.py index 04d6aa4778c3..b24ea2b803c6 100644 --- a/tests/models/multimodal/processing/test_minimax_m3.py +++ b/tests/models/multimodal/processing/test_minimax_m3.py @@ -1,138 +1,160 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for MiniMax-M3 VL ``max_long_side_pixel`` resize support. +"""Unit tests for MiniMax-M3 VL preprocessing (sglang-aligned size spec). -These exercise the vendored processor directly (no checkpoint / GPU needed), so -they validate the long-side resize spec and the resulting prompt-token counts -deterministically. +These exercise the vendored processor directly (no checkpoint / GPU needed) and +pin the ``compute_aligned_size`` / ``resolve_max_long_side_pixel`` contract +against the sglang reference algorithm +(``python/sglang/srt/multimodal/processors/minimax_m3_vl.py``). """ import pytest import torch from vllm.transformers_utils.processors.minimax_m3 import ( - IMAGE_MAX_TOTAL_PIXELS, MIN_SHORT_SIDE_PIXEL, - VIDEO_MAX_TOTAL_PIXELS, MiniMaxM3VLImageProcessor, MiniMaxM3VLVideoProcessor, - smart_resize, + compute_aligned_size, + resolve_max_long_side_pixel, ) -# Long sides are multiples of patch_size*merge_size (28) so the rounding is -# exact and the expected token counts are unambiguous. -LONG_SIDES = [252, 504, 1008] +FACTOR = 28 # patch_size * merge_size +PATCH_SIZE = 14 MERGE2 = 2**2 # merge_size ** 2 +# detail -> max_long_side_pixel, per the size spec. +IMAGE_DETAIL_CAPS = {"low": 672, "auto": 2016, "high": 3584} +VIDEO_DETAIL_CAPS = {"low": 504, "auto": 672, "high": 2188} -def _image_tokens(grid_thw) -> int: - g = list(grid_thw) - return int(g[0] * g[1] * g[2]) // MERGE2 + +def _merged_tokens(grid_thw) -> int: + return int(grid_thw.prod()) // MERGE2 # --------------------------------------------------------------------------- # -# smart_resize: the long-side spec (a) shrink / (b) enlarge / (c) hard cap +# compute_aligned_size: floor layers / ceil-align / long-side cap # --------------------------------------------------------------------------- # -def test_smart_resize_long_side_shrink(): - # (a) long side exceeds the cap -> shrink so the long side equals the cap. - h, w = smart_resize( - 2048, 1024, factor=28, max_long_side_pixel=1008, max_total_pixels=10**9 - ) - assert max(h, w) == 1008 - assert (h, w) == (1008, 504) # aspect ratio preserved +def test_compute_aligned_size_enlarges_short_side(): + # (200, 40) hits both floors: area 8000 < 168**2 scales to (376, 76); the + # short-side floor then scales by 168/76, where ceil(76 * 168/76) = 169 + # (float overshoot) aligns up to 196, and 832 aligns up to 840. + h, w = compute_aligned_size(200, 40, FACTOR, max_long_side_pixel=1008) + assert (h, w) == (840, 196) + assert h % FACTOR == 0 and w % FACTOR == 0 -def test_smart_resize_short_side_enlarge(): - # (b) long side within the cap but short side below the floor -> enlarge so - # the short side reaches min_short_side_pixel. - h, w = smart_resize( - 200, 40, factor=28, max_long_side_pixel=1008, max_total_pixels=10**9 - ) - assert min(h, w) == MIN_SHORT_SIDE_PIXEL # 112 +def test_compute_aligned_size_shrinks_over_cap_isotropically(): + # Long side 2048 ceil-aligns to 2072 > cap -> integer aspect-ratio + # scale-back to (1008, 504), already factor-aligned. + h, w = compute_aligned_size(2048, 1024, FACTOR, max_long_side_pixel=1008) + assert (h, w) == (1008, 504) -def test_smart_resize_total_pixels_raises(): - # (c) still over the area cap after resizing -> raise instead of inferring. - with pytest.raises(ValueError, match="max_total_pixels"): - smart_resize( - 5000, - 5000, - factor=28, - max_long_side_pixel=4000, - max_total_pixels=IMAGE_MAX_TOTAL_PIXELS, - ) +def test_compute_aligned_size_rejects_extreme_aspect_ratio(): + with pytest.raises(ValueError, match="aspect ratio"): + compute_aligned_size(20001, 100, FACTOR, max_long_side_pixel=1008) -def test_smart_resize_backward_compatible_area_bound(): - # Without max_long_side_pixel the original Qwen-style area bound is used. - assert smart_resize(2048, 2048, factor=28, max_pixels=451584) == (672, 672) +# --------------------------------------------------------------------------- # +# resolve_max_long_side_pixel: detail tiers and explicit override +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize( + ("detail", "expected"), + [("low", 672), ("auto", 2016), ("high", 3584), (None, 2016)], +) +def test_resolve_max_long_side_pixel_image(detail, expected): + assert resolve_max_long_side_pixel(detail, None, modality="image") == expected + + +@pytest.mark.parametrize( + ("detail", "expected"), + [("low", 504), ("auto", 672), ("high", 2188), (None, 672)], +) +def test_resolve_max_long_side_pixel_video(detail, expected): + assert resolve_max_long_side_pixel(detail, None, modality="video") == expected + + +@pytest.mark.parametrize("modality", ["image", "video"]) +def test_resolve_max_long_side_pixel_explicit_overrides_detail(modality): + assert resolve_max_long_side_pixel("high", 1024, modality=modality) == 1024 + assert resolve_max_long_side_pixel(None, 1024, modality=modality) == 1024 + + +def test_processor_default_caps(): + assert MIN_SHORT_SIDE_PIXEL == 168 + assert MiniMaxM3VLImageProcessor().max_long_side_pixel == 2016 + assert MiniMaxM3VLVideoProcessor().max_long_side_pixel == 672 # --------------------------------------------------------------------------- # -# Image processor: monotonic prompt-token counts for 252 < 504 < 1008 +# Image processor end to end: grid/patch-counter consistency and token counts # --------------------------------------------------------------------------- # -def test_image_tokens_increase_with_max_long_side_pixel(): +@pytest.fixture(scope="module") +def square_image() -> torch.Tensor: + # 4096x4096 so detail="high" actually reaches the 3584 cap (a 2048x2048 + # image only ceil-aligns to 2072 and never upscales). + g = torch.Generator().manual_seed(0) + return torch.randint(0, 255, (3, 4096, 4096), dtype=torch.uint8, generator=g) + + +@pytest.mark.parametrize("detail", ["low", "auto", "high"]) +def test_image_preprocess_grid_matches_patch_counter(detail, square_image): proc = MiniMaxM3VLImageProcessor() - counts = [] - for long_side in LONG_SIDES: - patches = proc.get_number_of_image_patches( - 2048, 2048, images_kwargs={"max_long_side_pixel": long_side} - ) - counts.append(patches // MERGE2) + cap = IMAGE_DETAIL_CAPS[detail] - assert counts == [81, 324, 1296] - assert counts[0] < counts[1] < counts[2] + out = proc.preprocess([square_image], detail=detail, return_tensors="pt") + grid_thw = out["image_grid_thw"][0] + patches = proc.get_number_of_image_patches( + 4096, 4096, images_kwargs={"max_long_side_pixel": cap} + ) + assert int(grid_thw.prod()) == patches -def test_image_processor_defaults_match_spec(): - proc = MiniMaxM3VLImageProcessor() - assert proc.max_long_side_pixel is None # opt-in - assert proc.min_short_side_pixel == MIN_SHORT_SIDE_PIXEL - assert proc.max_total_pixels == IMAGE_MAX_TOTAL_PIXELS + # Cross-check against compute_aligned_size: the square long side lands + # exactly on the cap, giving (cap / patch_size) ** 2 patches. + assert compute_aligned_size(4096, 4096, FACTOR, max_long_side_pixel=cap) == ( + cap, + cap, + ) + assert patches == (cap // PATCH_SIZE) ** 2 -def test_image_preprocess_pipeline_monotonic(): +def test_image_tokens_increase_with_detail(square_image): proc = MiniMaxM3VLImageProcessor() - image = torch.randint(0, 255, (3, 2048, 2048), dtype=torch.uint8) - counts = [] - for long_side in LONG_SIDES: - out = proc.preprocess( - [image], - do_resize=True, - max_long_side_pixel=long_side, - return_tensors="pt", + counts = [ + _merged_tokens( + proc.preprocess([square_image], detail=detail, return_tensors="pt")[ + "image_grid_thw" + ][0] ) - counts.append(_image_tokens(out["image_grid_thw"][0])) - assert counts == [81, 324, 1296] + for detail in ("low", "auto", "high") + ] + # (cap / patch_size) ** 2 / merge_size**2 for caps 672 / 2016 / 3584. + assert counts == [576, 5184, 16384] + assert counts[0] < counts[1] < counts[2] -# --------------------------------------------------------------------------- # -# Video processor: same monotonic behavior + volumetric (w*h*frames) cap -# --------------------------------------------------------------------------- # -def test_video_tokens_increase_with_max_long_side_pixel(): - proc = MiniMaxM3VLVideoProcessor() - assert proc.max_total_pixels == VIDEO_MAX_TOTAL_PIXELS - video = torch.randint(0, 255, (4, 3, 2048, 2048), dtype=torch.uint8) - counts = [] - for long_side in LONG_SIDES: - out = proc.preprocess( - videos=[video], - do_resize=True, - max_long_side_pixel=long_side, - return_tensors="pt", - ) - counts.append(_image_tokens(out["video_grid_thw"][0])) - assert counts[0] < counts[1] < counts[2] +def test_image_preprocess_explicit_cap_overrides_detail(square_image): + proc = MiniMaxM3VLImageProcessor() + out = proc.preprocess( + [square_image], + detail="high", + max_long_side_pixel=672, + return_tensors="pt", + ) + # The explicit cap wins over detail="high" -> same grid as detail="low". + assert _merged_tokens(out["image_grid_thw"][0]) == 576 -def test_video_volumetric_cap_raises(): +# --------------------------------------------------------------------------- # +# Video processor: volumetric (frames * h * w) cap raises unconditionally +# --------------------------------------------------------------------------- # +def test_video_volume_cap_raises_by_default(): proc = MiniMaxM3VLVideoProcessor() - # 400 frames at a 1008-long-side square: 1008*1008*400 >> 301,056,000. - video = torch.randint(0, 255, (400, 3, 2048, 2048), dtype=torch.uint8) + # Default cap 672: 700 frames * 672 * 672 = 316,108,800 > 301,056,000. + # The volumetric cap raises even without any detail/max_long_side_pixel + # override. + video = torch.randint(0, 255, (700, 3, 672, 672), dtype=torch.uint8) with pytest.raises(ValueError, match="max_total_pixels"): - proc.preprocess( - videos=[video], - do_resize=True, - max_long_side_pixel=1008, - return_tensors="pt", - ) + proc.preprocess(videos=[video], return_tensors="pt") diff --git a/tests/models/registry.py b/tests/models/registry.py index 2036c33df95b..c9ff9648cdf5 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1579,6 +1579,11 @@ def check_available_online( is_available_online=False, use_original_num_layers=True, # DSpark has >1 draft block ), + "DSparkMiniMaxDraftModel": _HfExamplesInfo( + "MiniMaxAI/MiniMax-M3", + is_available_online=False, + use_original_num_layers=True, + ), "K3DSparkModel": _HfExamplesInfo( "moonshotai/Kimi-K3", speculative_model="Inferact/Kimi-K3-DSpark", diff --git a/tests/models/test_minimax_m3_mega_moe.py b/tests/models/test_minimax_m3_mega_moe.py new file mode 100644 index 000000000000..e4b834e3f832 --- /dev/null +++ b/tests/models/test_minimax_m3_mega_moe.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3.1 on the DeepGEMM NVFP4 MegaMoE layer: gating, the checkpoint +adapter, and the fused block against a dequantized reference.""" + +import gc +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + DeepGemmMegaMoEExperts, +) +from vllm.models.minimax_m3.nvidia.model import ( + MEGA_MOE_FC2_INPUT_SCALE, + _mega_moe_expert_params_mapping, + _mega_moe_expert_weight, + _use_mega_moe, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="NVFP4 MegaMoE requires CUDA" +) + + +def _vllm_config(moe_backend: str, moe_quant_algo: str | None): + return SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend=moe_backend), + quant_config=SimpleNamespace(moe_quant_algo=moe_quant_algo), + parallel_config=SimpleNamespace(enable_expert_parallel=True), + ) + + +@pytest.mark.parametrize( + ("moe_backend", "moe_quant_algo", "expected"), + [ + ("auto", "NVFP4", False), + ("deep_gemm_mega_moe", "NVFP4", True), + ("deep_gemm_mega_moe", None, NotImplementedError), + ], +) +def test_mega_moe_is_gated_on_nvfp4_checkpoints(moe_backend, moe_quant_algo, expected): + """MegaMoE is opt-in and refuses checkpoints whose experts are not NVFP4.""" + vllm_config = _vllm_config(moe_backend, moe_quant_algo) + if expected is NotImplementedError: + with pytest.raises(NotImplementedError, match="NVFP4 routed experts"): + _use_mega_moe(vllm_config) + else: + assert _use_mega_moe(vllm_config) is expected + + +def test_mega_moe_expert_weight_follows_the_qat_recipe(): + """Compressed-tensors quantization scales become ModelOpt dequantization + scales; the FC1 input scale is dropped (per-token) and the FC2 one is the + fixed QAT scale whatever the checkpoint holds.""" + prefix = "layers.3.block_sparse_moe.experts.7" + packed = torch.randint(0, 255, (4, 8), dtype=torch.uint8) + block_scale = torch.rand(4, 1).to(torch.float8_e4m3fn) + + assert _mega_moe_expert_weight(f"{prefix}.w1.weight_packed", packed) == ( + f"{prefix}.w1.weight", + packed, + ) + assert _mega_moe_expert_weight(f"{prefix}.w3.weight_scale", block_scale) == ( + f"{prefix}.w3.weight_scale", + block_scale, + ) + name, scale = _mega_moe_expert_weight( + f"{prefix}.w2.weight_global_scale", torch.tensor(448.0) + ) + assert name == f"{prefix}.w2.weight_scale_2" + assert torch.equal(scale, torch.tensor(1 / 448, dtype=torch.float32)) + for shard_id in ("w1", "w3"): + assert ( + _mega_moe_expert_weight( + f"{prefix}.{shard_id}.input_global_scale", torch.tensor(1.0) + ) + is None + ) + name, scale = _mega_moe_expert_weight( + f"{prefix}.w2.input_global_scale", torch.tensor(1.0) + ) + assert name == f"{prefix}.w2.input_scale" + assert scale.item() == MEGA_MOE_FC2_INPUT_SCALE + + +def _load_checkpoint(moe, tensors: dict[str, torch.Tensor]) -> None: + """``MiniMaxM3Model.load_weights`` for one MegaMoE block's tensors.""" + experts = moe.experts + mapping = _mega_moe_expert_params_mapping(experts.num_experts) + for name, tensor in tensors.items(): + if ".shared_experts." in name: + name = name.replace(".gate_proj.", ".w1.").replace(".up_proj.", ".w3.") + experts.load_shared_expert_weight(name, tensor) + continue + converted = _mega_moe_expert_weight(name, tensor) + if converted is None: + continue + name, tensor = converted + param_name, weight_name, expert_id, shard_id = next( + entry for entry in mapping if entry[1] in name + ) + name = name.replace(weight_name, param_name) + param = getattr(experts, name.rpartition(".")[2]) + param.weight_loader(param, tensor, name, shard_id=shard_id, expert_id=expert_id) + + +def _mega_moe_block(num_experts=8, top_k=2, hidden_size=512, intermediate_size=256): + """An SM100 MiniMax-M3.1 MoE block on MegaMoE with an MXFP8 shared expert.""" + if not current_platform.is_device_capability_family(100): + pytest.skip("DeepGEMM NVFP4 MegaMoE requires SM100") + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.config.kernel import KernelConfig + from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config + from vllm.models.minimax_m3.nvidia.model import MiniMaxM3MoE + + config = SimpleNamespace( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=num_experts, + num_experts_per_tok=top_k, + n_shared_experts=1, + scoring_func="sigmoid", + routed_scaling_factor=2.0, + use_routing_bias=False, + hidden_act="swigluoai", + swiglu_limit=7.0, + swiglu_alpha=1.702, + swiglu_beta=1.0, + ) + quant_config = ModelOptMxFp8Config( + is_checkpoint_mxfp8_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + moe_quant_algo="NVFP4", + ) + vllm_config = VllmConfig( + kernel_config=KernelConfig(moe_backend="deep_gemm_mega_moe") + ) + vllm_config.scheduler_config.max_num_batched_tokens = 256 + with set_current_vllm_config(vllm_config), torch.device("cuda"): + moe = MiniMaxM3MoE( + config, + layer_id=0, + quant_config=quant_config, + prefix="layers.0.block_sparse_moe", + use_mega_moe=True, + ) + assert isinstance(moe.experts, DeepGemmMegaMoEExperts) + return moe, config + + +@pytest.mark.parametrize("fused", [True, False]) +def test_mega_moe_allocates_the_shared_expert_once(dist_init, monkeypatch, fused): + """The block builds a real shared MLP only when MegaMoE does not fuse it + (e.g. an older DeepGEMM); a fused layer owns the only copy, and building + the block leaves no discarded weights for the garbage collector.""" + if not fused: + monkeypatch.setattr( + DeepGemmMegaMoEExperts, + "_kernel_fuses_shared_experts", + lambda self, spec: False, + ) + gc.collect() + gc.disable() + try: + moe, _ = _mega_moe_block() + allocated = torch.accelerator.memory_allocated() + gc.collect() + finally: + gc.enable() + assert torch.accelerator.memory_allocated() == allocated + assert moe.experts.has_fused_shared_experts is fused + assert (moe.shared_experts is None) is fused + assert all(param.is_cuda for param in moe.parameters()) + + +@pytest.mark.parametrize("num_tokens", [1, 45, 200]) +def test_mega_moe_block_matches_the_qat_reference(dist_init, num_tokens): + """A MiniMax-M3.1 MoE block on MegaMoE fuses its MXFP8 shared expert and, + loaded from a compressed-tensors checkpoint, reproduces the QAT recipe: + per-token FC1 input scales and the fixed FC2 input scale on the unweighted + SwiGLU output. Routing weights scale the BF16 expert outputs in combine.""" + from vllm.third_party.deep_gemm.utils import ( + cast_back_from_fp8, + per_token_cast_to_fp8, + ) + from vllm.third_party.deep_gemm.utils.nvfp4 import ( + cast_back_from_nvfp4, + per_token_cast_to_nvfp4, + ) + + torch.manual_seed(0) + moe, config = _mega_moe_block() + assert moe.experts.has_fused_shared_experts and moe.shared_experts is None + num_experts, top_k = config.num_local_experts, config.num_experts_per_tok + hidden_size, intermediate_size = config.hidden_size, config.intermediate_size + clamp, alpha, beta = config.swiglu_limit, config.swiglu_alpha, config.swiglu_beta + + def nvfp4(n, k): + weight = torch.randn(n, k, device="cuda") * k**-0.5 + global_scale = weight.abs().max().item() / (6 * 448) + packed, sf = per_token_cast_to_nvfp4(weight, global_scale) + dequant = cast_back_from_nvfp4(packed, sf, global_scale) + return ( + packed.view(torch.uint8), + sf.view(torch.float8_e4m3fn), + torch.tensor(1 / global_scale), + dequant, + ) + + def mxfp8(n, k): + weight = torch.randn(n, k, device="cuda") * k**-0.5 + fp8, sf = per_token_cast_to_fp8(weight, use_ue8m0=True, gran_k=32) + ue8m0 = (sf.view(torch.int32) >> 23).to(torch.uint8) + return fp8, ue8m0, cast_back_from_fp8(fp8, sf, gran_k=32) + + prefix = "layers.0.block_sparse_moe" + checkpoint: dict[str, torch.Tensor] = {} + ref_w13, ref_w2 = [], [] + for e in range(num_experts): + w13 = nvfp4(2 * intermediate_size, hidden_size) + w2 = nvfp4(hidden_size, intermediate_size) + for shard_id, rows, (packed, sf, global_scale, _) in ( + ("w1", slice(0, intermediate_size), w13), + ("w3", slice(intermediate_size, None), w13), + ("w2", slice(None), w2), + ): + name = f"{prefix}.experts.{e}.{shard_id}" + checkpoint[f"{name}.weight_packed"] = packed[rows] + checkpoint[f"{name}.weight_scale"] = sf[rows] + checkpoint[f"{name}.weight_global_scale"] = global_scale + # A placeholder the block must ignore. + checkpoint[f"{name}.input_global_scale"] = torch.tensor(0.25) + ref_w13.append(w13[3]) + ref_w2.append(w2[3]) + shared_ref = {} + for proj, (n, k) in ( + ("gate_proj", (intermediate_size, hidden_size)), + ("up_proj", (intermediate_size, hidden_size)), + ("down_proj", (hidden_size, intermediate_size)), + ): + fp8, ue8m0, shared_ref[proj] = mxfp8(n, k) + checkpoint[f"{prefix}.shared_experts.{proj}.weight"] = fp8 + checkpoint[f"{prefix}.shared_experts.{proj}.weight_scale"] = ue8m0 + _load_checkpoint(moe, checkpoint) + + # Rows span six decades, which one tensor-wide input scale cannot cover. + x = ( + torch.randn(num_tokens, hidden_size, device="cuda") + * torch.logspace(-4, 2, num_tokens, device="cuda")[:, None] + ).to(torch.bfloat16) + scores = torch.randn(num_tokens, num_experts, device="cuda").sigmoid() + topk_weights, topk_ids = scores.topk(top_k, dim=-1) + topk_weights *= config.routed_scaling_factor / topk_weights.sum(-1, keepdim=True) + out = moe.experts(x, topk_weights, topk_ids) + + def activation(h): + gate, up = h.to(torch.bfloat16).float().chunk(2, dim=-1) + gate = gate.clamp(max=clamp) + up = up.clamp(-clamp, clamp) + return gate * torch.sigmoid(alpha * gate) * (up + beta) + + def fp8_round_trip(h): + return cast_back_from_fp8( + *per_token_cast_to_fp8(h, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True), + gran_k=32, + use_packed_ue8m0=True, + ) + + x_scales = x.float().abs().amax(-1, keepdim=True).clamp_min(1e-12) / 2688 + x_ref = cast_back_from_nvfp4(*per_token_cast_to_nvfp4(x.float() / x_scales)) + x_ref *= x_scales + ref = torch.zeros(num_tokens, hidden_size, device="cuda") + for e in range(num_experts): + rows, slots = torch.where(topk_ids == e) + if rows.numel() == 0: + continue + mid = activation(x_ref[rows] @ ref_w13[e].T) + mid = cast_back_from_nvfp4( + *per_token_cast_to_nvfp4(mid, MEGA_MOE_FC2_INPUT_SCALE), + MEGA_MOE_FC2_INPUT_SCALE, + ) + expert_out = (mid @ ref_w2[e].T).to(torch.bfloat16).float() + ref.index_add_(0, rows, topk_weights[rows, slots, None] * expert_out) + shared_w13 = torch.cat([shared_ref["gate_proj"], shared_ref["up_proj"]]) + shared_mid = fp8_round_trip(activation(fp8_round_trip(x) @ shared_w13.T)) + ref += (shared_mid @ shared_ref["down_proj"].T).to(torch.bfloat16).float() + + assert torch.isfinite(out).all() + # Per row: a missing per-token scale would only show in the small rows. + row_error = (out.float() - ref).norm(dim=-1) / ref.norm(dim=-1) + assert row_error.max() < 0.02, row_error diff --git a/tests/models/test_qwen3_dflash_mxfp8_context_kv.py b/tests/models/test_qwen3_dflash_mxfp8_context_kv.py new file mode 100644 index 000000000000..623c9ddb1ae4 --- /dev/null +++ b/tests/models/test_qwen3_dflash_mxfp8_context_kv.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DFlash/DSpark context K/V projection with MXFP8 QKV weights.""" + +import pytest +import torch +from torch import nn + +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + _mxfp8_e4m3_quantize_torch, +) +from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda() + or not current_platform.is_device_capability_family(100), + reason="Native MXFP8 kernels require SM100-family CUDA", +) + +HIDDEN, HEAD_DIM, NUM_HEADS, NUM_LAYERS, NUM_TOKENS = 256, 128, 2, 2, 7 +Q_SIZE, KV_SIZE = NUM_HEADS * HEAD_DIM, HEAD_DIM # one KV head + + +@pytest.fixture +def bf16_on_cuda(default_vllm_config, dist_init): + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + with torch.device("cuda"): + yield + torch.set_default_dtype(old_dtype) + + +def _build_model(mode: str, has_bias: bool) -> tuple[DFlashQwen3Model, list]: + # MiniMax-style MXFP8 config; "mixed" leaves layer 1 in BF16. + quant_config = ( + None + if mode == "bf16" + else ModelOptMxFp8Config.from_config( + { + "quant_method": "mxfp8", + "activation_scheme": "dynamic", + "weight_block_size": [1, 32], + "ignored_layers": ( + ["layers.1.self_attn.qkv_proj"] if mode == "mixed" else [] + ), + } + ) + ) + model = DFlashQwen3Model.__new__(DFlashQwen3Model) + nn.Module.__init__(model) + model.hidden_norm = RMSNorm(HIDDEN) + model._rms_norm_eps = 1e-6 + attentions = [] + for i in range(NUM_LAYERS): + attn = nn.Module() + attn.q_size, attn.kv_size = Q_SIZE, KV_SIZE + attn.k_norm = RMSNorm(HEAD_DIM) + attn.qkv_proj = QKVParallelLinear( + HIDDEN, + HEAD_DIM, + NUM_HEADS, + 1, + bias=has_bias, + quant_config=quant_config, + prefix=f"layers.{i}.self_attn.qkv_proj", + disable_tp=True, + ) + weight = torch.randn(Q_SIZE + 2 * KV_SIZE, HIDDEN) / 16 + proj = attn.qkv_proj + if isinstance(proj.quant_method, UnquantizedLinearMethod): + proj.weight.data.copy_(weight) + else: + weight, scale = _mxfp8_e4m3_quantize_torch(weight) + proj.weight.data.copy_(weight) + proj.weight_scale.data.copy_(scale) + if has_bias: + proj.bias.data.normal_(std=0.1) + attentions.append(attn) + return model, attentions + + +@pytest.mark.parametrize("mode", ["bf16", "mxfp8", "mixed"]) +@pytest.mark.parametrize("has_bias", [False, True]) +@torch.inference_mode() +def test_context_kv_matches_qkv_projection(bf16_on_cuda, mode, has_bias): + """The context K/V must equal the K/V slice of each layer's own QKV output + after kernel repacking, eagerly and under CUDA graph replay.""" + torch.manual_seed(0) + model, attentions = _build_model(mode, has_bias) + # Load-time order: context buffers are built before kernels repack weights. + model._build_context_kv_buffers(attentions, has_bias=has_bias) + assert (model._fused_kv_weight is not None) == (mode == "bf16") + assert (model._context_kv_proj is not None) == (mode == "mxfp8") + for module in [*model.modules(), *(m for a in attentions for m in a.modules())]: + method = getattr(module, "quant_method", None) + if method is not None: + method.process_weights_after_loading(module) + + x = torch.randn(NUM_TOKENS, HIDDEN) + normed = model.hidden_norm(x) + expected = torch.cat( + [a.qkv_proj(normed)[0][:, Q_SIZE:] for a in attentions], dim=-1 + ) + + def project() -> torch.Tensor: + k, v = model._project_context_kv(x, NUM_TOKENS, NUM_LAYERS, 1, HEAD_DIM) + return torch.stack([k, v]).permute(2, 1, 0, 3, 4).reshape(NUM_TOKENS, -1) + + torch.testing.assert_close(project(), expected, rtol=0.02, atol=0.02) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = project() + x.normal_() + graph.replay() + torch.testing.assert_close(captured, project(), rtol=0, atol=0) diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 16334fdb0dc7..bfb2756b7006 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -1287,12 +1287,13 @@ def unknown_fps(data) -> VideoSourceMetadata: @pytest.mark.parametrize( - "model_repo, expected_loader_cls, hf_sample_kwargs", + "model_repo, expected_loader_cls, hf_sample_kwargs, appends_last_frame", [ pytest.param( "allenai/Molmo2-4B", Molmo2VideoBackend, None, + False, marks=pytest.mark.skip( reason="Video processor not aligned, investigate later.", ), @@ -1302,18 +1303,21 @@ def unknown_fps(data) -> VideoSourceMetadata: "zai-org/GLM-4.1V-9B-Thinking", DynamicVideoBackend, None, + False, id="glm4v", ), pytest.param( "zai-org/GLM-4.6V-Flash", GLM46VVideoBackend, None, + False, id="glm46v", ), pytest.param( "Qwen/Qwen3-VL-4B-Instruct", Qwen3VLVideoBackend, None, + False, id="qwen3vl", ), # Qwen2-VL/Qwen2.5-VL ship no ``video_processor_type`` in their @@ -1325,18 +1329,23 @@ def unknown_fps(data) -> VideoSourceMetadata: "Qwen/Qwen2-VL-7B-Instruct", Qwen2VLVideoBackend, {"fps": 2}, + False, id="qwen2vl", ), pytest.param( "Qwen/Qwen2.5-VL-7B-Instruct", Qwen2VLVideoBackend, {"fps": 2}, + False, id="qwen2_5_vl", ), pytest.param( "MiniMaxAI/MiniMax-M3", MiniMaxM3VideoBackend, None, + # The MiniMax-M3 sampler also keeps the final frame, which the + # Hub ``sample_frames`` drops. + True, id="minimax_m3_vl", ), ], @@ -1345,6 +1354,7 @@ def test_video_processor_from_model_repo( model_repo: str, expected_loader_cls: type, hf_sample_kwargs: dict[str, int | float] | None, + appends_last_frame: bool, ): """Test that a model repo resolves to the correct video loader backend. @@ -1390,6 +1400,9 @@ def test_video_processor_from_model_repo( duration=vllm_meta["duration"], ) hf_indices = processor.sample_frames(hf_metadata, **(hf_sample_kwargs or {})) + last_frame = vllm_meta["total_num_frames"] - 1 + if appends_last_frame and hf_indices[-1] != last_frame: + hf_indices = np.append(hf_indices, last_frame) vllm_indices = np.array(vllm_meta["frames_indices"]) np.testing.assert_array_equal( hf_indices, diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index d489cb2b9acf..d22e030f65f9 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -154,6 +154,80 @@ def test_modelopt_mxfp8_preserves_per_row_checkpoint_scales(dist_init, monkeypat assert torch.equal(linear.weight_scale, scales) +_M31_QUANT_CONFIG = { + "quant_method": "mxfp8", + "ignored_layers": ["lm_head", "model.layers.3.block_sparse_moe.gate"], + "group_size": 16, + "moe_quant_algo": "NVFP4", + "moe_quant_format": "nvfp4-pack-quantized", +} + + +def test_modelopt_mxfp8_routes_nvfp4_experts_to_compressed_tensors(monkeypatch): + """MiniMax-M3.1 keeps MXFP8 linears but ships compressed-tensors NVFP4 + routed experts; loading them with the MXFP8 MoE method drops weights.""" + ct_method = Mock() + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.compressed_tensors." + "compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4." + "CompressedTensorsW4A4Nvfp4MoEMethod", + ct_method, + ) + config = ModelOptMxFp8Config.from_config(dict(_M31_QUANT_CONFIG)) + layer = Mock(spec=RoutedExperts) + layer.moe_config = Mock() + prefix = "model.layers.3.block_sparse_moe.experts" + + # MiniMax-M3.1 declares that it QAT'd the FC1 input per token, leaving a + # neutral 1.0 input_global_scale. A model that declares nothing keeps the + # checkpoint's per-tensor scale. + for declared, expected in ( + (kNvfp4DynamicToken, kNvfp4DynamicToken), + (None, kNvfp4Dynamic), + ): + ct_method.reset_mock() + layer.moe_config.qat_activation_key = declared + assert config.get_quant_method(layer, prefix) is ct_method.return_value + ct_method.assert_called_once_with( + layer.moe_config, layer_name=prefix, activation_key=expected + ) + + plain = ModelOptMxFp8Config.from_config({"quant_method": "mxfp8"}) + assert plain.moe_quant_algo is None + + +def test_per_token_nvfp4_moe_backends_match_the_experts_that_implement_it(): + """``make_nvfp4_moe_kernel`` dispatches unquantized inputs for a per-token + scheme and only forwards ``per_token_activation`` to this set, so it has to + be exactly the experts that take that argument. ``_supports_quant_scheme`` + cannot be used for the check: several NVFP4 backends match on the weight + key alone and accept any activation key.""" + import inspect + + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + PER_TOKEN_NVFP4_MOE_BACKENDS, + NvFp4MoeBackend, + backend_to_kernel_cls, + ) + + implementing = { + backend + for backend in NvFp4MoeBackend + for cls in backend_to_kernel_cls(backend) + if "per_token_activation" in inspect.signature(cls.__init__).parameters + } + assert implementing == PER_TOKEN_NVFP4_MOE_BACKENDS + + +@pytest.mark.parametrize( + "override", + [{"moe_quant_algo": "MXFP4"}, {"moe_quant_format": "float-quantized"}], +) +def test_modelopt_mxfp8_rejects_unknown_expert_format(override): + with pytest.raises(ValueError, match="MoE quantization"): + ModelOptMxFp8Config.from_config({**_M31_QUANT_CONFIG, **override}) + + def test_modelopt_mxfp8_pre_processed_weights_follow_kernel(monkeypatch): """The weight cache daemon can only serve MXFP8 layers whose kernel needs no post-load state beyond the parameters it exports.""" diff --git a/tests/reasoning/test_minimax_m3_reasoning_parser.py b/tests/reasoning/test_minimax_m3_reasoning_parser.py index 561e4020cb1f..3769b481f0b5 100644 --- a/tests/reasoning/test_minimax_m3_reasoning_parser.py +++ b/tests/reasoning/test_minimax_m3_reasoning_parser.py @@ -16,7 +16,7 @@ class MiniMaxM3Tokenizer: """Small tokenizer with MiniMax M3 reasoning tags as special tokens.""" - special_tokens = ("", "") + special_tokens = ("", "", "]~b]") def __init__(self): self._token_to_id: dict[str, int] = {} @@ -459,3 +459,38 @@ def test_token_id_helpers_enabled_mode(): assert parser.count_reasoning_tokens(open_reasoning_ids) == len( tokenizer.encode("abc") ) + + +# Rendered like the MiniMax-M3.1 chat template: the system message spells out +# both markers, and a prior assistant turn without reasoning is prefixed with +# the closer. +PROMPT = ( + "]~b]system\nWrap your reasoning in tags, or begin " + "directly after the prefix.[e~[\n" + "]~b]user\nCapital of France?[e~[\n" + "]~b]ai\nParis.[e~[\n" + "]~b]user\nAnd Germany?[e~[\n" + "]~b]ai\n" +) + + +def test_prompt_markers_do_not_end_reasoning(): + parser, tokenizer = make_parser() + prompt_ids = tokenizer.encode(PROMPT) + + assert not parser.is_reasoning_end(prompt_ids) + assert not parser.is_reasoning_end(tokenizer.encode(PROMPT + "ab")) + assert parser.is_reasoning_end(tokenizer.encode(PROMPT + "a")) + assert parser.is_reasoning_end(tokenizer.encode(PROMPT + "Berlin")) + + +def test_streaming_end_check_ignores_prompt_markers(): + parser, tokenizer = make_parser() + prompt_ids = tokenizer.encode(PROMPT) + open_ids = tokenizer.encode("ab") + closed_ids = tokenizer.encode("") + + assert not parser.is_reasoning_end_streaming(prompt_ids + open_ids, open_ids[-1:]) + assert parser.is_reasoning_end_streaming( + prompt_ids + open_ids + closed_ids, closed_ids + ) diff --git a/tests/v1/attention/test_flashinfer_ts.py b/tests/v1/attention/test_flashinfer_ts.py new file mode 100644 index 000000000000..f77186d4aa03 --- /dev/null +++ b/tests/v1/attention/test_flashinfer_ts.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FLASHINFER_TS: the non-causal draft-block decode path. + +Covers the two things that are easy to get wrong and silently produce plausible +output: the K/V plane layout the kernel validates, and the page-aligned window +restriction that stands in for a window the kernel cannot take with a dense +mask. +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.flashinfer_ts import ( + FlashInferTSMetadataBuilder, +) + +HEAD_DIM = 128 +PAGE = 128 +WINDOW = 4096 + + +def _sm100_family() -> bool: + cap = current_platform.get_device_capability() + return cap is not None and (cap.major, cap.minor) in ((10, 0), (10, 3), (10, 7)) + + +requires_kernel = pytest.mark.skipif( + not current_platform.is_cuda() or not _sm100_family(), + reason="prims_ts decode runs on SM100a/SM103a/SM107a only", +) + + +@pytest.mark.parametrize("num_kv_heads", [1, 4]) +def test_reads_vllms_existing_interleaved_page(num_kv_heads): + """No layout change: K/V are the halves of vLLM's [pages, Hkv, page, 2*D] + page, whose token stride is 2*D. The kernel takes its TMA strides from the + tensor, so this must stay a pure view with head_dim contiguous. + """ + page = torch.zeros(8, num_kv_heads, PAGE, 2 * HEAD_DIM, dtype=torch.uint8) + k, v = page[..., :HEAD_DIM], page[..., HEAD_DIM:] + + assert k.shape == (8, num_kv_heads, PAGE, HEAD_DIM) + assert k.stride(-1) == 1 and v.stride(-1) == 1 # the one hard requirement + assert k.stride(2) == 2 * HEAD_DIM # interleaved, not compact + assert k.data_ptr() == page.data_ptr() # views, never a repack + assert v.data_ptr() == page.data_ptr() + HEAD_DIM + + +@pytest.mark.parametrize("query_len", [1, 3, 7]) +def test_window_restriction_is_a_superset_of_the_true_window(query_len): + """The restricted page table must never drop a token the symmetric window + can reach, and must overshoot by less than one page. + """ + builder = FlashInferTSMetadataBuilder.__new__(FlashInferTSMetadataBuilder) + builder.sliding_window = WINDOW + builder.page_size = PAGE + builder.query_len = query_len + # The builder writes metadata into persistent buffers so a replayed CUDA + # graph keeps reading the same addresses; the window math is exercised + # through them, not around them. + pages = builder._window_pages(query_len) + builder._buf_seq_lens = torch.zeros(6, dtype=torch.int32) + builder._buf_block_table = torch.zeros((6, pages), dtype=torch.int32) + + seq_lens = torch.tensor( + [1, PAGE, WINDOW - 1, WINDOW + query_len, 8192, 120_000], dtype=torch.int32 + ) + block_table = torch.arange( + seq_lens.numel() * (int(seq_lens.max()) // PAGE + 1), dtype=torch.int32 + ).view(seq_lens.numel(), -1) + + new_lens, new_bt = FlashInferTSMetadataBuilder._restrict_to_window( + builder, seq_lens, block_table, query_len + ) + + span = WINDOW + query_len - 1 + reachable = torch.minimum(seq_lens, torch.full_like(seq_lens, span)) + assert torch.all(new_lens >= reachable), "window restriction dropped live tokens" + assert int((new_lens - reachable).max()) < PAGE, "overshoot exceeds one page" + # The kept pages must be the tail of each row, so the last token of the + # sequence is still addressable. + last_page = (seq_lens.long() - 1) // PAGE + assert torch.all(new_bt.long().max(dim=1).values >= last_page) + + +def test_impl_accepts_the_arguments_the_attention_layer_passes(): + """Attention() calls impl_cls positionally and forwards extra_impl_args + (``sinks`` among them); a signature that only matches AttentionImpl's + abstract __init__ raises TypeError at engine init, not at import. + """ + from vllm.v1.attention.backends.flashinfer_ts import FlashInferTSImpl + + impl = FlashInferTSImpl( + 16, + HEAD_DIM, + HEAD_DIM**-0.5, + 1, + None, + WINDOW, + "fp8", + None, + AttentionType.DECODER, + None, + sinks=None, + ) + # FP8 KV reads an FP8 query; the layer's own quant is skipped via + # query_pre_quantized because the fused QK-norm/RoPE kernel emits it. + assert impl.supports_quant_query_input + + with pytest.raises(NotImplementedError): + FlashInferTSImpl( + 16, + HEAD_DIM, + HEAD_DIM**-0.5, + 1, + None, + WINDOW, + "fp8", + None, + AttentionType.DECODER, + None, + sinks=torch.zeros(16), + ) + + +@requires_kernel +def test_dense_plan_matches_flash_attention_non_causal(): + """A dense prims_ts plan is the non-causal draft block FA4 produces.""" + from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd + + pytest.importorskip("flashinfer.attention.prims_ts") + from flashinfer.attention.prims_ts import BatchDecodePagedTSWrapper + + torch.manual_seed(0) + b, sq, hq, hkv, kv = 2, 3, 16, 1, PAGE + dev = torch.device("cuda") + scale = HEAD_DIM**-0.5 + pages = b * (kv // PAGE) + 1 + + page = torch.randn(pages, 2 * hkv, PAGE, HEAD_DIM, dtype=torch.bfloat16, device=dev) + k_plane, v_plane = page[:, :hkv].contiguous(), page[:, hkv:].contiguous() + interleaved = torch.cat([k_plane, v_plane], dim=-1).contiguous() + k_fa, v_fa = interleaved.transpose(1, 2).split(HEAD_DIM, dim=-1) + + q = torch.randn(b * sq, hq, HEAD_DIM, dtype=torch.bfloat16, device=dev) + bt = torch.arange(1, pages, dtype=torch.int32, device=dev).view(b, -1).contiguous() + lens = torch.full((b,), kv, dtype=torch.int32, device=dev) + cu_q = torch.arange(0, (b + 1) * sq, sq, dtype=torch.int32, device=dev) + + ref = torch.empty(b * sq, hq, HEAD_DIM, dtype=torch.bfloat16, device=dev) + _flash_attn_fwd( + q, + k_fa, + v_fa, + cu_seqlens_q=cu_q, + seqused_k=lens, + max_seqlen_q=sq, + max_seqlen_k=kv, + page_table=bt, + softmax_scale=scale, + causal=False, + window_size_left=WINDOW - 1, + window_size_right=WINDOW - 1, + num_splits=0, + out=ref, + ) + + wrapper = BatchDecodePagedTSWrapper(kv_layout="HND") + wrapper.plan( + device=dev, + batch_size=b, + num_qo_heads=hq, + num_kv_heads=hkv, + head_dim=HEAD_DIM, + page_size=PAGE, + max_kv_len=kv, + max_seq_len_q=sq, + packed_query=False, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + o_data_type=torch.bfloat16, + mask_type="dense", + window_left=-1, + ) + out = torch.empty(b, sq, hq, HEAD_DIM, dtype=torch.bfloat16, device=dev) + wrapper.run( + q=q.view(b, sq, hq, HEAD_DIM), + paged_kv_cache=(k_plane, v_plane), + seq_lens=lens, + block_tables=bt, + bmm1_scale=scale, + bmm2_scale=1.0, + out=out, + ) + + # Within one bf16 ulp: the kernels round differently on some SM100 parts. + torch.testing.assert_close( + out.reshape(b * sq, hq, HEAD_DIM), ref, rtol=1e-2, atol=2e-3 + ) + + +def test_metadata_tensors_keep_their_addresses_across_builds(): + """A captured CUDA graph replays the ADDRESS of seq_lens and the page table. + + Returning freshly allocated tensors per step leaves a replayed graph reading + memory the allocator has since handed to someone else, which shows up as an + engine that pins every GPU and never retires a step. + """ + builder = FlashInferTSMetadataBuilder.__new__(FlashInferTSMetadataBuilder) + builder.sliding_window = WINDOW + builder.page_size = PAGE + builder.query_len = 4 + pages = builder._window_pages(4) + builder._buf_seq_lens = torch.zeros(4, dtype=torch.int32) + builder._buf_block_table = torch.zeros((4, pages), dtype=torch.int32) + + block_table = torch.arange(4 * 1024, dtype=torch.int32).view(4, -1) + ptrs = set() + for step in range(3): + seq_lens = torch.tensor([512, 8192, 40_000 + step, 120_000], dtype=torch.int32) + lens, table = FlashInferTSMetadataBuilder._restrict_to_window( + builder, seq_lens, block_table, 4 + ) + ptrs.add((lens.data_ptr(), table.data_ptr(), tuple(table.shape))) + + assert len(ptrs) == 1, f"metadata moved between builds: {ptrs}" diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 53641f32f249..f56619f43e2c 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from unittest import mock import pytest @@ -766,6 +766,12 @@ def forward(self, hidden_states): hidden_norm=SimpleNamespace(weight=nn.Parameter(torch.ones(2))), _rms_norm_eps=1e-6, ) + model._normalize_context_states = MethodType( + qwen3_dflash.DFlashQwen3Model._normalize_context_states, model + ) + model._project_context_kv_token_major = MethodType( + qwen3_dflash.DFlashQwen3Model._project_context_kv_token_major, model + ) layers_attn = [ SimpleNamespace( qkv_proj=projection, diff --git a/tests/v1/spec_decode/test_draft_attention_backend_override.py b/tests/v1/spec_decode/test_draft_attention_backend_override.py index 26bcd68fea8f..33ec7f7b551c 100644 --- a/tests/v1/spec_decode/test_draft_attention_backend_override.py +++ b/tests/v1/spec_decode/test_draft_attention_backend_override.py @@ -12,13 +12,15 @@ import pytest -from vllm.config import LoadConfig, SpeculativeConfig +from vllm.config import CacheConfig, LoadConfig, SpeculativeConfig +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model @dataclass class _AttentionConfig: backend: str | None = None + use_non_causal: bool = False @dataclass @@ -98,3 +100,32 @@ def test_override_does_not_mutate_the_target_config(): cfg = _config("FLASHINFER", "TRITON_ATTN") _capture_draft_config(cfg) assert cfg.attention_config.backend == "FLASHINFER" + + +def test_draft_attn_builders_get_overrides_and_keep_resolved_cache_state(): + """The draft's metadata builders read the same kernel overrides as its + layers, without losing state the engine core resolved on the target, e.g. + an fp8 draft next to an NVFP4 target with a resolved KV cache layout.""" + cache_config = CacheConfig(cache_dtype="nvfp4") + cache_config.kv_cache_layout = "HND" + cfg = _VllmConfig( + attention_config=_AttentionConfig(backend="FLASHINFER"), + kernel_config=_KernelConfig(), + cache_config=cache_config, + speculative_config=_SpeculativeConfig( + attention_backend="TRITON_ATTN", kv_cache_dtype="fp8" + ), + ) + speculator = object.__new__(DFlashSpeculator) + speculator.vllm_config = cfg + speculator.speculative_config = cfg.speculative_config + speculator.requires_non_causal = True + + used = speculator.attn_vllm_config + + assert used.attention_config.backend == "TRITON_ATTN" + assert used.attention_config.use_non_causal + assert used.cache_config.cache_dtype == "fp8" + assert used.cache_config.kv_cache_layout == "HND" + assert cfg.attention_config.backend == "FLASHINFER" + assert cfg.cache_config.cache_dtype == "nvfp4" diff --git a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py index 89bfc6a53d50..952a6a380a1b 100644 --- a/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py +++ b/tests/v1/worker/test_eagle3_aux_hidden_states_pp.py @@ -57,29 +57,39 @@ def test_mimo_does_not_inherit_aux_hidden_state_pp_support(): class _PPAuxLayer(torch.nn.Module): - """Aliases and mutates residual in place like MiniMaxM3DecoderLayer.""" + """Aliases and mutates residual in place like MiniMaxM3DecoderLayer. - ffn_all_reduce_deferred = False + With ``deferred``, the FFN output is this rank's half of a 2-way TP sum, and + the next input norm completes the all-reduce. + """ - def __init__(self, index): + def __init__(self, index, deferred): super().__init__() self.index = index + self.ffn_all_reduce_deferred = deferred + self.fuse_input_allreduce = False - def forward(self, positions, hidden_states, residual): - total = hidden_states if residual is None else hidden_states + residual - residual = hidden_states if residual is None else residual - residual.copy_(total + self.index + 1) - return torch.full_like(hidden_states, 10 * (self.index + 1)), residual + def forward_input_norm(self, hidden_states, residual): + if residual is None: + return hidden_states, hidden_states + if self.fuse_input_allreduce: + hidden_states = hidden_states * 2 + residual.add_(hidden_states) + return hidden_states, residual + + def forward_after_input_norm(self, positions, hidden_states, residual): + residual.add_(self.index + 1) + output = torch.full_like(hidden_states, 10 * (self.index + 1)) + return output / 2 if self.ffn_all_reduce_deferred else output, residual class _PPAuxNorm(torch.nn.Module): - """Clobbers its inputs so aux taps that alias them would fail.""" + """Fused add + norm (identity scaling) that overwrites its inputs in place.""" def forward(self, hidden_states, residual): - output = hidden_states + residual - hidden_states.fill_(-101) - residual.fill_(-103) - return output, residual + residual.add_(hidden_states) + hidden_states.copy_(residual) + return hidden_states, residual @pytest.fixture @@ -90,6 +100,9 @@ def minimax_pp_stage(monkeypatch): monkeypatch.setattr(minimax_m3, "get_pp_group", lambda: group) monkeypatch.setattr(parallel_state, "get_pp_group", lambda: group) monkeypatch.setattr(parallel_state, "model_parallel_is_initialized", lambda: True) + monkeypatch.setattr( + minimax_m3.MiniMaxM3Model, "_init_fused_moe_finalize", lambda *_: None + ) monkeypatch.setattr( minimax_m3, "VocabParallelEmbedding", @@ -98,10 +111,18 @@ def minimax_pp_stage(monkeypatch): monkeypatch.setattr( minimax_m3, "MiniMAXGemmaRMSNorm", lambda *args, **kwargs: _PPAuxNorm() ) + deferred = [False] monkeypatch.setattr( minimax_m3, "MiniMaxM3DecoderLayer", - lambda **kwargs: _PPAuxLayer(int(kwargs["prefix"].rsplit(".", 1)[1])), + lambda **kwargs: _PPAuxLayer( + int(kwargs["prefix"].rsplit(".", 1)[1]), deferred[0] + ), + ) + monkeypatch.setattr( + minimax_m3, + "fused_allreduce_gemma_rms_norm", + lambda hidden_states, residual, norm: norm(hidden_states * 2, residual), ) def make_layers(num_layers, build, prefix): @@ -117,8 +138,9 @@ def make_layers(num_layers, build, prefix): monkeypatch.setattr(minimax_m3, "make_layers", make_layers) - def make_stage(start, end, taps, world_size=2): + def make_stage(start, end, taps, world_size=2, deferred_all_reduce=False): bounds[:] = [start, end] + deferred[0] = deferred_all_reduce group.world_size = world_size group.is_first_rank = start == 0 group.is_last_rank = end == 6 @@ -130,6 +152,7 @@ def make_stage(start, end, taps, world_size=2): model_config=SimpleNamespace(hf_text_config=config), quant_config=None, speculative_config=None, + kernel_config=SimpleNamespace(moe_backend="auto"), use_v2_model_runner=True, ) ) @@ -204,3 +227,15 @@ def test_minimax_pp2_without_aux_keeps_hidden_and_residual_transport(minimax_pp_ ) == {"hidden_states", "residual"} output = last(None, None, sent) torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6))) + + +@pytest.mark.parametrize("taps", [(0, 3, 6), (1, 2, 5)]) +def test_minimax_aux_taps_follow_deferred_tp_all_reduce(minimax_pp_stage, taps): + """Under TP each FFN output stays a partial sum until the next input norm or + the final norm all-reduces it, so aux taps must read the reduced stream.""" + model, _ = minimax_pp_stage(0, 6, taps, world_size=1, deferred_all_reduce=True) + assert model.fuse_final_norm_allreduce + output, aux = model(None, None, None, _pp_aux_state(7)) + for hidden, tap in zip(aux, taps, strict=True): + torch.testing.assert_close(hidden, _pp_aux_state(_expected_aux_value(tap))) + torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6))) diff --git a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py index 901a4d184c73..fd6333f267e3 100644 --- a/tests/v1/worker/test_spec_decode_embed_sharing_pp.py +++ b/tests/v1/worker/test_spec_decode_embed_sharing_pp.py @@ -162,7 +162,7 @@ def m3_config(): parallel_config=ParallelConfig(), quant_config=None, cache_config=None, - kernel_config=SimpleNamespace(monokernel=False), + kernel_config=SimpleNamespace(monokernel=False, moe_backend="auto"), ) @@ -269,6 +269,10 @@ def test_m3_target_allocates_embedding_on_required_stage( ) monkeypatch.setattr(model, "get_pp_group", pp) monkeypatch.setattr("vllm.distributed.parallel_state.get_pp_group", pp) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tp_group", + lambda: SimpleNamespace(world_size=1), + ) monkeypatch.setattr(model, "VocabParallelEmbedding", lambda *_, **kw: _embed()) monkeypatch.setattr(model, "make_layers", lambda *_, **kw: (0, 0, nn.ModuleList())) monkeypatch.setattr(model, "MiniMAXGemmaRMSNorm", lambda *_, **kw: nn.Identity()) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index c8574a53bccd..d56f825304ae 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2782,6 +2782,7 @@ def fused_minimax_m3_qknorm_rope_kv_insert( q_fp8_scale: float = 1.0, kv_k_scale: torch.Tensor | None = None, kv_v_scale: torch.Tensor | None = None, + write_qkv: bool = True, ) -> None: """Fused MiniMax-M3 attention pre-processing (in-place). @@ -2792,11 +2793,13 @@ def fused_minimax_m3_qknorm_rope_kv_insert( - sparse layer (``num_index_heads > 0``): ``[q | k | v | index_q | index_k]`` — the index branch is read straight out of ``qkv``. - When ``kv_cache`` is given (sparse serving), also scatter-inserts the - normed/roped k & v into the paged KV cache by ``slot_mapping`` and the - index key into ``index_cache`` by ``index_slot_mapping``. ``kv_cache_dtype`` - selects the cache storage/conversion path. If - ``index_slot_mapping`` is omitted, ``slot_mapping`` is used for both caches. + When ``kv_cache`` is given, also scatter-inserts the normed/roped k & v + into the paged KV cache by ``slot_mapping`` and, on sparse layers, the + index key into ``index_cache`` by ``index_slot_mapping``. + ``kv_cache_dtype`` selects the cache storage/conversion path; a dense + layer's insert stores exactly what ``reshape_and_cache_flash`` would with + unit k/v scales. If ``index_slot_mapping`` is omitted, ``slot_mapping`` is + used for both caches. If ``q_out`` / ``index_q_out`` (contiguous ``[N, nq*128]`` / ``[N, niq*128]``) are given, the normed/roped q / index_q are written there @@ -2818,6 +2821,11 @@ def fused_minimax_m3_qknorm_rope_kv_insert( ``[index_q | index_k]`` tail, but the kernel only processes the main q/k/v branches and main KV cache. This is used by MiniMax-M3 index-topk reuse layers that consume top-k block ids selected by an earlier sparse layer. + + ``write_qkv=False`` needs ``q_fp8_out`` without ``q_out``, for callers that + read only the FP8 q, ``index_q_out`` and the caches: k and index_k need not + be rewritten in ``qkv``. The persistent kernel (SM80+, ``rotary_dim`` 64) + skips those stores; the fallback kernel may still make them. """ torch.ops._C.fused_minimax_m3_qknorm_rope_kv_insert( qkv, @@ -2845,6 +2853,47 @@ def fused_minimax_m3_qknorm_rope_kv_insert( q_fp8_scale, kv_k_scale, kv_v_scale, + write_qkv, + ) + + +def fused_minimax_m3_dspark_context_kv( + all_kv: torch.Tensor, + k_norm_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache: torch.Tensor, + cache_offsets: torch.Tensor, + k_scale: torch.Tensor, + scale_offsets: torch.Tensor, + rotary_dim: int, + eps: float, + kv_cache_dtype: str, +) -> None: + """Store the MiniMax-M3 DSpark draft's context K/V for every layer. + + ``all_kv`` ``[N, L, 2, nkv, 128]`` holds each layer's K and V. K gets the + layer's Gemma RMSNorm (``k_norm_weight[L, 128]``) and partial NeoX RoPE; + both are stored like ``reshape_and_cache_flash`` into the layers' caches. + Layer ``l``'s cache is ``kv_cache`` (layer 0, ``[nb, nkv, bs, 2 * 128]``) + shifted by ``cache_offsets[l]`` elements; its fp8 k/v scales are + ``k_scale`` (layer 0's) shifted by ``scale_offsets[0, l]`` / + ``scale_offsets[1, l]`` floats. Slots < 0 are skipped. + """ + torch.ops._C.fused_minimax_m3_dspark_context_kv( + all_kv, + k_norm_weight, + cos_sin_cache, + positions, + slot_mapping, + kv_cache, + cache_offsets, + k_scale, + scale_offsets, + rotary_dim, + eps, + kv_cache_dtype, ) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 2ba6fdb6cb1b..a63cc3962631 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -40,8 +40,10 @@ class AttentionConfig: backend: AttentionBackendEnum | None = None """Attention backend to use. Use "auto" or None for automatic selection.""" - minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton" - """Sparse decode kernel used by the MiniMax M3 MSA backend.""" + minimax_m3_msa_decode_backend: MiniMaxM3MSADecodeBackend = "cutlass" + """Sparse decode kernel used by the MiniMax M3 MSA backend. "cutlass" falls + back to Triton where the CUTLASS kernel does not support the platform, KV + cache dtype, or head layout.""" backend_per_kind: dict[str, AttentionBackendEnum] = field(default_factory=dict) """Per-KV-cache-group attention backend overrides, keyed by @@ -95,7 +97,20 @@ class AttentionConfig: model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse indexer, mxfp4 for DeepSeek-V4 on ROCm CDNA4 with aiter). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel support in the - backend.""" + backend; MiniMax M3 supports nvfp4 on SM100/SM103/SM107 with at most 8 + decode queries per request.""" + + indexer_decode_cp: bool = False + """Split the sparse-attention indexer's decode scoring across the + tensor-parallel ranks by context instead of by index head: each rank scores + a 1/TP slice of the index-K pages for every head and the ranks exchange + their local top-k, so each reads 1/TP of the index-K cache per sparse + layer. The index-K cache stays replicated, so this is unrelated to + `--decode-context-parallel-size`, which shards the KV cache itself; the two + cannot be combined. Implemented for the MiniMax M3 NVFP4 indexer, which + needs `indexer_kv_dtype="nvfp4"` and TP 4 (one index head per rank); it is + ignored with a warning where it does not apply. It pays off from about 4 + concurrent decodes at 128K context and grows with batch and context.""" indexer_sparse_logits: bool | None = None """DeepSeek V4.1 two-level indexer: score only the candidate blocks with diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index fd2e3a397e13..b777fb6ea138 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -3,6 +3,7 @@ import copy import functools +import os from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Literal, Self, get_args @@ -16,6 +17,7 @@ from vllm.config.utils import config, replace from vllm.logger import init_logger from vllm.transformers_utils.config import get_hf_text_config, iter_rope_parameters +from vllm.transformers_utils.repo_utils import get_subfolder_path from vllm.utils.hashing import safe_hash from vllm.utils.import_utils import LazyLoader, has_arctic_inference from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -114,6 +116,31 @@ def _rope_only_sizes_cache(hf_config: PreTrainedConfig) -> bool: ) +_MINIMAX_M3_MODEL_TYPES = ("minimax_m3_vl", "minimax_m3_text") +_MINIMAX_M3_DSPARK_SUBFOLDER = "dspark" + + +def _get_minimax_m3_dspark_draft( + target_model_config: ModelConfig, cache_dir: str | None +) -> str: + """MiniMax-M3 checkpoints ship the DSpark draft as a standalone checkpoint + in their `dspark/` subfolder.""" + draft = get_subfolder_path( + target_model_config.model, + _MINIMAX_M3_DSPARK_SUBFOLDER, + revision=target_model_config.revision, + cache_dir=cache_dir, + ) + if not os.path.isfile(os.path.join(draft, "config.json")): + raise ValueError( + f"MiniMax-M3 DSpark found no draft checkpoint in the " + f"'{_MINIMAX_M3_DSPARK_SUBFOLDER}/' subfolder of " + f"{target_model_config.model}. Set speculative_config.model to the " + "draft checkpoint." + ) + return draft + + def _is_qwen3_omni_target(model_config: ModelConfig) -> bool: hf_config = model_config.hf_config architectures = set(getattr(model_config, "architectures", ()) or ()) @@ -1134,7 +1161,25 @@ def hf_config_override( text_config.num_kv_shared_layers = 0 hf_config.update({"n_predict": 1, "architectures": ["Gemma4MTPModel"]}) - if ( + if initial_architecture == "DSparkMiniMaxDraftModel": + quantization_config = getattr(hf_config, "quantization_config", None) + hf_config = getattr(hf_config, "text_config", hf_config) + if ( + quantization_config is not None + and getattr(hf_config, "quantization_config", None) is None + ): + hf_config.update({"quantization_config": quantization_config}) + dflash_config = dict(getattr(hf_config, "dflash_config", None) or {}) + dflash_config.setdefault("causal", False) + hf_config.model_type = "minimax_m3_dspark" + hf_config.update( + { + "architectures": ["DSparkMiniMaxDraftModel"], + "dflash_config": dflash_config, + "markov_rank": hf_config.dspark_markov_rank, + } + ) + elif ( hf_config.model_type == "minimax_m3_vl" or initial_architecture == "MiniMaxM3SparseForConditionalGeneration" ): @@ -1269,7 +1314,11 @@ def __post_init__(self): ) self.method = "mtp" - if self.model is None and self.num_speculative_tokens is not None: + # DSpark drafts carry their own width (block size), so the default + # draft is resolved even without num_speculative_tokens. + if self.model is None and ( + self.num_speculative_tokens is not None or self.method == "dspark" + ): if self.method == "mtp": if self.target_model_config is None: raise ValueError("target_model_config must be present for mtp") @@ -1287,12 +1336,24 @@ def __post_init__(self): if not self.quantization: self.quantization = self.target_model_config.quantization elif self.method == "dspark": - # DeepSeek DSpark can ship the weights inside the target checkpoint if self.target_model_config is None: raise ValueError("target_model_config must be present for dspark") - self.model = self.target_model_config.model - if not self.quantization: - self.quantization = self.target_model_config.quantization + if ( + self.target_model_config.hf_config.model_type + in _MINIMAX_M3_MODEL_TYPES + ): + self.model = _get_minimax_m3_dspark_draft( + self.target_model_config, + self.draft_load_config.download_dir + if self.draft_load_config + else None, + ) + else: + # DeepSeek DSpark can ship the weights inside the target + # checkpoint + self.model = self.target_model_config.model + if not self.quantization: + self.quantization = self.target_model_config.quantization elif self.method in ("ngram", "[ngram]"): self.model = "ngram" elif self.method == "ngram_gpu": @@ -1568,6 +1629,8 @@ def __post_init__(self): not in self.draft_model_config.architectures and "Gemma4DSparkModel" not in self.draft_model_config.architectures and "K3DSparkModel" not in self.draft_model_config.architectures + and "DSparkMiniMaxDraftModel" + not in self.draft_model_config.architectures ): # DeepSeek-V4(.1) DSpark reuses the full target config # and its weights ship in the target checkpoint. diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py index 033533a5c9a0..2cc4564efa26 100644 --- a/vllm/cute_utils/__init__.py +++ b/vllm/cute_utils/__init__.py @@ -60,11 +60,15 @@ def to_cta0_smem(ptr: cute.Pointer, *, loc=None, ip=None): ) -def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None): +def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None, mcast_mask=None): """A simple helper that wraps group_modes() and tma_partition() NOTE: this should be called WITHOUT cute.elect_one() """ - if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp): + g2s_ops = ( + cpasync.CopyBulkTensorTileG2SOp, + cpasync.CopyBulkTensorTileG2SMulticastOp, + ) + if isinstance(atom.op, g2s_ops): gmem = src smem = dst elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp): @@ -81,8 +85,13 @@ def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None): cute.group_modes(gmem, 0), ) - if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp): - cute.copy(atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy) + if isinstance(atom.op, g2s_ops): + kwargs = {} + if mcast_mask is not None: + kwargs["mcast_mask"] = mcast_mask + cute.copy( + atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy, **kwargs + ) elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp): cute.copy(atom, s_part, g_part, cache_policy=cache_policy) else: diff --git a/vllm/cute_utils/cvt.py b/vllm/cute_utils/cvt.py index 8bb4ad95198b..ce103ee741b8 100644 --- a/vllm/cute_utils/cvt.py +++ b/vllm/cute_utils/cvt.py @@ -38,6 +38,12 @@ def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]: ) elif isinstance(data, (cute.Tensor, cute.TensorSSA)): + # TensorSSA in, TensorSSA out (it can't be recast: stage it in rmem) + is_ssa = isinstance(data, cute.TensorSSA) + if is_ssa: + frag = cute.make_rmem_tensor(data.shape, data.dtype) + frag.store(data) + data = frag # recast to 32-bit registers if needed if data.element_type == BFloat16: data = cute.recast_tensor(data, Uint32) @@ -48,7 +54,7 @@ def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]: out = cute.make_rmem_tensor(size * 2, Float32) for i in range(size): out[i * 2], out[i * 2 + 1] = bf16x2_to_fp32x2(data[i]) - return out + return out.load() if is_ssa else out else: raise ValueError(f"Unsupported type {type(data)}") diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py index c3b99fc36ec3..f21d0bf37360 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py @@ -99,7 +99,7 @@ def apply_weights( out_dtype, input_shape = qa.orig_dtype, qa.orig_shape else: assert isinstance(x, torch.Tensor) - out_dtype, input_shape = x.dtype, x.shape + out_dtype, input_shape = layer.params_dtype, x.shape input_mxfp8, input_scale = mxfp8_e4m3_quantize( x.view(-1, K), is_sf_swizzled_layout=True ) @@ -177,7 +177,7 @@ def apply_weights( out_dtype, input_shape = qa.orig_dtype, qa.orig_shape else: assert isinstance(x, torch.Tensor) - out_dtype, input_shape = x.dtype, x.shape + out_dtype, input_shape = layer.params_dtype, x.shape input_mxfp8, input_scale = mxfp8_e4m3_quantize( x.view(-1, K), is_sf_swizzled_layout=True ) diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 93246d9160e3..800e0d1aa6cd 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -253,10 +253,12 @@ def _mxfp8_dot_scaled_gemm( # gfx950's async-copy LDS (default in the triton 3.8 this branch pins) leaves no # room for num_stages=3 at BLOCK_K=256; use 2 there. _BK256_STAGES = 3 +_ON_GFX950 = False if current_platform.is_rocm(): from vllm.platforms.rocm import on_gfx950 - if on_gfx950(): + _ON_GFX950 = on_gfx950() + if _ON_GFX950: _BK256_STAGES = 2 @@ -271,6 +273,11 @@ def _select_cfg(M, N, K): (the K-loop is unmasked), so every BLOCK_K below is guarded to be K-divisible (served K: 384/768/1024/2048/6144). """ + if _ON_GFX950 and M <= 16: + if K == 6144 and N in (1536, 2560): + return 16, 16, 1024, 2, 3 + if K == 768 and N == 6144: + return 16, 16, 256, 2, 3 if M <= 64: # decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The # lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 93958bb12782..98b7f6509303 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -256,6 +256,7 @@ def __init__( mm_prefix_clamp_sliding_window: bool = False, attn_backend: type[AttentionBackend] | None = None, head_size_v: int | None = None, + query_pre_quantized: bool = False, **extra_impl_args, ) -> None: """The KV cache is stored inside this class and is accessed via @@ -469,7 +470,11 @@ def __init__( # for attn backends supporting query quantization self.query_quant = None if ( - self.impl.supports_quant_query_input + # A model whose fused QK-norm/RoPE kernel already emits the + # quantized query (MiniMax-M3's DSpark draft) passes + # query_pre_quantized and skips this second quantization. + not query_pre_quantized + and self.impl.supports_quant_query_input and ( self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype.startswith("nvfp4") diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 83df5d16c16e..b0b6744c6c17 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -16,7 +16,10 @@ OCP_MX_DTYPES, OCP_MX_Scheme, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + QuantKey, +) from vllm.platforms import current_platform from vllm.utils.import_utils import get_triton_kernels_version from vllm.utils.math_utils import cdiv @@ -1339,6 +1342,12 @@ class FusedMoEConfig: tp_shard_with_padding: bool = False rocm_aiter_fmoe_enabled: bool = False is_fused_shared_expert_requested: bool = False + # Activation scheme the routed experts were trained (QAT) against, when the + # checkpoint does not record it. Quant methods use it in place of their + # default: MiniMax-M3.1 quantizes the FC1 input per token and ships a 1.0 + # ``input_global_scale`` placeholder, which is indistinguishable from a + # calibrated per-tensor scale of 1.0. + qat_activation_key: QuantKey | None = None # Whether to force MXFP4 (a4w4) MoE activations for DeepSeek V4.1 on # ROCm/AITER. Opt-in via VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1; rejected # for any other model type. Resolved here, not in the forward path, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 3b3e2aa39973..cbe6d0ce8938 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -40,6 +40,7 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey logger = init_logger(__name__) @@ -114,6 +115,7 @@ def FusedMoEFactory( swiglu_beta: float | None = None, activation_situ_beta: float | None = None, activation_situ_linear_beta: float | None = None, + qat_activation_key: QuantKey | None = None, e_score_correction_bias: torch.Tensor | None = None, apply_router_weight_on_input: bool = False, activation: str = "silu", @@ -186,6 +188,8 @@ def FusedMoEFactory( swiglu_limit: SwiGLU activation limit activation_situ_beta: SituGLU activation beta activation_situ_linear_beta: SituGLU linear beta + qat_activation_key: Activation quantization scheme the experts were + trained against, for checkpoints that do not record it e_score_correction_bias: Expert score correction bias tensor apply_router_weight_on_input: Whether to apply router weights on input activation: Activation function name ("silu", "gelu", etc.) @@ -387,6 +391,7 @@ def FusedMoEFactory( swiglu_beta=swiglu_beta, activation_situ_beta=activation_situ_beta, activation_situ_linear_beta=activation_situ_linear_beta, + qat_activation_key=qat_activation_key, max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size, skip_final_all_reduce=skip_final_all_reduce, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 62dceb99b324..39ed48f8488b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum +from typing import Any import torch @@ -55,6 +56,15 @@ class NvFp4MoeBackend(Enum): FLASHINFER_MOE_EP_CUTEDSL = "FLASHINFER_MOE_EP_CUTEDSL" +# Backends whose experts quantize the MoE input themselves, with one scale per +# token. ``make_nvfp4_moe_kernel`` only hands ``per_token_activation`` to these. +PER_TOKEN_NVFP4_MOE_BACKENDS = frozenset( + { + NvFp4MoeBackend.FLASHINFER_CUTEDSL, + NvFp4MoeBackend.FLASHINFER_TRTLLM, + } +) + FLASHINFER_NVFP4_MOE_BACKENDS = [ NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTLASS, @@ -616,6 +626,10 @@ def make_nvfp4_moe_kernel( routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, per_token_activation: bool = False, ) -> mk.FusedMoEKernel: + # Marlin consumes A16 and must not receive per-token A4 metadata. + if backend == NvFp4MoeBackend.MARLIN and moe_quant_config.use_nvfp4_w4a16: + per_token_activation = False + # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( moe=moe_config, @@ -629,15 +643,17 @@ def make_nvfp4_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) - extra_kwargs = {} - if ( - backend - in ( - NvFp4MoeBackend.FLASHINFER_TRTLLM, - NvFp4MoeBackend.FLASHINFER_CUTEDSL, - ) - and per_token_activation - ): + extra_kwargs: dict[str, Any] = {} + if per_token_activation: + # Prepare dispatched the input unquantized, so an expert that does + # not quantize it itself would read garbage. The backend selection + # does not catch this: several NVFP4 experts match on the weight key + # alone and accept any activation key. + if backend not in PER_TOKEN_NVFP4_MOE_BACKENDS: + raise ValueError( + f"Per-token NVFP4 activation scales are not supported by " + f"{backend.value}." + ) extra_kwargs["per_token_activation"] = True # Create Experts. diff --git a/vllm/model_executor/layers/fused_moe/prepare_megamoe.py b/vllm/model_executor/layers/fused_moe/prepare_megamoe.py index 5e8d8cde9a62..e1297e8380b6 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_megamoe.py +++ b/vllm/model_executor/layers/fused_moe/prepare_megamoe.py @@ -300,52 +300,75 @@ def _prepare_megamoe_nvfp4_inputs_kernel( hidden_size: tl.constexpr, top_k: tl.constexpr, num_local_experts: tl.constexpr, - BLOCK_M: tl.constexpr, + BLOCK_H: tl.constexpr, BLOCK_K: tl.constexpr, + NUM_SPLITS: tl.constexpr, BLOCK_TOPK: tl.constexpr, BLOCK_E: tl.constexpr, SHARED_BLOCK_M: tl.constexpr, SHARED_MXFP8: tl.constexpr, ) -> None: + # Grid (token, split): each program quantizes the BLOCK_K chunks split, + # split + NUM_SPLITS, ... of one token's row. + token = tl.program_id(0) + split = tl.program_id(1) if l2_alpha is not None: experts = tl.arange(0, BLOCK_E) - expert_mask = (experts < num_local_experts) & (tl.program_id(0) == 0) + expert_mask = (experts < num_local_experts) & (token == 0) & (split == 0) alpha = tl.load(w2_global_scale + experts, mask=expert_mask) tl.store(l2_alpha + experts, alpha * l2_activation_scale, mask=expert_mask) + if token >= num_tokens: + return - token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) - token_mask = token_id < num_tokens - k_offsets = tl.arange(0, BLOCK_K) - hidden_ptrs = hidden_states + token_id[:, None] * hidden_stride_m - - amax = tl.zeros((BLOCK_M,), tl.float32) - for k_block_id in range(hidden_size // BLOCK_K): - hidden = tl.load( - hidden_ptrs + (k_block_id * BLOCK_K + k_offsets[None, :]) * hidden_stride_k, - mask=token_mask[:, None], - other=0.0, - ).to(tl.float32) - amax = tl.maximum(amax, tl.max(tl.abs(hidden), axis=1)) - x_scale = tl.maximum(amax, 1.0e-12) * _INV_NVFP4_MAX - tl.store(x_scales + token_id, x_scale, mask=token_mask) - - for k_block_id in range(hidden_size // BLOCK_K): - k_ids = k_block_id * BLOCK_K + k_offsets - hidden = tl.load( - hidden_ptrs + k_ids[None, :] * hidden_stride_k, - mask=token_mask[:, None], - other=0.0, + if split == 0: + token_ids = token + tl.zeros([1], dtype=tl.int32) + _store_topk( + token_ids, + token_ids < num_tokens, + topk_ids, + topk_weights, + is_padding, + topk_idx_out, + topk_weights_out, + topk_ids_stride_m, + topk_ids_stride_k, + topk_weights_stride_m, + topk_weights_stride_k, + is_padding_stride_m, + topk_idx_stride_m, + topk_idx_stride_k, + topk_weights_out_stride_m, + topk_weights_out_stride_k, + top_k, + 1, + BLOCK_TOPK, ) + # Every program reads the whole row for the token scale. + row = hidden_states + token.to(tl.int64) * hidden_stride_m + cols = tl.arange(0, BLOCK_H) + hidden_row = tl.load( + row + cols * hidden_stride_k, mask=cols < hidden_size, other=0.0 + ) + amax = tl.max(tl.abs(hidden_row.to(tl.float32)), axis=0) + x_scale = tl.maximum(amax, 1.0e-12) * _INV_NVFP4_MAX + if split == 0: + tl.store(x_scales + token, x_scale) + + shared_sf_row = _shared_sf_row(token, SHARED_BLOCK_M) + for k0 in tl.range(split * BLOCK_K, hidden_size, NUM_SPLITS * BLOCK_K): + k_ids = k0 + tl.arange(0, BLOCK_K) + k_mask = k_ids < hidden_size + hidden = tl.load(row + k_ids * hidden_stride_k, mask=k_mask, other=0.0) + # E2M1 values with one E4M3 scale per 16, relative to the token scale. values = tl.reshape( - tl.div_rn(hidden.to(tl.float32), x_scale[:, None]), - [BLOCK_M, BLOCK_K // 16, 16], + tl.div_rn(hidden.to(tl.float32), x_scale), [BLOCK_K // 16, 16] ) - block_scale = tl.max(tl.abs(values), axis=2) * (1.0 / 6.0) + block_scale = tl.max(tl.abs(values), axis=1) * (1.0 / 6.0) block_scale = tl.minimum(tl.maximum(block_scale, 2.0**-9), 448.0) block_scale = block_scale.to(tl.float8e4nv) - values = tl.div_rn(values, block_scale.to(tl.float32)[:, :, None]) + values = tl.div_rn(values, block_scale.to(tl.float32)[:, None]) magnitude = tl.abs(values) # Round to nearest even on the grid {0, 0.5, 1, 1.5, 2, 3, 4, 6}. code = ( @@ -358,61 +381,45 @@ def _prepare_megamoe_nvfp4_inputs_kernel( + (magnitude > 5.0).to(tl.int32) ) code = tl.where((values < 0) & (code != 0), code | 8, code) - low, high = tl.split(tl.reshape(code, [BLOCK_M, BLOCK_K // 2, 2])) - packed_offsets = k_block_id * (BLOCK_K // 2) + tl.arange(0, BLOCK_K // 2) + low, high = tl.split(tl.reshape(code, [BLOCK_K // 2, 2])) + packed_ids = k0 // 2 + tl.arange(0, BLOCK_K // 2) tl.store( - x_out + token_id[:, None] * x_stride_m + packed_offsets[None, :], + x_out + token * x_stride_m + packed_ids, (low | (high << 4)).to(tl.uint8), - mask=token_mask[:, None], + mask=packed_ids < hidden_size // 2, ) - sf_offsets = k_block_id * (BLOCK_K // 64) + tl.arange(0, BLOCK_K // 64) + # Four E4M3 scales per little-endian int32 word, one word per 64 values. + sf_ids = k0 // 64 + tl.arange(0, BLOCK_K // 64) sf_bytes = tl.reshape( - block_scale.to(tl.uint8, bitcast=True), [BLOCK_M, BLOCK_K // 64, 4] + block_scale.to(tl.uint8, bitcast=True), [BLOCK_K // 64, 4] ) tl.store( - x_sf + token_id[:, None] * x_sf_stride_m + sf_offsets[None, :], + x_sf + token * x_sf_stride_m + sf_ids, _pack_bytes(sf_bytes, 4), - mask=token_mask[:, None], + mask=sf_ids < hidden_size // 64, ) if shared_x is not None: - shared_ptrs = shared_x + token_id[:, None] * shared_x_stride_m + shared_ptrs = shared_x + token * shared_x_stride_m + k_ids if SHARED_MXFP8: fp8, _, scale_exp = _quantize_fp8_groups( - hidden.to(tl.float32), BLOCK_M, BLOCK_K, 32, True + tl.reshape(hidden.to(tl.float32), (1, BLOCK_K)), + 1, + BLOCK_K, + 32, + True, ) - tl.store(shared_ptrs + k_ids[None, :], fp8, mask=token_mask[:, None]) + tl.store(shared_ptrs, tl.reshape(fp8, [BLOCK_K]), mask=k_mask) + word_ids = k0 // 128 + tl.arange(0, BLOCK_K // 128) tl.store( shared_x_sf - + _shared_sf_row(token_id, SHARED_BLOCK_M) * shared_x_sf_stride_m - + k_block_id * shared_x_sf_stride_k, - _pack_bytes(scale_exp, BLOCK_K // 32), - mask=token_mask, + + shared_sf_row * shared_x_sf_stride_m + + word_ids * shared_x_sf_stride_k, + _pack_bytes(tl.reshape(scale_exp, [BLOCK_K // 128, 4]), 4), + mask=word_ids < hidden_size // 128, ) else: - tl.store(shared_ptrs + k_ids[None, :], hidden, mask=token_mask[:, None]) - - _store_topk( - token_id, - token_mask, - topk_ids, - topk_weights, - is_padding, - topk_idx_out, - topk_weights_out, - topk_ids_stride_m, - topk_ids_stride_k, - topk_weights_stride_m, - topk_weights_stride_k, - is_padding_stride_m, - topk_idx_stride_m, - topk_idx_stride_k, - topk_weights_out_stride_m, - topk_weights_out_stride_k, - top_k, - BLOCK_M, - BLOCK_TOPK, - ) + tl.store(shared_ptrs, hidden, mask=k_mask) def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]: @@ -632,11 +639,19 @@ def prepare_megamoe_nvfp4_inputs( if tensor is not None and tensor.stride(-1) != 1: raise ValueError(f"MegaMoE NVFP4 staging requires a row-major {name}.") - block_m = 8 if num_tokens >= 64 else 1 top_k = topk_ids.shape[1] num_local_experts = 0 if l2_alpha is None else l2_alpha.numel() + # Up to 512 tokens each row is split over programs; larger batches fill the + # GPU with one program per row looping over wider chunks. + block_h = triton.next_power_of_2(hidden_size) + if num_tokens <= 512: + block_k = min(1024, block_h) + num_splits = triton.cdiv(hidden_size, block_k) + else: + block_k = min(2048, block_h) + num_splits = 1 # One program even without tokens, to write l2_alpha. - grid = (max(triton.cdiv(num_tokens, block_m), 1),) + grid = (max(num_tokens, 1), num_splits) _prepare_megamoe_nvfp4_inputs_kernel[grid]( hidden_states, x.view(torch.uint8), @@ -672,8 +687,9 @@ def prepare_megamoe_nvfp4_inputs( hidden_size, top_k, num_local_experts, - BLOCK_M=block_m, - BLOCK_K=128, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, BLOCK_TOPK=triton.next_power_of_2(top_k), BLOCK_E=triton.next_power_of_2(max(num_local_experts, 1)), SHARED_BLOCK_M=shared_block_m or 1, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index 86c7f8c7ba81..22a0f02848f4 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -26,10 +26,15 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 CompressedTensorsMoEMethod, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + trtllm_nvfp4_hidden_alignment, +) from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( dequantize_nvfp4_moe_weights, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + is_dynamic_per_token_outer_scale, kNvfp4Dynamic, kNvfp4Static, ) @@ -58,16 +63,24 @@ def __init__( moe: FusedMoEConfig, layer_name: str | None = None, use_a16: bool = False, + activation_key: QuantKey = kNvfp4Dynamic, ): super().__init__(moe) self.group_size = 16 self.use_a16 = use_a16 + # FC1 input scheme; a per-token one (``kNvfp4DynamicToken``) replaces + # the checkpoint's static input scales. Only the backends that quantize + # the input themselves accept it, so selection enforces the invariant. + self.activation_key: QuantKey | None = None if use_a16 else activation_key + self.per_token_activation = is_dynamic_per_token_outer_scale( + self.activation_key + ) # Select experts implementation. self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( config=self.moe, weight_key=kNvfp4Static, - activation_key=None if use_a16 else kNvfp4Dynamic, + activation_key=self.activation_key, ) self.use_global_sf = is_global_sf_supported_for_nvfp4_backend( self.nvfp4_backend @@ -224,6 +237,14 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous() + a13_scale = 1.0 / layer.w13_input_global_scale + a2_scale = 1.0 / layer.w2_input_global_scale + if self.per_token_activation: + # Per-token scales cover the FC1 input, and the per-token kernel + # requantizes the FC2 input per token: static scales are neutral. + a13_scale = torch.ones_like(a13_scale) + a2_scale = torch.ones_like(a2_scale) + # Shuffle weights into the NvFp4 kernel format. ( w13, @@ -240,13 +261,17 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w13=layer.w13_weight, w13_scale=layer.w13_weight_scale, w13_scale_2=(1.0 / w13_weight_global_scale), - a13_scale=(1.0 / layer.w13_input_global_scale), + a13_scale=a13_scale, w2=layer.w2_weight, w2_scale=layer.w2_weight_scale, w2_scale_2=(1.0 / layer.w2_weight_global_scale), - a2_scale=(1.0 / layer.w2_input_global_scale), + a2_scale=a2_scale, is_act_and_mul=self.moe.is_act_and_mul, use_a16=self.use_a16, + trtllm_hidden_alignment=trtllm_nvfp4_hidden_alignment( + per_token_activation=self.per_token_activation, + is_act_and_mul=self.moe.is_act_and_mul, + ), ) replace_parameter(layer, "w13_weight", w13) @@ -267,6 +292,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: experts_cls=self.experts_cls, backend=self.nvfp4_backend, routing_tables=layer._expert_routing_tables(), + per_token_activation=self.per_token_activation, ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index bc85d867233d..128a5ff06cf3 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1218,10 +1218,14 @@ def __init__( is_checkpoint_mxfp8_serialized: bool, kv_cache_quant_algo: str | None, exclude_modules: list[str], + moe_quant_algo: str | None = None, ) -> None: super().__init__(exclude_modules) self.quant_method = "MXFP8" self.is_checkpoint_mxfp8_serialized = is_checkpoint_mxfp8_serialized + # MiniMax-M3.1 checkpoints keep MXFP8 linears but store routed experts + # as compressed-tensors NVFP4 (`moe_quant_algo: NVFP4`). + self.moe_quant_algo = moe_quant_algo if not is_checkpoint_mxfp8_serialized: raise ValueError( @@ -1257,12 +1261,25 @@ def from_config(cls, config: dict[str, Any]) -> "ModelOptMxFp8Config": # (same on-disk format as ModelOpt MXFP8); normalize to the ModelOpt # schema and reuse the shared parser. if "quantization" not in config and not config.get("quant_algo"): + moe_quant_algo = (config.get("moe_quant_algo") or "").upper() or None + if moe_quant_algo is not None and ( + moe_quant_algo != "NVFP4" + or config.get("moe_quant_format") != "nvfp4-pack-quantized" + or config.get("group_size") != 16 + ): + raise ValueError( + "Unsupported MXFP8 checkpoint MoE quantization: " + f"moe_quant_algo={config.get('moe_quant_algo')!r}, " + f"moe_quant_format={config.get('moe_quant_format')!r}, " + f"group_size={config.get('group_size')!r}." + ) config = { "quant_method": "modelopt", "quantization": { "quant_algo": "MXFP8", "kv_cache_quant_algo": config.get("kv_cache_quant_algo"), "exclude_modules": config.get("ignored_layers", []) or [], + "moe_quant_algo": moe_quant_algo, }, } return cast("ModelOptMxFp8Config", super().from_config(config)) @@ -1296,8 +1313,41 @@ def _from_config( is_checkpoint_mxfp8_serialized, kv_cache_quant_method, exclude_modules, + moe_quant_algo=original_config.get("quantization", {}).get( + "moe_quant_algo" + ), ) + def routed_expert_weight_key(self) -> QuantKey: + """The weight format the routed experts are stored in. + + MXFP8, or NVFP4 (``moe_quant_algo: NVFP4``). + """ + if self.moe_quant_algo != "NVFP4": + return kMxfp8Static + return kNvfp4Static + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> "QuantizeMethodBase | None": + if ( + self.moe_quant_algo == "NVFP4" + and isinstance(layer, RoutedExperts) + and not self.is_layer_excluded(prefix) + ): + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import ( # noqa: E501 + CompressedTensorsW4A4Nvfp4MoEMethod, + ) + + # Nothing in the checkpoint says how its FC1 input was + # quantized during QAT, so the model declares it. + return CompressedTensorsW4A4Nvfp4MoEMethod( + layer.moe_config, + layer_name=prefix, + activation_key=(layer.moe_config.qat_activation_key or kNvfp4Dynamic), + ) + return super().get_quant_method(layer, prefix) + class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): """FlashInfer TRTLLM MXFP8 block-scale MoE for ModelOpt checkpoints.""" diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 8cdbb3e596ae..557034f3e3d7 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -217,6 +217,17 @@ def __str__(self): FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kDynamicTokenScale ) + +def is_dynamic_per_token_outer_scale(key: QuantKey | None) -> bool: + """Whether ``key``'s second-level scale is computed per token at runtime.""" + return ( + key is not None + and key.scale2 is not None + and not key.scale2.static + and key.scale2.group_shape == GroupShape.PER_TOKEN + ) + + kNvfp4StaticGroupScale = ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16)) kNvfp4Static = QuantKey( FP4_DTYPE, scale=kNvfp4StaticGroupScale, scale2=kStaticTensorScale diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 99d3b80eef3a..e9125baae425 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -86,6 +86,30 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: logger.info("Using bfloat16 kv-cache for DeepSeekV3.2") +class MiniMaxM3SparseConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + from vllm.v1.kv_cache_interface import get_kv_quant_mode + + if not get_kv_quant_mode(vllm_config.cache_config.cache_dtype).is_nvfp4: + return + # The NVFP4 KV attention kernels (MSA's CUTLASS Q8KV4 prefill) need 16 + # query heads per KV head on each rank; fail before the workers start. + text_config = vllm_config.model_config.hf_text_config + tp_size = vllm_config.parallel_config.tensor_parallel_size + num_heads = text_config.num_attention_heads // tp_size + num_kv_heads = max(1, text_config.num_key_value_heads // tp_size) + if num_heads != 16 * num_kv_heads: + raise ValueError( + "MiniMax-M3 does not support an NVFP4 KV cache with " + f"tensor_parallel_size={tp_size} yet: each rank holds {num_heads} " + f"query heads over {num_kv_heads} KV head(s), and the NVFP4 " + "attention kernels need 16 per KV head. Use a smaller " + "tensor_parallel_size (with data/expert parallelism) or another " + "kv_cache_dtype." + ) + + class GlmMoeDsaForCausalLM(VerifyAndUpdateConfig): @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: @@ -836,7 +860,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: cache_config.mamba_block_size = model_config.max_model_len -class MiniMaxM3SparseForConditionalGenerationConfig(VerifyAndUpdateConfig): +class MiniMaxM3SparseForConditionalGenerationConfig(MiniMaxM3SparseConfig): INDEX_CACHE_KEYS = ("use_index_cache", "index_topk_freq", "index_skip_topk_offset") @classmethod @@ -1293,6 +1317,7 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "LlamaBidirectionalModel": LlamaBidirectionalConfig, "LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig, "LlamaNemotronVLModel": LlamaNemotronVLConfig, + "MiniMaxM3SparseForCausalLM": MiniMaxM3SparseConfig, "Mamba2ForCausalLM": MambaModelConfig, "MambaForCausalLM": MambaModelConfig, "MiniMaxM3SparseForConditionalGeneration": ( diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index b7959cfc856b..e5eb0959ac97 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -27,12 +27,17 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.modelopt import ModelOptMxFp8Config from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import make_attention_sink +from vllm.model_executor.utils import ( + is_weights_pre_processed, + weights_already_processed, +) from vllm.multimodal.inputs import NestedTensors from vllm.platforms import current_platform from vllm.transformers_utils.config import set_default_rope_theta @@ -59,6 +64,12 @@ _SLIDING_ATTENTION = "sliding_attention" +def _is_mxfp8_linear(layer: nn.Module) -> bool: + return isinstance( + getattr(layer, "quant_config", None), ModelOptMxFp8Config + ) and not isinstance(layer.quant_method, UnquantizedLinearMethod) + + def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool: """Resolve explicit causality before falling back to legacy layer defaults.""" is_causal = getattr(config, "is_causal", None) @@ -182,6 +193,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", attn_type: str = AttentionType.DECODER, + query_pre_quantized: bool = False, ) -> None: super().__init__() self.layer_name = prefix @@ -242,6 +254,7 @@ def __init__( prefix=f"{prefix}.attn", attn_type=attn_type, sinks=self.attention_sink_bias, + query_pre_quantized=query_pre_quantized, ) self.causal = causal if current_platform.is_rocm(): @@ -408,6 +421,9 @@ class DFlashQwen3Model(nn.Module): ".up_proj": (".gate_up_proj", 1), }, ) + # Maps module paths in the draft checkpoint's quantization exclusions to + # this model's module paths. + quant_config_mapper: WeightsMapper | None = None def __init__( self, @@ -422,6 +438,8 @@ def __init__( self.config = speculative_config.draft_model_config.hf_config self.vocab_size = self.config.vocab_size self.quant_config = get_draft_quant_config(vllm_config) + if self.quant_config is not None and self.quant_config_mapper is not None: + self.quant_config.apply_vllm_mapper(self.quant_config_mapper) map_draft_quant_config( self.quant_config, self.hf_to_vllm_mapper, @@ -510,6 +528,7 @@ def _build_context_kv_buffers( self._hidden_norm_weight = self.hidden_norm.weight.data self._context_qkv_projs = [a.qkv_proj for a in layers_attn] self._context_q_sizes = [a.q_size for a in layers_attn] + self._context_kv_proj: ReplicatedLinear | None = None if all( isinstance(proj.quant_method, UnquantizedLinearMethod) @@ -528,6 +547,12 @@ def _build_context_kv_buffers( # consumed by F.linear. Run each projection through its quant method. self._fused_kv_weight = None self._fused_kv_bias = None + if not is_weights_pre_processed() and all( + _is_mxfp8_linear(proj) for proj in self._context_qkv_projs + ): + self._context_kv_proj = self._build_mxfp8_context_kv_proj( + layers_attn, has_bias + ) # K-norm weights stacked into one contiguous [num_layers, head_dim] # tensor so the per-layer K-norm runs as a single grouped kernel. @@ -535,6 +560,35 @@ def _build_context_kv_buffers( [a.k_norm.weight.data for a in layers_attn], dim=0 ).contiguous() + @staticmethod + def _build_mxfp8_context_kv_proj( + layers_attn: list[nn.Module], has_bias: bool + ) -> ReplicatedLinear: + """One MXFP8 GEMM over the K/V rows of every layer's QKV projection. + + MXFP8 block scales are per output row, so the K/V rows and their scales + are copied as-is. Must run before process_weights_after_loading repacks + the weights and scales into kernel layouts. + """ + first = layers_attn[0].qkv_proj + with torch.device(first.weight.device): + proj = ReplicatedLinear( + first.input_size, + sum(2 * a.kv_size for a in layers_attn), + bias=has_bias, + params_dtype=first.params_dtype, + quant_config=first.quant_config, + prefix=first.prefix, + return_bias=False, + disable_tp=True, # QKV weights are already sharded on this rank. + ) + names = ["weight", "weight_scale"] + (["bias"] if has_bias else []) + for name in names: + getattr(proj, name).data.copy_( + torch.cat([getattr(a.qkv_proj, name)[a.q_size :] for a in layers_attn]) + ) + return proj + def _build_fused_kv_buffers(self) -> None: """Build fused weight buffers for precompute_and_store_context_kv. @@ -578,15 +632,7 @@ def _build_fused_kv_buffers(self) -> None: # References to inner Attention layers for direct cache writes self._attn_layers = [layer.self_attn.attn for layer in self.layers] - def _project_context_kv( - self, - context_states: torch.Tensor, - num_ctx: int, - num_layers: int, - num_kv_heads: int, - head_dim: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - # --- Fused KV projection (one GEMM for all layers) --- + def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor: normed_context_states = torch.empty_like(context_states) ops.rms_norm( normed_context_states, @@ -594,11 +640,28 @@ def _project_context_kv( self._hidden_norm_weight, self._rms_norm_eps, ) + return normed_context_states + + def _project_context_kv_token_major( + self, + context_states: torch.Tensor, + num_ctx: int, + num_layers: int, + num_kv_heads: int, + head_dim: int, + ) -> torch.Tensor: + """Context K/V of every layer, ``[num_ctx, L, 2, nkv, hd]``.""" + # --- Fused KV projection (one GEMM for all layers) --- + normed_context_states = self._normalize_context_states(context_states) if self._fused_kv_weight is not None: all_kv_flat = F.linear( normed_context_states, self._fused_kv_weight, self._fused_kv_bias ) all_kv = all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim) + elif self._context_kv_proj is not None: + all_kv = self._context_kv_proj(normed_context_states).view( + num_ctx, num_layers, 2, num_kv_heads, head_dim + ) else: layer_kv = [] for proj, q_size in zip( @@ -609,7 +672,19 @@ def _project_context_kv( qkv[:, q_size:].view(num_ctx, 2, num_kv_heads, head_dim) ) all_kv = torch.stack(layer_kv, dim=1) + return all_kv + def _project_context_kv( + self, + context_states: torch.Tensor, + num_ctx: int, + num_layers: int, + num_kv_heads: int, + head_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + all_kv = self._project_context_kv_token_major( + context_states, num_ctx, num_layers, num_kv_heads, head_dim + ) # Single contiguous copy that separates K/V and transposes to # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous. # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V. @@ -653,7 +728,8 @@ def precompute_and_store_context_kv( "DFlash buffer initialization was skipped. If dummy weights are not " "in use, this may indicate an error in weight loading." ) - self._build_fused_kv_buffers() + with weights_already_processed(): + self._build_fused_kv_buffers() num_ctx = context_states.shape[0] L = self._num_attn_layers diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index 156f55c373ad..cfa06c63cb95 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -231,6 +231,8 @@ def __init__( class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): + model_cls = Qwen3DSparkModel + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: nn.Module.__init__(self) speculative_config = vllm_config.speculative_config @@ -240,7 +242,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: if getattr(self.config, "draft_vocab_size", None) is None: self.config.draft_vocab_size = getattr(self.config, "vocab_size", None) target_layer_num = vllm_config.model_config.get_total_num_hidden_layers() - self.model = Qwen3DSparkModel( + self.model = self.model_cls( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"), start_layer_id=target_layer_num, diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 0092feacd9ba..ccba81ce6197 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -680,6 +680,10 @@ "vllm.models.deepseek_v41", "DSparkDeepseekV4ForCausalLM", ), + "DSparkMiniMaxDraftModel": ( + "vllm.models.minimax_m3.common.dspark", + "MiniMaxM3DSparkForCausalLM", + ), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), "Qwen3OmniDSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), # Serving under method="dflash" means EAGLEConfig rewrites the draft diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index a51f3bb2373c..0e706a390541 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -162,9 +162,6 @@ def _autotune_kimi_k3_kda_qkvg(model: torch.nn.Module) -> None: def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): - from vllm.model_executor.warmup.minimax_m3_msa_warmup import ( - minimax_m3_msa_warmup, - ) from vllm.model_executor.warmup.step5_dsa_warmup import step5_dsa_warmup if not worker.use_v2_model_runner: @@ -249,7 +246,13 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): b12x_warmup(worker, cudagraph_capture_sizes) - minimax_m3_msa_warmup(worker) + if current_platform.is_cuda(): + # The MSA kernels, and the CuTe DSL their modules import, are CUDA-only. + from vllm.model_executor.warmup.minimax_m3_msa_warmup import ( + minimax_m3_msa_warmup, + ) + + minimax_m3_msa_warmup(worker) # Allocate the exact decode-sized workspace, autotune cache misses, and # resolve every CUDA Graph bucket before capture begins. diff --git a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py index 2c8996a93341..699cad10af08 100644 --- a/vllm/model_executor/warmup/minimax_m3_msa_warmup.py +++ b/vllm/model_executor/warmup/minimax_m3_msa_warmup.py @@ -1,21 +1,83 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any + +import torch from vllm.logger import init_logger from vllm.model_executor.warmup.attention_warmup import ( mixed_batch_attention_warmup, ) +from vllm.models.minimax_m3.nvidia.indexer_msa import MiniMaxM3IndexerMSAImpl +from vllm.models.minimax_m3.nvidia.indexer_nvfp4 import MiniMaxM3IndexerNvfp4Impl from vllm.models.minimax_m3.nvidia.model import MiniMaxM3SparseAttention +from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import ( + NVFP4_BLOCK_SCALE_SHIFT, + uses_q8kv4_decode, +) +from vllm.models.minimax_m3.nvidia.sparse_attention_msa import ( + MiniMaxM3SparseMSABaseImpl, +) from vllm.platforms import current_platform from vllm.tracing import instrument if TYPE_CHECKING: + from vllm.config import VllmConfig from vllm.v1.worker.gpu_worker import Worker logger = init_logger(__name__) +# KV-cache dtypes fmha_sm100.warmup plans attention kernels for. +_MSA_KV_CACHE_DTYPES = ("auto", "bfloat16", "fp8", "fp8_e4m3", "nvfp4") + + +def _msa_prebuild_kwargs( + vllm_config: "VllmConfig", module: MiniMaxM3SparseAttention +) -> dict[str, Any] | None: + """``fmha_sm100.warmup`` arguments for the kernels this rank can request.""" + impl = module.impl + indexer = module.indexer.impl + msa = ( + impl + if isinstance(impl, MiniMaxM3SparseMSABaseImpl) + and module.kv_cache_dtype in _MSA_KV_CACHE_DTYPES + else None + ) + attention = msa is not None + if isinstance(indexer, MiniMaxM3IndexerNvfp4Impl): + index_cache_dtype = "nvfp4" + elif isinstance(indexer, MiniMaxM3IndexerMSAImpl): + index_cache_dtype = module.indexer_kv_dtype + else: + index_cache_dtype = None + if not attention and index_cache_dtype is None: + return None + + spec_config = vllm_config.speculative_config + num_spec = spec_config.num_speculative_tokens if spec_config is not None else 0 + q8kv4_decode = attention and uses_q8kv4_decode( + kv_cache_dtype=module.kv_cache_dtype, + num_q_heads=module.num_heads, + num_kv_heads=module.num_kv_heads, + ) + return dict( + kv_cache_dtype=module.kv_cache_dtype if attention else "auto", + index_cache_dtype=index_cache_dtype, + num_q_heads=module.num_heads, + num_kv_heads=module.num_kv_heads, + topk=impl.topk_blocks, + decode_query_lens=range(1, num_spec + 2), + num_index_heads=indexer.num_index_heads, + page_size=impl.block_size, + decode_backend="q8kv4" if q8kv4_decode else "kv_mode3", + block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT, + sparse_decode=msa is not None and msa.use_cutlass_decode, + sparse_prefill=attention, + prefill_backend=msa.prefill_backend if msa is not None else "cute_dsl", + device=torch.device("cuda", torch.accelerator.current_device_index()), + ) + @instrument(span_name="MiniMax M3 MSA warmup") def minimax_m3_msa_warmup(worker: "Worker") -> None: @@ -36,5 +98,14 @@ def minimax_m3_msa_warmup(worker: "Worker") -> None: logger.info("Warming up MiniMax M3 MSA kernels.") - # Cover sparse prefill through the normal model path. + # Build every native kernel the configuration can request at once instead + # of compiling each one serially on first use. + kwargs = _msa_prebuild_kwargs(worker.vllm_config, sparse_module) + if kwargs is not None: + from vllm.third_party.fmha_sm100 import warmup + + logger.info("MiniMax M3 MSA kernels: %s", warmup(**kwargs)) + + # Cover sparse prefill (and its CuTe-DSL kernels) through the normal model + # path. mixed_batch_attention_warmup(worker) diff --git a/vllm/models/common/ops/moe_tail_fusion/__init__.py b/vllm/models/common/ops/moe_tail_fusion/__init__.py new file mode 100644 index 000000000000..3a391caf17df --- /dev/null +++ b/vllm/models/common/ops/moe_tail_fusion/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Building blocks shared by the CuTe DSL kernels that fuse a MoE layer's tail +(top-k finalize, TP all-reduce, norms) across models.""" diff --git a/vllm/models/common/ops/moe_tail_fusion/primitives.py b/vllm/models/common/ops/moe_tail_fusion/primitives.py new file mode 100644 index 000000000000..53a2384bf265 --- /dev/null +++ b/vllm/models/common/ops/moe_tail_fusion/primitives.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from FlashInfer's flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py +# and runtime.py (flashinfer-ai/flashinfer@139af6f6). + +"""CuTe DSL primitives shared by the fused MoE-tail kernels: PTX loads and +stores, NVLS multicast stores, Lamport mailbox helpers, BF16 packing, block +reductions and launch helpers. This module defines no kernel.""" + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor +from cutlass.cutlass_dsl import T, dsl_user_op + +WARP_SIZE = 32 +VEC_BF16 = 8 +QUAD_BF16 = 4 +PACKED_BYTES = 16 +NEGATIVE_ZERO_BF16_BITS = 0x8000 +NEGATIVE_ZERO_BF16_PAIR = 0x80008000 +NEGATIVE_ZERO_F32_BITS = 0x80000000 + +# Lamport mailboxes rotate through three generations; the stage-state words +# hold the next generation to publish into and the one being consumed. +LAMPORT_GENERATIONS = 3 +NEXT_STAGE = 0 +ACTIVE_STAGE = 1 + + +def _asm(result, operands, text, constraints, *, side_effects, loc, ip): + return llvm.inline_asm( + result, + operands, + text, + constraints, + has_side_effects=side_effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# --------------------------------------------------------------------------- +# Global memory +# --------------------------------------------------------------------------- + + +@dsl_user_op +def load_global_u32x4( + pointer: cute.Pointer, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +): + """Load one 128-bit fragment as four u32 registers. + + The volatile form is the Lamport polling load. Marking the asm + side-effecting prevents loop-invariant motion and common-subexpression + elimination across polling iterations. + """ + address = pointer.toint(loc=loc, ip=ip) + opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" + loaded = _asm( + llvm.StructType.get_literal([T.i32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + f"{opcode} {{$0, $1, $2, $3}}, [$4];", + "=r,=r,=r,=r,l", + side_effects=volatile, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(4) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None): + address = pointer.toint(loc=loc, ip=ip) + loaded = _asm( + llvm.StructType.get_literal([T.i32()] * 2), + [address.ir_value(loc=loc, ip=ip)], + "ld.global.v2.u32 {$0, $1}, [$2];", + "=r,=r,l", + side_effects=False, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([2], T.i32(), loc=loc), + [ + llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) + for index in range(2) + ], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def load_global_bf16_as_f32(address: Int64, *, loc=None, ip=None) -> Float32: + return Float32( + _asm( + T.f32(), + [address.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 bits;\n\tld.global.b16 bits, [$1];\n\t" + "cvt.f32.bf16 $0, bits;\n\t}", + "=f,l", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: + address = pointer.toint(loc=loc, ip=ip) + return Uint32( + _asm( + T.i32(), + [address.ir_value(loc=loc, ip=ip)], + "ld.volatile.global.u32 $0, [$1];", + "=r,l", + side_effects=True, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_global_u32x4( + address: Int64, + packed, + *, + volatile: cutlass.Constexpr[bool] = False, + loc=None, + ip=None, +) -> None: + """Store four packed words to an ordinary or NVLS multicast global VA.""" + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)] + opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32" + _asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + f"{opcode} [$0], {{$1, $2, $3, $4}};", + "l,r,r,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32x2(address: Int64, packed, *, loc=None, ip=None) -> None: + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(2)] + _asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "st.global.v2.u32 [$0], {$1, $2};", + "l,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_global_u32(address: Int64, word: Uint32, *, loc=None, ip=None) -> None: + _asm( + None, + [address.ir_value(loc=loc, ip=ip), word.ir_value(loc=loc, ip=ip)], + "st.global.u32 [$0], $1;", + "l,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_async_release_gpu_add_u32( + pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None +) -> None: + """The SM100 arrival primitive FlashInfer's LamportFlags use.""" + address = pointer.toint(loc=loc, ip=ip) + _asm( + None, + [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], + "red.async.release.global.gpu.add.u32 [$0], $1;", + "l,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +# --------------------------------------------------------------------------- +# NVLS multicast stores +# --------------------------------------------------------------------------- + + +@dsl_user_op +def stmc_bf16x2(address: Int64, packed: Uint32, *, loc=None, ip=None) -> None: + _asm( + None, + [address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)], + "multimem.st.relaxed.sys.global.bf16x2 [$0], $1;", + "l,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stmc_bf16x4(address: Int64, packed, *, loc=None, ip=None) -> None: + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(2)] + _asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};", + "l,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None: + """Publish eight BF16 values (four packed words) through NVLS multicast.""" + words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)] + _asm( + None, + [address.ir_value(loc=loc, ip=ip), *words], + "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +# --------------------------------------------------------------------------- +# Lamport sentinels +# +# A BF16 mailbox marks an empty slot with BF16 negative zero in every lane; an +# FP32 mailbox marks it with FP32 negative zero in every word. Senders replace +# real negative zeros with positive zeros so they never read as empty. +# --------------------------------------------------------------------------- + + +@dsl_user_op +def store_lamport_sentinel_u32x4(address: Int64, *, loc=None, ip=None) -> None: + """Reset one BF16 Lamport fragment to BF16 negative zeros.""" + sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip) + _asm( + None, + [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None: + """Reset one FP32 Lamport fragment to four FP32 negative-zero bit patterns.""" + address = pointer.toint(loc=loc, ip=ip) + sentinel = Uint32(NEGATIVE_ZERO_F32_BITS).ir_value(loc=loc, ip=ip) + _asm( + None, + [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@cute.jit +def sanitize_negative_zero_u32(word: Uint32) -> Uint32: + low = Uint16(word & Uint32(0xFFFF)) + high = Uint16(word >> Uint32(16)) + if low == Uint16(NEGATIVE_ZERO_BF16_BITS): + word = word & Uint32(0xFFFF0000) + if high == Uint16(NEGATIVE_ZERO_BF16_BITS): + word = word & Uint32(0x0000FFFF) + return word + + +@cute.jit +def sanitize_negative_zero_u32x2(packed): + sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) + for index in cutlass.range_constexpr(2): + sanitized[index] = sanitize_negative_zero_u32(packed[index]) + return sanitized.load() + + +@cute.jit +def sanitize_negative_zero_u32x4(packed): + """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.""" + sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for index in cutlass.range_constexpr(4): + sanitized[index] = sanitize_negative_zero_u32(packed[index]) + return sanitized.load() + + +@cute.jit +def fragment_has_negative_zero(packed): + """Whether a BF16 fragment still holds the BF16 sentinel in any lane.""" + dirty = False + for index in cutlass.range_constexpr(4): + word = packed[index] + dirty = ( + dirty + | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) + | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) + ) + return dirty + + +@cute.jit +def fragment_is_dirty(packed): + """Whether an FP32 fragment still holds the FP32 sentinel in any word.""" + dirty = packed[0] == Uint32(NEGATIVE_ZERO_F32_BITS) + for index in cutlass.range_constexpr(1, 4): + dirty = dirty | (packed[index] == Uint32(NEGATIVE_ZERO_F32_BITS)) + return dirty + + +# --------------------------------------------------------------------------- +# Shared and distributed shared memory +# --------------------------------------------------------------------------- + + +@dsl_user_op +def map_shared_to_peer( + smem_pointer: cute.Pointer, peer_rank: Int32, *, loc=None, ip=None +) -> Int32: + """Map a local shared-memory slot to the same slot in a peer CTA.""" + address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + return Int32( + _asm( + T.i32(), + [address, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def store_shared_cluster_f32( + remote_address: Int32, value: Float32, *, loc=None, ip=None +) -> None: + _asm( + None, + [remote_address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], + "st.shared::cluster.f32 [$0], $1;", + "r,f", + side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None): + """Load two aligned FP32 partials from local shared memory.""" + address = Int32(pointer.toint(loc=loc, ip=ip)) + loaded = _asm( + llvm.StructType.get_literal([T.f32()] * 2), + [address.ir_value(loc=loc, ip=ip)], + "ld.shared.v2.f32 {$0, $1}, [$2];", + "=f,=f,r", + side_effects=True, + loc=loc, + ip=ip, + ) + return tuple( + Float32(llvm.extractvalue(T.f32(), loaded, [index], loc=loc, ip=ip)) + for index in range(2) + ) + + +@dsl_user_op +def load_shared_f32x4(pointer: cute.Pointer, *, loc=None, ip=None): + """Load four aligned FP32 partials from local shared memory.""" + address = Int32(pointer.toint(loc=loc, ip=ip)) + loaded = _asm( + llvm.StructType.get_literal([T.f32()] * 4), + [address.ir_value(loc=loc, ip=ip)], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", + "=f,=f,=f,=f,r", + side_effects=True, + loc=loc, + ip=ip, + ) + return tuple( + Float32(llvm.extractvalue(T.f32(), loaded, [index], loc=loc, ip=ip)) + for index in range(4) + ) + + +# --------------------------------------------------------------------------- +# BF16 packing and arithmetic +# --------------------------------------------------------------------------- + + +@dsl_user_op +def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): + values = llvm.bitcast( + ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc), + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, VEC_BF16, BFloat16) + + +@dsl_user_op +def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None): + values = llvm.bitcast( + ir.VectorType.get([QUAD_BF16], BFloat16.mlir_type, loc=loc), + packed.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(values, QUAD_BF16, BFloat16) + + +@dsl_user_op +def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): + packed = llvm.bitcast( + ir.VectorType.get([4], T.i32(), loc=loc), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + +@dsl_user_op +def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): + packed = llvm.bitcast( + ir.VectorType.get([2], T.i32(), loc=loc), + values.ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 2, Uint32) + + +@dsl_user_op +def bf16x2_to_u32(values, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.bitcast(T.i32(), values.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + ) + + +@dsl_user_op +def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32: + return Uint32( + _asm( + T.i32(), + [value.ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 bits;\n\tcvt.rn.bf16.f32 bits, $1;\n\t" + "cvt.u32.u16 $0, bits;\n\t}", + "=r,f", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def fma_f32_bf16( + a: BFloat16, b: BFloat16, acc: Float32, *, loc=None, ip=None +) -> Float32: + a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + return Float32( + _asm( + T.f32(), + [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)], + "fma.rn.f32.bf16 $0, $1, $2, $3;", + "=f,h,h,f", + side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +# --------------------------------------------------------------------------- +# Warp and block reductions +# --------------------------------------------------------------------------- + + +@dsl_user_op +def shuffle_sync_idx_u32( + value: Uint32, source_lane: Int32, *, loc=None, ip=None +) -> Uint32: + return Uint32( + _asm( + T.i32(), + [value.ir_value(loc=loc, ip=ip), source_lane.ir_value(loc=loc, ip=ip)], + "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;", + "=r,r,r", + # Preserve full-warp execution across later divergent consumers. + side_effects=True, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def group_leader_block_sum( + value: Float32, + warp_sums: cute.Tensor, + warps: cutlass.Constexpr[int], + leader_stride: cutlass.Constexpr[int], +) -> Float32: + """Block sum of the values held by every ``leader_stride``-th lane.""" + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + for offset in cutlass.range_constexpr(1, WARP_SIZE): + if cutlass.const_expr(offset >= leader_stride and (offset & (offset - 1)) == 0): + value = value + cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + if lane == 0: + cute.arch.store((warp_sums + warp).llvm_ptr, value) + cute.arch.barrier() + + result = Float32(0.0) + if warp == 0: + if lane < Int32(warps): + result = cute.arch.load( + (warp_sums + lane).llvm_ptr, + Float32, + ) + result = cute.arch.warp_reduction_sum(result) + if lane == 0: + cute.arch.store(warp_sums.llvm_ptr, result) + cute.arch.barrier() + return cute.arch.load(warp_sums.llvm_ptr, Float32) + + +@cute.jit +def warp_sum_specialized( + value: Float32, + warp_idx: Int32, + lane: Int32, + warps: cutlass.Constexpr[int], + last_warp_lanes: cutlass.Constexpr[int], + last_warp_mask: cutlass.Constexpr[int], +) -> Float32: + """Warp sum supporting a compile-time partial final warp.""" + if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32): + for offset in cutlass.range_constexpr(16, 0, -1): + # range_constexpr does not provide powers-of-two stepping. + if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): + other = cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=last_warp_mask, + mask_and_clamp=31, + ) + if (lane ^ Int32(offset)) < Int32(last_warp_lanes): + value = value + other + else: + for offset in cutlass.range_constexpr(16, 0, -1): + if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): + value = value + cute.arch.shuffle_sync_bfly( + value, + offset=offset, + mask=-1, + mask_and_clamp=31, + ) + return value + + +# --------------------------------------------------------------------------- +# Launch helpers +# --------------------------------------------------------------------------- + + +class CUDAGraphCompatibleWrapper: + """DLPack view that does not synchronize with the producer stream.""" + + __slots__ = ("tensor",) + + def __init__(self, tensor: torch.Tensor) -> None: + self.tensor = tensor + + def __dlpack__(self, stream=None): + # stream=-1 skips producer sync; CuTe launches on the current stream, + # which is the capturing one under CUDA graphs. + return self.tensor.__dlpack__(stream=-1) + + def __dlpack_device__(self): + return self.tensor.__dlpack_device__() + + +def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor: + return from_dlpack( + CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align + ) + + +def to_cute_dynamic( + tensor: torch.Tensor, alignment: int, *, divisibility: int +) -> cute.Tensor: + """A 1-D view whose length is a runtime shape divisible by ``divisibility``.""" + return to_cute(tensor, alignment).mark_compact_shape_dynamic( + mode=0, divisibility=divisibility + ) + + +def to_cute_dynamic_m( + tensor: torch.Tensor, *, mode: int, assumed_align: int = 16 +) -> cute.Tensor: + """Expose exactly one compact tensor mode as a runtime shape. + + Model dimensions remain part of the compiled tensor type. Only the token + mode is symbolic, so changing M within an op's capacity reuses the same + compiled kernel. + """ + return to_cute(tensor, assumed_align).mark_compact_shape_dynamic( + mode=mode, stride_order=tensor.dim_order() + ) + + +def make_fake_dynamic_compact_tensor(dtype, *, alignment: int, divisibility: int): + return make_fake_compact_tensor( + dtype, (cute.sym_int32(divisibility=divisibility),), assumed_align=alignment + ) + + +def current_cu_stream() -> cuda.CUstream: + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py index 191a928fe2e9..3bbd13cf2f49 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/all_reduce_mhc.py @@ -34,10 +34,11 @@ from cutlass.cute.runtime import make_fake_compact_tensor from vllm.distributed import get_tp_group -from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes - -from .primitives import ( +from vllm.models.common.ops.moe_tail_fusion.primitives import ( + ACTIVE_STAGE, + LAMPORT_GENERATIONS, NEGATIVE_ZERO_BF16_BITS, + NEXT_STAGE, QUAD_BF16, VEC_BF16, WARP_SIZE, @@ -46,6 +47,7 @@ current_cu_stream, f32_to_bf16_bits, fragment_has_negative_zero, + group_leader_block_sum, load_global_bf16_as_f32, load_global_u32x2, load_global_u32x4, @@ -67,10 +69,8 @@ to_cute, to_cute_dynamic, ) +from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes -LAMPORT_GENERATIONS = 3 -NEXT_STAGE = 0 -ACTIVE_STAGE = 1 # One UE8M0 scale per 32 values: four adjacent threads' BF16 fragments. MXFP8_BLOCK = 32 FRAGMENTS_PER_MXFP8_BLOCK = MXFP8_BLOCK // VEC_BF16 @@ -79,41 +79,6 @@ SCALE_TILE_COLS = 4 -@cute.jit -def _group_leader_block_sum( - value: Float32, - warp_sums: cute.Tensor, - warps: cutlass.Constexpr[int], - leader_stride: cutlass.Constexpr[int], -) -> Float32: - lane = cute.arch.lane_idx() - warp = cute.arch.warp_idx() - for offset in cutlass.range_constexpr(1, WARP_SIZE): - if cutlass.const_expr(offset >= leader_stride and (offset & (offset - 1)) == 0): - value = value + cute.arch.shuffle_sync_bfly( - value, - offset=offset, - mask=-1, - mask_and_clamp=31, - ) - if lane == 0: - cute.arch.store((warp_sums + warp).llvm_ptr, value) - cute.arch.barrier() - - result = Float32(0.0) - if warp == 0: - if lane < Int32(warps): - result = cute.arch.load( - (warp_sums + lane).llvm_ptr, - Float32, - ) - result = cute.arch.warp_reduction_sum(result) - if lane == 0: - cute.arch.store(warp_sums.llvm_ptr, result) - cute.arch.barrier() - return cute.arch.load(warp_sums.llvm_ptr, Float32) - - class _SharedOnlyPublishDeviceKernel: def __init__( self, @@ -227,7 +192,7 @@ def kernel( if cutlass.const_expr(self.release_before_store): if block == 0 and tidx == 0: store_global_u32( - stage_state.iterator + ACTIVE_STAGE, + (stage_state.iterator + ACTIVE_STAGE).toint(), stage, ) cute.arch.barrier() @@ -270,7 +235,7 @@ def kernel( if cutlass.const_expr(not self.release_before_store): if block == 0 and tidx == 0: store_global_u32( - stage_state.iterator + ACTIVE_STAGE, + (stage_state.iterator + ACTIVE_STAGE).toint(), stage, ) cute.arch.barrier() @@ -523,7 +488,7 @@ def kernel( if block == 0 and tidx == 0: store_global_u32( - stage_state.iterator + ACTIVE_STAGE, + (stage_state.iterator + ACTIVE_STAGE).toint(), stage, ) cute.arch.barrier() @@ -826,7 +791,7 @@ def kernel( if token == 0 and cluster_rank == 0 and tidx == 0: store_global_u32( - stage_state.iterator + NEXT_STAGE, + (stage_state.iterator + NEXT_STAGE).toint(), (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS), ) if cutlass.const_expr(self.enable_pdl): @@ -861,7 +826,7 @@ def kernel( smem = cutlass.utils.SmemAllocator() warp_sums = smem.allocate_array(Float32, self.warps) cluster_sums = smem.allocate_array(Float32, self.cluster_size) - cta_sum = _group_leader_block_sum( + cta_sum = group_leader_block_sum( thread_sum, warp_sums, self.warps, diff --git a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py b/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py deleted file mode 100644 index 16d5f4721a33..000000000000 --- a/vllm/models/deepseek_v41/nvidia/ops/cute_dsl/primitives.py +++ /dev/null @@ -1,396 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from FlashInfer's flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py -# and runtime.py (flashinfer-ai/flashinfer@139af6f6). - -"""PTX building blocks and launch helpers for the DSV4.1 LL all-reduce.""" - -import cuda.bindings.driver as cuda -import cutlass -import cutlass.cute as cute -import torch -from cutlass import Float32, Int32, Int64, Uint16, Uint32 -from cutlass._mlir import ir -from cutlass._mlir.dialects import llvm, vector -from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor -from cutlass.cutlass_dsl import T, dsl_user_op - -WARP_SIZE = 32 -VEC_BF16 = 8 -QUAD_BF16 = 4 -NEGATIVE_ZERO_BF16_BITS = 0x8000 -NEGATIVE_ZERO_BF16_PAIR = 0x80008000 - - -def _asm(result, operands, text, constraints, *, side_effects, loc, ip): - return llvm.inline_asm( - result, - operands, - text, - constraints, - has_side_effects=side_effects, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def load_global_u32x4( - pointer: cute.Pointer, - *, - volatile: cutlass.Constexpr[bool] = False, - loc=None, - ip=None, -): - address = pointer.toint(loc=loc, ip=ip) - opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" - loaded = _asm( - llvm.StructType.get_literal([T.i32()] * 4), - [address.ir_value(loc=loc, ip=ip)], - f"{opcode} {{$0, $1, $2, $3}}, [$4];", - "=r,=r,=r,=r,l", - side_effects=volatile, - loc=loc, - ip=ip, - ) - packed = vector.from_elements( - ir.VectorType.get([4], T.i32(), loc=loc), - [ - llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) - for index in range(4) - ], - loc=loc, - ip=ip, - ) - return cute.TensorSSA(packed, 4, Uint32) - - -@dsl_user_op -def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None): - address = pointer.toint(loc=loc, ip=ip) - loaded = _asm( - llvm.StructType.get_literal([T.i32()] * 2), - [address.ir_value(loc=loc, ip=ip)], - "ld.global.v2.u32 {$0, $1}, [$2];", - "=r,=r,l", - side_effects=False, - loc=loc, - ip=ip, - ) - packed = vector.from_elements( - ir.VectorType.get([2], T.i32(), loc=loc), - [ - llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) - for index in range(2) - ], - loc=loc, - ip=ip, - ) - return cute.TensorSSA(packed, 2, Uint32) - - -@dsl_user_op -def load_global_bf16_as_f32(address: Int64, *, loc=None, ip=None) -> Float32: - return Float32( - _asm( - T.f32(), - [address.ir_value(loc=loc, ip=ip)], - "{\n\t.reg .b16 bits;\n\tld.global.b16 bits, [$1];\n\t" - "cvt.f32.bf16 $0, bits;\n\t}", - "=f,l", - side_effects=False, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def store_global_u32x4(address: Int64, packed, *, loc=None, ip=None) -> None: - words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)] - _asm( - None, - [address.ir_value(loc=loc, ip=ip), *words], - "st.global.v4.u32 [$0], {$1, $2, $3, $4};", - "l,r,r,r,r", - side_effects=True, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def store_lamport_sentinel_u32x4(address: Int64, *, loc=None, ip=None) -> None: - sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip) - _asm( - None, - [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel], - "st.global.v4.u32 [$0], {$1, $2, $3, $4};", - "l,r,r,r,r", - side_effects=True, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: - address = pointer.toint(loc=loc, ip=ip) - return Uint32( - _asm( - T.i32(), - [address.ir_value(loc=loc, ip=ip)], - "ld.volatile.global.u32 $0, [$1];", - "=r,l", - side_effects=True, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def store_global_u32( - pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None -) -> None: - address = pointer.toint(loc=loc, ip=ip) - _asm( - None, - [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], - "st.global.u32 [$0], $1;", - "l,r", - side_effects=True, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): - values = llvm.bitcast( - ir.VectorType.get([VEC_BF16], cutlass.BFloat16.mlir_type, loc=loc), - packed.ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - return cute.TensorSSA(values, VEC_BF16, cutlass.BFloat16) - - -@dsl_user_op -def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None): - values = llvm.bitcast( - ir.VectorType.get([QUAD_BF16], cutlass.BFloat16.mlir_type, loc=loc), - packed.ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - return cute.TensorSSA(values, QUAD_BF16, cutlass.BFloat16) - - -@dsl_user_op -def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): - packed = llvm.bitcast( - ir.VectorType.get([2], T.i32(), loc=loc), - values.ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - return cute.TensorSSA(packed, 2, Uint32) - - -@dsl_user_op -def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): - packed = llvm.bitcast( - ir.VectorType.get([4], T.i32(), loc=loc), - values.ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - return cute.TensorSSA(packed, 4, Uint32) - - -@cute.jit -def sanitize_negative_zero_u32(word: Uint32) -> Uint32: - low = Uint16(word & Uint32(0xFFFF)) - high = Uint16(word >> Uint32(16)) - if low == Uint16(NEGATIVE_ZERO_BF16_BITS): - word = word & Uint32(0xFFFF0000) - if high == Uint16(NEGATIVE_ZERO_BF16_BITS): - word = word & Uint32(0x0000FFFF) - return word - - -@cute.jit -def sanitize_negative_zero_u32x4(packed): - sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) - for index in cutlass.range_constexpr(4): - sanitized[index] = sanitize_negative_zero_u32(packed[index]) - return sanitized.load() - - -@cute.jit -def sanitize_negative_zero_u32x2(packed): - sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) - for index in cutlass.range_constexpr(2): - sanitized[index] = sanitize_negative_zero_u32(packed[index]) - return sanitized.load() - - -@cute.jit -def fragment_has_negative_zero(packed): - dirty = False - for index in cutlass.range_constexpr(4): - word = packed[index] - dirty = ( - dirty - | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) - | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) - ) - return dirty - - -@dsl_user_op -def map_shared_to_peer( - smem_pointer: cute.Pointer, peer_rank: Int32, *, loc=None, ip=None -) -> Int32: - address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) - return Int32( - _asm( - T.i32(), - [address, peer_rank.ir_value(loc=loc, ip=ip)], - "mapa.shared::cluster.u32 $0, $1, $2;", - "=r,r,r", - side_effects=False, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def store_shared_cluster_f32( - remote_address: Int32, value: Float32, *, loc=None, ip=None -) -> None: - _asm( - None, - [remote_address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], - "st.shared::cluster.f32 [$0], $1;", - "r,f", - side_effects=True, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32: - return Uint32( - _asm( - T.i32(), - [value.ir_value(loc=loc, ip=ip)], - "{\n\t.reg .b16 bits;\n\tcvt.rn.bf16.f32 bits, $1;\n\t" - "cvt.u32.u16 $0, bits;\n\t}", - "=r,f", - side_effects=False, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def shuffle_sync_idx_u32( - value: Uint32, source_lane: Int32, *, loc=None, ip=None -) -> Uint32: - return Uint32( - _asm( - T.i32(), - [value.ir_value(loc=loc, ip=ip), source_lane.ir_value(loc=loc, ip=ip)], - "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;", - "=r,r,r", - # Preserve full-warp execution across later divergent consumers. - side_effects=True, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def stmc_bf16x2(address: Int64, packed: Uint32, *, loc=None, ip=None) -> None: - _asm( - None, - [address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)], - "multimem.st.relaxed.sys.global.bf16x2 [$0], $1;", - "l,r", - side_effects=True, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def stmc_bf16x4(address: Int64, values, *, loc=None, ip=None) -> None: - words = [values[index].ir_value(loc=loc, ip=ip) for index in range(2)] - _asm( - None, - [address.ir_value(loc=loc, ip=ip), *words], - "multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};", - "l,r,r", - side_effects=True, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def stmc_bf16x8(address: Int64, values, *, loc=None, ip=None) -> None: - words = [values[index].ir_value(loc=loc, ip=ip) for index in range(4)] - _asm( - None, - [address.ir_value(loc=loc, ip=ip), *words], - "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};", - "l,r,r,r,r", - side_effects=True, - loc=loc, - ip=ip, - ) - - -class _GraphSafeDLPack: - __slots__ = ("tensor",) - - def __init__(self, tensor: torch.Tensor) -> None: - self.tensor = tensor - - def __dlpack__(self, stream=None): - # stream=-1 skips producer sync; CuTe launches on the current stream, - # which is the capturing one under CUDA graphs. - return self.tensor.__dlpack__(stream=-1) - - def __dlpack_device__(self): - return self.tensor.__dlpack_device__() - - -def to_cute(tensor: torch.Tensor, alignment: int) -> cute.Tensor: - return from_dlpack(_GraphSafeDLPack(tensor.detach()), assumed_align=alignment) - - -def to_cute_dynamic( - tensor: torch.Tensor, alignment: int, *, divisibility: int -) -> cute.Tensor: - return to_cute(tensor, alignment).mark_compact_shape_dynamic( - mode=0, divisibility=divisibility - ) - - -def make_fake_dynamic_compact_tensor(dtype, *, alignment: int, divisibility: int): - return make_fake_compact_tensor( - dtype, (cute.sym_int32(divisibility=divisibility),), assumed_align=alignment - ) - - -def current_cu_stream() -> cuda.CUstream: - return cuda.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py index 87e7e45fa099..3677a3432b3c 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/publishing_gemm.py @@ -101,10 +101,7 @@ QuantizedActivation, as_quantized_activation, ) -from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes - -from ...ops.cute_dsl.all_reduce_mhc import ACTIVE_STAGE, NEXT_STAGE -from ...ops.cute_dsl.primitives import ( +from vllm.models.common.ops.moe_tail_fusion.primitives import ( VEC_BF16, _asm, current_cu_stream, @@ -113,6 +110,9 @@ shuffle_sync_idx_u32, stmc_bf16x8, ) +from vllm.models.deepseek_v41.common.ops.query_quant import mxfp8_scale_bytes + +from ...ops.cute_dsl.all_reduce_mhc import ACTIVE_STAGE, NEXT_STAGE from .. import LOW_BATCH_MAX_ROWS from ..attention.mxfp8 import Tactic, mxfp8_eligible from .all_reduce import BoundaryAllReduce diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py index 99a87c22bcc7..871b3032ddef 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_front_cute.py @@ -74,7 +74,8 @@ from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import T, dsl_user_op -from ...ops.cute_dsl.primitives import load_global_u32x4 +from vllm.models.common.ops.moe_tail_fusion.primitives import load_global_u32x4 + from .routed_expert_route_cute import ( MAX_EXPERTS, WORDS, diff --git a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py index f7120d8948dc..590fe2675a2a 100644 --- a/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py +++ b/vllm/models/deepseek_v41/nvidia/specialized_decode/moe/routed_expert_gemm.py @@ -55,8 +55,8 @@ from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.models.common.ops.moe_tail_fusion.primitives import current_cu_stream -from ...ops.cute_dsl.primitives import current_cu_stream from .routed_expert_front_cute import RoutedExpertFront from .routed_expert_gemm_cute import ( RoutedExpertGemm1, diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py index 67d5cfdf081d..41b10f7d1dc9 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py @@ -15,15 +15,16 @@ import torch.distributed as dist import torch.distributed._symmetric_memory as symm_mem from cutlass import BFloat16, Float32, Int32, Int64, Uint32 +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, vector +from cutlass.cutlass_dsl import T, dsl_user_op from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput - -from .primitives import ( - NUM_LAMPORT_BUFFERS, +from vllm.models.common.ops.moe_tail_fusion.primitives import ( + LAMPORT_GENERATIONS, PACKED_BYTES, VEC_BF16, bf16x8_to_packed_u32x4, - finalize_top16_bf16, fragment_is_dirty, load_global_u32x4, load_shared_f32x2, @@ -32,7 +33,7 @@ map_shared_to_peer, packed_u32x4_to_bf16x8, red_async_release_gpu_add_u32, - sanitize_negative_zero, + sanitize_negative_zero_u32x4, stmc_bf16x8, store_global_u32x4, store_lamport_sentinel_128, @@ -45,6 +46,103 @@ _SEVEN_CTA_MAX_M = 4 +def _make_top16_bf16_finalize_asm() -> str: + lines = [ + "{", + ".reg .pred _valid<16>;", + ".reg .s32 _index<16>;", + ".reg .b32 _weight_words<8>;", + ".reg .b16 _weights<16>;", + ".reg .b32 _data<64>;", + ".reg .b16 _lo, _hi;", + ".reg .f32 _acc<8>;", + ".reg .u64 _address<16>;", + "ld.global.v4.u32 {_index0, _index1, _index2, _index3}, [$5];", + "ld.global.v4.u32 {_index4, _index5, _index6, _index7}, [$5+16];", + "ld.global.v4.u32 {_index8, _index9, _index10, _index11}, [$5+32];", + "ld.global.v4.u32 {_index12, _index13, _index14, _index15}, [$5+48];", + "ld.global.v4.u32 " + "{_weight_words0, _weight_words1, _weight_words2, " + "_weight_words3}, [$6];", + "ld.global.v4.u32 " + "{_weight_words4, _weight_words5, _weight_words6, " + "_weight_words7}, [$6+16];", + ] + for pair in range(8): + lines.append( + f"mov.b32 {{_weights{2 * pair}, _weights{2 * pair + 1}}}, " + f"_weight_words{pair};" + ) + for element in range(8): + lines.append(f"mov.f32 _acc{element}, 0f00000000;") + for route in range(16): + data = 4 * route + lines.extend( + [ + f"setp.ge.s32 _valid{route}, _index{route}, 0;", + f"@_valid{route} mad.wide.s32 _address{route}, " + f"_index{route}, 7168, $4;", + f"@_valid{route} ld.global.v4.u32 " + f"{{_data{data}, _data{data + 1}, _data{data + 2}, " + f"_data{data + 3}}}, [_address{route}];", + ] + ) + for route in range(16): + for pair in range(4): + data = 4 * route + pair + element = 2 * pair + lines.extend( + [ + f"@_valid{route} mov.b32 {{_lo, _hi}}, _data{data};", + f"@_valid{route} fma.rn.f32.bf16 _acc{element}, _lo, " + f"_weights{route}, _acc{element};", + f"@_valid{route} fma.rn.f32.bf16 _acc{element + 1}, _hi, " + f"_weights{route}, _acc{element + 1};", + ] + ) + for pair in range(4): + lines.append(f"cvt.rn.bf16x2.f32 ${pair}, _acc{2 * pair + 1}, _acc{2 * pair};") + lines.append("}") + return "\n".join(lines) + + +_TOP16_BF16_FINALIZE_ASM = _make_top16_bf16_finalize_asm() + + +@dsl_user_op +def finalize_top16_bf16( + gemm2_vector: cute.Pointer, + route_indices: cute.Pointer, + route_weights: cute.Pointer, + *, + loc=None, + ip=None, +): + """Finalize one Kimi K3 top-16 vector with wide metadata loads.""" + addresses = [ + pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + for pointer in (gemm2_vector, route_indices, route_weights) + ] + out = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 4), + addresses, + _TOP16_BF16_FINALIZE_ASM, + "=r,=r,=r,=r,l,l,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + packed = vector.from_elements( + ir.VectorType.get([4], T.i32(), loc=loc), + [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(packed, 4, Uint32) + + def _mapping( tp_size: int, latent_dim: int, @@ -367,7 +465,7 @@ def _token_device( cute.AddressSpace.gmem, assumed_align=16, ) - local_packed = sanitize_negative_zero( + local_packed = sanitize_negative_zero_u32x4( finalize_top16_bf16( gemm2_vector, route_indices, @@ -402,7 +500,7 @@ def _token_device( local_values[element] + values[element].to(Float32) * weight ) - local_packed = sanitize_negative_zero( + local_packed = sanitize_negative_zero_u32x4( bf16x8_to_packed_u32x4(local_values.load().to(BFloat16)) ) else: @@ -412,7 +510,7 @@ def _token_device( cute.AddressSpace.gmem, assumed_align=16, ) - local_packed = sanitize_negative_zero( + local_packed = sanitize_negative_zero_u32x4( load_global_u32x4(local_ptr, volatile=False) ) multicast_offset = ( @@ -592,7 +690,7 @@ def _token_device( arrived = load_volatile_u32(access_counter) while arrived < Uint32(m): arrived = load_volatile_u32(access_counter) - next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS) + next_index = (current_index + Uint32(1)) % Uint32(LAMPORT_GENERATIONS) actual_bytes = Uint32(m) * Uint32(self.tp_size * self.latent_dim * 2) cute.arch.store((latent_flags.iterator + 0).llvm_ptr, next_index) cute.arch.store((latent_flags.iterator + 1).llvm_ptr, current_index) @@ -653,7 +751,7 @@ def _token_device( cute.AddressSpace.gmem, assumed_align=16, ) - local_packed = sanitize_negative_zero( + local_packed = sanitize_negative_zero_u32x4( load_global_u32x4(source_ptr, volatile=False) ) destination_element = current_elements + ( @@ -771,7 +869,7 @@ def _token_device( target = Uint32(m) * Uint32(self.shared_roles) while arrived < target: arrived = load_volatile_u32(access_counter) - next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS) + next_index = (current_index + Uint32(1)) % Uint32(LAMPORT_GENERATIONS) actual_bytes = Uint32(m) * Uint32(self.tp_size * self.shard_dim * 2) cute.arch.store((shared_flags.iterator + 0).llvm_ptr, next_index) cute.arch.store((shared_flags.iterator + 1).llvm_ptr, current_index) @@ -1071,7 +1169,7 @@ def __init__( ) bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2 - routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer + routed_bytes = LAMPORT_GENERATIONS * bytes_per_routed_buffer self._routed_workspace = symm_mem.empty( routed_bytes // 4, dtype=torch.float32, @@ -1080,7 +1178,7 @@ def __init__( self._routed_symm_mem = symm_mem.rendezvous(self._routed_workspace, group) self._routed_workspace.fill_(-0.0) actual_bytes_per_buffer = ( - self._routed_symm_mem.buffer_size // NUM_LAMPORT_BUFFERS // 16 * 16 + self._routed_symm_mem.buffer_size // LAMPORT_GENERATIONS // 16 * 16 ) if actual_bytes_per_buffer < bytes_per_routed_buffer: raise RuntimeError("routed symmetric workspace is too small") @@ -1105,7 +1203,7 @@ def __init__( self._shared_shard = self._shared_output[:, shard_start:shard_end] self._shared_workspace = symm_mem.empty( - (NUM_LAMPORT_BUFFERS, max_m, tp_size, self.shard_dim), + (LAMPORT_GENERATIONS, max_m, tp_size, self.shard_dim), dtype=torch.bfloat16, device=device, ) diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py index ea9d86d6fa85..fd91fad1f398 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_gemm.py @@ -52,10 +52,11 @@ from cutlass.cute.runtime import from_dlpack from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from vllm.models.common.ops.moe_tail_fusion.primitives import CUDAGraphCompatibleWrapper + from .fused_add_multicast_skinny_gemm import ( FusedAddMulticastSkinnyGemmKernel, ) -from .primitives import CUDAGraphCompatibleWrapper def _as_cute(tensor: torch.Tensor, *, dynamic_m: bool = False): diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py index 21b27b278f60..635939e7fcee 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py @@ -14,15 +14,15 @@ from cutlass import BFloat16, Float32, Int64, const_expr from cutlass.cute.runtime import from_dlpack -from .primitives import ( +from vllm.models.common.ops.moe_tail_fusion.primitives import ( CUDAGraphCompatibleWrapper, bf16x2_to_u32, bf16x4_to_packed_u32x2, bf16x8_to_packed_u32x4, fma_f32_bf16, - sanitize_negative_zero, sanitize_negative_zero_u32, sanitize_negative_zero_u32x2, + sanitize_negative_zero_u32x4, store_global_u32, store_global_u32x2, store_global_u32x4, @@ -283,7 +283,7 @@ def kernel( packed, ) else: - packed = sanitize_negative_zero( + packed = sanitize_negative_zero_u32x4( bf16x8_to_packed_u32x4(fused.load()) ) store_global_u32x4( diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py index b2a9100b711c..0eccdd9b8740 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py @@ -13,7 +13,7 @@ import torch from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream -from .primitives import ( +from vllm.models.common.ops.moe_tail_fusion.primitives import ( VEC_BF16, fragment_is_dirty, load_global_u32x4, 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 deleted file mode 100644 index 3840caf24a27..000000000000 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py +++ /dev/null @@ -1,576 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Shared CuTe DSL primitives; this module does not define a CUDA kernel.""" - -from __future__ import annotations - -import cutlass -import cutlass.cute as cute -import torch -from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32 -from cutlass._mlir import ir -from cutlass._mlir.dialects import llvm, vector -from cutlass.cute.runtime import from_dlpack -from cutlass.cutlass_dsl import T, dsl_user_op - -VEC_BF16 = 8 -PACKED_BYTES = 16 -NUM_LAMPORT_BUFFERS = 3 -NEG_ZERO_F32_BITS = 0x80000000 -NEG_ZERO_BF16_BITS = 0x8000 - - -class CUDAGraphCompatibleWrapper: - """DLPack view that does not synchronize with the producer stream.""" - - def __init__(self, tensor: torch.Tensor): - self.tensor = tensor - - def __dlpack__(self, stream=None): - return self.tensor.__dlpack__(stream=-1) - - def __dlpack_device__(self): - return self.tensor.__dlpack_device__() - - -def to_cute(tensor: torch.Tensor, assumed_align: int = 16) -> cute.Tensor: - return from_dlpack( - CUDAGraphCompatibleWrapper(tensor.detach()), assumed_align=assumed_align - ) - - -def to_cute_dynamic_m( - tensor: torch.Tensor, - *, - mode: int, - assumed_align: int = 16, -) -> cute.Tensor: - """Expose exactly one compact tensor mode as a runtime shape. - - Model dimensions remain part of the compiled tensor type. Only the token - mode is symbolic, so changing M within an Op's capacity reuses the same - compiled kernel. - """ - return to_cute(tensor, assumed_align).mark_compact_shape_dynamic( - mode=mode, - stride_order=tensor.dim_order(), - ) - - -@dsl_user_op -def fma_f32_bf16( - a: BFloat16, - b: BFloat16, - acc: Float32, - *, - loc=None, - ip=None, -) -> Float32: - a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) - b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) - result = llvm.inline_asm( - T.f32(), - [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)], - "fma.rn.f32.bf16 $0, $1, $2, $3;", - "=f,h,h,f", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - return Float32(result) - - -@dsl_user_op -def load_global_u32x4( - pointer: cute.Pointer, - *, - volatile: cutlass.Constexpr[bool] = False, - loc=None, - ip=None, -): - """Load one 128-bit fragment as four u32 registers. - - The volatile form is the Lamport polling load. Marking the asm - side-effecting prevents loop-invariant motion and common-subexpression - elimination across polling iterations. - """ - address = pointer.toint(loc=loc, ip=ip) - opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" - out = llvm.inline_asm( - llvm.StructType.get_literal([T.i32()] * 4), - [address.ir_value(loc=loc, ip=ip)], - f"{opcode} {{$0, $1, $2, $3}}, [$4];", - "=r,=r,=r,=r,l", - has_side_effects=volatile, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - packed = vector.from_elements( - ir.VectorType.get([4], T.i32(), loc=loc), - [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)], - loc=loc, - ip=ip, - ) - return cute.TensorSSA(packed, 4, Uint32) - - -def _make_top16_bf16_finalize_asm() -> str: - lines = [ - "{", - ".reg .pred _valid<16>;", - ".reg .s32 _index<16>;", - ".reg .b32 _weight_words<8>;", - ".reg .b16 _weights<16>;", - ".reg .b32 _data<64>;", - ".reg .b16 _lo, _hi;", - ".reg .f32 _acc<8>;", - ".reg .u64 _address<16>;", - "ld.global.v4.u32 {_index0, _index1, _index2, _index3}, [$5];", - "ld.global.v4.u32 {_index4, _index5, _index6, _index7}, [$5+16];", - "ld.global.v4.u32 {_index8, _index9, _index10, _index11}, [$5+32];", - "ld.global.v4.u32 {_index12, _index13, _index14, _index15}, [$5+48];", - "ld.global.v4.u32 " - "{_weight_words0, _weight_words1, _weight_words2, " - "_weight_words3}, [$6];", - "ld.global.v4.u32 " - "{_weight_words4, _weight_words5, _weight_words6, " - "_weight_words7}, [$6+16];", - ] - for pair in range(8): - lines.append( - f"mov.b32 {{_weights{2 * pair}, _weights{2 * pair + 1}}}, " - f"_weight_words{pair};" - ) - for element in range(8): - lines.append(f"mov.f32 _acc{element}, 0f00000000;") - for route in range(16): - data = 4 * route - lines.extend( - [ - f"setp.ge.s32 _valid{route}, _index{route}, 0;", - f"@_valid{route} mad.wide.s32 _address{route}, " - f"_index{route}, 7168, $4;", - f"@_valid{route} ld.global.v4.u32 " - f"{{_data{data}, _data{data + 1}, _data{data + 2}, " - f"_data{data + 3}}}, [_address{route}];", - ] - ) - for route in range(16): - for pair in range(4): - data = 4 * route + pair - element = 2 * pair - lines.extend( - [ - f"@_valid{route} mov.b32 {{_lo, _hi}}, _data{data};", - f"@_valid{route} fma.rn.f32.bf16 _acc{element}, _lo, " - f"_weights{route}, _acc{element};", - f"@_valid{route} fma.rn.f32.bf16 _acc{element + 1}, _hi, " - f"_weights{route}, _acc{element + 1};", - ] - ) - for pair in range(4): - lines.append(f"cvt.rn.bf16x2.f32 ${pair}, _acc{2 * pair + 1}, _acc{2 * pair};") - lines.append("}") - return "\n".join(lines) - - -_TOP16_BF16_FINALIZE_ASM = _make_top16_bf16_finalize_asm() - - -@dsl_user_op -def finalize_top16_bf16( - gemm2_vector: cute.Pointer, - route_indices: cute.Pointer, - route_weights: cute.Pointer, - *, - loc=None, - ip=None, -): - """Finalize one Kimi K3 top-16 vector with wide metadata loads.""" - addresses = [ - pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) - for pointer in (gemm2_vector, route_indices, route_weights) - ] - out = llvm.inline_asm( - llvm.StructType.get_literal([T.i32()] * 4), - addresses, - _TOP16_BF16_FINALIZE_ASM, - "=r,=r,=r,=r,l,l,l", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - packed = vector.from_elements( - ir.VectorType.get([4], T.i32(), loc=loc), - [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)], - loc=loc, - ip=ip, - ) - return cute.TensorSSA(packed, 4, Uint32) - - -@dsl_user_op -def store_global_u32x4( - address: Int64, - packed, - *, - volatile: cutlass.Constexpr[bool] = False, - loc=None, - ip=None, -) -> None: - """Store four packed words to an ordinary or NVLS multicast global VA.""" - words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)] - opcode = "st.volatile.global.v4.u32" if volatile else "st.global.v4.u32" - llvm.inline_asm( - None, - [address.ir_value(loc=loc, ip=ip), *words], - f"{opcode} [$0], {{$1, $2, $3, $4}};", - "l,r,r,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def store_global_u32x2( - address: Int64, - packed, - *, - loc=None, - ip=None, -) -> None: - words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(2)] - llvm.inline_asm( - None, - [address.ir_value(loc=loc, ip=ip), *words], - "st.global.v2.u32 [$0], {$1, $2};", - "l,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def store_global_u32( - address: Int64, - word: Uint32, - *, - loc=None, - ip=None, -) -> None: - llvm.inline_asm( - None, - [ - address.ir_value(loc=loc, ip=ip), - word.ir_value(loc=loc, ip=ip), - ], - "st.global.u32 [$0], $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None: - """Publish eight BF16 values through an NVLS multicast mapping.""" - words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)] - llvm.inline_asm( - None, - [address.ir_value(loc=loc, ip=ip), *words], - "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};", - "l,r,r,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None: - """Reset one Lamport fragment to four FP32 negative-zero bit patterns.""" - address = pointer.toint(loc=loc, ip=ip) - value = Uint32(NEG_ZERO_F32_BITS).ir_value(loc=loc, ip=ip) - llvm.inline_asm( - None, - [address.ir_value(loc=loc, ip=ip), value, value, value, value], - "st.global.v4.u32 [$0], {$1, $2, $3, $4};", - "l,r,r,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def red_async_release_gpu_add_u32( - pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None -) -> None: - """The exact SM100 arrival primitive used by upstream LamportFlags.""" - address = pointer.toint(loc=loc, ip=ip) - llvm.inline_asm( - None, - [ - address.ir_value(loc=loc, ip=ip), - value.ir_value(loc=loc, ip=ip), - ], - "red.async.release.global.gpu.add.u32 [$0], $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: - address = pointer.toint(loc=loc, ip=ip) - return Uint32( - llvm.inline_asm( - T.i32(), - [address.ir_value(loc=loc, ip=ip)], - "ld.volatile.global.u32 $0, [$1];", - "=r,l", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def map_shared_to_peer( - smem_ptr: cute.Pointer, - peer_rank: Int32, - *, - loc=None, - ip=None, -) -> Int32: - """Map a local shared-memory slot to the same slot in a peer CTA.""" - smem_address = smem_ptr.toint(loc=loc, ip=ip).ir_value() - return Int32( - llvm.inline_asm( - T.i32(), - [smem_address, peer_rank.ir_value(loc=loc, ip=ip)], - "mapa.shared::cluster.u32 $0, $1, $2;", - "=r,r,r", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def store_shared_cluster_f32( - remote_address: Int32, - value: Float32, - *, - loc=None, - ip=None, -) -> None: - llvm.inline_asm( - None, - [ - remote_address.ir_value(loc=loc, ip=ip), - value.ir_value(loc=loc, ip=ip), - ], - "st.shared::cluster.f32 [$0], $1;", - "r,f", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None): - """Load two aligned FP32 DSM partials from local shared memory.""" - address = Int32(pointer.toint(loc=loc, ip=ip)) - out = llvm.inline_asm( - llvm.StructType.get_literal([T.f32()] * 2), - [address.ir_value(loc=loc, ip=ip)], - "ld.shared.v2.f32 {$0, $1}, [$2];", - "=f,=f,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - return tuple( - Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(2) - ) - - -@dsl_user_op -def load_shared_f32x4(pointer: cute.Pointer, *, loc=None, ip=None): - """Load four aligned FP32 DSM partials from local shared memory.""" - address = Int32(pointer.toint(loc=loc, ip=ip)) - out = llvm.inline_asm( - llvm.StructType.get_literal([T.f32()] * 4), - [address.ir_value(loc=loc, ip=ip)], - "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", - "=f,=f,=f,=f,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - return tuple( - Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) - ) - - -@dsl_user_op -def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): - target_type = ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc) - values = llvm.bitcast( - target_type, - packed.ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - return cute.TensorSSA(values, VEC_BF16, BFloat16) - - -@dsl_user_op -def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): - target_type = ir.VectorType.get([4], T.i32(), loc=loc) - packed = llvm.bitcast( - target_type, - values.ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - return cute.TensorSSA(packed, 4, Uint32) - - -@dsl_user_op -def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): - target_type = ir.VectorType.get([2], T.i32(), loc=loc) - packed = llvm.bitcast( - target_type, - values.ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - return cute.TensorSSA(packed, 2, Uint32) - - -@dsl_user_op -def bf16x2_to_u32(values, *, loc=None, ip=None): - packed = llvm.bitcast( - T.i32(), - values.ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) - return Uint32(packed) - - -@cute.jit -def sanitize_negative_zero_u32(word): - low = Uint16(word & Uint32(0xFFFF)) - high = Uint16(word >> Uint32(16)) - if low == Uint16(NEG_ZERO_BF16_BITS): - word = word & Uint32(0xFFFF0000) - if high == Uint16(NEG_ZERO_BF16_BITS): - word = word & Uint32(0x0000FFFF) - return word - - -@cute.jit -def sanitize_negative_zero_u32x2(packed): - result = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) - for i in cutlass.range_constexpr(2): - result[i] = sanitize_negative_zero_u32(packed[i]) - return result.load() - - -@cute.jit -def sanitize_negative_zero(packed): - """Turn real BF16 -0 into +0 so it cannot equal the empty sentinel.""" - result = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) - for i in cutlass.range_constexpr(4): - word = packed[i] - low = Uint16(word & Uint32(0xFFFF)) - high = Uint16(word >> Uint32(16)) - if low == Uint16(NEG_ZERO_BF16_BITS): - word = word & Uint32(0xFFFF0000) - if high == Uint16(NEG_ZERO_BF16_BITS): - word = word & Uint32(0x0000FFFF) - result[i] = word - return result.load() - - -@cute.jit -def fragment_is_dirty(packed): - """Bit-exact upstream sentinel check: one comparison per 32-bit word.""" - dirty = packed[0] == Uint32(NEG_ZERO_F32_BITS) - for i in cutlass.range_constexpr(1, 4): - dirty = dirty | (packed[i] == Uint32(NEG_ZERO_F32_BITS)) - return dirty - - -@cute.jit -def warp_sum_specialized( - value: Float32, - warp_idx: Int32, - lane: Int32, - warps: cutlass.Constexpr[int], - last_warp_lanes: cutlass.Constexpr[int], - last_warp_mask: cutlass.Constexpr[int], -) -> Float32: - """Warp sum supporting a compile-time partial final warp.""" - if warp_idx == Int32(warps - 1) and cutlass.const_expr(last_warp_lanes < 32): - for offset in cutlass.range_constexpr(16, 0, -1): - # range_constexpr does not provide powers-of-two stepping. - if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): - other = cute.arch.shuffle_sync_bfly( - value, - offset=offset, - mask=last_warp_mask, - mask_and_clamp=31, - ) - if (lane ^ Int32(offset)) < Int32(last_warp_lanes): - value = value + other - else: - for offset in cutlass.range_constexpr(16, 0, -1): - if cutlass.const_expr(offset in (16, 8, 4, 2, 1)): - value = value + cute.arch.shuffle_sync_bfly( - value, - offset=offset, - mask=-1, - mask_and_clamp=31, - ) - return value diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py index 2833c3a37d13..c6efad7cc3d3 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/common/host.py @@ -11,7 +11,7 @@ import torch from cutlass.cute.runtime import make_fake_tensor -from vllm.models.kimi_k3.nvidia.ops.cute_dsl.latent_moe_tail.primitives import to_cute +from vllm.models.common.ops.moe_tail_fusion.primitives import to_cute _NAMES = ("BFloat16", "Float32", "Uint8", "Uint32", "Int32", "Int64") _DTYPES = {getattr(torch, name.lower()): getattr(cutlass, name) for name in _NAMES} diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py index 7e80d430488e..292734d746da 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/cutedsl/moe/small_batch/front_tp16.py @@ -12,7 +12,8 @@ from cutlass import BFloat16, Float32, Int32, Uint8, Uint32, Uint64 from cutlass.cute.runtime import make_fake_stream, make_fake_tensor -from .....ops.cute_dsl.latent_moe_tail.primitives import to_cute +from vllm.models.common.ops.moe_tail_fusion.primitives import to_cute + from ...common.primitives import ( alloc, asm, diff --git a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py index 8ca3e1e5277c..d739e39bb58c 100644 --- a/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py +++ b/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/prefetch.py @@ -36,7 +36,8 @@ import cuda.bindings.driver as cuda import torch -from ...ops.cute_dsl.latent_moe_tail.primitives import to_cute +from vllm.models.common.ops.moe_tail_fusion.primitives import to_cute + from ..cutedsl.common.host import _checked, resolve_kernel from ..cutedsl.common.policy import batch_table from ..cutedsl.common.prefetch import compile_kernel diff --git a/vllm/models/minimax_m3/amd/indexer_nvfp4.py b/vllm/models/minimax_m3/amd/indexer_nvfp4.py new file mode 100644 index 000000000000..47d67244f3ff --- /dev/null +++ b/vllm/models/minimax_m3/amd/indexer_nvfp4.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax M3 lightning indexer over NVFP4 index-K pages on ROCm (gfx950). + +Keys score as ``e4m3_rn_satfinite(e2m1 * scale)`` with AITER's init / local +sentinels; the top-k writes logical block ids for the NVFP4 attend, no page +table. With ``indexer_decode_cp`` at TP4, decode steps of 4 rows per request +and at most 256 requests score context-parallel (``ops/index_cp``). +""" + +import torch + +from vllm.config import get_current_vllm_config +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.models.minimax_m3.amd.indexer_msa import ( + MSA_INDEX_HEAD_DIM, + MSA_SCORE_TYPE, + MSA_SPARSE_BLOCK_SIZE, + MSA_TOPK_BLOCKS, + MiniMaxM3IndexerMSABackend, + MiniMaxM3IndexerMSAMetadata, + _score_width, +) +from vllm.models.minimax_m3.common.indexer import MiniMaxM3IndexerImpl +from vllm.platforms import current_platform +from vllm.v1.worker.ubatching import dbo_enabled + +logger = init_logger(__name__) + +# Index heads per rank (TP4 / TP2 / TP1). +NUM_INDEX_HEADS = (1, 2, 4) +# Decode CP: rows per request (DSpark 3) and the request cap. +CP_QUERY_LEN = 4 +CP_MAX_REQS = 256 + + +def rocm_nvfp4_indexer_unsupported_reason( + *, + topk_blocks: int, + sparse_block_size: int, + num_index_heads: int, + index_head_dim: int, + score_type: str, +) -> str | None: + """Why the ROCm NVFP4 indexer cannot serve this config, None when it can.""" + if not current_platform.is_rocm(): + return "needs ROCm" + from vllm.platforms.rocm import on_gfx950 + + if not on_gfx950(): + return "needs gfx950" + if (topk_blocks, sparse_block_size, index_head_dim, score_type) != ( + MSA_TOPK_BLOCKS, + MSA_SPARSE_BLOCK_SIZE, + MSA_INDEX_HEAD_DIM, + MSA_SCORE_TYPE, + ): + return ( + f"needs topk_blocks={MSA_TOPK_BLOCKS}, block_size=" + f"{MSA_SPARSE_BLOCK_SIZE}, index_head_dim={MSA_INDEX_HEAD_DIM} and " + f"score_type={MSA_SCORE_TYPE!r}" + ) + if num_index_heads not in NUM_INDEX_HEADS: + return f"needs num_index_heads in {NUM_INDEX_HEADS}, got {num_index_heads}" + return None + + +def init_rocm_nvfp4_indexer_decode_cp(max_model_len: int) -> bool: + """Set up decode CP if configured (a TP collective); returns whether it is on.""" + from vllm.distributed import get_tensor_model_parallel_world_size + + from .ops.index_cp import host as index_cp + + if not get_current_vllm_config().attention_config.indexer_decode_cp: + return False + if get_tensor_model_parallel_world_size() != index_cp.WORLD: + logger.warning_once( + "indexer_decode_cp needs TP %d on ROCm; scoring decode locally.", + index_cp.WORLD, + ) + return False + index_cp.init( + torch.device("cuda", torch.accelerator.current_device_index()), + _score_width(max_model_len, MSA_SPARSE_BLOCK_SIZE), + ) + if index_cp.get() is None: + return False + logger.info_once("MiniMax M3 NVFP4 indexer: decode context parallelism on (TP4)") + return True + + +class MiniMaxM3IndexerRocmNvfp4Impl(MiniMaxM3IndexerImpl): + """FlyDSL NVFP4 index scores + top-k block ids for prefill and decode.""" + + # Its builder materializes the prefill rows' block counts. + indexer_backend_cls = MiniMaxM3IndexerMSABackend + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + if self.indexer_kv_dtype != "nvfp4": + raise ValueError( + f"{type(self).__name__} requires indexer_kv_dtype='nvfp4', got " + f"{self.indexer_kv_dtype!r}" + ) + reason = rocm_nvfp4_indexer_unsupported_reason( + topk_blocks=self.topk_blocks, + sparse_block_size=self.block_size, + num_index_heads=self.num_index_heads, + index_head_dim=self.index_head_dim, + score_type=self.score_type, + ) + if reason is not None: + raise ValueError(f"MiniMax M3 NVFP4 indexer on ROCm {reason}") + + def _new_score(self, rows: int, max_seq_len: int) -> torch.Tensor: + """Uninitialized: the scorer writes every block the top-k reads.""" + return torch.empty( + ( + self.num_index_heads, + rows, + _score_width(max_seq_len, self.block_size), + ), + dtype=torch.float32, + device=self.index_cache.kv_cache.device, + ) + + def _decode_cp(self, num_decodes: int, query_len: int): + from .ops.index_cp import host as index_cp + + cp = index_cp.get() + if ( + cp is None + # The push has no back-pressure; under DBO no all-reduce orders it. + or dbo_enabled() + or self.num_index_heads != 1 + or query_len != CP_QUERY_LEN + or num_decodes > CP_MAX_REQS + or num_decodes * query_len > cp.max_rows + ): + return None + return cp + + def forward( + self, + index_query: torch.Tensor, + **_aiter_table_args, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + from .ops.index_topk_decode.host import ( + minimax_m3_index_topk_decode_indices, + minimax_m3_index_topk_prefill_indices, + ) + from .ops.nvfp4_index_decode.host import minimax_m3_nvfp4_index_score_decode + from .ops.nvfp4_index_prefill.host import ( + minimax_m3_nvfp4_index_score_prefill, + ) + + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return None, None # profiling run; caches unbound + md = attn_metadata[self.index_cache.prefix] + assert isinstance(md, MiniMaxM3IndexerMSAMetadata) + num_tokens = md.num_actual_tokens + nd = md.num_decode_tokens + iq = index_query[:num_tokens].view( + -1, self.num_index_heads, self.index_head_dim + ) + kv = self.index_cache.kv_cache + buf = self.topk_indices_buffer + if buf is None: + buf = torch.empty( + (self.num_index_heads, num_tokens, self.topk_blocks), + dtype=torch.int32, + device=iq.device, + ) + + decode_topk: torch.Tensor | None = None + prefill_topk: torch.Tensor | None = None + if md.num_decodes > 0: + d = md.decode + assert d is not None + cp = self._decode_cp(md.num_decodes, d.decode_query_len) + if cp is not None: + # A view of the receive buffer, valid until the next call. + score = cp.score_push( + iq[:nd, 0], + kv, + d.block_table, + d.seq_lens, + self.init_blocks, + self.local_blocks, + _score_width(d.max_seq_len, self.block_size), + ) + else: + score = self._new_score(nd, d.max_seq_len) + minimax_m3_nvfp4_index_score_decode( + iq[:nd], + kv, + score, + d.block_table, + d.seq_lens, + init_blocks=self.init_blocks, + local_blocks=self.local_blocks, + query_len=d.decode_query_len, + max_seq_len=d.max_seq_len, + ) + decode_topk = buf[:, :nd, :] + minimax_m3_index_topk_decode_indices( + score, + decode_topk, + d.seq_lens, + max_seq_len=d.max_seq_len, + query_len=d.decode_query_len, + ) + + if md.num_prefills > 0: + p = md.prefill + assert p is not None and md.prefill_num_valid_pages is not None + score = self._new_score(num_tokens - nd, p.max_seq_len) + minimax_m3_nvfp4_index_score_prefill( + iq[nd:], + kv, + score, + p.block_table, + p.cu_seqlens_q, + p.seq_lens, + init_blocks=self.init_blocks, + local_blocks=self.local_blocks, + max_query_len=p.max_query_len, + max_seq_len=p.max_seq_len, + ) + prefill_topk = buf[:, nd:num_tokens, :] + minimax_m3_index_topk_prefill_indices( + score, + prefill_topk, + max_seq_len=p.max_seq_len, + num_valid_pages=md.prefill_num_valid_pages, + ) + return decode_topk, prefill_topk diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 424b24e51a87..45ea80f51eee 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -102,6 +102,11 @@ msa_indexer_cp_enabled, select_msa_indexer_impl_cls, ) +from vllm.models.minimax_m3.amd.indexer_nvfp4 import ( + MiniMaxM3IndexerRocmNvfp4Impl, + init_rocm_nvfp4_indexer_decode_cp, +) +from vllm.models.minimax_m3.amd.mono.envs import envs as mono_envs from vllm.models.minimax_m3.amd.ops import ( gemma_fused_add_rmsnorm, gemma_fused_add_rmsnorm_mxfp8, @@ -140,8 +145,12 @@ ) from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms.rocm import on_gfx942 from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype +from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( + triton_reshape_and_cache_flash, +) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheSpec, @@ -154,6 +163,9 @@ logger = init_logger(__name__) +# The fused QK-norm/RoPE/KV-insert HIP op segfaults on gfx942. +_USE_GFX942_NATIVE_QKNORM_ROPE = on_gfx942() + def _sparse_attention_layer_ids(config: PreTrainedConfig) -> set[int]: """Layer ids whose attention runs the extra sparse "index" branch.""" @@ -559,6 +571,19 @@ def forward( hidden_states: torch.Tensor | QuantizedActivation, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) + if _USE_GFX942_NATIVE_QKNORM_ROPE: + num_tokens = qkv.shape[0] + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q = self.q_norm(q.view(num_tokens, self.num_heads, self.head_dim)) + k = self.k_norm(k.view(num_tokens, self.num_kv_heads, self.head_dim)) + q, k = self.rotary_emb.forward_native(positions, q, k) + assert k is not None + q = q.reshape(num_tokens, self.q_size) + k = k.reshape(num_tokens, self.kv_size) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + # Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place (dense # mode: no index branch, no KV-cache insert). Matches nvidia/model.py and # replaces the unfused split -> q_norm/k_norm -> rotary_emb chain; verified @@ -735,6 +760,9 @@ def __init__( # Shared top-k buffer: the indexer writes the selected blocks into it and # the attend impl reads them back (no Python value crosses the break). self.topk_indices_buffer = topk_indices_buffer + # One fused kernel inserts both caches, so NVFP4 K/V needs an NVFP4 + # index cache. + self.use_nvfp4_kv = get_kv_quant_mode(self.kv_cache_dtype).is_nvfp4 # Indexer and main attention are separate impls. This indexer is # ROCm-only, so the model selects it rather than the neutral # MiniMaxM3Indexer, which knows nothing about it and would pick its own. @@ -765,8 +793,10 @@ def __init__( topk_blocks=sparse_cfg["sparse_topk_blocks"], sparse_block_size=sparse_cfg["sparse_block_size"], ) - self.use_aiter_sparse_pa = minimax_m3_use_aiter_sparse_pa( - self.num_kv_heads, emits_sparse_block_table=self.indexer_emits_table + self.use_aiter_sparse_pa = not self.use_nvfp4_kv and ( + minimax_m3_use_aiter_sparse_pa( + self.num_kv_heads, emits_sparse_block_table=self.indexer_emits_table + ) ) self.kv_cache_k = torch.tensor([]) self.kv_cache_v = torch.tensor([]) @@ -793,7 +823,11 @@ def __init__( self.sparse_ctx_buffer = sparse_ctx_buffer # Same condition the forward dispatches on, so the indexer built here # and the branch that calls it cannot disagree about which one this is. - if msa_indexer_impl_cls is None: + if self.use_nvfp4_kv: + self.indexer = MiniMaxM3MSAIndexer( + impl_cls=MiniMaxM3IndexerRocmNvfp4Impl, **indexer_kwargs + ) + elif msa_indexer_impl_cls is None: # Self-contained nn.Module: owns its side cache, selects its impl. self.indexer = MiniMaxM3Indexer(**indexer_kwargs) else: @@ -964,6 +998,118 @@ def _insert_aiter_sparse_pa_kv( return minimax_m3_insert_index_cache(index_k, index_cache, index_slot_mapping) + def _native_gfx942_qknorm_rope_cache( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + fwd_slot_mapping: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run QK norm, RoPE, and cache insertion without the fused HIP op.""" + from vllm.models.minimax_m3.amd.ops.sparse_pa import ( + minimax_m3_insert_index_cache, + ) + + num_tokens = qkv.shape[0] + q, k, v, index_q_raw, index_k_raw = qkv.split( + [ + self.q_size, + self.kv_size, + self.kv_size, + self.index_q_size, + self.idx_head_dim, + ], + dim=-1, + ) + + q = self.q_norm(q.view(num_tokens, self.num_heads, self.head_dim)) + k = self.k_norm(k.view(num_tokens, self.num_kv_heads, self.head_dim)) + q, k = self.rotary_emb.forward_native(positions, q, k) + assert k is not None + q = q.reshape(num_tokens, self.q_size) + + main_slot_mapping = fwd_slot_mapping[self.layer_name] + if self.kv_cache.numel() != 0: + # [block, head, token, 2 * dim] -> [block, token, head, dim]. + key_cache, value_cache = self.kv_cache.transpose(1, 2).split( + self.head_dim, dim=-1 + ) + triton_reshape_and_cache_flash( + k, + v.view(num_tokens, self.num_kv_heads, self.head_dim), + key_cache, + value_cache, + main_slot_mapping, + self.kv_cache_dtype, + self._k_scale, + self._v_scale, + ) + + if self.skip_index_topk: + return q, None + + index_q = self.index_q_norm( + index_q_raw.view(num_tokens, self.num_idx_heads, self.idx_head_dim) + ) + index_k = self.index_k_norm(index_k_raw.view(num_tokens, self.idx_head_dim)) + index_q, index_k = self.index_rotary_emb.forward_native( + positions, index_q, index_k + ) + assert index_k is not None + index_q = index_q.reshape(num_tokens, self.index_q_size).to( + self.indexer.index_cache.dtype + ) + + index_cache = self.indexer.index_cache.kv_cache + if index_cache.numel() != 0: + index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] + minimax_m3_insert_index_cache( + index_k.reshape(num_tokens, self.idx_head_dim), + index_cache, + index_slot_mapping, + ) + return q, index_q + + def _norm_rope_nvfp4_insert( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + index_slot_mapping: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fused norm + RoPE + NVFP4 K/V and index-K insert; returns q, index-q.""" + from vllm.models.minimax_m3.amd.ops.norm_rope_nvfp4_insert.host import ( + minimax_m3_norm_rope_nvfp4_insert, + ) + + num_tokens = qkv.shape[0] + q = qkv.new_empty((num_tokens, self.q_size)) + index_q = qkv.new_empty( + (num_tokens, self.index_q_size), + dtype=self.indexer.index_cache.query_dtype, + ) + minimax_m3_norm_rope_nvfp4_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.rotary_emb.cos_sin_cache, + positions, + slot_mapping, + index_slot_mapping, + self.kv_cache, + self.indexer.index_cache.kv_cache, + self._k_scale, + self._v_scale, + q, + index_q, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + num_index_heads=self.num_idx_heads, + eps=self.q_norm.variance_epsilon, + ) + return q, index_q + def forward( self, positions: torch.Tensor, @@ -993,7 +1139,29 @@ def forward( # Memory-profiling run: caches not yet bound, slot_mapping is empty. return qkv.new_zeros((num_tokens, self.hidden_size)) + if _USE_GFX942_NATIVE_QKNORM_ROPE: + q, index_q = self._native_gfx942_qknorm_rope_cache( + qkv, positions, fwd_slot_mapping + ) + output = torch.empty_like(q) + attn_output = self._run_attention(q, index_q, output) + output, _ = self.o_proj(attn_output) + return output + main_slot_mapping = fwd_slot_mapping[self.layer_name] + if self.use_nvfp4_kv: + q, index_q = self._norm_rope_nvfp4_insert( + qkv, + positions, + main_slot_mapping, + fwd_slot_mapping[self.indexer.index_cache.prefix], + ) + output = torch.empty_like(q) + attn_output = self._run_attention( + q, None if self.skip_index_topk else index_q, output + ) + output, _ = self.o_proj(attn_output) + return output q = qkv.new_empty((num_tokens, self.q_size)) key_cache = value_cache = page16_slot_mapping = None k_scale = v_scale = None @@ -1403,6 +1571,7 @@ def forward( class MiniMaxM3Model(nn.Module, EagleModelMixin): fall_back_to_pt_during_load = False + supports_aux_hidden_states_over_pp = True def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1451,6 +1620,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype( "bf16" ) + use_nvfp4_kv = get_kv_quant_mode( + vllm_config.cache_config.cache_dtype + ).is_nvfp4 # The only MSA indexer probe. It has to live here rather than in # the layer because the buffers it gates are one allocation shared # by all of them, so the answer is needed before any layer exists, @@ -1460,13 +1632,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # this indexer and there is no such choice until it exists; the # probe owes nothing to CP in return, since what it checks is the # platform and the kernels' shape contract, neither of which a - # head count enters. - impl_cls = select_msa_indexer_impl_cls( - topk_blocks=sparse_cfg["sparse_topk_blocks"], - sparse_block_size=sparse_cfg["sparse_block_size"], - index_head_dim=sparse_cfg["sparse_index_dim"], - indexer_kv_dtype=indexer_kv_dtype, - score_type=sparse_cfg.get("sparse_score_type", "max"), + # head count enters. NVFP4 K/V runs its own indexer, which emits + # no page table. + impl_cls = ( + None + if use_nvfp4_kv + else select_msa_indexer_impl_cls( + topk_blocks=sparse_cfg["sparse_topk_blocks"], + sparse_block_size=sparse_cfg["sparse_block_size"], + index_head_dim=sparse_cfg["sparse_index_dim"], + indexer_kv_dtype=indexer_kv_dtype, + score_type=sparse_cfg.get("sparse_score_type", "max"), + ) ) if impl_cls is not None and minimax_m3_use_aiter_sparse_pa( num_kv_heads, emits_sparse_block_table=True @@ -1498,6 +1675,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config, msa_indexer_selected=self.msa_indexer_impl_cls is not None, ) + if use_nvfp4_kv: + # A TP collective (IPC buffers), outside graph capture. + init_rocm_nvfp4_indexer_decode_cp( + vllm_config.model_config.max_model_len + ) else: self.topk_indices_buffer = None @@ -1529,6 +1711,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) + # VLLM_ROCM_MINIMAX_M3_MONO=1: supported decode steps run one FlyDSL + # launch per layer (mono kernels) + self._mono = None + if mono_envs.MONO_ENABLE: + from vllm.models.minimax_m3.amd.mono.dispatch import MonoDecode + + self._mono = MonoDecode(self, vllm_config) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -1548,6 +1737,10 @@ def forward( intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: + if self._mono is not None and self._mono.supports( + input_ids, positions, intermediate_tensors, inputs_embeds + ): + return self._mono.forward(input_ids, positions, inputs_embeds) if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -1560,10 +1753,15 @@ def forward( residual = intermediate_tensors["residual"] decode = self.specialized_decode - # EAGLE3 is not yet compatible with pipeline parallel - aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) + aux_hidden_states: list[torch.Tensor] = [] + if get_pp_group().is_first_rank: + self._maybe_add_hidden_state( + aux_hidden_states, self.start_layer, hidden_states, residual + ) + layers = self.layers[self.start_layer : self.end_layer] if decode is None: - for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]): + for idx, layer in enumerate(layers, start=self.start_layer): hidden_states, residual = layer(positions, hidden_states, residual) self._maybe_add_hidden_state( aux_hidden_states, idx + 1, hidden_states, residual @@ -1572,8 +1770,7 @@ def forward( step = decode.plan(positions, intermediate_tensors) failure: Exception | None = None try: - layers = self.layers[self.start_layer : self.end_layer] - for idx, layer in enumerate(layers): + for idx, layer in enumerate(layers, start=self.start_layer): kernel = step.layer(layer.layer_id) if step else None hidden_states, residual = (kernel or layer)( positions, hidden_states, residual @@ -1587,11 +1784,16 @@ def forward( if not get_pp_group().is_last_rank: return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} + { + "hidden_states": hidden_states, + "residual": residual, + **self.pack_local_aux_hidden_states(aux_hidden_states), + } ) hidden_states, _ = self.norm(hidden_states, residual) + aux_hidden_states = remote_aux + aux_hidden_states if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/models/minimax_m3/amd/mono/__init__.py b/vllm/models/minimax_m3/amd/mono/__init__.py new file mode 100644 index 000000000000..7a79e4b7c0da --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3 fused per-layer decode of up to ``config.MAX_TOKENS`` tokens; entered +through ``dispatch.MonoDecode`` (``VLLM_ROCM_MINIMAX_M3_MONO=1``).""" diff --git a/vllm/models/minimax_m3/amd/mono/check.py b/vllm/models/minimax_m3/amd/mono/check.py new file mode 100644 index 000000000000..bef5dbe7d736 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/check.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Debug tooling for the mono path. + +``forward_checked`` (``VLLM_ROCM_MINIMAX_M3_MONO_CHECK=1``, meant for +``--enforce-eager``) runs every sparse layer twice on the same (ar, res): first +the mono layer kernel, then the original decoder layer. Both insert the same cache +slots (the kernel writes the same bytes up to the re-quantized attention weights), +so running both is side-effect free. Rank 0 logs, per layer and token, the +difference of the residual stream and of the reduced MoE output. +""" + +import torch + +from vllm.logger import init_logger +from vllm.models.minimax_m3.amd.mono.config import HIDDEN + +logger = init_logger(__name__) + + +def _diff(a: torch.Tensor, b: torch.Tensor) -> str: + a, b = a.float().flatten(), b.float().flatten() + d = (a - b).abs() + rel = (d.norm() / b.norm().clamp(min=1e-30)).item() + return f"max {d.max().item():.3e} rel {rel:.3e}" + + +def forward_checked( + runner, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None, +) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: + """``runner.forward`` with the original decoder layer run beside every mono + layer; the mono chain's output is returned.""" + aux_hidden_states: list[torch.Tensor] | None = ( + [] if runner.model.aux_hidden_state_layers else None + ) + res = runner.start(input_ids, positions, inputs_embeds, aux_hidden_states) + n = res.shape[0] + tables = runner.step_tables(n) + ars = runner.ars + for i, lw in enumerate(runner.sparse): + layer = runner.model.layers[lw.layer_id] + h_in = ars[i % 2][:n].clone() + res_in = res.clone() + res = runner.run_sparse_layer(i, lw, tables, positions, res) + ref_h, ref_res = layer(positions, h_in, res_in) + if aux_hidden_states is not None: + runner.model._maybe_add_hidden_state( + aux_hidden_states, + lw.layer_id + 1, + ars[(i + 1) % 2][:n], + res, + ) + if runner.rank == 0: + for k in range(n): + logger.info( + "mono check layer %d token %d/%d pos %d seq_len %d: residual %s | " + "reduced out %s", + lw.layer_id, k, n, int(positions[k]), int(tables.seq_lens[k]), + _diff(res[k], ref_res.view(n, HIDDEN)[k]), + _diff(ars[(i + 1) % 2][k], ref_h.view(n, HIDDEN)[k]), + ) # fmt: skip + hidden = runner.run_final_norm(res) + if aux_hidden_states is not None: + return hidden, aux_hidden_states + return hidden diff --git a/vllm/models/minimax_m3/amd/mono/config.py b/vllm/models/minimax_m3/amd/mono/config.py new file mode 100644 index 000000000000..9ed7777f3cef --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/config.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Shapes the mono kernels are compiled for: one MiniMax-M3 TP4 rank. + +Every value here is checked against the loaded model in ``weights.py``; a +model that does not match is never routed to the mono path. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +# One CTA per MI355X CU; the kernels rely on every CTA being co-resident. +BLOCKS = 256 +THREADS = 512 +WAVES = THREADS // 64 + +HIDDEN = 6144 +HEAD_DIM = 128 +ROTARY_DIM = 64 +LOCAL_Q_HEADS = 16 + + +def qkv_rows(idx_heads: int) -> int: + """Rows of a rank's fused projection, q | k | v | index_q | index_k: 16 q + heads, one k and v head, ``idx_heads`` index q heads (the rank's one, or every + one under indexer context parallelism) and the index k head.""" + return (LOCAL_Q_HEADS + 3 + idx_heads) * HEAD_DIM + + +O_K = LOCAL_Q_HEADS * HEAD_DIM + +N_ROUTED = 128 +TOP_K = 4 +SHARED_EXPERT = N_ROUTED # the fused shared expert's id +MOE_SLOTS = TOP_K + 1 +INTER = 768 # expert intermediate per rank + +SPARSE_BLOCK = 128 +TOPK_BLOCKS = 16 +MAX_SPARSE_KEYS = SPARSE_BLOCK * TOPK_BLOCKS +# the longest context served: the score region spans it +MAX_CONTEXT = 1 << 20 +MAX_INDEX_BLOCKS = MAX_CONTEXT // SPARSE_BLOCK +PAGE16 = 16 + +MAX_TOKENS = 16 # tokens one mono step serves (the MFMA B operand holds 16) +TP = 4 +# indexer context parallelism: a rank computes every index q head +MAX_QKV_ROWS = qkv_rows(TP) +# indexer context parallelism serves a step's requests past this many index +# blocks: its selection costs a fixed ~7 us (S = 1: ~43.5 us/layer at any length) +# that the one-head path's O(n^2) ranking passes between 258 and 297 blocks +# (S = 1 and S = 16 alike; logs/cp_threshold*.log) +INDEX_CP_FROM_BLOCKS = 288 + + +@dataclass(frozen=True) +class IndexHeads: + """The index q heads in a rank's fused projection, a build parameter: + ``count`` of them (1, or all TP in head order under indexer context + parallelism) and ``own``, the one this rank's selection scores.""" + + count: int = 1 + own: int = 0 + + def __post_init__(self): + assert self.count in (1, TP) and 0 <= self.own < self.count + + @property + def rows(self) -> int: + return qkv_rows(self.count) + + @property + def iq_off(self) -> int: + return (LOCAL_Q_HEADS + 2 + self.own) * HEAD_DIM + + @property + def ik_off(self) -> int: + return (LOCAL_Q_HEADS + 2 + self.count) * HEAD_DIM + + +# without indexer context parallelism: the rank's own index q head only +ONE_INDEX_HEAD = IndexHeads() + +LAYER_SLOTS = 128 # mailbox epochs: step * LAYER_SLOTS + layer + 1 + +# host K/V cache geometry, set by the runner before the kernels are built: the +# page-16 ids one sparse block spans (8: K and V as separate planes; 16: a block +# interleaves K|V and V is reached from the same id at +8 pages, as vLLM lays +# it out) and +# whether fp8 K/V carry the host's fixed unit scale instead of per-token scales +KV_PAGE_STRIDE = SPARSE_BLOCK // PAGE16 +STATIC_KV_SCALE = False + + +def configure(*, kv_page_stride: int, static_kv_scale: bool) -> None: + global KV_PAGE_STRIDE, STATIC_KV_SCALE + assert kv_page_stride in (SPARSE_BLOCK // PAGE16, 2 * SPARSE_BLOCK // PAGE16) + KV_PAGE_STRIDE = kv_page_stride + STATIC_KV_SCALE = static_kv_scale + + +class MonoUnsupported(Exception): + """The loaded model or runtime configuration is outside what mono serves.""" diff --git a/vllm/models/minimax_m3/amd/mono/dispatch.py b/vllm/models/minimax_m3/amd/mono/dispatch.py new file mode 100644 index 000000000000..386f9d023bc9 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/dispatch.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Which MiniMax-M3 forwards take the mono path (``VLLM_ROCM_MINIMAX_M3_MONO=1``). + +The deployment is checked once (``_config_refusal``); a step is routed to the +kernels only when it is a pure decode of 1..MAX_TOKENS tokens whose tables the +kernels can read (``_step_refusal``). Every per-step condition is decided from +TP-replicated state: the kernels all-reduce in-kernel, so ranks that disagreed +would wait on each other forever. + +The runner is built on the first eligible step, which under cudagraphs is the +warmup run before the first capture (its allocations and IPC handshake cannot be +captured), and rebuilt when vLLM re-binds the KV cache after its memory-profiling +capture round. +""" + +from typing import Any, cast + +import torch + +from vllm.config import VllmConfig +from vllm.forward_context import ForwardContext, get_forward_context +from vllm.logger import init_logger +from vllm.models.minimax_m3.amd.mono.check import forward_checked +from vllm.models.minimax_m3.amd.mono.config import ( + MAX_CONTEXT, + MAX_TOKENS, + SPARSE_BLOCK, + TP, + MonoUnsupported, +) +from vllm.models.minimax_m3.amd.mono.envs import envs +from vllm.models.minimax_m3.amd.mono.runner import MonoDecodeRunner, dense_prefix + +logger = init_logger(__name__) + + +def _dense_prefix_ok(layers) -> bool: + try: + dense_prefix(layers) + except MonoUnsupported: + return False + return True + + +def _separate_shared_ok(layers) -> bool: + """Every MoE layer keeps its shared expert as a separate MLP (the kernels then + run it from that MLP's MXFP8 weights, ``weights.mxfp8_a_tiles``).""" + moes = [lyr.block_sparse_moe for lyr in layers if lyr.is_moe_layer] + return bool(moes) and all( + not m.is_fused_shared_expert_enabled + and getattr(m, "shared_experts", None) is not None + for m in moes + ) + + +def _config_refusal(vllm_config: VllmConfig, model) -> str | None: + """Why this deployment cannot use the kernels at all, or None.""" + pc = vllm_config.parallel_config + cc = vllm_config.cache_config + cg = vllm_config.compilation_config.cudagraph_mode + max_len = vllm_config.model_config.max_model_len + index_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("bf16") + checks = ( + (pc.tensor_parallel_size == TP, f"TP {pc.tensor_parallel_size} != {TP}"), + (pc.data_parallel_size == 1, "DP > 1"), + (pc.pipeline_parallel_size == 1, "PP > 1"), + (not pc.enable_expert_parallel, "EP"), + (str(cc.cache_dtype).startswith("fp8"), f"kv cache {cc.cache_dtype}"), + (index_dtype == "fp8", f"index cache {index_dtype}"), + (cc.block_size == SPARSE_BLOCK, f"block size {cc.block_size}"), + (max_len <= MAX_CONTEXT, f"max_model_len {max_len} > {MAX_CONTEXT}"), + (cg is None or not cg.requires_piecewise_compilation(), "piecewise cudagraph"), + ( + getattr(model, "is_fused_shared_expert_enabled", False) + or _separate_shared_ok(model.layers), + "shared expert neither fused nor a separate MLP", + ), + (_dense_prefix_ok(model.layers), "layer mix"), + ) + for ok, why in checks: + if not ok: + return why + return None + + +def _step_refusal( + fwd: ForwardContext, layer_name: str, n: int, positions +) -> str | None: + """Why this step cannot run on the kernels, or None when it can.""" + if not 1 <= n <= MAX_TOKENS: + return f"{n} tokens" + if positions.dtype != torch.int64: # read as an int64's low word + return f"positions {positions.dtype}" + if fwd.ubatch_slices is not None: + return "microbatched" + if not isinstance(fwd.attn_metadata, dict) or not isinstance( + fwd.slot_mapping, dict + ): + return "no per-layer metadata" + md = cast(Any, fwd.attn_metadata.get(layer_name)) + if md is None or md.decode is None: + return "no decode metadata" + if md.num_prefills: + return "prefill in the step" + q = md.decode.decode_query_len + if q < 1 or n % q: + return f"{n} tokens for {q} per request" + bt, seq_lens = md.decode.block_table, md.decode.seq_lens + if bt.dtype != torch.int32 or seq_lens.dtype != torch.int32: + return "tables not int32" + if bt.shape[0] < n // q or seq_lens.shape[0] < n // q: + return "tables shorter than the batch" + if bt.stride(0) != bt.shape[1] or bt.stride(1) != 1: # rows bt_width apart + return "block table not row-contiguous" + slots = fwd.slot_mapping.get(layer_name) + if slots is None or slots.dtype != torch.int64 or slots.numel() < n: + return "slot mapping" + return None + + +class MonoDecode: + """Routes supported decode steps of one ``MiniMaxM3Model`` to the kernels.""" + + def __init__(self, model, vllm_config: VllmConfig) -> None: + self._model = model + self._config = vllm_config.model_config.hf_text_config + self._runner: MonoDecodeRunner | None = None + self._enabled = False + self._slots_verified = False + if not envs.MONO_ENABLE: + return + why = _config_refusal(vllm_config, model) + if why is not None: + logger.info("MiniMax-M3 mono decode off: %s", why) + return + attn0 = model.layers[dense_prefix(model.layers)].self_attn + self._layer_name = attn0.layer_name + self._index_prefix = attn0.indexer.index_cache.prefix + self._enabled = True + + def _disable(self, why: str) -> None: + logger.warning("MiniMax-M3 mono decode off: %s", why) + self._enabled = False + + def _ensure_runner(self) -> bool: + """Build (or rebuild after a KV cache re-bind) the runner outside capture.""" + if self._runner is not None and self._runner.stale(): + if torch.cuda.is_current_stream_capturing(): + return False + self._runner.close() + self._runner = None + self._slots_verified = False + logger.info("MiniMax-M3 mono decode: KV cache re-bound, rebuilding") + if self._runner is not None: + return True + if torch.cuda.is_current_stream_capturing(): + return False + try: + self._runner = MonoDecodeRunner(self._model, self._config) + except MonoUnsupported as why: + self._disable(str(why)) + return False + except Exception: + # TP-uniform: every rank fails the same way, so all fall back + logger.exception("MiniMax-M3 mono decode: runner build failed") + self._enabled = False + return False + if self._runner.shared_fp8: + copy = sum(lw.shared_bytes() for lw in self._runner.sparse) + how = f"separate MXFP8 ({copy / 2**30:.2f} GiB tiled copy per rank)" + else: + how = "fused into the MXFP4 expert tables" + logger.info("MiniMax-M3 mono decode on: shared expert %s", how) + return True + + def _slots_shared(self, fwd: ForwardContext, n: int) -> bool: + """The kernels insert K/V and the index key through one page-128 slot + mapping, so the index cache must share the main cache's block ids. + Checked once (a device sync).""" + if self._slots_verified: + return True + slot_mapping = cast(dict[str, Any], fwd.slot_mapping) + slots = slot_mapping[self._layer_name] + index_slots = slot_mapping.get(self._index_prefix) + if index_slots is None or not torch.equal(index_slots[:n], slots[:n]): + self._disable("index cache does not share the KV cache's block ids") + return False + self._slots_verified = True + return True + + def supports( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors, + inputs_embeds: torch.Tensor | None, + ) -> bool: + if not self._enabled or intermediate_tensors is not None: + return False + fwd = get_forward_context() + n = positions.numel() + if _step_refusal(fwd, self._layer_name, n, positions) is not None: + return False + return self._ensure_runner() and self._slots_shared(fwd, n) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None, + ) -> torch.Tensor: + assert self._runner is not None + if envs.MONO_CHECK: + return forward_checked(self._runner, input_ids, positions, inputs_embeds) + return self._runner.forward(input_ids, positions, inputs_embeds) diff --git a/vllm/models/minimax_m3/amd/mono/envs.py b/vllm/models/minimax_m3/amd/mono/envs.py new file mode 100644 index 000000000000..0ef17463c786 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/envs.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Environment switches of the mono decode path (read once at import).""" + +import os + + +def _flag(name: str) -> bool: + return os.getenv(name, "0") in ("1", "true", "True") + + +class _Envs: + # route supported decode steps to the mono kernels + MONO_ENABLE = _flag("VLLM_ROCM_MINIMAX_M3_MONO") + # run every sparse layer twice (kernels, then the original modules) and log the + # per-token differences; the layer-parity self-check, meant for --enforce-eager + MONO_CHECK = _flag("VLLM_ROCM_MINIMAX_M3_MONO_CHECK") + + +envs = _Envs() diff --git a/vllm/models/minimax_m3/amd/mono/kernels/__init__.py b/vllm/models/minimax_m3/amd/mono/kernels/__init__.py new file mode 100644 index 000000000000..c8834edb71b7 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/kernels/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""FlyDSL kernels of the mono decode path: K1 ``pre_attn`` and K4 ``post_attn``.""" diff --git a/vllm/models/minimax_m3/amd/mono/kernels/common.py b/vllm/models/minimax_m3/amd/mono/kernels/common.py new file mode 100644 index 000000000000..b42f1e375e3a --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/kernels/common.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Device-side building blocks shared by the mono kernels. + +Ported from the FlyDSL GLM-5 layer kernel (ROCm/FlyDSL branch +``glm5-mla-moe-layer-monokernel`` @ 1e28fd8a, +``kernels/mla_moe_layer/shared_reuse_moe_kernel.py``): wave reductions, tagged +mailboxes and the MFMA GEMV loop. Only published flydsl APIs and aiter's vendored +``buffer_ops`` / ``dpp_utils`` are used, so nothing here needs the FlyDSL repo. + +Every function but ``kernel_symbol`` is traced inside a ``@flyc.kernel`` body. + +Raw boundaries the kernels keep, for want of an equivalent on the pinned FlyDSL: +``buffer_ops`` loads / stores (the kernels take raw Int64 addresses, no tensors +for ``make_buffer_tensor``); raw MFMA intrinsics (hand-laid operands: permlane +transposes, fp8 packing, MXFP4 scales); the LLVM intrinsics here (``expect``, +``s_memrealtime``, the approximate rsq / rcp / exp2) and the struct unpacking of +``permlane*_swap``; ``inttoptr`` for the timeline's global stores. +""" + +from __future__ import annotations + +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.dpp_utils import update_dpp_i32 +from aiter.ops.flydsl.kernels.kernels_common import kernel_signature +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.compiler.ast_rewriter import ASTRewriter +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import T, as_ir_value + +from vllm.models.minimax_m3.amd.mono.config import LAYER_SLOTS, WAVES + +# gfx95x cache policy bits (LLVM CPol): SC0 = 1, SC1 = 16. SC1 alone is device +# scope (past the per-XCD caches); SC0 | SC1 is system scope (peers over XGMI). +CM_DEV = 16 +CM_SYS = 17 +POLL_MAX = 12 # mailbox specs polled per batch (bounds live registers) + + +def kernel_symbol(stem, **params): + """Kernel symbol that profilers demangle to ``vllm::_``. The + amdhsa assembler rejects ``::`` in a symbol, so it is Itanium-mangled, as + aiter's ``_ZN5aiter...E`` kernels are.""" + name = f"{stem}_{kernel_signature(**params)}" + return f"_ZN4vllm{len(name)}{name}E" + + +def rsrc(addr, nbytes=None): + return bo.create_buffer_resource_from_addr(addr, num_records_bytes=nbytes) + + +def permlane_swap(off, x, y): + """v_permlane{32,16}_swap on registers (x, y): off 32 trades x's upper-half + lanes with y's lower-half lanes; off 16 trades x's odd 16-lane rows with y's + even rows (each within its half). -> (x', y').""" + swap = rocdl.permlane32_swap if off == 32 else rocdl.permlane16_swap + pr = swap( + llvm.StructType.get_literal([T.i32, T.i32]), + as_ir_value(fx.Int32(x)), + as_ir_value(fx.Int32(y)), + False, + False, + ) + return tuple(fx.Int32(llvm.extractvalue(T.i32, pr, [j])) for j in range(2)) + + +def lane_gather(v, src_lane): + """Lane ``src_lane``'s value of v, a source lane a lane (ds_bpermute).""" + return fx.Int32( + rocdl.ds_bpermute(T.i32, fx.Int32(src_lane) * 4, fx.Int32(v).ir_value()) + ) + + +def readlane(v, src_lane): + """Lane ``src_lane``'s value of v (wave-uniform).""" + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + r = fx.Int32(rocdl.readlane(T.i32, bits, fx.Int32(src_lane))) + return r.bitcast(fx.Float32) if isinstance(v, fx.Float32) else r + + +def memrealtime(): + """s_memrealtime: the 100 MHz device clock (no FlyDSL wrapper).""" + return fx.Int64(llvm.call_intrinsic(T.i64, "llvm.amdgcn.s.memrealtime", [], [], [])) + + +def unlikely(cond): + """``cond`` marked cold (llvm.expect): LLVM lays the branch out of the hot + path. A rare path's code inline among the hot stages costs them even when it + never runs (the long-context selection: S = 16, 3k context +1.4 us).""" + return fx.Boolean( + llvm.call_intrinsic( + ir.IntegerType.get_signless(1), + "llvm.expect.i1", + [as_ir_value(cond), as_ir_value(fx.Boolean(False))], + [], + [], + ) + ) + + +def uniform(v): + return fx.Int32(rocdl.readfirstlane(T.i32, fx.Int32(v).ir_value())) + + +def _hw_f32(name, x): + return fx.Float32( + llvm.call_intrinsic(T.f32, name, [fx.Float32(x).ir_value()], [], []) + ) + + +def hw_rsq(x): + """v_rsq_f32 (~1 ulp, no libm range fixups).""" + return _hw_f32("llvm.amdgcn.rsq.f32", x) + + +def hw_rcp(x): + return _hw_f32("llvm.amdgcn.rcp.f32", x) + + +def hw_exp2(x): + return _hw_f32("llvm.amdgcn.exp2.f32", x) + + +def div_rn(x, d, r): + """``x / d`` rounded to nearest, bit-identical to the IEEE division, given + ``r = 1.0 / d`` from one true division per divisor (Markstein: a quotient + within 1 ulp plus its exact fma residual times RN(1/d) rounds correctly).""" + q = x * r + return fx.Float32(fmath.fma(fx.Float32(fmath.fma(-q, d, x)), r, q)) + + +def xshfl(v, off): + """Value of lane ``lane ^ off``: offsets 32 / 16 via v_permlane*_swap, in-row + offsets via DPP (VALU latency, not LDS).""" + if off >= 16: + return v.shuffle_xor(off, 64) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + if off == 8: + y = fx.Int32(update_dpp_i32(x, x, 0x118, 0xF, 0xC, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x108, 0xF, 0x3, False)) + elif off == 4: + y = fx.Int32(update_dpp_i32(x, x, 0x114, 0xF, 0xA, False)) + y = fx.Int32(update_dpp_i32(y, x, 0x104, 0xF, 0x5, False)) + elif off == 2: + y = fx.Int32(update_dpp_i32(x, x, 0x4E, 0xF, 0xF, False)) + else: + y = fx.Int32(update_dpp_i32(x, x, 0xB1, 0xF, 0xF, False)) + return y.bitcast(fx.Float32) if is_f else y + + +def xred(v, off, op): + """op(v, lane ^ off) for a symmetric op; 32 / 16 take both halves of one swap.""" + if off < 16: + return op(v, xshfl(v, off)) + is_f = isinstance(v, fx.Float32) + x = v.bitcast(fx.Int32) if is_f else fx.Int32(v) + a, b = permlane_swap(off, x, x) + if is_f: + return op(a.bitcast(fx.Float32), b.bitcast(fx.Float32)) + return op(type(v)(a), type(v)(b)) + + +def _add(a, b): + return a + b + + +def butterfly(v, offsets, op=_add): + """Butterfly reduction over the given xor offsets, in the given order. + + The order is the numerical contract: the aiter per-head norm reduces 32 + lanes with offsets 16, 8, 4, 2, 1, and matching it keeps rstd bit-exact.""" + for off in offsets: + v = xred(v, off, op) + return v + + +def wave_sum(v): + return butterfly(v, (32, 16, 8, 4, 2, 1)) + + +def traced(fn): + """Apply flydsl's AST rewriting to a helper, so its ``while`` / ``if`` over + traced values lower to scf ops. Kernels rewrite only their own body.""" + return ASTRewriter.transform(fn) + + +def wave_max(v): + return butterfly(v, (32, 16, 8, 4, 2, 1), fx.max) + + +@traced +def block_max(v, lane, wave, red): + """Max of v over the CTA's WAVES waves; ``red`` needs WAVES words.""" + wv = wave_max(v) + gpu.barrier() + if lane == 0: + fx.ptr_store(wv, red + wave) + gpu.barrier() + t = fx.ptr_load(red + 0) + for w in range_constexpr(1, WAVES): + t = fx.max(t, fx.ptr_load(red + w)) + gpu.barrier() + return t + + +def mfma_bf16(a, b, c): + return fx.Vector(rocdl.mfma_f32_16x16x32_bf16(T.vec(4, T.f32), [a, b, c])) + + +def fp8x8_bf16_pk(words): + """Eight fp8 of two dwords -> bf16 MFMA operand (exact) by + v_cvt_scalef32_pk_bf16_fp8 (scale 1), two values an instruction.""" + out = [] + for w in words: + for hi in (False, True): + p = rocdl.cvt_scalef32_pk_bf16_fp8( + T.bf16x2, as_ir_value(fx.Int32(w)), as_ir_value(fx.Float32(1.0)), hi + ) + out.append(fx.Vector(p).bitcast(fx.Int32)[0]) + return fx.Vector.from_elements(out, fx.Int32).bitcast(fx.BFloat16) + + +def bf16_round(a): + return fx.Float32(fx.Float32(a).to(fx.BFloat16)) + + +def bf16_pair(a, b): + """Two f32 -> one f32-typed word holding (bf16(a), bf16(b)).""" + return ( + fx.Vector.from_elements([a, b], fx.Float32) + .to(fx.BFloat16) + .bitcast(fx.Float32)[0] + ) + + +def bf2_f32(w): + """Packed bf16 pair word -> (f32 low, f32 high).""" + return (w << 16).bitcast(fx.Float32), (w & fx.Int32(-65536)).bitcast(fx.Float32) + + +def fp8_pack4(a, b, c, d): + """Four f32 (already in range) -> one dword of four E4M3 bytes, RNE.""" + lo = fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, a, b, fx.Int32(0), False)) + return fx.Int32(rocdl.cvt_pk_fp8_f32(T.i32, c, d, lo, True)) + + +@traced +def _spin(load_all, pending): + """Re-issue ``load_all`` until ``pending`` clears; a side-effecting op in the + loop keeps the loads from being hoisted.""" + v = load_all() + while pending(v): + rocdl.s_sleep(0) + v = load_all() + return v + + +class Mailbox: + """Tagged-pair hand-off between CTAs (device scope) and GPUs (system scope). + + Every 32-bit value is stored next to this launch's epoch tag; a consumer + polls the payload until all tags match, so a hand-off costs one round trip + (no store drain, no separate flag). The epoch is ``step * LAYER_SLOTS + + layer + 1`` with ``step`` a device counter, so the sequence is graph-safe. + """ + + def __init__(self, scratch, step_ptr, layer): + self.scratch = scratch + self.tag = ( + uniform(bo.buffer_load(rsrc(step_ptr), 0, vec_width=1, dtype=T.i32)) + * LAYER_SLOTS + + layer + + 1 + ) + + def addr(self, offset): + return self.scratch + fx.Int64(offset) + + def put(self, base_addr, i, v, cm=CM_DEV): + bits = v.bitcast(fx.Int32) if isinstance(v, fx.Float32) else fx.Int32(v) + bo.buffer_store( + fx.Vector.from_elements([bits, self.tag], fx.Int32), + rsrc(base_addr), + i * 2, + cache_modifier=cm, + ) + + def put_bf(self, base_addr, i, vs, cm=CM_DEV): + """Store 2 or 4 elements as aligned packed bf16 pairs.""" + words = [] + for j in range_constexpr(len(vs) // 2): + words += [bf16_pair(vs[2 * j], vs[2 * j + 1]).bitcast(fx.Int32), self.tag] + bo.buffer_store( + fx.Vector.from_elements(words, fx.Int32), + rsrc(base_addr), + i, + cache_modifier=cm, + ) + + def put_words(self, base_addr, i, words, cm=CM_DEV): + """Raw 32-bit words (1 or 2) at pairs i.. as (word, tag). Integer payloads + (fp8 words) must come this way: through an f32 value a NaN bit pattern is + not kept.""" + vals = [] + for w in words: + vals += [fx.Int32(w), self.tag] + bo.buffer_store( + fx.Vector.from_elements(vals, fx.Int32), + rsrc(base_addr), + i * 2, + cache_modifier=cm, + ) + + def poll(self, specs, scope="agent", batch=POLL_MAX): + """Poll ``(base_addr, pair index, npairs)`` specs in one batch. + + All pairs are re-loaded together while any tag is stale, so a batch costs + one round trip after its last producer lands. ``scope``: "agent" for pairs + another CTA of this GPU wrote, "system" for pairs a peer GPU wrote. Returns + the value words.""" + if const_expr(len(specs) == 0): + return [] + if const_expr(len(specs) > batch): + return self.poll(specs[:batch], scope, batch) + self.poll( + specs[batch:], scope, batch + ) + assert scope in ("agent", "system"), scope + cm = CM_DEV if const_expr(scope == "agent") else CM_SYS + tag = self.tag + + def load_all(): + words = [] + for b, i, n in specs: + w = fx.Vector( + bo.buffer_load( + rsrc(b), + fx.Int32(i) * 2, + vec_width=2 * n, + dtype=T.i32, + cache_modifier=cm, + ) + ) + words += [w[e] for e in range(2 * n)] + return fx.Vector.from_elements(words, fx.Int32) + + nw = sum(2 * n for _, _, n in specs) + + def pending(v): + bad = v[1] != tag + for e in range_constexpr(3, nw, 2): + bad = bad | (v[e] != tag) + return bad + + v = _spin(load_all, pending) + outs, e = [], 0 + for _, _, n in specs: + outs.append([v[e + 2 * q] for q in range(n)]) + e += 2 * n + return outs + + +def ar_pack_sumsq(vals): + """One logical thread of the custom all-reduce epilogue: its 8 elements' + squares accumulated in order (hipcc contracts ``acc += v * v`` into fma).""" + acc = fx.Float32(0.0) + for v in vals: + acc = fmath.fma(v, v, acc) + return acc + + +@traced +def ar_block_sum(p0, p1, lane, wave, red): + """Sum of the 768 per-pack partials of a 6144-wide bf16 row in the order the + 1-stage fused all-reduce + RMSNorm reduces them: 24 warps of 32, a 16..1 + butterfly per warp, then the 24 warp sums butterflied again in one warp. + + Logical warp ``3 * wave + lane // 32`` holds ``p0``; logical warp + ``3 * wave + 2`` holds ``p1`` in lanes < 32 (``WAVES`` = 8). ``red`` needs + 24 words. Matching this order keeps rstd bit-exact with the original path. + """ + return ar_block_sums([(p0, p1)], lane, wave, red)[0] + + +@traced +def ar_block_sums(parts, lane, wave, red): + """``ar_block_sum`` of several rows ``parts = [(p0, p1), ...]`` behind one set + of barriers; ``red`` needs 24 words per row.""" + bs = [ + (butterfly(p0, (16, 8, 4, 2, 1)), butterfly(p1, (16, 8, 4, 2, 1))) + for p0, p1 in parts + ] + gpu.barrier() + for k in range_constexpr(len(bs)): + b0, b1 = bs[k] + if lane % 32 == 0: + fx.ptr_store(b0, red + (24 * k + wave * 3 + lane // 32)) + if lane == 0: + fx.ptr_store(b1, red + (24 * k + wave * 3 + 2)) + gpu.barrier() + li = lane % 32 + tots = [ + butterfly( + (li < 24).select( + fx.ptr_load(red + (24 * k + fx.min(li, 23))), fx.Float32(0.0) + ), + (16, 8, 4, 2, 1), + ) + for k in range(len(parts)) + ] + gpu.barrier() + return tots diff --git a/vllm/models/minimax_m3/amd/mono/kernels/index_score.py b/vllm/models/minimax_m3/amd/mono/kernels/index_score.py new file mode 100644 index 000000000000..aaad81f44796 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/kernels/index_score.py @@ -0,0 +1,474 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Indexer block scores for K4's sparse selection (``select_blocks`` polls them). + +The Triton decode scorer's math, one 128-key block of index-K per task: bf16(fp8 +K) x bf16 index q on the MFMA, max over the causal keys, times fp32(sm_scale * +log2 e). A context of <= TOPK_BLOCKS blocks keeps every block and is not scored; +nor are the blocks the selector pins (the first ``init_blocks``, the last +``local_blocks``): their keys are forced. A scored block waits for the index keys +K1 inserts this step only when it may hold one (``wait_new_keys``). + +A task is one block of one request and runs on one wave: the request's q tokens +(a speculative verify's rows) are the MFMA's B columns, so the block's index-K is +read once for all of them, and no task needs the CTA (no barrier). + +Emitted at the end of K1 (``emit_index_scores``), which makes index q: K4's +selector then finds the scores in place instead of scoring after its launch. +``build_index_score_kernel`` runs the same tasks alone, for K4's harnesses. +""" + +import struct +from typing import NamedTuple + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.kernels_common import LOG2E +from flydsl.expr import const_expr, gpu, range_constexpr +from flydsl.expr.typing import Int32, Int64, T + +from vllm.models.minimax_m3.amd.mono.config import ( + BLOCKS, + HEAD_DIM, + INDEX_CP_FROM_BLOCKS, + MAX_INDEX_BLOCKS, + MAX_TOKENS, + ONE_INDEX_HEAD, + SPARSE_BLOCK, + THREADS, + TOPK_BLOCKS, + TP, + WAVES, + IndexHeads, +) +from vllm.models.minimax_m3.amd.mono.kernels.common import ( + Mailbox, + butterfly, + fp8x8_bf16_pk, + kernel_symbol, + lane_gather, + mfma_bf16, + permlane_swap, + readlane, + rsrc, + traced, + uniform, + unlikely, +) + +INDEX_BLOCK_BYTES = SPARSE_BLOCK * HEAD_DIM # one fp8 index-K block +KEY_TILES = SPARSE_BLOCK // 16 # MFMA row tiles of a block +_CURRENT_STREAM = fx.Stream(None) + + +def _f32(x: float) -> float: + return struct.unpack("f", struct.pack("f", x))[0] + + +def index_scale_log2e(sm_scale: float) -> float: + """fp32(sm_scale * log2 e), the scorer's scale.""" + return _f32(_f32(sm_scale) * _f32(LOG2E)) + + +def scored_count(n_blk, init_blocks: int, local_blocks: int): + """How many blocks (from block init_blocks on) a context of n_blk blocks scores.""" + return (n_blk > TOPK_BLOCKS).select(n_blk - (init_blocks + local_blocks), 0) + + +class StepRows(NamedTuple): + """A decode step's rows as the indexer sees them (``step_rows``).""" + + seq: list # each row's context, >= 1 (a graph's pad rows carry 0) + n_blks: list # each row's index blocks + long: list # each row's request is long + + def any_long(self): + v = self.long[0] + for lg in self.long[1:]: + v = v | lg + return v + + +def step_rows(seq_lens, tokens: int, q_len, heads: IndexHeads) -> StepRows: + """The one place a request is decided long: past INDEX_CP_FROM_BLOCKS index + blocks with every index head (its scores take the context-parallel layout, + ``emit_index_scores``), past THREADS with one (more blocks than + ``rank_blocks`` ranks); either way K4 selects it past the split stage. + + Decided by the request's last row (a request is ``q_len`` consecutive rows): + a speculative verify's rows can sit either side of the threshold, and a + request's rows must all read the layout its scores were written in.""" + seq = [ + fx.max(uniform(bo.buffer_load(rsrc(seq_lens), k, vec_width=1, dtype=T.i32)), 1) + for k in range(tokens) + ] + n_blks = [(s + SPARSE_BLOCK - 1) // SPARSE_BLOCK for s in seq] + req_blks = list(n_blks) + for k in range(tokens - 2, -1, -1): + req_blks[k] = (fx.Int32(k) % q_len == q_len - 1).select( + n_blks[k], req_blks[k + 1] + ) + past = INDEX_CP_FROM_BLOCKS if heads.count > 1 else THREADS + return StepRows(seq, n_blks, [nb > past for nb in req_blks]) + + +def _pick(vals, k): + """vals[k] for a traced index k.""" + v = vals[0] + for j in range_constexpr(1, len(vals)): + v = (k == j).select(vals[j], v) + return v + + +@traced +def emit_index_scores( + slot, lane, wave, red, mb_put, iscore, q_frag, tokens, q_len, init_blocks, + local_blocks, scale, index_cache, block_table, seq_lens, bt_width, + wait_new_keys=None, prepare=None, heads=ONE_INDEX_HEAD, q_frag_head=None, + prepare_heads=None, +): # fmt: skip + """This CTA's (``slot``) share of the step's score tasks: up to BLOCKS tasks, + task u runs on CTA u, else on wave (u // BLOCKS) % WAVES of CTA u % BLOCKS + (``red``: WAVES * 16 floats of LDS). The tokens form requests of ``q_len`` + consecutive rows sharing a block table; the last one has the longest context + and lists the request's scored blocks. ``q_frag(k, s)``: this lane's B operand + for token k, k-chunk s (bf16x8 of index q). Pair k * MAX_INDEX_BLOCKS + b of + mailbox ``iscore`` <- the score of token k's block b. + + ``wait_new_keys()`` (K1 fused into the layer kernel; wave-local) waits for the + index keys this launch inserts: a block within MAX_TOKENS keys of a token's + end may hold an earlier token of its request (speculative decode rows). + Without it, every key this launch inserts must sit in a pinned block + (one-token requests). ``prepare(first, last)`` (with CTA barriers) runs on a + CTA with tasks before any index q is read: K1 makes rows first..last's + there. + + Indexer context parallelism (``heads.count`` > 1): a long request + (``step_rows``) scores only this rank's blocks, b = heads.own + (mod TP) -- every index q head of them, a task's key block read once for all: + pair k * MAX_INDEX_BLOCKS + h * (MAX_INDEX_BLOCKS / TP) + b / TP <- the score + of token k's block b for head h. ``q_frag_head(k, s, h)``: head h's B operand + (every lane its own h); ``prepare_heads()`` (with CTA barriers, on every CTA) + makes those requests' rows' heads first. The other requests score as above.""" + assert local_blocks >= 1, "the block of the current token must be pinned" + l16 = lane % 16 + g4 = lane // 16 + r_ic = rsrc(index_cache) + rows = step_rows(seq_lens, tokens, q_len, heads) + seq = rows.seq + # every request's scored blocks in one list, a task per wave across the grid; + # context-parallel requests' (this rank's blocks) in a second + cp = heads.count > 1 + first_own = init_blocks + (heads.own - init_blocks) % TP + pre = [fx.Int32(0)] + pre_cp = [fx.Int32(0)] + for k in range_constexpr(tokens): + last = fx.Int32(k) % q_len == q_len - 1 + n_scored = scored_count(rows.n_blks[k], init_blocks, local_blocks) + if const_expr(cp): + long = rows.long[k] + n_own = fx.max(init_blocks + n_scored - first_own + TP - 1, 0) // TP + pre.append(pre[-1] + (last & ~long).select(n_scored, fx.Int32(0))) + pre_cp.append(pre_cp[-1] + (last & long).select(n_own, fx.Int32(0))) + else: + pre.append(pre[-1] + last.select(n_scored, fx.Int32(0))) + total = pre[tokens] + total_cp = pre_cp[-1] # 0 without context parallelism + + r_bt = rsrc(block_table) + + def task_block(u, pre=pre, first=init_blocks, stride=1): + """Task u of list ``pre`` -> (the request's last token, block b); per lane + if u is.""" + tk = fx.Int32(0) + for j in range_constexpr(1, tokens): + tk = tk + (u >= pre[j]).select(fx.Int32(1), fx.Int32(0)) + return tk, first + stride * (u - _pick(pre[:tokens], tk)) + + def page_of(tk, b): + return fx.Int32( + bo.buffer_load(r_bt, tk * bt_width + b, vec_width=1, dtype=T.i32) + ) + + def new_keys_in(tk, b): + """Block b of the request ending at token tk may hold keys this launch + inserts (a speculative verify's earlier rows).""" + return (b + 1) * SPARSE_BLOCK > _pick(seq, tk) - MAX_TOKENS + + def cols_of(tk): + """B column l16's token (the request's token l16, the last one repeated) + and its context.""" + col = tk - (q_len - 1) + fx.min(l16, q_len - 1) + return col, _pick(seq, col) + + def load_raw(page, tiles): + """Key tiles ``tiles`` (16 keys each) of a block, in flight: lane g4 = + lane / 16 reads its key's bytes 32 g4 .. 32 g4 + 32 with two 16 B loads + (half the load instructions of eight-byte operand reads, scattered over + the rows: 1M context, S = 16, scores 125 -> 100 us); ``split_tiles`` + makes them MFMA operands.""" + words = [] + for i in tiles: + row = (page * INDEX_BLOCK_BYTES + (16 * i + l16) * HEAD_DIM + 32 * g4) // 4 + for h in range_constexpr(2): + w = fx.Vector( + bo.buffer_load(r_ic, row + 4 * h, vec_width=4, dtype=T.i32) + ) + words += [w[e] for e in range(4)] + return fx.Vector.from_elements(words, fx.Int32) + + def split_tiles(raw, tiles): + """``load_raw``'s words -> per tile the MFMA A operands: tile i holds keys + 16 i + lane % 16, k = 32 s + 8 (lane / 16). A 4 x 4 transpose of the 8 B + parts across the key's four lanes (permlane swaps) gives each lane part + g4 of every chunk s.""" + out = [] + for n in range_constexpr(len(tiles)): + d = [raw[8 * n + e] for e in range(8)] + # part j = dwords 2 j, 2 j + 1; lanes g4 and g4 ^ 2, then g4 ^ 1, trade + for off, pairs in ((32, ((0, 2), (1, 3))), (16, ((0, 1), (2, 3)))): + for a, b in pairs: + for w in range_constexpr(2): + d[2 * a + w], d[2 * b + w] = permlane_swap( + off, d[2 * a + w], d[2 * b + w] + ) + out.append( + [ + fx.Vector.from_elements([d[2 * s], d[2 * s + 1]], fx.Int32) + for s in range(HEAD_DIM // 32) + ] + ) + return out + + def tiles_max(b, qb, col_seq, tiles, kw): + """Max over key tiles ``tiles`` of block b (``kw``: ``split_tiles``) of this + lane's column's causal scores (its B operand ``qb``, per k-chunk; the + column's context ``col_seq``), the same in every lane of the column.""" + v = fx.Float32(float("-inf")) + for n, i in enumerate(tiles): + c = fx.Vector.filled(4, 0.0, fx.Float32) + for s in range_constexpr(HEAD_DIM // 32): + c = mfma_bf16(fp8x8_bf16_pk(kw[n][s]), qb[s], c) + # lane holds keys 16 i + 4 (lane / 16) + e of column lane % 16 + key0 = b * SPARSE_BLOCK + 16 * i + 4 * g4 + for e in range_constexpr(4): + v = fx.max(v, (key0 + e < col_seq).select(c[e], float("-inf"))) + return butterfly(v, (32, 16), fx.max) + + def put(b, col, v): + if (lane < 16) & (l16 < q_len): + mb_put(iscore, col * MAX_INDEX_BLOCKS + b, v * scale) + + def emit_cp_tasks(): + """The context-parallel requests' tasks, a task a wave: B column f of pass + p (f = 16 p + lane % 16) is head f / q_len's row f % q_len of the request, + the key block's operands kept for every pass.""" + if const_expr(prepare_heads is not None): + prepare_heads() + tiles = range(KEY_TILES) + seq_lane = _pick(seq, fx.min(lane, tokens - 1)) # lane k: token k's context + n_cols = TP * q_len + + def score_cp(tk, b, page): + """Block b of the request ending at token tk, at ``page``: every head.""" + if const_expr(wait_new_keys is not None): # noqa: SIM102 + if (b + 1) * SPARSE_BLOCK > readlane(seq_lane, tk) - MAX_TOKENS: + wait_new_keys() + kw = split_tiles(load_raw(page, tiles), tiles) + # every pass unrolled, the unused ones skipped: a runtime pass loop + # slowed the whole kernel's short contexts (S = 1, 3k: +0.6 us) + for p in range_constexpr(TP * MAX_TOKENS // 16): + if p * 16 < n_cols: + f = p * 16 + l16 + head = fx.min(f // q_len, TP - 1) + col = tk - (q_len - 1) + f % q_len + qb = [q_frag_head(col, s, head) for s in range(HEAD_DIM // 32)] + v = tiles_max(b, qb, lane_gather(seq_lane, col), tiles, kw) + if (lane < 16) & (f < n_cols): + mb_put( + iscore, + col * MAX_INDEX_BLOCKS + + head * (MAX_INDEX_BLOCKS // TP) + + b // TP, + v * scale, + ) + + # as the wave mode below: task u on wave (u // BLOCKS) % WAVES of CTA + # u % BLOCKS, lane j decoding task c0 + step j + step = BLOCKS * WAVES + for c0 in range(slot + BLOCKS * wave, total_cp, step * 64): + c0 = fx.Int32(c0) + tk_j, b_j = task_block( + fx.min(c0 + step * lane, total_cp - 1), pre_cp, first_own, TP + ) + pages = page_of(tk_j, b_j) + for i in range(fx.min((total_cp - c0 + step - 1) // step, 64)): + i = fx.Int32(i) + score_cp(readlane(tk_j, i), readlane(b_j, i), readlane(pages, i)) + + # Up to BLOCKS tasks: one a CTA, a key tile a wave, the columns' maxima met in + # LDS; the keys go out before prepare() (only the task's request's rows), so + # their latency hides under the index q's: the scores are on the step's + # critical path (keys after: S = 1, 3k context +1 us). More: a task a wave, + # every row's index q first (a wave's task count differs: no barrier inside), + # the pages of a wave's next 64 tasks fetched at once, a lane each. + if total <= BLOCKS: + for u in range(slot, total, BLOCKS): + tk, b = task_block(fx.Int32(u)) + page = uniform(page_of(tk, b)) + if const_expr(wait_new_keys is not None): # noqa: SIM102 + if new_keys_in(tk, b): + wait_new_keys() + raw = load_raw(page, [wave]) + if const_expr(prepare is not None): + prepare(tk - (q_len - 1), tk) + col, col_seq = cols_of(tk) + qb = [q_frag(col, s) for s in range(HEAD_DIM // 32)] + v = tiles_max(b, qb, col_seq, [wave], split_tiles(raw, [wave])) + if lane < 16: + fx.ptr_store(v, red + (wave * 16 + l16)) + gpu.barrier() + if wave == 0: + for w in range_constexpr(1, WAVES): + v = fx.max(v, fx.ptr_load(red + (w * 16 + l16))) + put(b, col, v) + gpu.barrier() + if total > BLOCKS: + if const_expr(prepare is not None): + prepare(fx.Int32(0), fx.Int32(tokens - 1)) + tiles = range(KEY_TILES) + seq_lane = _pick(seq, fx.min(lane, tokens - 1)) # lane k: token k's context + + def score_task(tk, b, page): + """Block b of the request ending at token tk, at ``page``.""" + if const_expr(wait_new_keys is not None): # noqa: SIM102 + if (b + 1) * SPARSE_BLOCK > readlane(seq_lane, tk) - MAX_TOKENS: + wait_new_keys() + raw = load_raw(page, tiles) + col = tk - (q_len - 1) + fx.min(l16, q_len - 1) + col_seq = lane_gather(seq_lane, col) + qb = [q_frag(col, s) for s in range(HEAD_DIM // 32)] + put(b, col, tiles_max(b, qb, col_seq, tiles, split_tiles(raw, tiles))) + + # Every wave the same task count: task u on wave + # (u // BLOCKS) % WAVES of CTA u % BLOCKS; lane j decodes task c0 + step j + # once -- token, block, page -- read back per task (a task's select chains + # sat between one task's keys and the next: S = 16, +8 us). + step = BLOCKS * WAVES + for c0 in range(slot + BLOCKS * wave, total, step * 64): + c0 = fx.Int32(c0) + tk_j, b_j = task_block(fx.min(c0 + step * lane, total - 1)) + pages = page_of(tk_j, b_j) + for i in range(fx.min((total - c0 + step - 1) // step, 64)): + i = fx.Int32(i) + score_task(readlane(tk_j, i), readlane(b_j, i), readlane(pages, i)) + if const_expr(cp): # noqa: SIM102 + # cold: a step without context-parallel requests runs none of it + if unlikely(total_cp > 0): + emit_cp_tasks() + + +def q_frag_from(iq, heads=ONE_INDEX_HEAD): + """(``q_frag``, ``q_frag_head``) reading index q from memory: [token][head] + [HEAD_DIM] bf16, ``heads.count`` heads a token.""" + + def q_frag_head(tk, s, head): + lane = fx.thread_idx.x % 64 + return fx.Vector( + bo.buffer_load( + rsrc(iq), + ((tk * heads.count + head) * HEAD_DIM + 32 * s + 8 * (lane // 16)) // 2, + vec_width=4, + dtype=T.i32, + ) + ).bitcast(fx.BFloat16) + + def q_frag(tk, s): + return q_frag_head(tk, s, heads.own) + + return q_frag, q_frag_head + + +def build_index_score_kernel( + sm_scale: float, + init_blocks: int, + local_blocks: int, + tokens: int = 1, + heads: IndexHeads = ONE_INDEX_HEAD, +): + """``@flyc.jit`` launcher of the score tasks alone (index q from memory, + ``heads.count`` heads a token), for K4's harnesses: scores into mailbox + ``iscore`` with K4's (step, layer) tag.""" + assert 1 <= tokens <= MAX_TOKENS + scale = index_scale_log2e(sm_scale) + # the JIT cache key holds scalar closure values only: the heads go in as ints + ih_count, ih_own = heads.count, heads.own + + @fx.struct + class Smem: + red: fx.Array[fx.Float32, WAVES * 16, 16] + + kernel_name = kernel_symbol( + "minimax_m3_index_score", + s=tokens, + ib=init_blocks, + lb=local_blocks, + ih=heads.count, + io=heads.own, + ) + + @flyc.kernel(name=kernel_name, known_block_size=[THREADS, 1, 1]) + def index_score_kernel( + iq: Int64, + index_cache: Int64, + block_table: Int64, + seq_lens: Int64, + bt_width: Int32, + q_len: Int32, + iscore: Int64, + step: Int64, + layer: Int32, + ): + tid = fx.thread_idx.x + lds = fx.SharedAllocator().allocate(Smem).peek() + red = lds.red.ptr + mbox = Mailbox(iscore, step, layer) + mb_put = mbox.put + k_heads = IndexHeads(ih_count, ih_own) + q_frag, q_frag_head = q_frag_from(iq, k_heads) + emit_index_scores( + fx.block_idx.x, tid % 64, uniform(tid // 64), red, mb_put, iscore, + q_frag, tokens, q_len, init_blocks, local_blocks, scale, index_cache, + block_table, seq_lens, bt_width, heads=k_heads, q_frag_head=q_frag_head, + ) # fmt: skip + + @flyc.jit + def launch_index_score( + iq: Int64, + index_cache: Int64, + block_table: Int64, + seq_lens: Int64, + bt_width: Int32, + q_len: Int32, + iscore: Int64, + step: Int64, + layer: Int32, + stream: fx.Stream = _CURRENT_STREAM, + ): + index_score_kernel( + iq, + index_cache, + block_table, + seq_lens, + bt_width, + q_len, + iscore, + step, + layer, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch_index_score diff --git a/vllm/models/minimax_m3/amd/mono/kernels/post_attn.py b/vllm/models/minimax_m3/amd/mono/kernels/post_attn.py new file mode 100644 index 000000000000..7824d15b8030 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/kernels/post_attn.py @@ -0,0 +1,2963 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""K4 ``mono_post``: one MiniMax-M3 sparse MoE layer, attention through FFN reduce. + +One launch per layer and rank, ``BLOCKS`` CTAs x ``THREADS`` threads, for a step of +S = 1..MAX_TOKENS tokens (each token's stage tasks side by side); for one token:: + + index : K1, fused in, scored every index-K block (the Triton decode scorer's + block score; with indexer context parallelism a long context's + blocks of this rank, for every index head); every split CTA then + runs the Triton selector's top-k and page-16 table emission itself + (ranking, no sort; for a long request the split tasks' shares' + best TOPK_BLOCKS are ranked, gathered from the peers under CP) + split : one 256-key context partition x 16 heads per task (the gluon decode's + partitioning) over that page-16 table, FP8 QK / PV with its scale + placement + o : combine the partitions (the gluon reduce), per-token FP8 (standalone + quant formula) -> o_proj GEMV + -> push bf16 partial rows to every peer -> rank-ordered sum + -> acc = bf16(sum) + h (h_mid = bf16(acc), mailbox ``a`` = acc) + router: gemma RMSNorm(acc) in the fused all-reduce's reduction order -> bf16 + gate GEMV -> logits; publishes its slice of the normalized input + moe : sigmoid + bias top-4 (+ the fused shared expert) -> MXFP4 a16w4 + up/gate on aiter's shuffled layout -> swiglu -> mid (bf16); + the CTA's down task reuses that routing, its weights prefetched + under the up/gate GEMV: route-weighted down GEMV -> push bf16 + partial rows -> rank-ordered sum -> ar_out = bf16(sum) + +From ``WIDE_FROM`` tokens the MoE groups its work by expert (``wide_moe``): the +router tasks publish xn as MXFP8, each up / gate task pair runs one distinct +expert's 32 rows for every token routed to it (a8w4 MFMA, a B column per token) +and publishes MXFP8 mid rows, and a down task sums over the step's (expert, +token) mids. From ``AG_RS_FROM`` tokens both all-reduces are a reduce-scatter to +the row group's owner rank, which sums in rank order and all-gathers the result. + +The next layer's K1 takes (ar_out, h_mid) exactly as the original path's fused +all-reduce + RMSNorm takes (partials, residual). + +Every stage is a task list; task t runs on CTA (base + t) % BLOCKS and every CTA +walks the stages in order. Dependencies only point to earlier stages and all +CTAs are co-resident, so every spin wait makes progress. +""" + +import struct + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from aiter.ops.flydsl.kernels.act import swiglu_mul_batch +from aiter.ops.flydsl.kernels.kernels_common import LOG2E +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import Int32, Int64, T, as_ir_value + +from vllm.models.minimax_m3.amd.mono import config as mono_config +from vllm.models.minimax_m3.amd.mono.config import ( + BLOCKS, + HEAD_DIM, + HIDDEN, + INDEX_CP_FROM_BLOCKS, + INTER, + MAX_INDEX_BLOCKS, + MAX_TOKENS, + MOE_SLOTS, + N_ROUTED, + ONE_INDEX_HEAD, + PAGE16, + SHARED_EXPERT, + SPARSE_BLOCK, + THREADS, + TOP_K, + TOPK_BLOCKS, + TP, + WAVES, + IndexHeads, +) +from vllm.models.minimax_m3.amd.mono.kernels.common import ( + CM_DEV, + CM_SYS, + Mailbox, + bf2_f32, + bf16_pair, + bf16_round, + butterfly, + fp8_pack4, + hw_exp2, + hw_rcp, + hw_rsq, + kernel_symbol, + memrealtime, + rsrc, + uniform, + unlikely, + wave_max, + wave_sum, + xshfl, +) +from vllm.models.minimax_m3.amd.mono.kernels.common import ( + block_max as block_max_of, +) +from vllm.models.minimax_m3.amd.mono.kernels.index_score import ( + index_scale_log2e, + step_rows, +) +from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import ( + K1_ARGS, + emit_k1, +) +from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import ( + N_HEAD_TASKS as K1_HEAD_TASKS, +) +from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import ( + SCRATCH_HDONE as K1_SCRATCH_HDONE, +) +from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import ( + SCRATCH_RDONE as K1_SCRATCH_RDONE, +) +from vllm.models.minimax_m3.amd.mono.layout import ( + AG_RS_FROM, + DN_ROWS, + MID_ROW, + MIDSC_ROW, + N_DN, + N_O, + N_ROUTER, + N_SPLIT, + O_K, + O_ROW, + O_ROWS, + PAGES_PER_BLOCK, + SCORE_COPIES, + SCRATCH, + SH_ROWS, + SH_TASKS, + SPARSE_TABLE_WORDS, + SPLIT_KEYS, + UG_PER_SLOT, + UG_ROWS, + XN8_ROW, + XSC_ROW, + H, + pool_words, + shared_after_routing, + stage_bases, + sym_layout, + ug_ctas_per_token, + ug_tasks_of, + wide_moe, +) + +FP8_MAX = 448.0 +NEG = -3.4e38 # the gluon decode's masked-score value +PAGE_BYTES = PAGE16 * HEAD_DIM +XN_SLICE = HIDDEN // N_ROUTER # normalized-input elements each router task publishes +W13_ROWS = 2 * INTER # gate rows then up rows (GGUU) +W13_BYTES = W13_ROWS * HIDDEN // 2 # one expert, fp4 +W2_BYTES = HIDDEN * INTER // 2 +W13_SCALE_COLS = HIDDEN // 32 +W2_SCALE_COLS = INTER // 32 +# the separate MXFP8 shared expert (weights.mxfp8_a_tiles): 16-row x 128-k tiles, +# its four tables reached through the ``sh_args`` pointer table (SH_ARGS order) +F8_TILE = 16 * 128 +SH8_W13_BYTES = W13_ROWS * HIDDEN +SH8_W2_BYTES = HIDDEN * INTER +SH8_S13_BYTES = W13_ROWS * W13_SCALE_COLS +SH8_S2_BYTES = HIDDEN * W2_SCALE_COLS +SH_ARGS = ("sh_w13", "sh_s13", "sh_w2", "sh_s2") +SH_W13, SH_S13, SH_W2, SH_S2 = range(len(SH_ARGS)) + +# The launch goes on the caller's current stream (graph capture included). +_CURRENT_STREAM = fx.Stream(None) + + +def _f32(x: float) -> float: + return struct.unpack("f", struct.pack("f", x))[0] + + +TL_POINTS = 32 # timeline stamps per CTA +K1_STAMPS = {1: 27, 2: 28, 3: 29, 4: 30, 6: 31} # fuse_k1: K1's points -> K4's +RANK_SPLIT = 8 # threads ranking one indexer block +RANK_LANES = (4, 2, 1) # butterfly over those RANK_SPLIT lanes +KEY_NONE = -(2**31) # the sort key below every block's +# blocks a thread keys in its split task's eighth of a context (<= MAX_INDEX_BLOCKS) +CAND_BATCH = 2 +assert N_SPLIT * THREADS * CAND_BATCH >= MAX_INDEX_BLOCKS +# Indexer context parallelism: a long context's split task p scans index head +# p // 2's scores of half p % 2 of this rank's blocks -- as many as a TP split task +# scans -- for the head's rank; a request is long past INDEX_CP_FROM_BLOCKS blocks +# (``step_rows``; every long request context-parallel: no TP long selection in that +# build), its rows ending less than a block apart are >= INDEX_CP_FROM_BLOCKS blocks +# each, and a split task then scans >= TOPK_BLOCKS blocks +assert N_SPLIT == 2 * TP +assert N_SPLIT * TOPK_BLOCKS <= INDEX_CP_FROM_BLOCKS <= THREADS +assert MAX_TOKENS <= SPARSE_BLOCK + + +def scale_index(row, col, cols): + """Byte of e8m0 scale (row, col) in aiter's ``shuffle_scale`` layout: + ``view(rows/32, 2, 16, cols/8, 2, 4).permute(0, 3, 5, 2, 4, 1)``.""" + r32, a, b = row // 32, (row // 16) % 2, row % 16 + c8, d, e = col // 8, (col // 4) % 2, col % 4 + return ((((r32 * (cols // 8) + c8) * 4 + e) * 16 + b) * 2 + d) * 2 + a + + +def build_post_attn_kernel( + npes: int, + sm_scale: float, + eps: float, + route_scale: float, + shared_weight: float, + swiglu_limit: float, + init_blocks: int, + local_blocks: int, + tokens: int = 1, + timeline: bool = False, + fuse_k1: bool = False, + heads: IndexHeads = ONE_INDEX_HEAD, + shared_fp8: bool = False, +): + """``@flyc.jit`` launcher of K4 for one rank of an ``npes``-way TP group and a + decode step of ``tokens`` (<= MAX_TOKENS) rows, ``q_len`` consecutive rows a + request (a speculative verify's), with contexts up to MAX_CONTEXT. + ``fuse_k1``: the layer's K1 runs first in the same launch (``K1_ARGS``, + ``positions``, ``slot_mapping``, ``res``; 0 otherwise), for a fused projection + of index q ``heads`` (``IndexHeads``). ``shared_fp8``: the expert tables hold + the routed experts only and the shared expert is MXFP8, its tables behind + the ``sh_args`` pointer table (``SH_ARGS``; 0 otherwise), its GEMVs fp8 x fp8 + MFMAs.""" + assert 1 <= tokens <= MAX_TOKENS + N_E = N_ROUTED if shared_fp8 else N_ROUTED + 1 # experts in w13 .. s2 + E_TOP = N_E - 1 # the bound a routed expert id is clamped to + S13_BYTES = N_E * W13_ROWS * W13_SCALE_COLS # e8m0, every expert + S2_BYTES = N_E * HIDDEN * W2_SCALE_COLS + W = npes + G = BLOCKS + BASE = stage_bases(tokens) + SY = sym_layout(W) + # fuse_k1: K1's outputs are read device-coherent (K1 stores them so) + CM_K1 = CM_DEV if fuse_k1 else 0 + WIDE = wide_moe(tokens) + POOL = pool_words(tokens) + U_MAX = TOP_K * tokens + 1 # distinct routed experts at most, then the shared + + @fx.struct + class Smem: + # activations, bf16 pairs / fp8 quads: every token's xn, o_proj input or + # mid; wide: also the split's PV partials (the stages take turns) + x: fx.Array[fx.Int32, POOL, 16] + red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] + p8: fx.Array[fx.Int32, H * SPLIT_KEYS // 4, 16] + q8: fx.Array[fx.Int32, O_K // 4, 16] + opart: fx.Array[fx.Float32, 1 if WIDE else WAVES * O_K, 16] # PV partials + # wide MoE: the step's distinct experts (the shared one last), token masks + uexp: fx.Array[fx.Int32, U_MAX + 2, 16] + umask: fx.Array[fx.Int32, U_MAX + 2, 16] + # and per (expert, B column) the down stage's mid row and route weight + urow: fx.Array[fx.Int32, U_MAX * 16 if WIDE else 1, 16] + uwt: fx.Array[fx.Float32, U_MAX * 16 if WIDE else 1, 16] + hst: fx.Array[ + fx.Float32, 3 * WAVES * H + WAVES, 16 + ] # head max / sum, max, v_scale max + misc: fx.Array[fx.Float32, O_ROWS * MAX_TOKENS, 16] # a task's output rows + route: fx.Array[fx.Int32, 8 * MAX_TOKENS, 16] # token k's slots at 8 k + rwt: fx.Array[fx.Float32, 8 * MAX_TOKENS, 16] # and their routing weights + # indexer sort keys, a block a thread (a long request: the waves' + # candidate (key, block) pairs) and their ranks + keys: fx.Array[fx.Int32, THREADS, 16] + ranks: fx.Array[fx.Int32, THREADS, 16] + blk: fx.Array[fx.Int32, 32, 16] # this split's 16 pages, n_ctx, tail flag + tls: fx.Array[fx.Int64, TL_POINTS if timeline else 1, 16] # timeline stamps + # fuse_k1: index q (every head's: in the pool) + qs: fx.Array[fx.Float32, HEAD_DIM // 2 * tokens if heads.count == 1 else 1, 16] + + kernel_name = kernel_symbol( + "minimax_m3_mono_layer" if fuse_k1 else "minimax_m3_post_attn", + tp=npes, + s=tokens, + ib=init_blocks, + lb=local_blocks, + ih=heads.count, + io=heads.own, + tl=timeline, + kps=mono_config.KV_PAGE_STRIDE, + sks=mono_config.STATIC_KV_SCALE, + sf8=shared_fp8, + ) + + # the JIT cache key holds the kernel's scalar closure values, not objects: the + # index heads go in as ints (a shared cache otherwise serves one rank's build + # to every rank) + ih_count, ih_own = heads.count, heads.own + # fuse_k1 with every index head: K1's index q of every head (qs) after its x8 + # rows in the pool (K1 is done with the pool before K4's stages use it) + QS_CP = tokens * HIDDEN // 4 + assert not (fuse_k1 and heads.count > 1) or ( + QS_CP + HEAD_DIM // 2 * tokens * heads.count <= POOL + ) + + @flyc.kernel(name=kernel_name, known_block_size=[THREADS, 1, 1]) + def post_attn_kernel( + h_in: Int64, + q: Int64, + block_table: Int64, + seq_lens: Int64, + k_cache: Int64, + v_cache: Int64, + k_scale: Int64, + v_scale: Int64, + w_o: Int64, + s_o: Int64, + g_post: Int64, + w_gate: Int64, + bias: Int64, + w13: Int64, + s13: Int64, + w2: Int64, + s2: Int64, + sh_args: Int64, + h_mid: Int64, + ar_out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + step: Int64, + rank: Int32, + layer: Int32, + bt_width: Int32, + q_len: Int32, + tl: Int64, + k1_args: Int64, + positions: Int64, + slot_mapping: Int64, + res: Int64, + ): + tid = fx.thread_idx.x + bid = fx.block_idx.x + lane = tid % 64 + wave = tid // 64 + g4 = lane // 16 + l16 = lane % 16 + lds = fx.SharedAllocator().allocate(Smem).peek() + xs = lds.x.ptr + red = lds.red.ptr + p8 = lds.p8.ptr + q8 = lds.q8.ptr + # the pool read and written as f32 (wide: the split's PV partials live there) + fpool = fx.recast_iter(fx.Float32, xs) + opart = fpool if WIDE else lds.opart.ptr + uexp = lds.uexp.ptr + umask = lds.umask.ptr + urow = lds.urow.ptr + uwt = lds.uwt.ptr + hst = lds.hst.ptr + misc = lds.misc.ptr + route = lds.route.ptr + rwt = lds.rwt.ptr + keys = lds.keys.ptr + ranks = lds.ranks.ptr + blk = lds.blk.ptr + tls = lds.tls.ptr + v4f = fx.Vector.make_type(4, fx.Float32) + v2i = fx.Vector.make_type(2, fx.Int32) + + mbox = Mailbox(scratch, step, layer) + # Plain names, not method calls: flydsl's AST rewriter carries every local + # whose method is called inside a dynamic if/for as loop state. + mb_put = mbox.put + mb_put_bf = mbox.put_bf + mb_put_words = mbox.put_words + mb_poll = mbox.poll + + def mb(name): + return scratch + fx.Int64(SCRATCH[name]) + + def acquire(): + fx.memory_fence( + syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Acquire + ) + + def scores_copy(c, k=0): + """Copy c of token k's router outputs: expert e's (routing key, sigmoid) + at pairs 2 e, 2 e + 1.""" + return mb("scores") + (fx.Int64(c) * MAX_TOKENS + fx.Int64(k)) * ( + N_ROUTED * 16 + ) + + def route_key(v, e): + """Signed-orderable i32 of score ``v`` with its low 7 bits replaced by + 127 - expert ``e``: a pick is one integer max, a tie goes to the lower + expert.""" + b = v.bitcast(fx.Int32) + b = (b < 0).select(b ^ fx.Int32(0x7FFFFFFF), b) + return (b & fx.Int32(-128)) | (127 - e) + + # each wave pushes to one peer + pv = fx.Vector( + bo.buffer_load( + rsrc(peers), fx.min(wave, W - 1) * 2, vec_width=2, dtype=T.i32 + ) + ) + peer_dst = (fx.Int64(uniform(pv[1])) << 32) | fx.Int64( + fx.Uint32(uniform(pv[0])) + ) + AG_RS = tokens >= AG_RS_FROM + + def peer_addr(w): + pw = fx.Vector(bo.buffer_load(rsrc(peers), w * 2, vec_width=2, dtype=T.i32)) + return (fx.Int64(uniform(pw[1])) << 32) | fx.Int64( + fx.Uint32(uniform(pw[0])) + ) + + def start(name): + return (bid + (G - BASE[name])) % G + + def stamp(k): + """``timeline``: s_memrealtime (100 MHz) of this CTA passing point k, kept + in LDS until the kernel ends -- a global store per stamp would put its + write latency on every later ``s_waitcnt vmcnt(0)``.""" + # The compile-time and traced conditions cannot share one ``and``. + if const_expr(timeline): # noqa: SIM102 + if tid == 0: + fx.ptr_store(memrealtime(), tls + k) + + def flush_stamps(): + if const_expr(timeline): # noqa: SIM102 + if tid < TL_POINTS: + fx.generic_store( + fx.inttoptr( + fx.PointerType.get( + fx.Int64.ir_type, fx.AddressSpace.Global, 8 + ), + tl + fx.Int64((bid * TL_POINTS + tid) * 8), + ), + fx.ptr_load(tls + tid), + ) + + def ld_bf16x4(r, k, cm=0): + w = fx.Vector( + bo.buffer_load(r, k // 2, vec_width=2, dtype=T.i32, cache_modifier=cm) + ) + v = w.bitcast(fx.BFloat16).to(fx.Float32) + return [v[j] for j in range(4)] + + def mfma_fp8(a64, b64, c): + return fx.Vector( + rocdl.mfma_f32_16x16x32_fp8_fp8(T.vec(4, T.f32), [a64, b64, c, 0, 0, 0]) + ) + + def block_max(v): + return block_max_of(v, lane, wave, red) + + def push_rows(region, row0, n_rows, only=None): + """This task's bf16 partial rows (misc: token k's n_rows from k n_rows) to + the peers' ``region``, a pair a lane, wave w -> peer w (itself too); + ``only``: to that peer alone.""" + half = n_rows // 2 + to = (wave < W) if only is None else (wave == only) + for c in range_constexpr((half * tokens + 63) // 64): + i = lane + 64 * c + if to & (i < half * tokens): + mb_put_bf( + peer_dst + fx.Int64(SY[region]), + (rank * MAX_TOKENS + i // half) * HIDDEN + + row0 + + (i % half) * 2, + [fx.ptr_load(misc + i * 2), fx.ptr_load(misc + (i * 2 + 1))], + CM_SYS, + ) + + if const_expr(timeline): # noqa: SIM102 + if tid < TL_POINTS: # same wave as the stamping thread: ordered + fx.ptr_store(fx.Int64(0), tls + tid) + stamp(0) + + if const_expr(fuse_k1): + # this layer's K1 first: its GEMV / head / norm tasks and the scores + def k1_ptr(i): + v = fx.Vector( + bo.buffer_load(rsrc(k1_args), 2 * i, vec_width=2, dtype=T.i32) + ) + return (fx.Int64(uniform(v[1])) << 32) | fx.Int64( + fx.Uint32(uniform(v[0])) + ) + + def k1_stamp(k): + """K1's timeline points 1-4 and 6 on K4's free stamps 27-31.""" + if const_expr(k in K1_STAMPS): + stamp(K1_STAMPS[k]) + + ( + ar_in, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, cos_sin, index_cache, + iq_out, scratch1, + ) = [k1_ptr(i) for i in range(len(K1_ARGS))] # fmt: skip + qs_k1 = lds.qs.ptr + if const_expr(ih_count > 1): + qs_k1 = fpool + QS_CP + emit_k1( + tid, bid, lane, wave, xs, red, qs_k1, k1_stamp, step, layer, + eps, index_scale_log2e(sm_scale), tokens, q_len, init_blocks, + local_blocks, IndexHeads(ih_count, ih_own), + ( + ar_in, res, h_in, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, + cos_sin, positions, slot_mapping, k_cache, v_cache, k_scale, + v_scale, index_cache, q, iq_out, block_table, seq_lens, + mb("iscore"), scratch1, + ), + bt_width, + True, + ) # fmt: skip + gpu.barrier() + k1_mbox = Mailbox(scratch1, step, layer) + k1_poll = k1_mbox.poll + hdone_mb = k1_mbox.addr(K1_SCRATCH_HDONE) + rdone_mb = k1_mbox.addr(K1_SCRATCH_RDONE) + + # ============================================================ 0. index + # Indexer: K1 scored every block the selection does not pin (index_score); + # here only the selection runs. + # a graph's pad rows carry seq_len 0 and an all-zero block-table row: as one + # key of page 0 they stay finite and select a real page + seq_lens_k, n_blks, long_rows = step_rows( + seq_lens, tokens, q_len, IndexHeads(ih_count, ih_own) + ) + # Up to TOPK_BLOCKS blocks the top-k keeps every block and the scores would + # only order them. The original selector's own contract calls either order + # equally correct (index_topk._pack_score_key), so a short context is not + # scored and takes the blocks in id order -- the tail block is last either + # way -- instead of putting the scoring chain in front of the attention. + scored = [nb > TOPK_BLOCKS for nb in n_blks] + + def bt_row(k): + """Token k's block-table row (rows are ``bt_width`` apart).""" + return block_table + fx.Int64(k) * fx.Int64(bt_width) * 4 + + def pick(vals, k): + """vals[k] for a traced token index k.""" + v = vals[0] + for i in range_constexpr(1, tokens): + v = (k == i).select(vals[i], v) + return v + + def index_key(b, word, n_blk): + """Sort key of block b's raw score word: forced blocks pinned, then a + signed int ordered like the float (-0 == 0, NaN below -inf).""" + sc = word.bitcast(fx.Float32) + sc = (b < init_blocks).select(fx.Float32(1e30), sc) + sc = (b >= n_blk - local_blocks).select(fx.Float32(1e29), sc) + bits = sc.bitcast(fx.Int32) + bits = (bits == fx.Int32(-(2**31))).select(fx.Int32(0), bits) + key = (bits < 0).select(bits ^ 0x7FFFFFFF, bits) + return ((bits & 0x7FFFFFFF) > 0x7F800000).select(fx.Int32(-(2**31)), key) + + def rank_blocks(k, iscore, n_blk, r_bt, scored_k): + """Up to THREADS blocks: thread b keys block b, and a block's rank + counts the blocks that beat it. -> this thread's (key, block, rank, + live, the tail block's key, page). ``scored_k``: the scores are read + (else the blocks rank by id).""" + b = tid + live = b < n_blk + bb = fx.min(b, n_blk - 1) + page = fx.Int32(bo.buffer_load(r_bt, bb, vec_width=1, dtype=T.i32)) + # K1 (or the harnesses' scorer) wrote every score before this launch, so + # a plain load reads them: it goes out before the context length is in. + # A pinned block's slot holds no score (index_key forces its key); past + # n_blk the keys are never read + if const_expr(fuse_k1): + # the scorers are CTAs of this launch: poll (a pinned block's slot + # polls a scored block's) + word = fx.Int32(0) + if scored_k: + word = mb_poll( + [ + ( + iscore, + fx.min( + fx.max(tid, init_blocks), + n_blk - 1 - local_blocks, + ), + 1, + ) + ] + )[0][0] + else: + word = fx.Int32( + bo.buffer_load(rsrc(iscore), 2 * tid, vec_width=1, dtype=T.i32) + ) + fx.ptr_store(scored_k.select(index_key(tid, word, n_blk), -tid), keys + tid) + stamp(18) + gpu.barrier() + stamp(20) + # rank = the keys that beat it. RANK_SPLIT threads share a block, each + # comparing every RANK_SPLIT-th key: one lane walking all 61 keys of a + # 7.8k context took 1.9 us (a serial compare -> SGPR mask chain) + part = tid % RANK_SPLIT + for b0 in range(0, n_blk, THREADS // RANK_SPLIT): + rb = fx.Int32(b0) + tid // RANK_SPLIT + rbc = fx.min(rb, n_blk - 1) + kb = fx.ptr_load(keys + rbc) + cnt = fx.Int32(0) + for j0 in range(0, n_blk, RANK_SPLIT): + j = fx.Int32(j0) + part + kj = fx.ptr_load(keys + fx.min(j, n_blk - 1)) + beats = (kj > kb) | ((kj == kb) & (j > rbc)) + cnt = cnt + (beats & (j < n_blk)).select(fx.Int32(1), fx.Int32(0)) + cnt = butterfly(cnt, RANK_LANES) + if (part == 0) & (rb < n_blk): + fx.ptr_store(cnt, ranks + rb) + gpu.barrier() + key = fx.ptr_load(keys + bb) + rank = fx.ptr_load(ranks + bb) + stamp(19) + kt = fx.ptr_load(keys + (n_blk - 1)) + return key, bb, rank, live, kt, page + + def ballot(pred): + return fx.Int64(rocdl.ballot(T.i64, pred)) + + def beats(ah, al, bh, bl): + """(key, block) pair a orders before b.""" + return (ah > bh) | ((ah == bh) & (al > bl)) + + def choose_pair(key, block, peer_key, peer_block, take_high): + partner_high = beats(peer_key, peer_block, key, block) + high_h = partner_high.select(peer_key, key) + high_l = partner_high.select(peer_block, block) + low_h = partner_high.select(key, peer_key) + low_l = partner_high.select(block, peer_block) + return take_high.select(high_h, low_h), take_high.select(high_l, low_l) + + def merge16_best(ah, al, bh, bl): + """Return the top half of two descending 16-item lists.""" + b_hi = beats(bh, bl, ah, al) + key, block = b_hi.select(bh, ah), b_hi.select(bl, al) + for off in (8, 4, 2, 1): + ph, pl = xshfl(key, off), xshfl(block, off) + p_hi = beats(ph, pl, key, block) + hi_h, hi_l = p_hi.select(ph, key), p_hi.select(pl, block) + lo_h, lo_l = p_hi.select(key, ph), p_hi.select(block, pl) + take_hi = (l16 & off) == 0 + key = take_hi.select(hi_h, lo_h) + block = take_hi.select(hi_l, lo_l) + return key, block + + def sort16_desc(key, block): + for size, strides in ( + (2, (1,)), + (4, (2, 1)), + (8, (4, 2, 1)), + (16, (8, 4, 2, 1)), + ): + for stride in strides: + ph, pl = xshfl(key, stride), xshfl(block, stride) + asc = (l16 & size) != 0 + low = (l16 & stride) == 0 + key, block = choose_pair(key, block, ph, pl, low ^ asc) + return key, block + + def wave_best(h0, l0, h1, l1): + """Return the exact top 16 of 128 pairs in lanes 0..15.""" + h0, l0 = sort16_desc(h0, l0) + h1, l1 = sort16_desc(h1, l1) + h0, l0 = merge16_best(h0, l0, xshfl(h0, 31), xshfl(l0, 31)) + h1, l1 = merge16_best(h1, l1, xshfl(h1, 31), xshfl(l1, 31)) + h0, l0 = merge16_best(h0, l0, xshfl(h0, 47), xshfl(l0, 47)) + h1, l1 = merge16_best(h1, l1, xshfl(h1, 47), xshfl(l1, 47)) + h0, l0 = merge16_best( + h0, + l0, + h1.shuffle_xor(15, 64), + l1.shuffle_xor(15, 64), + ) + return h0, l0 + + def rank_cands(): + """Merge eight sorted candidate rows to the exact top 16.""" + key, block = fx.Int32(KEY_NONE), fx.Int32(-1) + if wave == 0: + group = lane // TOPK_BLOCKS + index = l16 + ai = (2 * group) * TOPK_BLOCKS + index + bi = (2 * group + 1) * TOPK_BLOCKS + (TOPK_BLOCKS - 1 - index) + key, block = merge16_best( + fx.ptr_load(keys + 2 * ai), + fx.ptr_load(keys + (2 * ai + 1)), + fx.ptr_load(keys + 2 * bi), + fx.ptr_load(keys + (2 * bi + 1)), + ) + key, block = merge16_best(key, block, xshfl(key, 31), xshfl(block, 31)) + key, block = merge16_best(key, block, xshfl(key, 47), xshfl(block, 47)) + return key, block, tid, (tid < TOPK_BLOCKS) & (block >= 0) + + def place(key, bb, rank, live, kt, n_blk, seq_len, active): + """The ranked blocks' slots, as the Triton selector emits them: full + blocks by rank, the tail block (the one holding the current token) + last. -> (selected, slot, sparse context length); ``active``: this + CTA places (a uniform flag).""" + n_sel = fx.min(n_blk, TOPK_BLOCKS) + tail = n_blk - 1 + tail_first = (kt > key) | ((kt == key) & (tail > bb)) + sel = live & (rank < n_sel) & active + is_tail = sel & (bb == tail) + if is_tail: + fx.ptr_store(fx.Int32(1), blk + 17) + gpu.barrier() + tail_sel = fx.ptr_load(blk + 17) != 0 + n_full = n_sel - tail_sel.select(fx.Int32(1), fx.Int32(0)) + slot = is_tail.select( + n_full, rank - tail_first.select(fx.Int32(1), fx.Int32(0)) + ) + n_ctx = tail_sel.select( + n_full * SPARSE_BLOCK + seq_len - tail * SPARSE_BLOCK, + fx.min(n_sel * SPARSE_BLOCK, seq_len), + ) + return sel, slot, n_ctx + + def table_of(k): + return rsrc(mb("sparse_table") + fx.Int64(k) * (SPARSE_TABLE_WORDS * 4)) + + def select_blocks(k, t): + """blk[0:16] := the pages of selection slots 2t, 2t+1 and blk[16] := the + sparse context length, exactly as the Triton selector emits them: top + TOPK_BLOCKS by (score, block id) descending with the init blocks pinned + to 1e30 and the local ones to 1e29, full blocks in that order and the + tail block (the one holding the current token) last. A block's slot + follows from its rank, the count of blocks that beat it. k: token. + For a long request ``stage_select_long`` left blk: a one-block + stand-in is ranked, placing nothing.""" + iscore = mb("iscore") + fx.Int64(k) * (MAX_INDEX_BLOCKS * 8) + short = ~pick(long_rows, k) + n_blk = short.select(pick(n_blks, k), fx.Int32(1)) + seq_len = pick(seq_lens_k, k) + r_bt = rsrc(bt_row(k)) + table = table_of(k) + if (tid < 2 * PAGES_PER_BLOCK + 2) & short: + fx.ptr_store(fx.Int32(0), blk + tid) + key, bb, rank, live, kt, page = rank_blocks( + k, iscore, n_blk, r_bt, pick(scored, k) & short + ) + sel, slot, n_ctx = place(key, bb, rank, live, kt, n_blk, seq_len, short) + if sel & (slot // 2 == t): + for j in range_constexpr(PAGES_PER_BLOCK): + fx.ptr_store( + page * mono_config.KV_PAGE_STRIDE + j, + blk + ((slot % 2) * PAGES_PER_BLOCK + j), + ) + if sel & (t == 0): + for j in range_constexpr(PAGES_PER_BLOCK): + bo.buffer_store( + page * mono_config.KV_PAGE_STRIDE + j, + table, + slot * PAGES_PER_BLOCK + j, + ) + if (tid == 0) & short: + fx.ptr_store(n_ctx, blk + 16) + if t == 0: + bo.buffer_store(n_ctx, table, TOPK_BLOCKS * PAGES_PER_BLOCK) + gpu.barrier() + + def select_long(k, part): + """Split task (k, part) of a long request's token, before the + split stage. Its share of the blocks (a thread keys CAND_BATCH) -> each + wave's TOPK_BLOCKS best -> the CTA's, to a candidate mailbox; then the + N_SPLIT shares' candidates (the context's TOPK_BLOCKS best are among + them) are ranked and placed -> blk as ``select_blocks`` leaves it, and + (part 0) the sparse table. The share: an eighth of the blocks, to this + rank's sel_long; with every index head (indexer context parallelism) + head part // 2's scores of half part % 2 of this rank's blocks, to the + head's rank's sel_cp -- the shares its own head's candidates come in.""" + n_blk = pick(n_blks, k) + iscore = mb("iscore") + fx.Int64(k) * (MAX_INDEX_BLOCKS * 8) + first = k * (N_SPLIT * TOPK_BLOCKS) # token k's candidates (pair pairs) + table = table_of(k) + if tid == 0: # the tail flag ``place`` reads: LDS, not zeroed on this path + fx.ptr_store(fx.Int32(0), blk + 17) + if const_expr(ih_count > 1): + own0 = init_blocks + (ih_own - init_blocks) % TP # this rank's first + n_own = (n_blk - own0 + TP - 1) // TP + # the last scored one: the local blocks' slots hold no score + last = own0 + TP * ((n_blk - local_blocks - own0 + TP - 1) // TP - 1) + span = (n_own + 1) // 2 + lo = (part % 2) * span + hi = fx.min(lo + span, n_own) + js = [lo + tid + THREADS * i for i in range(CAND_BATCH)] + blocks = [own0 + TP * j for j in js] + live = [j < hi for j in js] + src = iscore + fx.Int64(part // 2) * (MAX_INDEX_BLOCKS // TP * 8) + slots = [fx.min(b, last) // TP for b in blocks] + out, out_slot = peer_addr(part // 2) + fx.Int64(SY["sel_cp"]), 2 * rank + out_slot = out_slot + part % 2 + got_from, got_scope = sym + fx.Int64(SY["sel_cp"]), "system" + else: + span = (n_blk + N_SPLIT - 1) // N_SPLIT + lo = part * span + hi = fx.min(lo + span, n_blk) + blocks = [lo + tid + THREADS * i for i in range(CAND_BATCH)] + live = [b < hi for b in blocks] + src = iscore + slots = [ + fx.min(fx.max(b, init_blocks), n_blk - 1 - local_blocks) + for b in blocks + ] + out, out_slot = mb("sel_long"), part + got_from, got_scope = out, "agent" + if const_expr(fuse_k1): + words = [w[0] for w in mb_poll([(src, sl, 1) for sl in slots])] + else: + words = [ + fx.Int32( + bo.buffer_load(rsrc(src), 2 * sl, vec_width=1, dtype=T.i32) + ) + for sl in slots + ] + hs = [ + lv.select(index_key(b, w, n_blk), fx.Int32(KEY_NONE)) + for b, w, lv in zip(blocks, words, live) + ] + ls = [lv.select(b, fx.Int32(-1)) for b, lv in zip(blocks, live)] + stamp(18) + rh, rl = wave_best(hs[0], ls[0], hs[1], ls[1]) + if lane < TOPK_BLOCKS: + fx.ptr_store(rh, keys + 2 * (wave * TOPK_BLOCKS + lane)) + fx.ptr_store(rl, keys + (2 * (wave * TOPK_BLOCKS + lane) + 1)) + gpu.barrier() + ch, cl, cnt, held = rank_cands() + # a share holds >= TOPK_BLOCKS blocks: its 16 best are blocks of + # distinct ranks + if held & (cnt < TOPK_BLOCKS): + mb_put_words( + out, + 2 * (first + out_slot * TOPK_BLOCKS + cnt), + [ch, cl], + CM_SYS if ih_count > 1 else CM_DEV, + ) + gpu.barrier() + stamp(20) + if tid < N_SPLIT * TOPK_BLOCKS: + got = mb_poll([(got_from, 2 * (first + tid), 2)], got_scope)[0] + fx.ptr_store(got[0], keys + 2 * tid) + fx.ptr_store(got[1], keys + (2 * tid + 1)) + gpu.barrier() + ch, cl, cnt, held = rank_cands() + page = fx.Int32( + bo.buffer_load( + rsrc(bt_row(k), MAX_INDEX_BLOCKS * 4), + cl, + vec_width=1, + dtype=T.i32, + mask=held, + ) + ) + stamp(19) + # the tail block is a local one: its key is forced + kt = index_key(n_blk - 1, fx.Int32(0), n_blk) + sel, slot, n_ctx = place( + ch, cl, cnt, held, kt, n_blk, pick(seq_lens_k, k), pick(long_rows, k) + ) + if sel & (slot // 2 == part): + for j in range_constexpr(PAGES_PER_BLOCK): + fx.ptr_store( + page * mono_config.KV_PAGE_STRIDE + j, + blk + ((slot % 2) * PAGES_PER_BLOCK + j), + ) + if sel & (part == 0): + for j in range_constexpr(PAGES_PER_BLOCK): + bo.buffer_store( + page * mono_config.KV_PAGE_STRIDE + j, + table, + slot * PAGES_PER_BLOCK + j, + ) + if tid == 0: + fx.ptr_store(n_ctx, blk + 16) + if part == 0: + bo.buffer_store(n_ctx, table, TOPK_BLOCKS * PAGES_PER_BLOCK) + gpu.barrier() + + def stage_select_long(): + """The split tasks of a long request's token select before the + split stage (``select_long``), each leaving its blk for its split task + (a CTA runs at most one). Out of the split stage and cold: inline + there, this code slowed the split for every context (64k: +1 us).""" + ts = start("split") + k = ts // N_SPLIT + if unlikely((ts < N_SPLIT * tokens) & pick(long_rows, k)): + select_long(k, ts % N_SPLIT) + + # Each stage is its own function: flydsl carries a local reassigned inside a + # dynamic if / for as region state, so names must not be shared across stages. + # ============================================================ 1. split + def stage_split(): + r_k, r_v = rsrc(k_cache), rsrc(v_cache) + r_ks, r_vs = rsrc(k_scale), rsrc(v_scale) + PPW = SPLIT_KEYS // WAVES // PAGE16 # pages per wave + for ts in range(start("split"), N_SPLIT * tokens, G): + ts = fx.Int32(ts) # token tok's partition t: ts = tok N_SPLIT + t + tok = ts // N_SPLIT + t = ts % N_SPLIT + select_blocks(tok, t) + stamp(17) + if const_expr(fuse_k1): + # this token's q heads and every token's new K / V: a request's + # tokens of one step (speculative decode) attend to each other's + if wave == 0: + kv = lane - H + k1_poll( + [ + ( + hdone_mb, + (lane < H).select( + tok * K1_HEAD_TASKS + lane, + fx.min(kv // 2, tokens - 1) * K1_HEAD_TASKS + + H + + kv % 2, + ), + 1, + ) + ] + ) + gpu.barrier() + acquire() + # Every independent load goes out first: this wave's pages and q. + # Pages past the context read junk masked below. + n_ctx = uniform(fx.ptr_load(blk + 16)) + pages = [ + uniform(fx.ptr_load(blk + (wave * PPW + j))) for j in range(PPW) + ] + qv = ld_bf16x4(rsrc(q), tok * O_K + tid * 4, CM_K1) + # K (QK B operand, key = lane % 16 of page j), V (PV B operand: dims + # 16 jd + lane % 16, keys 8 (lane / 16) .. + 8), per-token scales + kw = [] + for j in range_constexpr(PPW): + for s in range_constexpr(HEAD_DIM // 32): + kw.append( + fx.Vector( + bo.buffer_load( + r_k, + ( + pages[j] * PAGE_BYTES + + (2 * s + g4 // 2) * 256 + + l16 * 16 + + (g4 % 2) * 8 + ) + // 4, + vec_width=2, + dtype=T.i32, + cache_modifier=CM_K1, + ) + ).bitcast(fx.Int64)[0] + ) + vpg = (g4 // 2 == 0).select(pages[0], pages[1]) + vw = [ + fx.Vector( + bo.buffer_load( + r_v, + (vpg * PAGE_BYTES + (16 * jd + l16) * PAGE16 + (g4 % 2) * 8) + // 4, + vec_width=2, + dtype=T.i32, + cache_modifier=CM_K1, + ) + ).bitcast(fx.Int64)[0] + for jd in range(HEAD_DIM // 16) + ] + kss = [ + fx.Float32( + bo.buffer_load( + r_ks, + pages[j] * PAGE16 + l16, + vec_width=1, + dtype=T.f32, + cache_modifier=CM_K1, + ) + ) + for j in range(PPW) + ] + vss = [ + fx.Float32( + bo.buffer_load( + r_vs, + pages[j] * PAGE16 + l16, + vec_width=1, + dtype=T.f32, + cache_modifier=CM_K1, + ) + ) + for j in range(PPW) + ] + fx.ptr_store(fp8_pack4(qv[0], qv[1], qv[2], qv[3]), q8 + tid) + gpu.barrier() + stamp(12) + # QK per page: A=q (heads), B=K (keys); a lane holds heads + # 4*g4+e and key l16. + key0 = t * SPLIT_KEYS + wave * (PPW * PAGE16) + valid = [key0 + j * PAGE16 + l16 < n_ctx for j in range(PPW)] + sc = [] + for j in range_constexpr(PPW): + c = fx.Vector.filled(4, 0.0, fx.Float32) + for s in range_constexpr(HEAD_DIM // 32): + qw = fx.Vector( + fx.ptr_load( + q8 + (l16 * 32 + 8 * s + 2 * g4), result_type=v2i + ) + ) + c = mfma_fp8( + qw.bitcast(fx.Int64)[0], kw[j * (HEAD_DIM // 32) + s], c + ) + qk = sm_scale * kss[j] + sc.append( + [valid[j].select(qk * c[e], fx.Float32(NEG)) for e in range(4)] + ) + # partition max per head and max v_scale: in-wave butterfly, then LDS + vloc = fx.max( + valid[0].select(vss[0], fx.Float32(0.0)), + valid[1].select(vss[1], fx.Float32(0.0)), + ) + vmax_w = wave_max(vloc) + for e in range_constexpr(4): + mh = butterfly(fx.max(sc[0][e], sc[1][e]), (8, 4, 2, 1), fx.max) + if l16 == 0: + fx.ptr_store(mh, hst + (wave * H + g4 * 4 + e)) + if lane == 0: + fx.ptr_store(vmax_w, hst + (3 * WAVES * H + wave)) + gpu.barrier() + stamp(13) + # every wave partial is read before the first p8 store below: LDS + # stores in between serialize each read behind its own wait + vmax = fx.ptr_load(hst + 3 * WAVES * H) + for w in range_constexpr(1, WAVES): + vmax = fx.max(vmax, fx.ptr_load(hst + (3 * WAVES * H + w))) + mxs = [] + for e in range_constexpr(4): + mx = fx.ptr_load(hst + (g4 * 4 + e)) + for w in range_constexpr(1, WAVES): + mx = fx.max(mx, fx.ptr_load(hst + (w * H + g4 * 4 + e))) + mxs.append(mx) + fscale = FP8_MAX * hw_rcp(vmax + 1e-8) + pscale = vmax * (1.0 / FP8_MAX) + for e in range_constexpr(4): + head = g4 * 4 + e + mx = mxs[e] + ls = fx.Float32(0.0) + for j in range_constexpr(PPW): + pj = hw_exp2((sc[j][e] - mx) * LOG2E) + ls = ls + pj + b8 = ( + fp8_pack4( + (valid[j].select(vss[j], fx.Float32(0.0)) * fscale) + * pj, + 0.0, + 0.0, + 0.0, + ) + & 0xFF + ) + w8 = b8 | (xshfl(b8, 1) << 8) + w8 = w8 | (xshfl(w8, 2) << 16) + if l16 % 4 == 0: + kloc = wave * (PPW * PAGE16) + j * PAGE16 + l16 + fx.ptr_store( + w8, p8 + (head * (SPLIT_KEYS // 4) + kloc // 4) + ) + ls = butterfly(ls, (8, 4, 2, 1)) + if l16 == 0: + fx.ptr_store(ls, hst + (WAVES * H + wave * H + head)) + fx.ptr_store(mx, hst + (2 * WAVES * H + head)) + gpu.barrier() + stamp(14) + # PV over this wave's 32 keys: A = P (heads), B = V (dims) + pw = fx.Vector( + fx.ptr_load( + p8 + + ( + l16 * (SPLIT_KEYS // 4) + + (wave * (PPW * PAGE16) + 8 * g4) // 4 + ), + result_type=v2i, + ) + ).bitcast(fx.Int64)[0] + # keys past the context carry P = 0, but a NaN V byte would still poison + nv = n_ctx - (key0 + 8 * g4) + vmask = (nv >= 8).select( + fx.Int64(-1), + (nv <= 0).select( + fx.Int64(0), (fx.Int64(1) << (fx.Int64(nv) * 8)) - 1 + ), + ) + for jd in range_constexpr(HEAD_DIM // 16): + cv = mfma_fp8( + pw, vw[jd] & vmask, fx.Vector.filled(4, 0.0, fx.Float32) + ) + for e in range_constexpr(4): + fx.ptr_store( + cv[e], + opart + + (wave * O_K + (g4 * 4 + e) * HEAD_DIM + 16 * jd + l16), + ) + gpu.barrier() + stamp(15) + # Sum wave partials: acc=prob_scale*PV, then /exp_sum, bf16. + e0 = tid * 4 + head = e0 // HEAD_DIM + lsum = fx.ptr_load(hst + (WAVES * H + head)) + for w in range_constexpr(1, WAVES): + lsum = lsum + fx.ptr_load(hst + (WAVES * H + w * H + head)) + inv_l = 1.0 / (lsum > 0.0).select(lsum, fx.Float32(1.0)) + ov = [] + for k in range_constexpr(4): + acc = fx.ptr_load(opart + (e0 + k)) + for w in range_constexpr(1, WAVES): + acc = acc + fx.ptr_load(opart + (w * O_K + e0 + k)) + ov.append((pscale * acc) * inv_l) + mb_put_bf(mb("sp_o"), ts * O_K + e0, ov) + if tid < H: + mb_put( + mb("sp_m"), + ts * H + tid, + fx.ptr_load(hst + (2 * WAVES * H + tid)), + ) + lt = fx.ptr_load(hst + (WAVES * H + tid)) + for w in range_constexpr(1, WAVES): + lt = lt + fx.ptr_load(hst + (WAVES * H + w * H + tid)) + mb_put(mb("sp_l"), ts * H + tid, lt) + gpu.barrier() + + stage_select_long() + stage_split() + stamp(1) + + def merge_token(k): + """Token k's partitions merged (the gluon decode's reduce; thread = 4 dims + of one head) and quantized per token (standalone quant: scale = amax * + (1 / FP8_MAX), x * rcp(scale)) -> (this thread's fp8 word, scale).""" + hm = tid // (HEAD_DIM // 4) + e0 = tid * 4 + sp0 = k * N_SPLIT # token k's partitions + # wave 0 waits on every partition's l, then the CTA polls once: 512 + # threads x 24 loads spinning through the whole split load the memory + # system the split itself reads (-0.4 us) + if wave == 0: + mb_poll( + [ + (mb("sp_l"), sp0 * H + lane + 64 * i, 1) + for i in range(N_SPLIT * H // 64) + ] + ) + gpu.barrier() + # one batch: every partition's (m, l) and this thread's 4 outputs + got = mb_poll( + [(mb("sp_m"), (sp0 + i) * H + hm, 1) for i in range(N_SPLIT)] + + [(mb("sp_l"), (sp0 + i) * H + hm, 1) for i in range(N_SPLIT)] + + [ + (mb("sp_o"), ((sp0 + i) * O_K + e0) // 2, 2) for i in range(N_SPLIT) + ], + batch=3 * N_SPLIT, + ) + mls = got[: 2 * N_SPLIT] + ovs = got[2 * N_SPLIT :] + stamp(3) + m_i = [mls[i][0].bitcast(fx.Float32) for i in range(N_SPLIT)] + l_i = [mls[N_SPLIT + i][0].bitcast(fx.Float32) for i in range(N_SPLIT)] + gm = m_i[0] + for i in range_constexpr(1, N_SPLIT): + gm = fx.max(gm, m_i[i]) + sl = [l_i[i] * hw_exp2((m_i[i] - gm) * LOG2E) for i in range(N_SPLIT)] + gs = sl[0] + for i in range_constexpr(1, N_SPLIT): + gs = gs + sl[i] + gs = (gs > 0.0).select(gs, fx.Float32(1.0)) + av = [fx.Float32(0.0) for _ in range(4)] + for i in range_constexpr(N_SPLIT): + wi = sl[i] / gs + o01 = bf2_f32(ovs[i][0]) + o23 = bf2_f32(ovs[i][1]) + av = [ + av[0] + wi * o01[0], + av[1] + wi * o01[1], + av[2] + wi * o23[0], + av[3] + wi * o23[1], + ] + a0, a1, a2, a3 = (bf16_round(x) for x in av) + amax = block_max( + fx.max( + fx.max(fmath.absf(a0), fmath.absf(a1)), + fx.max(fmath.absf(a2), fmath.absf(a3)), + ) + ) + x_scale = amax * (1.0 / FP8_MAX) + # an all-zero row (a pad row) quantizes to zeros, not 0 * inf + inv = (amax == 0.0).select(fx.Float32(0.0), hw_rcp(x_scale)) + return fp8_pack4(a0 * inv, a1 * inv, a2 * inv, a3 * inv), x_scale + + # ======================================= 2b. attention merge (S > 1) + # One task per token, in parallel on its own CTA: merged in the o tasks, the + # tokens' merges ran one after another in every one of them (+14 us at S = 4) + def stage_merge(): + if const_expr(tokens > 1): + for k in range(start("merge"), tokens, G): + k = fx.Int32(k) + word, x_scale = merge_token(k) + # plain words, then the scale: its pair is the flag + bo.buffer_store( + word, + rsrc(mb("attnq")), + k * (O_K // 4) + tid, + cache_modifier=CM_DEV, + ) + fx.memory_fence( + syncscope=rocdl.SyncScope.Workgroup, + ordering=fx.AtomicOrdering.Release, + ) + gpu.barrier() + if tid == 0: + mb_put(mb("attnq_s"), k, x_scale) + + stage_merge() + + # ================================================ 3. o_proj + attn reduce + def stage_o(): + r_wo = rsrc(w_o) + O_CPW = O_K // 64 // 4 # 64-k chunks per wave: 4 waves per 16-row group + for t in range(start("o"), N_O, G): + t = fx.Int32(t) + rg = t * 2 + wave // 4 + wts = [] + for cc in range_constexpr(O_CPW): + kc = (wave % 4) * O_CPW + cc + wts.append( + fx.Vector( + bo.buffer_load( + r_wo, + ((rg * (O_K // 32) + kc * 2) * 512 + lane * 16) // 4, + vec_width=4, + dtype=T.i32, + ) + ) + ) + # this task's row scales, before the wait + ws = fx.Float32( + bo.buffer_load( + rsrc(s_o), t * O_ROWS + tid % O_ROWS, vec_width=1, dtype=T.f32 + ) + ) + + if const_expr(tokens == 1): + # the merge folded in: the partitions -> this CTA's fp8 input + word, x_scale = merge_token(0) + x_scales = [x_scale] + fx.ptr_store(word, xs + tid) + else: + # the merge tasks' fp8 inputs: every token's scale (published + # after its plain words), then the words, 16 B a load + got = mb_poll([(mb("attnq_s"), k, 1) for k in range(tokens)]) + stamp(3) + x_scales = [got[k][0].bitcast(fx.Float32) for k in range(tokens)] + acquire() + n_v = tokens * O_K // 16 + for i in range_constexpr((n_v + THREADS - 1) // THREADS): + v = tid + THREADS * i + if v < n_v: + fx.ptr_store( + fx.Vector( + bo.buffer_load( + rsrc(mb("attnq")), + v * 4, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ), + xs + (v * 4 // (O_K // 4) * O_ROW + v * 4 % (O_K // 4)), + ) + gpu.barrier() + # B column l % 16 is token l % 16 (token 0 past the last): C column k + # is token k's rows, the weight stream is the same for any token count + col_tok = fx.min(l16, tokens - 1) + c = fx.Vector.filled(4, 0.0, fx.Float32) + # chunks kc, kc + 1 in one 16x16x128 FP8 MFMA (unit scales): its lane + # holds k 16 (l / 16) .. + 16 of each 64-k half -- the chunks' packing + for p in range_constexpr(O_CPW // 2): + xb = [ + fx.Vector( + fx.ptr_load( + xs + + ( + col_tok * O_ROW + + ((wave % 4) * O_CPW + 2 * p + h) * 16 + + g4 * 4 + ), + result_type=v4f, + ) + ).bitcast(fx.Int32) + for h in range(2) + ] + c = fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), + [ + fx.Vector.from_elements( + [ + wts[2 * p + h][e] + for h in range(2) + for e in range(4) + ], + fx.Int32, + ), + fx.Vector.from_elements( + [xb[h][e] for h in range(2) for e in range(4)], + fx.Int32, + ), + c, + 0, + 0, + 0, + fx.Int32(127), + 0, + fx.Int32(127), + ], + ) + ) + fx.ptr_store(c, red + (wave * 64 + lane) * 4) + gpu.barrier() + if tid < O_ROWS * tokens: + tk = tid // O_ROWS + grp = (tid % O_ROWS) // 16 + rr = tid % 16 + tot = fx.Float32(0.0) + for w in range_constexpr(4): + tot = tot + fx.ptr_load( + red + + (((grp * 4 + w) * 64 + 16 * (rr // 4) + tk) * 4 + rr % 4) + ) + fx.ptr_store(bf16_round(tot * pick(x_scales, tk) * ws), misc + tid) + gpu.barrier() + # push bf16 partial rows, wave w -> peer w (to itself too); with AG_RS + # to the row group's owner rank only + push_rows("attn", t * O_ROWS, O_ROWS, t % W if AG_RS else None) + gpu.barrier() + if const_expr(AG_RS): # noqa: SIM102 + # the owner sums every peer's partials in rank order and gathers + # the sums to every rank's a_ag (h_mid: the router tasks) + if (t % W == rank) & (tid < O_ROWS // 2 * tokens): + ptk = tid // (O_ROWS // 2) + row = t * O_ROWS + (tid % (O_ROWS // 2)) * 2 + own = sym + fx.Int64(SY["attn"]) + prt = mb_poll( + [ + (own, ((src * MAX_TOKENS + ptk) * HIDDEN + row) // 2, 1) + for src in range(W) + ], + "system", + ) + t0 = fx.Float32(0.0) + t1 = fx.Float32(0.0) + for src in range_constexpr(W): + f0, f1 = bf2_f32(prt[src][0]) + t0 = t0 + f0 + t1 = t1 + f1 + for w in range_constexpr(W): + mb_put_bf( + peer_addr(w) + fx.Int64(SY["a_ag"]), + ptk * HIDDEN + row, + [t0, t1], + CM_SYS, + ) + if const_expr(not AG_RS): # noqa: SIM102 + # this task's rows of the all-reduce, finished here (GLM's + # peer_reduce): sum peer partials in rank order, bf16, then + # add the residual into h_mid and the 'a' mailbox. A router + # reads 24 KB device-scope, not 4 peers' 48 KB system-scope. + if tid < O_ROWS // 2 * tokens: + ptk = tid // (O_ROWS // 2) + row = t * O_ROWS + (tid % (O_ROWS // 2)) * 2 + if const_expr(tokens != 1): + if const_expr(fuse_k1): + k1_poll([(rdone_mb, ptk, 1)]) + acquire() + hr = fx.Int32( + bo.buffer_load( + rsrc(h_in), + (ptk * HIDDEN + row) // 2, + vec_width=1, + dtype=T.i32, + cache_modifier=CM_K1, + ) + ) + own = sym + fx.Int64(SY["attn"]) + prt = mb_poll( + [ + (own, ((src * MAX_TOKENS + ptk) * HIDDEN + row) // 2, 1) + for src in range(W) + ], + "system", + ) + t0 = fx.Float32(0.0) + t1 = fx.Float32(0.0) + for src in range_constexpr(W): + f0, f1 = bf2_f32(prt[src][0]) + t0 = t0 + f0 + t1 = t1 + f1 + # the rank sum only: the router adds the residual in f32, so the + # norm reads bf16(sum) plus the unrounded residual. + mb_put_bf(mb("a"), ptk * HIDDEN + row, [t0, t1]) + if const_expr(tokens != 1): + r0, r1 = bf2_f32(hr) + a0 = bf16_round(t0) + r0 + a1 = bf16_round(t1) + r1 + bo.buffer_store( + fx.Vector.from_elements([a0, a1], fx.Float32).to( + fx.BFloat16 + ), + rsrc(h_mid), + ptk * HIDDEN + row, + ) + + stage_o() + stamp(4) + + # ======================================================= MoE helpers + def e8m0_load(r_s, byte_idx): + """An e8m0 scale byte, still in flight (the MFMA takes it as is).""" + return fx.Int32(bo.buffer_load(r_s, byte_idx, vec_width=1, dtype=T.i8)) + + def group_scale_byte(expert, rows, rg, kc, cols): + """scale_index(expert rows + 16 rg + l16, 4 kc + g4, cols): this lane's + byte for row group rg, 128-k chunk kc, folded to wave-uniform terms plus + the lane (the general form's per-lane arithmetic cost the S = 16 up / + gate loop ~8 us).""" + base = (expert * (rows // 32) + rg // 2) * (cols // 8) + kc // 2 + return (base * 64 + lane) * 4 + (kc % 2) * 2 + rg % 2 + + def route_top4(k, k0, sc0, k1, sc1, sig, wave_mask=None): + """route[8k : 8k+4] / rwt[..] := token k's sigmoid top-k gating from the + router tasks' routing keys ``k0`` / ``k1`` (experts ``lane`` / ``lane + + 64``, sigmoid + bias, see route_key) and unbiased sigmoids (wave k % WAVES + computes it, looking the sigmoids up in LDS at ``sig``): a pick is one + wave max; the weights are the picks' sigmoids renormalized and scaled by + route_scale. Slot 4 is the fused shared expert. ``wave_mask``: this + wave's expert -> token mask, the picks' bit k set (the picks are + distinct, and the wave's tokens run in program order). The caller + syncs.""" + if wave == k % WAVES: + # the unbiased weights, looked up by expert once the picks are known + fx.ptr_store(sc0, sig + (k * N_ROUTED + lane)) + fx.ptr_store(sc1, sig + (k * N_ROUTED + 64 + lane)) + pid = fx.Int32(0) + for p in range_constexpr(TOP_K): + m = butterfly(fx.max(k0, k1), (32, 16, 8, 4, 2, 1), fx.max) + hit0 = k0 == m + k0 = hit0.select(fx.Int32(-(2**31)), k0) + k1 = ((k1 == m) & (k0 != m)).select(fx.Int32(-(2**31)), k1) + pid = (lane == p).select(127 - (m & 127), pid) + if lane < TOP_K: + w = fx.ptr_load(sig + (k * N_ROUTED + pid)) + tot = butterfly(w, (1, 2)) + fx.ptr_store(pid, route + (8 * k + lane)) + fx.ptr_store( + w * (route_scale / fx.max(tot, 1e-20)), rwt + (8 * k + lane) + ) + if const_expr(wave_mask is not None): + was = fx.Int32(fx.ptr_load(wave_mask + pid)) + fx.ptr_store(was | (fx.Int32(1) << k), wave_mask + pid) + if lane == 0: + fx.ptr_store(fx.Int32(SHARED_EXPERT), route + (8 * k + TOP_K)) + fx.ptr_store(fx.Float32(shared_weight), rwt + (8 * k + TOP_K)) + + # a8w4 operands in xs (Int32 words, never through f32 arithmetic): the up / + # gate stage holds every token's MXFP8 xn then its E8M0 scales; the down + # stage every token's MXFP8 mid then its scales + XN8_WORDS = HIDDEN // 4 + XSC_WORDS = HIDDEN // 32 + XSC_LDS = tokens * XN8_ROW + MID16_WORDS = MOE_SLOTS * INTER // 2 # bf16 pairs, as the mids land + MID8_WORDS = MOE_SLOTS * INTER // 4 + MID8_LDS = tokens * MID16_WORDS + MIDSC_WORDS = MOE_SLOTS * INTER // 32 + MIDSC_LDS = MID8_LDS + tokens * MID8_WORDS + + def mx8_scale(vals, offsets): + """E8M0 byte (>= 1) of the 32-block the lanes xor-``offsets`` apart share, + aiter's MXFP8 RoundUp rule ceil_pow2(amax / 448), and its reciprocal.""" + amax = fx.max(vals[0], -vals[0]) + for v in vals[1:]: + amax = fx.max(amax, fx.max(v, -v)) + amax = butterfly(amax, offsets, fx.max) + bits = (amax * fx.Float32(1.0 / FP8_MAX)).bitcast(fx.Int32) + e = (bits >> 23) & 0xFF + e = e + ((bits & 0x7FFFFF) != 0).select(fx.Int32(1), fx.Int32(0)) + e = fx.max(fx.min(e, fx.Int32(254)), fx.Int32(1)) + return e, hw_rcp((e << 23).bitcast(fx.Float32)) + + def mx8_q(v, inv): + return fx.max(fx.min(v * inv, fx.Float32(FP8_MAX)), fx.Float32(-FP8_MAX)) + + def lds_b8(chunk_word): + """This lane's B operand of the 128-k chunk at LDS word ``chunk_word``: + fp8 k 16 g .. + 16 then 64 + 16 g .. + 16 (g = lane // 16), the + f8f6f4 MFMA's B layout (harness/mfma_f4f8_probe.py).""" + h = [ + fx.Vector( + fx.ptr_load(xs + (chunk_word + 16 * j + 4 * g4), result_type=v4f) + ).bitcast(fx.Int32) + for j in range(2) + ] + return fx.Vector.from_elements( + [h[j][e] for j in range(2) for e in range(4)], fx.Int32 + ) + + def bf16x2(word): + return fx.Vector.from_elements([word], fx.Int32).bitcast(fx.BFloat16) + + def lds_i32(word): + return fx.ptr_load(xs + word).bitcast(fx.Int32) + + def mfma_f4f8(w4, sa, b8, sb, c, sa_byte=None): + """16x16x128 scaled MFMA: A = this lane's 32 fp4 (row l % 16, k block + l / 16) with E8M0 ``sa``, B = its 32 fp8 with E8M0 ``sb``. ``sa_byte``: + ``sa`` is a scale dword, the MFMA takes that byte of it.""" + if const_expr(sa_byte is None): + sa, sa_byte = sa & 0xFF, 0 + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [w4, b8, c, 4, 0, sa_byte, sa, 0, sb] + ) + ) + + def mfma_f8f8(w8, sa, b8, sb, c, sa_byte=None): + """``mfma_f4f8`` with an fp8 A: this lane's 32 fp8 of row l % 16 in the B + operand's k order (``f8_tile_load``); the scales are placed alike.""" + if const_expr(sa_byte is None): + sa, sa_byte = sa & 0xFF, 0 + return fx.Vector( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + T.vec(4, T.f32), [w8, b8, c, 0, 0, sa_byte, sa, 0, sb] + ) + ) + + def f8_tile_load(r_w, tile, slot): + """Lane ``slot``'s 32 fp8 of 16 x 128 tile ``tile`` (weights.mxfp8_a_tiles) + as the f8f6f4 MFMA's 8-dword A operand.""" + word = (tile * F8_TILE + slot * 32) // 4 + h = [ + fx.Vector(bo.buffer_load(r_w, word + 4 * j, vec_width=4, dtype=T.i32)) + for j in range(2) + ] + return fx.Vector.from_elements( + [h[j][e] for j in range(2) for e in range(4)], fx.Int32 + ) + + def sh_ptr(i): + """The MXFP8 shared expert's table ``SH_ARGS[i]``, read at use from the + ``sh_args`` pointer table (as K1's are): as four more kernel arguments + they widened the kernarg pointer tuple, which ROCm LLVM's partial- + register rewrite spilled incompletely (a fault at S = 2).""" + v = fx.Vector( + bo.buffer_load(rsrc(sh_args), 2 * i, vec_width=2, dtype=T.i32) + ) + return (fx.Int64(uniform(v[1])) << 32) | fx.Int64(fx.Uint32(uniform(v[0]))) + + def w_tile(r_w, rg, kc, slot, k, fp8): + """Lane ``slot``'s A operand of row group rg, 128-k chunk kc of a table + of k columns: fp4 (two 16 x 64 tiles of 512 B, shuffle_weight(16, 16)) + or MXFP8 (``f8_tile_load``).""" + if const_expr(fp8): + return f8_tile_load(r_w, rg * (k // 128) + kc, slot) + return fx.Vector( + bo.buffer_load( + r_w, + ((rg * (k // 64) + kc * 2) * 512 + slot * 16) // 4, + vec_width=4, + dtype=T.i32, + ) + ) + + mfma_sh = mfma_f8f8 if shared_fp8 else mfma_f4f8 # the shared expert's + + # ============================================================ 4. router + def stage_router(): + r_g = rsrc(g_post) + r_gate = rsrc(w_gate) + # this thread's 12 consecutive elements of the row (512 x 12 = HIDDEN): + # the poll is 3 specs per rank, 12 in one batch, every thread busy + RE = HIDDEN // THREADS + rw0 = tid * (RE // 2) + + def ld_bf16x12(r, word, cm=0): + """12 bf16 from ``word`` (a 16 B then an 8 B load).""" + a8 = fx.Vector( + bo.buffer_load(r, word, vec_width=4, dtype=T.i32, cache_modifier=cm) + ).bitcast(fx.BFloat16) + a4 = fx.Vector( + bo.buffer_load( + r, word + 4, vec_width=2, dtype=T.i32, cache_modifier=cm + ) + ).bitcast(fx.BFloat16) + return [a8[e] for e in range(8)] + [a4[e] for e in range(4)] + + # one task per (token, 8 experts): the tokens' norms and dot products run + # on different CTAs, not one after another in each (S = 4: 74.7 -> 64.1 us) + for tt in range(start("router"), N_ROUTER * tokens, G): + tt = fx.Int32(tt) + k = tt // N_ROUTER + t = tt % N_ROUTER + e_r = t * WAVES + wave + bias_r = fx.Float32( + bo.buffer_load(rsrc(bias), e_r, vec_width=1, dtype=T.f32) + ) + gw = [ # bf16 pairs + fx.Vector( + bo.buffer_load( + r_gate, + (e_r * HIDDEN + (i * 64 + lane) * 8) // 2, + vec_width=4, + dtype=T.i32, + ) + ) + for i in range(HIDDEN // 512) + ] + gp = ld_bf16x12(r_g, rw0) + if const_expr(fuse_k1): + # the residual (K1's norm / GEMV task) + if wave == 0: + k1_poll([(rdone_mb, k, 1)]) + gpu.barrier() + acquire() + hres = ld_bf16x12(rsrc(h_in), k * (HIDDEN // 2) + rw0, CM_K1) + # the all-reduced post-attention row (the o tasks finished it; with + # AG_RS the owner ranks gathered it into this rank's a_ag) + a_src, a_scope = ( + (sym + fx.Int64(SY["a_ag"]), "system") + if AG_RS + else (mb("a"), "agent") + ) + arow = mb_poll( + [ + (a_src, (k * HIDDEN + tid * RE) // 2 + 2 * h, 2) + for h in range(RE // 4) + ], + a_scope, + ) + stamp(22) + v = [ + f + fx.Float32(hres[4 * h + 2 * qq + i]) + for h in range(RE // 4) + for qq in range(2) + for i, f in enumerate(bf2_f32(arow[h][qq])) + ] + if const_expr(AG_RS or tokens == 1): # noqa: SIM102 + # h_mid = bf16(bf16(sum) + residual): task t writes slice t + if tid // (XN_SLICE // RE) == t: + hw = (k * HIDDEN + tid * RE) // 2 + hb = fx.Vector.from_elements(v, fx.Float32).to(fx.BFloat16) + hi = hb.bitcast(fx.Int32) + bo.buffer_store( + fx.Vector.from_elements( + [hi[e] for e in range(4)], fx.Int32 + ), + rsrc(h_mid), + hw, + ) + bo.buffer_store( + fx.Vector.from_elements([hi[4], hi[5]], fx.Int32), + rsrc(h_mid), + hw + 4, + ) + # the row's sum of squares: a wave sum each, one barrier, every lane + # adds the 8 + ssw = fx.Float32(0.0) + for e in range_constexpr(RE): + ssw = fmath.fma(v[e], v[e], ssw) + ssw = wave_sum(ssw) + if lane == 0: + fx.ptr_store(ssw, red + wave) + gpu.barrier() + ss = fx.ptr_load(red) + for w in range_constexpr(1, WAVES): + ss = ss + fx.ptr_load(red + w) + stamp(24) + rstd = hw_rsq(ss / float(HIDDEN) + eps) + for j in range_constexpr(RE // 2): + fx.ptr_store( + bf16_pair( + v[2 * j] * rstd * (fx.Float32(gp[2 * j]) + 1.0), + v[2 * j + 1] * rstd * (fx.Float32(gp[2 * j + 1]) + 1.0), + ), + xs + (rw0 + j), + ) + gpu.barrier() + stamp(25) + if tid < XN_SLICE // 4: # 4 elements a lane, 8 lanes a 32-block + w0 = t * (XN_SLICE // 2) + tid * 2 + f = [ + v + for j in range(2) + for v in bf2_f32(fx.ptr_load(xs + (w0 + j)).bitcast(fx.Int32)) + ] + e, inv = mx8_scale(f, (1, 2, 4)) + word = fp8_pack4(*[mx8_q(v, inv) for v in f]) + xw = k * XN8_WORDS + t * (XN_SLICE // 4) + tid + xsw = k * XSC_WORDS + t * (XN_SLICE // 32) + tid // 8 + if const_expr(WIDE): + # plain words, flagged once the task is done (below) + bo.buffer_store( + word, rsrc(mb("xn8p")), xw, cache_modifier=CM_DEV + ) + if tid % 8 == 0: + bo.buffer_store( + e, rsrc(mb("xscp")), xsw, cache_modifier=CM_DEV + ) + else: + mb_put_words(mb("xn8"), xw, [word]) + if tid % 8 == 0: + mb_put_words(mb("xsc"), xsw, [e]) + stamp(26) + # bf16 pair dot products, one accumulator per pair position: four + # 12-long chains instead of one 96-long FMA chain behind 192 widenings + accs = [fx.Float32(0.0) for _ in range(4)] + for i in range_constexpr(HIDDEN // 512): + xv = fx.Vector( + fx.ptr_load(xs + (i * 64 + lane) * 4, result_type=v4f) + ).bitcast(fx.Int32) + for j in range_constexpr(4): + accs[j] = fx.Float32( + rocdl.fdot2_f32_bf16( + T.f32, + as_ir_value(bf16x2(gw[i][j])), + as_ir_value(bf16x2(xv[j])), + as_ir_value(accs[j]), + clamp=False, + ).result + ) + logit = wave_sum((accs[0] + accs[1]) + (accs[2] + accs[3])) + stamp(23) + # the routing key and sigmoid go out, not the logit: the consumers' + # routing starts at the max rounds + sc = hw_rcp(1.0 + hw_exp2(-bf16_round(logit) * LOG2E)) + rk = route_key(sc + bias_r, e_r) + if lane < SCORE_COPIES: + mb_put_words( + scores_copy(lane, k), e_r * 2, [rk, sc.bitcast(fx.Int32)] + ) + if const_expr(WIDE): + # this task's xn slice is out: its flag (after the logits, not + # in front of them) + fx.memory_fence( + syncscope=rocdl.SyncScope.Workgroup, + ordering=fx.AtomicOrdering.Release, + ) + gpu.barrier() + if tid == 0: + mb_put(mb("xndone"), tt, fx.Int32(1)) + + stage_router() + stamp(5) + + # ================================ 5. experts: up / gate -> down -> FFN reduce + UG_KC = HIDDEN // 128 # 128-k chunks + UG_CPW = UG_KC // WAVES # 128-k chunks per wave: each wave a k eighth + DN_KC = INTER // 128 # 128-k chunks of one slot + DN_UNITS = TOP_K * DN_KC # routed (slot, 128-k chunk) units per row group + DN_UPW = (DN_UNITS + 3) // 4 # 4 waves per row group + SH_UPW = (DN_KC + 3) // 4 # the shared expert's chunks, once for all tokens + # routed (token, slot, row group) tasks t < ROUTED_TASKS; the shared expert's + # row groups (one GEMV with a column per token) are tasks ROUTED_TASKS + j, + # run as task j of stage "shared" while the routing is still on its way, so + # the routed tasks alone set the rounds (S = 4: 768 = 3 x 256) + ROUTED_TASKS = TOP_K * UG_PER_SLOT * tokens + TPT = TOP_K * UG_PER_SLOT # routed tasks per token + CPT = ug_ctas_per_token(tokens) + UG_ITERS = max(ug_tasks_of(c, tokens) for c in range(G)) # per CTA, at most + + def ug_task_id(it): + """(this CTA's it-th routed task, clamped; whether it exists).""" + if const_expr(CPT is not None): + u = bid % CPT + it * CPT + return (bid // CPT) * TPT + fx.min(u, TPT - 1), u < TPT + t = start("ug") + it * G + return fx.min(t, ROUTED_TASKS - 1), t < ROUTED_TASKS + + def ug_task(t): + """(shared expert?, token, slot, row group) of up / gate task t.""" + shared = t >= ROUTED_TASKS + r = fx.min(t, ROUTED_TASKS - 1) + tk = r // (TOP_K * UG_PER_SLOT) + slot = shared.select( + fx.Int32(TOP_K), (r % (TOP_K * UG_PER_SLOT)) // UG_PER_SLOT + ) + return shared, tk, slot, t % UG_PER_SLOT + + def ug_load(t): + """Up / gate weights and scales of task t.""" + shared, tk, slot, cg = ug_task(t) + # the bound keeps a corrupt route from addressing past the expert table + # (a shared task never gets here: ug_task_id clamps to the routed tasks + # and sh_load runs the shared expert) + e_u = shared.select( + fx.Int32(min(SHARED_EXPERT, E_TOP)), + fx.min(uniform(fx.ptr_load(route + (8 * tk + slot))), E_TOP), + ) + r_w = rsrc(w13 + fx.Int64(e_u) * W13_BYTES) + r_s = rsrc(s13) + wts: list[list] = [] # [gate, up][chunk] + scs: list[list] = [] + for gu in range_constexpr(2): + rg = gu * UG_PER_SLOT + cg + wts.append([]) + scs.append([]) + for i in range_constexpr(UG_CPW): + kc = wave * UG_CPW + i + wts[gu].append( + fx.Vector( + bo.buffer_load( + r_w, + ((rg * (HIDDEN // 64) + kc * 2) * 512 + lane * 16) // 4, + vec_width=4, + dtype=T.i32, + ) + ) + ) + scs[gu].append( + e8m0_load( + r_s, + group_scale_byte(e_u, W13_ROWS, rg, kc, W13_SCALE_COLS), + ) + ) + return wts, scs + + def ug_compute(t, wts, scs): + """Up / gate GEMV of task t -> swiglu -> its 16 mid values (of every token + for the shared expert: B column l % 16 is token l % 16).""" + shared, tk, slot, cg = ug_task(t) + col = shared.select(fx.min(l16, tokens - 1), tk) + cg_ = fx.Vector.filled(4, 0.0, fx.Float32) + cu_ = fx.Vector.filled(4, 0.0, fx.Float32) + # every chunk's B scale up front: read at use, each MFMA waited out an LDS + # latency (S = 4 up / gate 3.5 us over the unscaled floor) + sbs = [ + lds_i32(XSC_LDS + col * XSC_ROW + (wave * UG_CPW + i) * 4 + g4) + for i in range(UG_CPW) + ] + for i in range_constexpr(UG_CPW): + kc = wave * UG_CPW + i + xb = lds_b8(col * XN8_ROW + kc * 32) + cg_ = mfma_f4f8(wts[0][i], scs[0][i], xb, sbs[i], cg_) + cu_ = mfma_f4f8(wts[1][i], scs[1][i], xb, sbs[i], cu_) + stamp(9) + # gate partials in red, up partials in opart (free once the split is done) + fx.ptr_store(cg_, red + (wave * 64 + lane) * 4) + fx.ptr_store(cu_, opart + (wave * 64 + lane) * 4) + gpu.barrier() + # row tid % 16; a shared task's thread tid reads token tid / 16's column + if tid < shared.select(fx.Int32(UG_ROWS * tokens), fx.Int32(UG_ROWS)): + r = tid % UG_ROWS + col = tid // UG_ROWS + kk = shared.select(col, tk) + gsum = fx.Float32(0.0) + usum = fx.Float32(0.0) + for w in range_constexpr(WAVES): + gsum = gsum + fx.ptr_load( + red + ((w * 64 + 16 * (r // 4) + col) * 4 + r % 4) + ) + usum = usum + fx.ptr_load( + opart + ((w * 64 + 16 * (r // 4) + col) * 4 + r % 4) + ) + y = swiglu_mul_batch([gsum], [usum], fx.Float32(-swiglu_limit))[0] + mb_put( + mb("mid"), + (kk * MOE_SLOTS + slot) * INTER + cg * UG_ROWS + r, + bf16_round(y), + ) + gpu.barrier() + + def sh_load(j, has=None): + """The shared expert's half task j: gate rows 8 j .. + 8 on lanes l16 < 8, + the up rows alike on the rest -- one MFMA tile; the waves split k.""" + row = (l16 < 8).select(SH_ROWS * j + l16, INTER + SH_ROWS * j + l16 - 8) + rg = row // 16 + if const_expr(shared_fp8): + w_base, s_base = sh_ptr(SH_W13), sh_ptr(SH_S13) + w_bytes, s_bytes, e_rows = SH8_W13_BYTES, SH8_S13_BYTES, 0 + else: + w_base, s_base = w13 + fx.Int64(SHARED_EXPERT) * W13_BYTES, s13 + w_bytes, s_bytes = W13_BYTES, S13_BYTES + e_rows = SHARED_EXPERT * W13_ROWS + if const_expr(has is None): + r_w = rsrc(w_base) + r_s = rsrc(s_base) + else: # a CTA without the task reads through zero-sized buffers + r_w = rsrc(w_base, has.select(fx.Int32(w_bytes), fx.Int32(0))) + r_s = rsrc(s_base, has.select(fx.Int32(s_bytes), fx.Int32(0))) + wts = [] + scs = [] + # Row offset; each chunk adds 256 * (kc // 2) + 2 * (kc % 2). + sc_row = scale_index(e_rows + row, g4, W13_SCALE_COLS) + for i in range_constexpr(UG_CPW): + kc = wave * UG_CPW + i + wts.append(w_tile(r_w, rg, kc, g4 * 16 + row % 16, HIDDEN, shared_fp8)) + scs.append(e8m0_load(r_s, sc_row + kc // 2 * 256 + kc % 2 * 2)) + return wts, scs + + def sh_compute(j, wts, scs): + """Half task j's GEMV (B column l % 16 = token l % 16) -> swiglu -> the + shared slot's mid rows 8 j .. + 8 of every token.""" + col = fx.min(l16, tokens - 1) + c = fx.Vector.filled(4, 0.0, fx.Float32) + sbs = [ + lds_i32(XSC_LDS + col * XSC_ROW + (wave * UG_CPW + i) * 4 + g4) + for i in range(UG_CPW) + ] + for i in range_constexpr(UG_CPW): + kc = wave * UG_CPW + i + c = mfma_sh(wts[i], scs[i], lds_b8(col * XN8_ROW + kc * 32), sbs[i], c) + fx.ptr_store(c, red + (wave * 64 + lane) * 4) + gpu.barrier() + # mid row r = tid % 8 of token tid / 8: gate in tile row r, up in r + 8 + if tid < SH_ROWS * tokens: + r = tid % SH_ROWS + col = tid // SH_ROWS + gsum = fx.Float32(0.0) + usum = fx.Float32(0.0) + for w in range_constexpr(WAVES): + gsum = gsum + fx.ptr_load( + red + ((w * 64 + 16 * (r // 4) + col) * 4 + r % 4) + ) + usum = usum + fx.ptr_load( + red + ((w * 64 + 16 * (r // 4 + 2) + col) * 4 + r % 4) + ) + y = swiglu_mul_batch([gsum], [usum], fx.Float32(-swiglu_limit))[0] + mb_put( + mb("mid"), + (col * MOE_SLOTS + TOP_K) * INTER + j * SH_ROWS + r, + bf16_round(y), + ) + gpu.barrier() + + def down_load(k, td, has_dn): + """Token k's route-weighted down units of row group td; a CTA without a + down task reads through zero-sized buffers.""" + dn_bytes = has_dn.select(fx.Int32(W2_BYTES), fx.Int32(0)) + dn_scale_bytes = has_dn.select(fx.Int32(S2_BYTES), fx.Int32(0)) + rg_d = td * 2 + wave // 4 + units = [] + for i in range_constexpr(DN_UPW): + j = (wave % 4) + 4 * i + jj = fx.min(j, DN_UNITS - 1) + sl = jj // DN_KC + kc = jj % DN_KC + e_d = fx.min(uniform(fx.ptr_load(route + (8 * k + sl))), E_TOP) + wv = fx.Vector( + bo.buffer_load( + rsrc(w2 + fx.Int64(e_d) * W2_BYTES, dn_bytes), + ((rg_d * (INTER // 64) + kc * 2) * 512 + lane * 16) // 4, + vec_width=4, + dtype=T.i32, + ) + ) + sc = e8m0_load( + rsrc(s2, dn_scale_bytes), + group_scale_byte(e_d, HIDDEN, rg_d, kc, W2_SCALE_COLS), + ) + wgt = (j < DN_UNITS).select( + fx.ptr_load(rwt + (8 * k + sl)), fx.Float32(0.0) + ) + units.append((wv, sc, sl, kc, wgt)) + return units + + def down_load_shared(td, has_dn): + """The shared expert's down units of row group td (every token's).""" + rg_d = td * 2 + wave // 4 + if const_expr(shared_fp8): + w_base, s_base, w_bytes, s_bytes = ( + sh_ptr(SH_W2), sh_ptr(SH_S2), SH8_W2_BYTES, SH8_S2_BYTES, + ) # fmt: skip + e = 0 + else: + w_base, s_base, w_bytes, s_bytes = ( + w2 + fx.Int64(SHARED_EXPERT) * W2_BYTES, s2, W2_BYTES, S2_BYTES, + ) # fmt: skip + e = SHARED_EXPERT + r_w = rsrc(w_base, has_dn.select(fx.Int32(w_bytes), fx.Int32(0))) + r_s = rsrc(s_base, has_dn.select(fx.Int32(s_bytes), fx.Int32(0))) + units = [] + for i in range_constexpr(SH_UPW): + j = (wave % 4) + 4 * i + kc = fx.min(j, DN_KC - 1) + wv = w_tile(r_w, rg_d, kc, lane, INTER, shared_fp8) + sc = e8m0_load( + r_s, group_scale_byte(e, HIDDEN, rg_d, kc, W2_SCALE_COLS) + ) + wgt = (j < DN_KC).select(fx.Float32(shared_weight), fx.Float32(0.0)) + units.append((wv, sc, kc, wgt)) + return units + + def ffn_finish(t): + """Misc (row group t's bf16 partial rows, token k's DN_ROWS from k DN_ROWS) + -> every peer -> rank-ordered sum -> ar_out.""" + # every token's rows: reduce threads take token tid // 16, pair tid % 16 + row_tok = tid // (DN_ROWS // 2) + row_pair = tid % (DN_ROWS // 2) + if const_expr(W > 1): # push bf16 partial rows: wave w -> peer w + # (with AG_RS to the row group's owner rank only) + push_rows("ffn", t * DN_ROWS, DN_ROWS, t % W if AG_RS else None) + gpu.barrier() + if tid < DN_ROWS // 2 * tokens: + row = t * DN_ROWS + row_pair * 2 + if const_expr(W == 1): + t0 = fx.ptr_load(misc + tid * 2) + t1 = fx.ptr_load(misc + (tid * 2 + 1)) + else: + own = sym + fx.Int64(SY["ffn"]) + + def rank_sum(): + parts = mb_poll( + [ + ( + own, + ((src * MAX_TOKENS + row_tok) * HIDDEN + row) // 2, + 1, + ) + for src in range(W) + ], + "system", + ) + s0 = fx.Float32(0.0) + s1 = fx.Float32(0.0) + for src in range_constexpr(W): + f0, f1 = bf2_f32(parts[src][0]) + s0 = s0 + f0 + s1 = s1 + f1 + return s0, s1 + + if const_expr(AG_RS): + t0 = fx.Float32(0.0) + t1 = fx.Float32(0.0) + if t % W == rank: # the owner: sum, gather to every rank + t0, t1 = rank_sum() + for w in range_constexpr(W): + mb_put_bf( + peer_addr(w) + fx.Int64(SY["ffn_ag"]), + row_tok * HIDDEN + row, + [t0, t1], + CM_SYS, + ) + else: # the owner's sums (bf16, as ar_out stores them) + got = mb_poll( + [ + ( + sym + fx.Int64(SY["ffn_ag"]), + (row_tok * HIDDEN + row) // 2, + 1, + ) + ], + "system", + ) + t0, t1 = bf2_f32(got[0][0]) + else: + t0, t1 = rank_sum() + bo.buffer_store( + fx.Vector.from_elements([t0, t1], fx.Float32).to(fx.BFloat16), + rsrc(ar_out), + row_tok * HIDDEN + row, + ) + + # ======================================== 5w. experts, wide (S > 4) + # LDS pool words: MXFP8 xn, its scales, the up partials, router sigmoids; the + # down stage reuses the pool for MXFP8 mid rows (token k's slot s at row + # k MOE_SLOTS + s, a zero row last) and their scales + W_UPP = tokens * (XN8_ROW + XSC_ROW) + W_SIG = W_UPP + WAVES * 64 * 4 + MID_ROWS = MOE_SLOTS * tokens # the zero row's index + MID_ROW_WORDS = INTER // 4 + MIDSC_ROW_WORDS = INTER // 32 + W_MIDSC = (MID_ROWS + 1) * MID_ROW + # up / gate tasks go in pairs, a CTA runs both: an MXFP8 block is 32 rows + UG_PAIRS = UG_PER_SLOT // 2 + UGW_ITERS = 2 * ((U_MAX * UG_PAIRS + G - 1) // G) # up / gate tasks a CTA runs + # distinct experts the expert-table loops run (shared_fp8: the shared + # expert's MXFP8 tasks / units follow them, compiled apart) + U_TBL = U_MAX - 1 if shared_fp8 else U_MAX + # an MXFP8 shared task runs as SH_PARTS k parts, a part's weights in flight + # as many as a table task's: the whole task's and the next's were 192 VGPRs + # (256 with spills, and the table's tasks ran slower too) + SH_PARTS = 2 + assert UG_CPW % SH_PARTS == 0 + + def sh_ks(part): + n = UG_CPW // SH_PARTS + return tuple(range(part * n, (part + 1) * n)) + + DN_WITERS = (U_TBL * DN_KC + 3) // 4 # (expert, chunk) units a wave runs + DN_DEPTH = 12 # units whose weights are in flight + ROUTE_ITERS = (tokens + WAVES - 1) // WAVES # tokens a wave routes + + def popc64(x): + return fx.Int32(fx.ctpop(x)) + + def slot_of(k, e): + """Token k's slot routed to expert e (TOP_K: the shared expert).""" + s = fx.Int32(TOP_K) + for j in range_constexpr(TOP_K): + s = (fx.Int32(fx.ptr_load(route + (8 * k + j))) == e).select( + fx.Int32(j), s + ) + return s + + def build_experts(): + """Uexp / umask[0 .. n_u) := the tokens' distinct routed experts in + expert order and their token masks, [n_u] := the shared expert (every + token), uexp[U_MAX + 1] := n_u. The waves' masks must be in (synced).""" + m = fx.Int32(0) + if tid < N_ROUTED: + for w in range_constexpr(WAVES): + m = m | fx.Int32(fx.ptr_load(xs + (W_UPP + w * N_ROUTED + tid))) + bal = ballot(m != 0) + below = popc64(bal & ((fx.Int64(1) << fx.Int64(lane)) - 1)) + if tid == 0: + fx.ptr_store(popc64(bal), umask + (U_MAX + 1)) + gpu.barrier() + if tid < N_ROUTED: + off = (wave == 0).select( + fx.Int32(0), fx.Int32(fx.ptr_load(umask + (U_MAX + 1))) + ) + rank_ = off + below + if m != 0: + fx.ptr_store(fx.Int32(tid), uexp + rank_) + fx.ptr_store(m, umask + rank_) + if tid == N_ROUTED - 1: + nu_ = off + popc64(bal) + fx.ptr_store(nu_, uexp + (U_MAX + 1)) + fx.ptr_store(fx.Int32(SHARED_EXPERT), uexp + nu_) + fx.ptr_store(fx.Int32((1 << tokens) - 1), umask + nu_) + gpu.barrier() + + UG_KS = tuple(range(UG_CPW)) + + def wug_load(t, has, scales=True, fp8=False, ks=UG_KS): + """Up / gate weights (this wave's 128-k chunks ``ks``; and ``scales``: + the scale dwords, which hold the pair's both halves) of task t = + (distinct expert t / UG_PER_SLOT, row group t % UG_PER_SLOT); without + the task, reads go through zero-sized buffers. ``fp8``: the MXFP8 + shared expert's task (weights.mxfp8_a_tiles; the scales -- every + chunk's -- in the routed tables' order).""" + u = t // UG_PER_SLOT + cg = t % UG_PER_SLOT + if const_expr(fp8): + e_u = 0 + r_w = rsrc( + sh_ptr(SH_W13), has.select(fx.Int32(SH8_W13_BYTES), fx.Int32(0)) + ) + r_s = rsrc( + sh_ptr(SH_S13), has.select(fx.Int32(SH8_S13_BYTES), fx.Int32(0)) + ) + else: + e_u = fx.min(uniform(fx.ptr_load(uexp + u)), E_TOP) + r_w = rsrc( + w13 + fx.Int64(e_u) * W13_BYTES, + has.select(fx.Int32(W13_BYTES), fx.Int32(0)), + ) + r_s = rsrc(s13, has.select(fx.Int32(S13_BYTES), fx.Int32(0))) + wts: list[list] = [] + scs: list[list] = [] + for gu in range_constexpr(2): + rg = gu * UG_PER_SLOT + cg + wts.append([]) + scs.append([]) + for i in range_constexpr(UG_CPW): + kc = wave * UG_CPW + i + if const_expr(i in ks): + wts[gu].append(w_tile(r_w, rg, kc, lane, HIDDEN, fp8)) + if const_expr(scales and i % 2 == 0): + # the dword holding this lane's scales of chunks kc, kc + 1 + # of both 16-row halves of the 32-row block (bytes 2 d + a) + scs[gu].append( + fx.Int32( + bo.buffer_load( + r_s, + group_scale_byte( + e_u, W13_ROWS, rg, kc, W13_SCALE_COLS + ) + // 4, + vec_width=1, + dtype=T.i32, + ) + ) + ) + return wts, scs + + def wug_acc(wts, scs, half, cg_, cu_, ks=UG_KS, fp8=False): + """A task's up / gate MFMAs over its 128-k chunks ``ks`` of this wave + (wts[gu][j]: chunk ks[j]) for every token (B column l % 16 = token + l % 16), into (cg_, cu_). ``half``: the task's 16-row half of its + 32-row scale block (compile time); ``fp8``: an MXFP8 task + (``wug_load(fp8=True)``).""" + mfma_ = mfma_f8f8 if fp8 else mfma_f4f8 + col = fx.min(l16, tokens - 1) + sbs = [ + lds_i32(XSC_LDS + col * XSC_ROW + (wave * UG_CPW + i) * 4 + g4) + for i in ks + ] + for j in range_constexpr(len(ks)): + i = ks[j] + kc = wave * UG_CPW + i + xb = lds_b8(col * XN8_ROW + kc * 32) + sel = (i % 2) * 2 + half + cg_ = mfma_(wts[0][j], scs[0][i // 2], xb, sbs[j], cg_, sel) + cu_ = mfma_(wts[1][j], scs[1][i // 2], xb, sbs[j], cu_, sel) + return cg_, cu_ + + def wug_y(wts, scs, half): + """A task's up / gate GEMV (``wug_acc``) -> swiglu: thread (row r, + token k)'s mid value.""" + zero = fx.Vector.filled(4, 0.0, fx.Float32) + return wug_swiglu(*wug_acc(wts, scs, half, zero, zero)) + + def wug_swiglu(cg_, cu_): + """The waves' up / gate partials -> swiglu (``wug_y``).""" + fx.ptr_store(cg_, red + (wave * 64 + lane) * 4) + fx.ptr_store(cu_, fpool + (W_UPP + (wave * 64 + lane) * 4)) + gpu.barrier() + r = tid % UG_ROWS + k = fx.min(tid // UG_ROWS, tokens - 1) + gsum = fx.Float32(0.0) + usum = fx.Float32(0.0) + for w in range_constexpr(WAVES): + gsum = gsum + fx.ptr_load( + red + ((w * 64 + 16 * (r // 4) + k) * 4 + r % 4) + ) + usum = usum + fx.ptr_load( + fpool + (W_UPP + (w * 64 + 16 * (r // 4) + k) * 4 + r % 4) + ) + gpu.barrier() # red and the up partials are free for the next task + return swiglu_mul_batch([gsum], [usum], fx.Float32(-swiglu_limit))[0] + + def fp8_lanes4(q): + """Lane r (r % 4 == 0): the E4M3 dword of lanes r .. r + 3's values.""" + lo = fx.Int32( + rocdl.cvt_pk_fp8_f32(T.i32, q, xshfl(q, 1), fx.Int32(0), False) + ) + return lo | (xshfl(lo, 2) << 16) + + def wug_publish(pr, y_lo, y_hi): + """Task pair pr = (distinct expert pr / UG_PAIRS, 32-row block + pr % UG_PAIRS): its mid rows (thread (r, k) holds rows r and 16 + r of + the block) as the down stage's MXFP8 -- the rounding its bf16 staging + did, one 32-block a 16-lane group -- then the pair's flag.""" + u = pr // UG_PAIRS + blk_ = pr % UG_PAIRS + r = tid % UG_ROWS + k = fx.min(tid // UG_ROWS, tokens - 1) + lo = bf16_round(y_lo) + hi = bf16_round(y_hi) + sc, inv = mx8_scale([lo, hi], (1, 2, 4, 8)) + words = [fp8_lanes4(mx8_q(v, inv)) for v in (lo, hi)] + e_u = fx.Int32(fx.ptr_load(uexp + u)) + m_u = fx.Int32(fx.ptr_load(umask + u)) + row = k * MOE_SLOTS + slot_of(k, e_u) + if (tid < UG_ROWS * tokens) & (((m_u >> k) & 1) != 0): + if r % 4 == 0: + for h in range_constexpr(2): + bo.buffer_store( + words[h], + rsrc(mb("mid8p")), + row * MID_ROW_WORDS + blk_ * 8 + h * 4 + r // 4, + cache_modifier=CM_DEV, + ) + if r == 0: + bo.buffer_store( + sc, + rsrc(mb("midsp")), + row * MIDSC_ROW_WORDS + blk_, + cache_modifier=CM_DEV, + ) + fx.memory_fence( + syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Release + ) + gpu.barrier() + if tid == 0: + mb_put(mb("ugdone"), pr, fx.Int32(1)) + + def down_tables(): + """Urow / uwt := per (distinct expert, B column) the down stage's mid row + (the zero row where the column's token is not routed there) and route + weight. Routing only: built while the first up / gate weights land.""" + nu = uniform(fx.ptr_load(uexp + (U_MAX + 1))) + for i in range_constexpr((U_MAX * 16 + THREADS - 1) // THREADS): + ix = tid + THREADS * i + if ix < U_MAX * 16: + xu = ix // 16 + xk = ix % 16 + xe = fx.Int32(fx.ptr_load(uexp + fx.min(xu, nu))) + xm = fx.Int32(fx.ptr_load(umask + fx.min(xu, nu))) + xt = fx.min(xk, tokens - 1) + xs_ = slot_of(xt, xe) + xh = (((xm >> xk) & 1) != 0) & (xu <= nu) + fx.ptr_store( + xh.select(xt * MOE_SLOTS + xs_, fx.Int32(MID_ROWS)), urow + ix + ) + fx.ptr_store( + xh.select(fx.ptr_load(rwt + (8 * xt + xs_)), fx.Float32(0.0)), + uwt + ix, + ) + + def stage_moe_wide(): + td = start("down") # this CTA's down task (the split tasks come first) + has_dn = td < N_DN + td = fx.min(td, N_DN - 1) + # every token's routing: wave w routes tokens w, w + WAVES, .., their + # scores polled in one batch + mine = bid % SCORE_COPIES + # wave w's expert -> mask of the tokens it routed there, in the up + # partials' words (free until the first up / gate task) + wave_mask = xs + (W_UPP + wave * N_ROUTED) + fx.ptr_store(fx.Int32(0), wave_mask + lane) + fx.ptr_store(fx.Int32(0), wave_mask + (lane + 64)) + wave_toks = [ + fx.min(wave + WAVES * j, tokens - 1) for j in range(ROUTE_ITERS) + ] + sg = mb_poll( + [ + (scores_copy(mine, k), (lane + 64 * h) * 2, 2) + for k in wave_toks + for h in range(2) + ] + ) + for j in range_constexpr(ROUTE_ITERS): + if wave + WAVES * j < tokens: + route_top4( + wave_toks[j], + sg[2 * j][0], + sg[2 * j][1].bitcast(fx.Float32), + sg[2 * j + 1][0], + sg[2 * j + 1][1].bitcast(fx.Float32), + fpool + W_SIG, + wave_mask, + ) + stamp(2) + gpu.barrier() + build_experts() + stamp(6) + if (bid == BASE["down"]) & (tid < MOE_SLOTS * tokens): + sk = tid // MOE_SLOTS + sj = tid % MOE_SLOTS + mb_put( + mb("sel"), + sk * 2 * MOE_SLOTS + sj, + fx.ptr_load(route + (8 * sk + sj)), + ) + mb_put( + mb("sel"), + sk * 2 * MOE_SLOTS + MOE_SLOTS + sj, + fx.ptr_load(rwt + (8 * sk + sj)), + ) + nu_ug = uniform(fx.ptr_load(uexp + (U_MAX + 1))) + n_pairs = (nu_ug if shared_fp8 else nu_ug + 1) * UG_PAIRS + # shared_fp8: shared pair j runs after the table's pairs on CTA + # (sh0 + j) % G. It costs two of theirs: from a full round of pairs on, + # the table's skip slots hole0 .. hole0 + UG_PAIRS - 1 (slot s: CTA s % + # G's (s / G)-th pair), the shared CTAs' last full round, so that those + # carry one pair fewer (otherwise the up / gate stage ran a pair longer + # than the fused build's at every S; S = 8: 25.8 -> 31.6 us slowest CTA) + sh0 = n_pairs + if const_expr(shared_fp8): + hole0 = n_pairs + (UG_PAIRS - G) + holes = hole0 >= 0 + sh0 = n_pairs + holes.select(fx.Int32(UG_PAIRS), fx.Int32(0)) + sh_j = (bid + (G - sh0 % G)) % G + has_shp = sh_j < UG_PAIRS + sh_pr = n_pairs + fx.min(sh_j, UG_PAIRS - 1) + UG_ALL = UGW_ITERS + (2 * SH_PARTS if shared_fp8 else 0) + + def task(it): + """Iteration it's task: slot bid + (it / 2) G's pair, half it % 2 + (past UGW_ITERS: the shared pair's, half (it - UGW_ITERS) / SH_PARTS + a k part an iteration).""" + if const_expr(it >= UGW_ITERS): + return sh_pr * 2 + (it - UGW_ITERS) // SH_PARTS, has_shp + pr = fx.Int32(bid + it // 2 * G) + if const_expr(shared_fp8): + # a hole runs nothing (pair n_pairs); past them, UG_PAIRS lower + pr = (holes & (pr >= hole0)).select( + (pr < hole0 + UG_PAIRS).select(n_pairs, pr - UG_PAIRS), pr + ) + return fx.min(pr, n_pairs - 1) * 2 + it % 2, pr < n_pairs + + def task_load(it, t, has, scales=True): + if const_expr(it >= UGW_ITERS): + # the shared pair's scales come with its first part + q = it - UGW_ITERS + return wug_load( + t, has, scales=q == 0, fp8=True, ks=sh_ks(q % SH_PARTS) + ) + return wug_load(t, has, scales) + + t0, h0 = task(0) + cur = task_load(0, t0, h0) # task 0: a lo half, with the pair's scales + down_tables() + stamp(7) + # the first task's weights are in flight: now the xn + # every token's MXFP8 xn and scales -> the pool, once every router task + # flagged its slice (plain words, 16 B a load) + n_rt = N_ROUTER * tokens + if wave == 0: + mb_poll( + [ + (mb("xndone"), fx.min(lane + 64 * c, n_rt - 1), 1) + for c in range((n_rt + 63) // 64) + ] + ) + gpu.barrier() + acquire() + for base_w, n_row, lds_w, lds_row in ( + (mb("xn8p"), XN8_WORDS, 0, XN8_ROW), + (mb("xscp"), XSC_WORDS, XSC_LDS, XSC_ROW), + ): + n_v = tokens * n_row // 4 + for i in range_constexpr((n_v + THREADS - 1) // THREADS): + v = tid + THREADS * i + if v < n_v: + fx.ptr_store( + fx.Vector( + bo.buffer_load( + rsrc(base_w), + v * 4, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ), + xs + (lds_w + v * 4 // n_row * lds_row + v * 4 % n_row), + ) + gpu.barrier() + y_lo = fx.Float32(0.0) # the pair's first half, until the second's + if const_expr(shared_fp8): + # a shared half's up / gate accumulators, part to part + sh_cg = fx.Vector.filled(4, 0.0, fx.Float32) + sh_cu = fx.Vector.filled(4, 0.0, fx.Float32) + for it in range_constexpr(UG_ALL): + t, has_t = task(it) + if const_expr(it + 1 < UG_ALL): + tn, hn = task(it + 1) + # a hi half reads no scales: its lo half's dwords hold them + nxt = task_load(it + 1, tn, hn, scales=(it + 1) % 2 == 0) + if const_expr(it < UGW_ITERS): + if const_expr(it % 2 == 0): + pair_scs = cur[1] + if has_t: + # a pair: row groups cg, cg + 1 + y = wug_y(cur[0], pair_scs, it % 2) + if const_expr(it % 2 == 0): + y_lo = y + else: + wug_publish(t // 2, y_lo, y) + else: + half, part = divmod(it - UGW_ITERS, SH_PARTS) + if const_expr(half == 0 and part == 0): + pair_scs = cur[1] + if has_t: + if const_expr(part == 0): + sh_cg = fx.Vector.filled(4, 0.0, fx.Float32) + sh_cu = fx.Vector.filled(4, 0.0, fx.Float32) + sh_cg, sh_cu = wug_acc( + cur[0], pair_scs, half, sh_cg, sh_cu, sh_ks(part), fp8=True + ) + if const_expr(part == SH_PARTS - 1): + y = wug_swiglu(sh_cg, sh_cu) + if const_expr(half == 0): + y_lo = y + else: + wug_publish(t // 2, y_lo, y) + if const_expr(it + 1 < UG_ALL): + cur = nxt + stamp(8) + if has_dn: + wide_down(td) + stamp(11) + + def wide_down(td): + """Row group td's down GEMV over every distinct expert (B column l % 16 = + token l % 16's mid for it, the zero row where the token is not routed + there), route-weighted per column -> ffn_finish.""" + nu = uniform(fx.ptr_load(uexp + (U_MAX + 1))) + rg_d = td * 2 + wave // 4 + n_units = (nu if shared_fp8 else nu + 1) * DN_KC + + def unit_load(i): + """(weights, scale, unit) of this wave's i-th (expert, 128-k chunk) + unit, zero-sized reads past the last.""" + j = wave % 4 + 4 * i + ok = j < n_units + # clamped into the table: a unit past the last reads a real or zero + # row, never past urow (a garbage row is garbage B, and 0 * NaN) + du = fx.min(j // DN_KC, U_MAX - 1) + dkc = j % DN_KC + de = fx.min(uniform(fx.ptr_load(uexp + du)), E_TOP) + dwv = fx.Vector( + bo.buffer_load( + rsrc( + w2 + fx.Int64(de) * W2_BYTES, + ok.select(fx.Int32(W2_BYTES), fx.Int32(0)), + ), + ((rg_d * (INTER // 64) + dkc * 2) * 512 + lane * 16) // 4, + vec_width=4, + dtype=T.i32, + ) + ) + dsc = e8m0_load( + rsrc(s2, ok.select(fx.Int32(S2_BYTES), fx.Int32(0))), + group_scale_byte(de, HIDDEN, rg_d, dkc, W2_SCALE_COLS), + ) + return dwv, dsc, du, dkc, ok + + def with_rows(u): + """A loaded unit plus its B row and route weight (the tables are in), + read ahead so the unit's B loads wait on one LDS trip.""" + dwv, dsc, du, dkc, ok = u + drow = fx.Int32(fx.ptr_load(urow + (du * 16 + l16))) + dw = ok.select(fx.ptr_load(uwt + (du * 16 + l16)), fx.Float32(0.0)) + return dwv, dsc, drow, dkc, dw, ok + + # software pipeline: the first units' weights go out before the mids are + # in; unit i + DN_DEPTH's before unit i's MFMA + acc = [fx.Float32(0.0) for _ in range(4)] + if const_expr(shared_fp8): + # the MXFP8 shared expert's units (wave % 4 + 4 i): no routing needed, + # their weights go out first + r_shared_w = rsrc(sh_ptr(SH_W2)) + r_shared_s = rsrc(sh_ptr(SH_S2)) + sh_dn = [] + for i in range_constexpr(SH_UPW): + skc = fx.min(wave % 4 + 4 * i, DN_KC - 1) + sh_dn.append( + ( + f8_tile_load(r_shared_w, rg_d * (INTER // 128) + skc, lane), + e8m0_load( + r_shared_s, + group_scale_byte(0, HIDDEN, rg_d, skc, W2_SCALE_COLS), + ), + skc, + wave % 4 + 4 * i < DN_KC, + ) + ) + inflight = [unit_load(i) for i in range(min(DN_DEPTH, DN_WITERS))] + # every up / gate task pair's flag, then the MXFP8 mid rows and their + # scales, the pool's layout (plain words, 16 B a load) + n_ug = (nu + 1) * UG_PAIRS + mb_poll( + [ + (mb("ugdone"), fx.min(tid + THREADS * i, n_ug - 1), 1) + for i in range((U_MAX * UG_PAIRS + THREADS - 1) // THREADS) + ] + ) + gpu.barrier() + acquire() + for base_w, n_row, lds_w, lds_row in ( + (mb("mid8p"), MID_ROW_WORDS, 0, MID_ROW), + (mb("midsp"), MIDSC_ROW_WORDS, W_MIDSC, MIDSC_ROW), + ): + n_v = MID_ROWS * n_row // 4 + for i in range_constexpr((n_v + THREADS - 1) // THREADS): + v = tid + THREADS * i + if v < n_v: + fx.ptr_store( + fx.Vector( + bo.buffer_load( + rsrc(base_w), + v * 4, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ), + xs + (lds_w + v * 4 // n_row * lds_row + v * 4 % n_row), + ) + if tid < MID_ROW_WORDS: + fx.ptr_store(fx.Int32(0), xs + (MID_ROWS * MID_ROW + tid)) + if tid < MIDSC_ROW_WORDS: + fx.ptr_store(fx.Int32(0), xs + (W_MIDSC + MID_ROWS * MIDSC_ROW + tid)) + gpu.barrier() + stamp(10) + for i in range_constexpr(DN_WITERS): + if const_expr(i == 0): # the tables are in: the early units' rows + inflight = [with_rows(u) for u in inflight] + dwv, dsc, drow, dkc, dw, ok = inflight.pop(0) + if const_expr(i + DN_DEPTH < DN_WITERS): + inflight.append(with_rows(unit_load(i + DN_DEPTH))) + dcu = mfma_f4f8( + dwv, + dsc, + lds_b8(drow * MID_ROW + dkc * 32), + lds_i32(W_MIDSC + drow * MIDSC_ROW + dkc * 4 + g4), + fx.Vector.filled(4, 0.0, fx.Float32), + ) + # a unit past the last adds nothing, whatever its MFMA gave + acc = [ + acc[e] + ok.select(dcu[e] * dw, fx.Float32(0.0)) for e in range(4) + ] + if const_expr(shared_fp8): + # the shared expert's B column l16: its mid row (slot TOP_K) and weight + srow = fx.Int32(fx.ptr_load(urow + (nu * 16 + l16))) + sw = fx.ptr_load(uwt + (nu * 16 + l16)) + for swv, ssc, skc, sok in sh_dn: + dcu = mfma_f8f8( + swv, + ssc, + lds_b8(srow * MID_ROW + skc * 32), + lds_i32(W_MIDSC + srow * MIDSC_ROW + skc * 4 + g4), + fx.Vector.filled(4, 0.0, fx.Float32), + ) + acc = [ + acc[e] + sok.select(dcu[e] * sw, fx.Float32(0.0)) + for e in range(4) + ] + a0, a1, a2, a3 = acc + fx.ptr_store( + fx.Vector.from_elements([a0, a1, a2, a3], fx.Float32), + red + (wave * 64 + lane) * 4, + ) + gpu.barrier() + if tid < DN_ROWS * tokens: + k = tid // DN_ROWS + grp = (tid % DN_ROWS) // 16 + rr = tid % 16 + tot = fx.Float32(0.0) + for w in range_constexpr(4): + tot = tot + fx.ptr_load( + red + (((grp * 4 + w) * 64 + 16 * (rr // 4) + k) * 4 + rr % 4) + ) + fx.ptr_store(bf16_round(tot), misc + tid) + gpu.barrier() + ffn_finish(td) + + def stage_moe(): + td = bid - BASE["down"] # this CTA's down task + has_dn = (td >= 0) & (td < N_DN) + td = fx.min(fx.max(td, 0), N_DN - 1) + has_sh = start("shared") < SH_TASKS + n_xw = XN8_WORDS // THREADS + + def xn_specs(k): + return [ + (mb("xn8"), k * XN8_WORDS + tid + i * THREADS, 1) + for i in range(n_xw) + ] + [(mb("xsc"), k * XSC_WORDS + fx.min(tid, XSC_WORDS - 1), 1)] + + def stage_xn(k, got): + for i in range_constexpr(n_xw): + fx.ptr_store(got[i][0], xs + (k * XN8_ROW + tid + i * THREADS)) + if tid < XSC_WORDS: + fx.ptr_store(got[n_xw][0], xs + (XSC_LDS + k * XSC_ROW + tid)) + + def fetch_xn(tok=None): + """Token tok's (default: every token's) MXFP8 xn and scales -> xs. Wave + 0 first polls the last word of each router task's slice: all 512 + threads of 256 CTAs spinning on the whole of it (S = 4) held the + shared-expert CTAs' xn back 3 us.""" + gate_tok = fx.min(lane // N_ROUTER, tokens - 1) if tok is None else tok + if wave == 0: + mb_poll( + [ + ( + mb("xn8"), + gate_tok * XN8_WORDS + + (lane % N_ROUTER + 1) * (XN_SLICE // 4) + - 1, + 1, + ) + ] + ) + gpu.barrier() + if tok is None: + got = mb_poll([sp for k in range(tokens) for sp in xn_specs(k)]) + for k in range_constexpr(tokens): + stage_xn(k, got[k * (n_xw + 1) : (k + 1) * (n_xw + 1)]) + else: + stage_xn(tok, mb_poll(xn_specs(tok))) + + sh_late = shared_after_routing(tokens) + if const_expr(sh_late): + sh_j = fx.min(start("shared"), SH_TASKS - 1) + sh_w = sh_load(sh_j, has_sh) + if has_sh: + fetch_xn() + elif has_sh: # the shared expert needs xn, not the routing + sh_j0 = start("shared") + sh_w0 = sh_load(sh_j0) + fetch_xn() + gpu.barrier() + sh_compute(sh_j0, sh_w0[0], sh_w0[1]) + if ug_task_id(0)[1] | has_dn: + # xn lands before the logits: stage it while waiting (a shared-expert + # CTA has it already); a CTA's routed tasks are one token's + if start("shared") >= SH_TASKS: + fetch_xn(bid // CPT if CPT is not None else None) + # wave k polls token k's router logits and routes it right away + mine = bid % SCORE_COPIES + for k in range_constexpr(tokens): + if wave == k: + got = mb_poll( + [ + (scores_copy(mine, k), (lane + 64 * h) * 2, 2) + for h in range(2) + ] + ) + stamp(21) + route_top4( + k, + got[0][0], + got[0][1].bitcast(fx.Float32), + got[1][0], + got[1][1].bitcast(fx.Float32), + opart, + ) + stamp(2) + gpu.barrier() + stamp(6) + if (bid == BASE["down"]) & (tid < MOE_SLOTS * tokens): + sk = tid // MOE_SLOTS + sj = tid % MOE_SLOTS + mb_put( + mb("sel"), + sk * 2 * MOE_SLOTS + sj, + fx.ptr_load(route + (8 * sk + sj)), + ) + mb_put( + mb("sel"), + sk * 2 * MOE_SLOTS + MOE_SLOTS + sj, + fx.ptr_load(rwt + (8 * sk + sj)), + ) + cur = ug_load(ug_task_id(0)[0]) + stamp(7) + if const_expr(sh_late): # noqa: SIM102 + if has_sh: # while the first routed task's weights are in flight + sh_compute(sh_j, sh_w[0], sh_w[1]) + # task i + 1's weights go out before task i's GEMV (software pipeline); + # past the last task the load index is clamped and the result dropped + for it in range_constexpr(UG_ITERS): + t, has_t = ug_task_id(it) + if const_expr(it + 1 < UG_ITERS): + nxt = ug_load(ug_task_id(it + 1)[0]) + if has_t: + ug_compute(t, cur[0], cur[1]) + if const_expr(it == 0): + # token 0's down weights stream under the mid exchange; issued + # with the first task's up / gate weights they share the CU's + # bandwidth and delay them (K4 on 4 GPUs 34.3 -> 32.8 us) + dn_units = {0: down_load(0, td, has_dn)} + sh_units = down_load_shared(td, has_dn) + if const_expr(tokens > 1 and it == UG_ITERS - 1): + # two tokens in flight when the down GEMVs start + dn_units[1] = down_load(1, td, has_dn) + if const_expr(it + 1 < UG_ITERS): + cur = nxt + stamp(8) + if has_dn: + moe_down(td, has_dn, dn_units, sh_units) + stamp(11) + + def moe_down(t, has_dn, dn_units, sh_units): + """Every token's route-weighted down GEMV of row group ``t`` -> push bf16 + partial rows -> rank-ordered sum -> ar_out.""" + + def stage_mids(k): + """Token k's mid rows (every slot) -> xs, as bf16 pairs.""" + mids = mb_poll( + [ + ( + mb("mid"), + k * MOE_SLOTS * INTER + + fx.min(tid * 2 + i * 2 * THREADS, MOE_SLOTS * INTER - 2), + 2, + ) + for i in range(4) + ] + ) + for i in range_constexpr(4): + e = tid * 2 + i * 2 * THREADS + if e < MOE_SLOTS * INTER: + fx.ptr_store( + bf16_pair( + mids[i][0].bitcast(fx.Float32), + mids[i][1].bitcast(fx.Float32), + ), + xs + (k * (MOE_SLOTS * INTER // 2) + e // 2), + ) + + # every token's mid in one go: they land together (S = 4: -1.1 us) + for k in range_constexpr(tokens): + stage_mids(k) + stamp(10) + gpu.barrier() + # then MXFP8 in a pass of its own: quantizing inside stage_mids (after + # each poll) cost S = 1 / 4 1.75 / 6 us, though the arithmetic is tiny + for k in range_constexpr(tokens): + if tid < MID8_WORDS // 2: # 8 elements a lane, 4 lanes a 32-block + w = fx.Vector( + fx.ptr_load(xs + (k * MID16_WORDS + tid * 4), result_type=v4f) + ).bitcast(fx.Int32) + f = [v for j in range(4) for v in bf2_f32(w[j])] + sc, inv = mx8_scale(f, (1, 2)) + qv = [mx8_q(v, inv) for v in f] + fx.ptr_store( + fp8_pack4(*qv[:4]), xs + (MID8_LDS + k * MID8_WORDS + tid * 2) + ) + fx.ptr_store( + fp8_pack4(*qv[4:]), + xs + (MID8_LDS + k * MID8_WORDS + tid * 2 + 1), + ) + if tid % 4 == 0: + fx.ptr_store(sc, xs + (MIDSC_LDS + k * MIDSC_WORDS + tid // 4)) + gpu.barrier() + # the shared expert once, B column l % 16 = token l % 16's mid -> opart + # (free once the split is done); each token's reduce adds its column + col_tok = fx.min(l16, tokens - 1) + sh = [fx.Float32(0.0) for _ in range(4)] + for wv, sc, kc, wgt in sh_units: + cu = mfma_sh( + wv, + sc, + lds_b8( + MID8_LDS + + col_tok * MID8_WORDS + + (TOP_K * INTER + kc * 128) // 4 + ), + lds_i32( + MIDSC_LDS + + col_tok * MIDSC_WORDS + + (TOP_K * INTER + kc * 128) // 32 + + g4 + ), + fx.Vector.filled(4, 0.0, fx.Float32), + ) + sh = [sh[j4] + cu[j4] * wgt for j4 in range(4)] + fx.ptr_store( + fx.Vector.from_elements(sh, fx.Float32), opart + (wave * 64 + lane) * 4 + ) + # every token's GEMV first, partials in opart past the shared ones: one + # barrier and one reduce for all tokens, not two barriers per token + for k in range_constexpr(tokens): + # token k + 2's weights go out before token k's GEMV + if const_expr(k + 2 < tokens): + dn_units[k + 2] = down_load(k + 2, t, has_dn) + acc = [fx.Float32(0.0) for _ in range(4)] + for wv, sc, sl, kc, wgt in dn_units.pop(k): + cu = mfma_f4f8( + wv, + sc, + lds_b8( + MID8_LDS + k * MID8_WORDS + (sl * INTER + kc * 128) // 4 + ), + lds_i32( + MIDSC_LDS + + k * MIDSC_WORDS + + (sl * INTER + kc * 128) // 32 + + g4 + ), + fx.Vector.filled(4, 0.0, fx.Float32), + ) + acc = [acc[j4] + cu[j4] * wgt for j4 in range(4)] + fx.ptr_store( + fx.Vector.from_elements(acc, fx.Float32), + opart + ((k + 1) * WAVES * 64 + wave * 64 + lane) * 4, + ) + gpu.barrier() + if tid < DN_ROWS * tokens: + k = tid // DN_ROWS + grp = (tid % DN_ROWS) // 16 + rr = tid % 16 + tot = fx.Float32(0.0) + for w in range_constexpr(4): + tot = tot + fx.ptr_load( + opart + + ( + ((k + 1) * WAVES * 64 + (grp * 4 + w) * 64 + 16 * (rr // 4)) + * 4 + + rr % 4 + ) + ) + for w in range_constexpr(4): + tot = tot + fx.ptr_load( + opart + (((grp * 4 + w) * 64 + 16 * (rr // 4) + k) * 4 + rr % 4) + ) + fx.ptr_store(bf16_round(tot), misc + tid) + gpu.barrier() + ffn_finish(t) + + (stage_moe_wide if WIDE else stage_moe)() + flush_stamps() + + @flyc.jit + def launch_post_attn( + h_in: Int64, + q: Int64, + block_table: Int64, + seq_lens: Int64, + k_cache: Int64, + v_cache: Int64, + k_scale: Int64, + v_scale: Int64, + w_o: Int64, + s_o: Int64, + g_post: Int64, + w_gate: Int64, + bias: Int64, + w13: Int64, + s13: Int64, + w2: Int64, + s2: Int64, + sh_args: Int64, + h_mid: Int64, + ar_out: Int64, + scratch: Int64, + sym: Int64, + peers: Int64, + step: Int64, + rank: Int32, + layer: Int32, + bt_width: Int32, + q_len: Int32, + tl: Int64, + k1_args: Int64, + positions: Int64, + slot_mapping: Int64, + res: Int64, + stream: fx.Stream = _CURRENT_STREAM, + ): + post_attn_kernel( + h_in, + q, + block_table, + seq_lens, + k_cache, + v_cache, + k_scale, + v_scale, + w_o, + s_o, + g_post, + w_gate, + bias, + w13, + s13, + w2, + s2, + sh_args, + h_mid, + ar_out, + scratch, + sym, + peers, + step, + rank, + layer, + bt_width, + q_len, + tl, + k1_args, + positions, + slot_mapping, + res, + ).launch(grid=(G,), block=(THREADS,), stream=stream) + + return launch_post_attn diff --git a/vllm/models/minimax_m3/amd/mono/kernels/pre_attn.py b/vllm/models/minimax_m3/amd/mono/kernels/pre_attn.py new file mode 100644 index 000000000000..b8b9b2133e39 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/kernels/pre_attn.py @@ -0,0 +1,902 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""K1 ``mono_pre``: one MiniMax-M3 sparse layer, input norm through cache insert. + +One launch per layer and rank, ``BLOCKS`` CTAs x ``THREADS`` threads, for S = +``tokens`` decode rows, ``q_len`` consecutive rows a request (each row below is +per token):: + + acc = ar + res (f32) -> residual out bf16(acc) + -> gemma RMSNorm -> per-token FP8 (amax / FP8_MAX, true divide) + -> qkv | index_q | index_k GEMV on the preshuffled ptpc-FP8 weight + -> per-head gemma norm + partial NeoX RoPE + -> q, index_q out (bf16); K / V -> page-16 SHUFFLE cache with per-token + scales; index_k -> index cache (E4M3, unit scale) + +It reads exactly what the original path leaves in memory (the ``ar`` / ``res`` +pair of the custom all-reduce, the loaded weights, the bound caches) and writes +what the original ``fused_qknorm_idxrqknorm`` writes, so either path can run any +layer. The numerical contracts (reduction orders, where values round to bf16) +follow the aiter kernels the original path uses; see each stage. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops as bo +from flydsl.expr import const_expr, gpu, range_constexpr, rocdl +from flydsl.expr import math as fmath +from flydsl.expr.typing import Int32, Int64, T + +from vllm.models.minimax_m3.amd.mono import config as mono_config +from vllm.models.minimax_m3.amd.mono.config import ( + BLOCKS, + HEAD_DIM, + HIDDEN, + LOCAL_Q_HEADS, + MAX_QKV_ROWS, + MAX_TOKENS, + ONE_INDEX_HEAD, + PAGE16, + ROTARY_DIM, + THREADS, + TP, + WAVES, + IndexHeads, + qkv_rows, +) +from vllm.models.minimax_m3.amd.mono.kernels.common import ( + CM_DEV, + Mailbox, + ar_block_sums, + ar_pack_sumsq, + bf16_pair, + bf16_round, + butterfly, + div_rn, + fp8_pack4, + hw_rsq, + kernel_symbol, + memrealtime, + rsrc, + traced, + uniform, + wave_max, +) +from vllm.models.minimax_m3.amd.mono.kernels.index_score import ( + emit_index_scores, + index_scale_log2e, + step_rows, +) + +FP8_MAX = 448.0 +ROWS_PER_TASK = 16 +K_CHUNKS = HIDDEN // 64 # 64-k chunks: one dwordx4 of weight per lane +CHUNKS_PER_WAVE = K_CHUNKS // WAVES +Q_OFF, K_OFF, V_OFF = 0, LOCAL_Q_HEADS * HEAD_DIM, (LOCAL_Q_HEADS + 1) * HEAD_DIM +N_HEAD_TASKS = LOCAL_Q_HEADS + 4 # q heads, k, v, the rank's index_q, index_k +IK_TASK = LOCAL_Q_HEADS + 3 +HEAD_ROW_GROUPS = HEAD_DIM // ROWS_PER_TASK # GEMV tasks a head +ELEMS = 4 # per lane in a head task: 32 lanes x 4 = 128, the aiter kernel's split +PAGE_BYTES = PAGE16 * HEAD_DIM # one page-16 of one kv head (fp8) + + +def page16_of(slot): + """Page-16 id of a page-128 slot under the host's K/V block page stride.""" + kps = mono_config.KV_PAGE_STRIDE + if kps == mono_config.SPARSE_BLOCK // PAGE16: + return slot // PAGE16 + blk, t = slot // mono_config.SPARSE_BLOCK, slot % mono_config.SPARSE_BLOCK + return blk * kps + t // PAGE16 + + +# GEMV tasks, head tasks (every token of a head on one CTA) and the tokens' norm +# tasks each own a CTA, in this order; a head CTA's half waves hold the tokens. +# The rank's own rows (q | k | v | its index q | index k) are the first ONE_GEMV +# tasks under indexer context parallelism too, and its other index q heads' the +# OTHER_GEMV after the norm tasks, run only in a step that needs them: the CTA +# layout of a step without long requests is the one-head build's. +ONE_GEMV = qkv_rows(1) // ROWS_PER_TASK +OTHER_GEMV = (TP - 1) * HEAD_ROW_GROUPS +HT0 = ONE_GEMV +NT0 = HT0 + N_HEAD_TASKS +XT0 = NT0 + MAX_TOKENS +assert XT0 + OTHER_GEMV <= BLOCKS +assert MAX_TOKENS * 32 <= THREADS +assert HIDDEN // 4 % THREADS == 0 + +# The launch goes on the caller's current stream (graph capture included). +_CURRENT_STREAM = fx.Stream(None) + +SCRATCH_QKV = 0 # (value, tag) pairs, one per qkv row +SCRATCH_X8 = MAX_TOKENS * MAX_QKV_ROWS * 8 # S > 2: the norm tasks' x8 rows, words +SCRATCH_X8S = SCRATCH_X8 + MAX_TOKENS * HIDDEN // 4 * 8 # and their scales +# fused layer kernel only: a head task's / the residual's done flags +SCRATCH_HDONE = SCRATCH_X8S + MAX_TOKENS * 8 +SCRATCH_RDONE = SCRATCH_HDONE + MAX_TOKENS * N_HEAD_TASKS * 8 +SCRATCH_BYTES = SCRATCH_RDONE + MAX_TOKENS * 8 + + +def _ld_bf16x4(r, k): + w = fx.Vector(bo.buffer_load(r, k // 2, vec_width=2, dtype=T.i32)) + v = w.bitcast(fx.BFloat16).to(fx.Float32) + return [v[j] for j in range(4)] + + +def ld_raw8(r, lt): + """Logical thread lt's 8 bf16 of a row, as loaded (for ``bf16x8_f32``).""" + return fx.Vector(bo.buffer_load(r, lt * 4, vec_width=4, dtype=T.i32)) + + +def bf16x8_f32(raw): + v = raw.bitcast(fx.BFloat16) + return [fx.Float32(v[j]) for j in range(8)] + + +def head_rope_inputs(tk, gamma, d0, cos_sin, positions): + """Token tk's head gamma and cos / sin at dims d0 .. d0 + ELEMS (as f32).""" + # positions are int64: this token's low word + pos = fx.Int32(bo.buffer_load(rsrc(positions), 2 * tk, vec_width=1, dtype=T.i32)) + half = ROTARY_DIM // 2 + ib = (d0 < half).select(d0, d0 - half) + r_cs = rsrc(cos_sin) + return ( + _ld_bf16x4(rsrc(gamma), d0), + _ld_bf16x4(r_cs, pos * ROTARY_DIM + fx.min(ib, half - ELEMS)), + _ld_bf16x4(r_cs, pos * ROTARY_DIM + half + fx.min(ib, half - ELEMS)), + ) + + +def head_norm_rope(e, lane, gw4, cs, sn, eps): + """One head's gemma RMSNorm + partial NeoX RoPE on a half wave, lane l of it + holding dims 4 (l % 32) .. + 4 (the aiter kernel's split and reduction order).""" + d0 = (lane % 32) * ELEMS + half = ROTARY_DIM // 2 + first = d0 < half + ss = fx.Float32(0.0) + for j in range_constexpr(ELEMS): + ss = ss + e[j] * e[j] + ss = butterfly(ss, (16, 8, 4, 2, 1)) + rr = hw_rsq(ss / float(HEAD_DIM) + eps) + n = [e[j] * rr * (gw4[j] + 1.0) for j in range(ELEMS)] + # partial NeoX RoPE over the first ROTARY_DIM dims: the partner is in this half + in_rope = d0 < ROTARY_DIM + partner_lane = first.select(lane + half // ELEMS, lane - half // ELEMS) + out = [] + for j in range_constexpr(ELEMS): + p = fx.Int32( + rocdl.ds_bpermute( + T.i32, + (fx.min(fx.max(partner_lane, 0), 63) * 4).ir_value(), + n[j].bitcast(fx.Int32).ir_value(), + ) + ).bitcast(fx.Float32) + # hipcc contracts ``x*c - p*s`` into fma(x, c, -(p*s)); match it + ps = p * sn[j] + rot = fmath.fma(n[j], cs[j], first.select(-ps, ps)) + out.append(in_rope.select(rot, n[j])) + return out + + +@traced +def emit_pre_attn( + tid, bid, lane, wave, x8, red, mb_put, mb_put_words, mb_poll, qkv_mb, x8_mb, + x8s_mb, stamp, eps, tokens, heads, ptrs, signal=None, long_step=None, +): # fmt: skip + """K1's body: the GEMV tasks (ROWS_PER_TASK of the ``heads`` projection's rows + each: the rank's own on CTAs 0 .. ONE_GEMV, its other index q heads' on CTAs + XT0 .., in a step where ``long_step()``), the head tasks (CTAs HT0 .., a token + per half wave) and, with S > 2, one norm task per token (CTAs NT0 ..) feeding + the GEMV tasks.""" + rows = heads.rows + ( + ar, + res, + h_out, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, cos_sin, positions, + slot_mapping, k_cache, v_cache, k_scale, v_scale, index_cache, q_out, iq_out, + ) = ptrs # fmt: skip + + # stores K4 reads in the same launch (signal): device-coherent like the + # mailboxes, so the done flags need no L2 writeback + cm_out = CM_DEV if signal is not None else 0 + + def publish_done(mb_addr, idx, who=None): + """The CTA's stores done, then flag idx of ``mb_addr`` (from thread 0, or + from every thread where ``who``).""" + if const_expr(signal is not None): + fx.memory_fence( + syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Release + ) + gpu.barrier() + if (tid == 0) if who is None else who: + mb_put(mb_addr, idx, fx.Int32(1)) + + # acc = bf16(sum of partials) + residual, f32; the norm reads acc unrounded. + # Elements are laid out as the fused all-reduce's logical threads (8 each): this + # lane holds lt0 and, in lanes < 32, lt1, so rstd reduces in its order. + lts = [wave * 96 + lane, wave * 96 + 64 + fx.min(lane, 31)] + + def load_norm_inputs(toks): + """(gamma, [token][lt] -> (all-reduce, residual)) as loaded.""" + raw_g = [ld_raw8(rsrc(g_in), lt) for lt in lts] + raw_ar = [ + [ + ( + ld_raw8(rsrc(ar), tk * (HIDDEN // 8) + lt), + ld_raw8(rsrc(res), tk * (HIDDEN // 8) + lt), + ) + for lt in lts + ] + for tk in toks + ] + return raw_g, raw_ar + + def norm_quant(toks, raw_g, raw_ar, store_residual): + """Tokens toks: the residual out (where ``store_residual(j)``), gemma RMSNorm + and per-token FP8 into x8 row j (the j-th of toks) -> the x scales.""" + gs = [bf16x8_f32(raw) for raw in raw_g] + accs = [ # [j][lt][8] + [ + [a + r for a, r in zip(bf16x8_f32(ra), bf16x8_f32(rr))] + for ra, rr in raw_ar[j] + ] + for j in range(len(toks)) + ] + for j in range_constexpr(len(toks)): + if store_residual(j): + r_h = rsrc(h_out) + bo.buffer_store( + fx.Vector.from_elements(accs[j][0], fx.Float32).to(fx.BFloat16), + r_h, + toks[j] * HIDDEN + lts[0] * 8, + cache_modifier=cm_out, + ) + if lane < 32: + bo.buffer_store( + fx.Vector.from_elements(accs[j][1], fx.Float32).to(fx.BFloat16), + r_h, + toks[j] * HIDDEN + lts[1] * 8, + cache_modifier=cm_out, + ) + # every token's sum of squares behind one set of barriers + tots = ar_block_sums( + [ + ( + ar_pack_sumsq(accs[j][0]), + (lane < 32).select(ar_pack_sumsq(accs[j][1]), fx.Float32(0.0)), + ) + for j in range(len(toks)) + ], + lane, + wave, + red, + ) + stamp(1) + normed = [] + for j in range_constexpr(len(toks)): + rstd = hw_rsq(tots[j] / float(HIDDEN) + eps) + normed.append( + [ + [accs[j][i][e] * rstd * (gs[i][e] + 1.0) for e in range(8)] + for i in range(2) + ] + ) + for j in range_constexpr(len(toks)): + amax = fx.Float32(0.0) + for e in range_constexpr(8): + amax = fx.max(amax, fmath.absf(normed[j][0][e])) + amax = fx.max( + amax, + (lane < 32).select(fmath.absf(normed[j][1][e]), fx.Float32(0.0)), + ) + w_max = wave_max(amax) + if lane == 0: + fx.ptr_store(w_max, red + (j * WAVES + wave)) + gpu.barrier() + x_scales = [] + for j in range_constexpr(len(toks)): + amax = fx.ptr_load(red + j * WAVES) + for w in range_constexpr(1, WAVES): + amax = fx.max(amax, fx.ptr_load(red + (j * WAVES + w))) + x_scale = (amax == 0.0).select(fx.Float32(1.0), amax / FP8_MAX) + x_scales.append(x_scale) + x_rcp = 1.0 / x_scale + for i in range_constexpr(2): + q = [div_rn(normed[j][i][e], x_scale, x_rcp) for e in range(8)] + w01 = [ + fp8_pack4(q[0], q[1], q[2], q[3]), + fp8_pack4(q[4], q[5], q[6], q[7]), + ] + xo = x8 + (j * (HIDDEN // 4) + lts[i] * 2) + if (i == 0) | (lane < 32): + fx.ptr_store(w01[0].bitcast(fx.Float32), xo) + fx.ptr_store(w01[1].bitcast(fx.Float32), xo + 1) + gpu.barrier() + stamp(2) + return x_scales + + # ============================================ norm tasks (S > 2) + # (S <= 2 the GEMV tasks normalize every token themselves: S = 2 fused layer + # 43.8 -> 43.6 us; S = 4 57.5 -> 59.4 us the other way round) + if const_expr(tokens > 2): # noqa: SIM102 + if (bid >= NT0) & (bid < NT0 + tokens): + tk = bid - NT0 + raw_g, raw_ar = load_norm_inputs([tk]) + x_scale = norm_quant([tk], raw_g, raw_ar, lambda j: True)[0] + if const_expr(signal is not None): + publish_done(signal[1], tk) # the residual + # the x8 row as plain words, then its scale: the scale's pair is the flag + for i in range_constexpr(2): + if (i == 0) | (lane < 32): + bo.buffer_store( + fx.Vector( + fx.ptr_load( + x8 + lts[i] * 2, + result_type=fx.Vector.make_type(2, fx.Int32), + ) + ), + rsrc(x8_mb), + tk * (HIDDEN // 4) + lts[i] * 2, + cache_modifier=CM_DEV, + ) + fx.memory_fence( + syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Release + ) + gpu.barrier() + if tid == 0: + mb_put(x8s_mb, tk, x_scale) + + # ============================================ GEMV tasks (one per CTA) + # CTA bid < ONE_GEMV: row group bid of q | k | v, then the own index q head's + # and index k's (the identity in the one-head build) + own_q = (LOCAL_Q_HEADS + 2) * HEAD_ROW_GROUPS + rg = (bid < own_q).select( + bid, + (bid < own_q + HEAD_ROW_GROUPS).select( + heads.iq_off // ROWS_PER_TASK + bid - own_q, + heads.ik_off // ROWS_PER_TASK + bid - own_q - HEAD_ROW_GROUPS, + ), + ) + gemv = bid < ONE_GEMV + if const_expr(heads.count > 1): + xt = bid - XT0 + if (xt >= 0) & (xt < OTHER_GEMV): + other = xt // HEAD_ROW_GROUPS # the other heads in head order, own skipped + head = other + (other >= heads.own).select(1, 0) + rg = (LOCAL_Q_HEADS + 2 + head) * HEAD_ROW_GROUPS + xt % HEAD_ROW_GROUPS + gemv = long_step() + if gemv: + if const_expr(tokens <= 2): + raw_g, raw_ar = load_norm_inputs(list(range(tokens))) + # weight stream after the norm's inputs: the norm waits on its loads in issue + # order (vmcnt), so it would otherwise wait for the 96 KB of weights too; the + # scheduling barrier keeps the compiler from interleaving them again (nothing + # before it uses a loaded value, so it does not wait either) + rocdl.sched_barrier(0) + r_w = rsrc(w_qkv) + # 64-k chunk kc of row group rg is 1 KB at ((rg*K/32 + 2kc)*512); lane l + # holds row l%16, k = 64kc + 16(l/16) + [0, 16). + wts = [] + for c in range_constexpr(CHUNKS_PER_WAVE): + kc = wave * CHUNKS_PER_WAVE + c + wts.append( + fx.Vector( + bo.buffer_load( + r_w, + ((rg * (HIDDEN // 32) + kc * 2) * 512 + lane * 16) // 4, + vec_width=4, + dtype=T.i32, + ) + ) + ) + ws = fx.Float32( # row tid % 16's scale, sampled below by tid < 16 tokens + bo.buffer_load( + rsrc(s_qkv), + rg * ROWS_PER_TASK + tid % ROWS_PER_TASK, + vec_width=1, + dtype=T.f32, + ) + ) + if const_expr(tokens <= 2): + x_scales = norm_quant( + list(range(tokens)), raw_g, raw_ar, lambda j: bid == j + ) + if const_expr(signal is not None): # noqa: SIM102 + if bid < tokens: + publish_done(signal[1], bid) # the residual + else: + # the norm tasks' scales (each published after its row's plain x8 words), + # then every row, 16 B a load + got = mb_poll([(x8s_mb, tk, 1) for tk in range(tokens)]) + x_scales = [got[tk][0].bitcast(fx.Float32) for tk in range(tokens)] + fx.memory_fence( + syncscope=rocdl.SyncScope.Workgroup, ordering=fx.AtomicOrdering.Acquire + ) + n_v = tokens * HIDDEN // 16 # 16 B loads + for i in range_constexpr((n_v + THREADS - 1) // THREADS): + v = tid + THREADS * i + if v < n_v: + fx.ptr_store( + fx.Vector( + bo.buffer_load( + rsrc(x8_mb), + v * 4, + vec_width=4, + dtype=T.i32, + cache_modifier=CM_DEV, + ) + ), + x8 + v * 4, + ) + gpu.barrier() + stamp(2) + # 12 chunks per wave: two 16x16x32 FP8 MFMAs each (k halves of 8 bytes). + # B column l % 16 is token l % 16 (token 0 past the last): C column t is + # token t's rows, the weight stream is the same for any token count. + col_tok = fx.min(lane % 16, tokens - 1) + c = fx.Vector.filled(4, 0.0, fx.Float32) + for cc in range_constexpr(CHUNKS_PER_WAVE): + kc = wave * CHUNKS_PER_WAVE + cc + xb = fx.Vector( + fx.ptr_load( + x8 + (col_tok * (HIDDEN // 4) + kc * 16 + (lane // 16) * 4), + result_type=fx.Vector.make_type(4, fx.Float32), + ) + ).bitcast(fx.Int64) + wa = wts[cc].bitcast(fx.Int64) + for h in range_constexpr(2): + c = fx.Vector( + rocdl.mfma_f32_16x16x32_fp8_fp8( + T.vec(4, T.f32), [wa[h], xb[h], c, 0, 0, 0] + ) + ) + fx.ptr_store(c, red + (wave * 64 + lane) * 4) + gpu.barrier() + # sample column t of C: lanes t, 16 + t, 32 + t, 48 + t hold rows 4 (l/16) + e + if tid < ROWS_PER_TASK * tokens: + r = tid % ROWS_PER_TASK + tk = tid // ROWS_PER_TASK + tot_r = fx.Float32(0.0) + for w in range_constexpr(WAVES): + tot_r = tot_r + fx.ptr_load( + red + ((w * 64 + 16 * (r // 4) + tk) * 4 + r % 4) + ) + x_scale = x_scales[0] + for k in range_constexpr(1, tokens): + x_scale = (tk == k).select(x_scales[k], x_scale) + mb_put( + qkv_mb, + tk * rows + rg * ROWS_PER_TASK + r, + bf16_round(tot_r * x_scale * ws), + ) + + stamp(3) + + # ======================================= head tasks: norm / rope / cache + # one CTA per head, token k on the half wave of threads 32 k .. 32 k + 31 + if (bid >= HT0) & (bid < HT0 + N_HEAD_TASKS): + ht = bid - HT0 + tk = tid // 32 + l32 = tid % 32 + is_q = ht < LOCAL_Q_HEADS + is_k = ht == LOCAL_Q_HEADS + is_v = ht == LOCAL_Q_HEADS + 1 + is_iq = ht == LOCAL_Q_HEADS + 2 + base = tk * rows + is_q.select( + ht * HEAD_DIM, + is_k.select( + K_OFF, is_v.select(V_OFF, is_iq.select(heads.iq_off, heads.ik_off)) + ), + ) + if tk < tokens: + d0 = l32 * ELEMS + # slot_mapping is int64: this token's low word + slot = fx.Int32( + bo.buffer_load(rsrc(slot_mapping), 2 * tk, vec_width=1, dtype=T.i32) + ) + # gamma and this position's cos / sin go out before the wait + gw = is_q.select(g_q, is_k.select(g_k, is_iq.select(g_iq, g_ik))) + gw4, cs, sn = head_rope_inputs(tk, gw, d0, cos_sin, positions) + got = mb_poll([(qkv_mb, base + d0, 2), (qkv_mb, base + d0 + 2, 2)]) + stamp(4) + e = [ + got[0][0].bitcast(fx.Float32), + got[0][1].bitcast(fx.Float32), + got[1][0].bitcast(fx.Float32), + got[1][1].bitcast(fx.Float32), + ] + if is_v: + # raw V: per-token amax over the head, scale = amax / FP8_MAX + if const_expr(mono_config.STATIC_KV_SCALE): + vs = fx.Float32(1.0) + else: + am = fx.Float32(0.0) + for j in range_constexpr(ELEMS): + am = fx.max(am, fmath.absf(e[j])) + am = butterfly(am, (16, 8, 4, 2, 1), fx.max) + vs = (am > 0.0).select(am / FP8_MAX, fx.Float32(1.0)) + if slot >= 0: + b = page16_of(slot) + t = slot % PAGE16 + if l32 == 0: + bo.buffer_store( + vs, rsrc(v_scale), b * PAGE16 + t, cache_modifier=cm_out + ) + r_v = rsrc(v_cache) + for j in range_constexpr(ELEMS): + q8 = fp8_pack4(e[j] / vs, 0.0, 0.0, 0.0) & 0xFF + bo.buffer_store( + fx.Int8(q8), + r_v, + b * PAGE_BYTES + (d0 + j) * PAGE16 + t, + cache_modifier=cm_out, + ) + else: + out = head_norm_rope(e, lane, gw4, cs, sn, eps) + if is_q: + bo.buffer_store( + fx.Vector.from_elements(out, fx.Float32).to(fx.BFloat16), + rsrc(q_out), + (tk * LOCAL_Q_HEADS + ht) * HEAD_DIM + d0, + cache_modifier=cm_out, + ) + if is_iq: + bo.buffer_store( + fx.Vector.from_elements(out, fx.Float32).to(fx.BFloat16), + rsrc(iq_out), + tk * HEAD_DIM + d0, + ) + if is_k: + if const_expr(mono_config.STATIC_KV_SCALE): + ksc = fx.Float32(1.0) + else: + am = fx.Float32(0.0) + for j in range_constexpr(ELEMS): + am = fx.max(am, fmath.absf(out[j])) + am = butterfly(am, (16, 8, 4, 2, 1), fx.max) + ksc = (am > 0.0).select(am / FP8_MAX, fx.Float32(1.0)) + if slot >= 0: + b = page16_of(slot) + t = slot % PAGE16 + if l32 == 0: + bo.buffer_store( + ksc, + rsrc(k_scale), + b * PAGE16 + t, + cache_modifier=cm_out, + ) + qk = [bf16_round(out[j]) / ksc for j in range(ELEMS)] + bo.buffer_store( + fp8_pack4(qk[0], qk[1], qk[2], qk[3]), + rsrc(k_cache), + ( + b * PAGE_BYTES + + (d0 // 16) * (PAGE16 * 16) + + t * 16 + + d0 % 16 + ) + // 4, + cache_modifier=cm_out, + ) + if (ht == IK_TASK) & (slot >= 0): + qi = [bf16_round(out[j]) for j in range(ELEMS)] + bo.buffer_store( + fp8_pack4(qi[0], qi[1], qi[2], qi[3]), + rsrc(index_cache), + (slot * HEAD_DIM + d0) // 4, + cache_modifier=cm_out, + ) + + if const_expr(signal is not None): + # every token's flag once the whole CTA's stores are out + publish_done(signal[0], tid * N_HEAD_TASKS + ht, tid < tokens) + + +# the fused layer kernel's K1 pointers that do not change per step, in its device +# array (the rest are K4's own, or kernel arguments) +K1_ARGS = ( + "ar", "g_in", "w_qkv", "s_qkv", "g_q", "g_k", "g_iq", "g_ik", "cos_sin", + "index_cache", "iq_out", "scratch", +) # fmt: skip + + +@traced +def emit_k1( + tid, bid, lane, wave, x8, red, qs, stamp, step, layer, eps, index_scale, tokens, + q_len, init_blocks, local_blocks, heads, ptrs, bt_width, signal=False, +): # fmt: skip + """K1 whole: the GEMV / head / norm tasks, then every CTA's indexer scores. + ``x8``: tokens * HIDDEN / 4 words of LDS, ``red``: WAVES * 256 floats, ``qs``: + tokens * heads.count * HEAD_DIM / 2 floats. ``q_len``: tokens per request + (speculative verify rows). ``heads``: the index q heads of the fused projection + (``IndexHeads``). ``signal`` (the fused layer kernel): the head tasks and + residual writers publish done flags (SCRATCH_HDONE / SCRATCH_RDONE).""" + ( + ar, res, h_out, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, cos_sin, positions, + slot_mapping, k_cache, v_cache, k_scale, v_scale, index_cache, q_out, iq_out, + block_table, seq_lens, iscore, scratch, + ) = ptrs # fmt: skip + mbox = Mailbox(scratch, step, layer) + qkv_mb = mbox.addr(SCRATCH_QKV) + x8_mb = mbox.addr(SCRATCH_X8) + x8s_mb = mbox.addr(SCRATCH_X8S) + # Plain names, not method calls: flydsl's AST rewriter carries every local whose + # method is called inside a dynamic if/for as loop state. + mb_put = mbox.put + mb_put_words = mbox.put_words + mb_poll = mbox.poll + + def index_q_row(tk, head): + """Token tk's index q of the projection's index head ``head`` into qs (row + tk * heads.count + head, HEAD_DIM / 2 words), made from the GEMV's raw rows + with the index-q head's math, by this wave's lanes < 32.""" + d0 = lane * ELEMS + gw4, cs, sn = head_rope_inputs(tk, g_iq, d0, cos_sin, positions) + base = tk * heads.rows + IndexHeads(heads.count, 0).iq_off + head * HEAD_DIM + got = mb_poll([(qkv_mb, base + d0, 2), (qkv_mb, base + d0 + 2, 2)]) + e = [got[h][j].bitcast(fx.Float32) for h in range(2) for j in range(2)] + out = head_norm_rope(e, lane, gw4, cs, sn, eps) + row = qs + (tk * heads.count + head) * (HEAD_DIM // 2) + fx.ptr_store(bf16_pair(out[0], out[1]), row + lane * 2) + fx.ptr_store(bf16_pair(out[2], out[3]), row + (lane * 2 + 1)) + + def make_index_q(first, last): + """Tokens first..last's index q of this rank's head into qs: wave w makes + tokens w, w + WAVES, ..; run by a CTA with score tasks only.""" + for k0 in range_constexpr(0, tokens, WAVES): + tk = wave + k0 + if (tk < tokens) & (tk >= first) & (tk <= last) & (lane < 32): + index_q_row(tk, heads.own) + gpu.barrier() + + def make_index_q_heads(): + """Every token's index q of every head into qs (indexer context + parallelism): (token, head) job j on wave j % WAVES.""" + for j0 in range_constexpr(0, tokens * heads.count, WAVES): + j = wave + j0 + if (j < tokens * heads.count) & (lane < 32): + index_q_row(j // heads.count, j % heads.count) + gpu.barrier() + + def q_frag_head(tk, s, head): + """This lane's MFMA B operand of token tk's index q of head ``head``, + k-chunk s.""" + row = (tk * heads.count + head) * (HEAD_DIM // 2) + return fx.Vector( + fx.ptr_load( + qs + (row + (32 * s + 8 * (lane // 16)) // 2), + result_type=fx.Vector.make_type(4, fx.Float32), + ) + ).bitcast(fx.BFloat16) + + def q_frag(tk, s): + """This lane's MFMA B operand of token tk's index q (this rank's head), + k-chunk s.""" + return q_frag_head(tk, s, heads.own) + + def long_step(): + """A request of the step is long (``step_rows``): it scores every index q + head (indexer context parallelism).""" + return step_rows(seq_lens, tokens, q_len, heads).any_long() + + emit_pre_attn( + tid, bid, lane, wave, x8, red, mb_put, mb_put_words, mb_poll, qkv_mb, x8_mb, + x8s_mb, stamp, eps, tokens, heads, + ( + ar, res, h_out, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, + cos_sin, positions, slot_mapping, k_cache, v_cache, k_scale, + v_scale, index_cache, q_out, iq_out, + ), + (mbox.addr(SCRATCH_HDONE), mbox.addr(SCRATCH_RDONE)) if signal else None, + long_step, + ) # fmt: skip + stamp(5) + # the scorers: the idle CTAs first (the norm tasks' come first but finish + # early), then the GEMV CTAs, the head CTAs last + score_base = ONE_GEMV + N_HEAD_TASKS + hdone_mb = mbox.addr(SCRATCH_HDONE) + + def wait_new_keys(): + """Every token's index_k head task done, its key visible to this wave.""" + mb_poll([(hdone_mb, fx.min(lane, tokens - 1) * N_HEAD_TASKS + IK_TASK, 1)]) + fx.memory_fence( + syncscope=rocdl.SyncScope.Agent, ordering=fx.AtomicOrdering.Acquire + ) + + emit_index_scores( + (bid + (BLOCKS - score_base)) % BLOCKS, lane, uniform(wave), red, + mb_put, iscore, q_frag, tokens, q_len, init_blocks, local_blocks, + index_scale, index_cache, block_table, seq_lens, bt_width, + wait_new_keys if signal else None, + make_index_q, heads, q_frag_head, make_index_q_heads, + ) # fmt: skip + stamp(6) + + +TL_POINTS = 8 # timeline stamps per CTA + + +def build_pre_attn_kernel( + eps: float, + sm_scale: float, + init_blocks: int, + local_blocks: int, + tokens: int = 1, + timeline: bool = False, + heads: IndexHeads = ONE_INDEX_HEAD, +): + """``@flyc.jit`` launcher of K1 for this model's RMSNorm epsilon, attention + scale and sparse pinned blocks, a decode batch of ``tokens`` (<= MAX_TOKENS) + rows and the fused projection's index q ``heads``.""" + assert 1 <= tokens <= MAX_TOKENS + index_scale = index_scale_log2e(sm_scale) + + @fx.struct + class Smem: + # FP8 activations, natural k order, one row per token + x8: fx.Array[fx.Int32, tokens * HIDDEN // 4, 16] + red: fx.Array[fx.Float32, WAVES * 64 * 4, 16] + tls: fx.Array[fx.Int64, TL_POINTS if timeline else 1, 16] # timeline stamps + # every token's index q for the scorers, bf16 pairs + qs: fx.Array[fx.Float32, HEAD_DIM // 2 * tokens * heads.count, 16] + + kernel_name = kernel_symbol( + "minimax_m3_pre_attn", + s=tokens, + ib=init_blocks, + lb=local_blocks, + ih=heads.count, + io=heads.own, + tl=timeline, + kps=mono_config.KV_PAGE_STRIDE, + sks=mono_config.STATIC_KV_SCALE, + ) + + # the JIT cache key holds the kernel's scalar closure values, not objects: the + # index heads go in as ints (a shared cache otherwise serves one rank's build + # to every rank) + ih_count, ih_own = heads.count, heads.own + + @flyc.kernel(name=kernel_name, known_block_size=[THREADS, 1, 1]) + def pre_attn_kernel( + ar: Int64, + res: Int64, + h_out: Int64, + g_in: Int64, + w_qkv: Int64, + s_qkv: Int64, + g_q: Int64, + g_k: Int64, + g_iq: Int64, + g_ik: Int64, + cos_sin: Int64, + positions: Int64, + slot_mapping: Int64, + k_cache: Int64, + v_cache: Int64, + k_scale: Int64, + v_scale: Int64, + index_cache: Int64, + q_out: Int64, + iq_out: Int64, + block_table: Int64, + seq_lens: Int64, + bt_width: Int32, + q_len: Int32, + iscore: Int64, + scratch: Int64, + step: Int64, + layer: Int32, + tl: Int64, + ): + tid = fx.thread_idx.x + bid = fx.block_idx.x + lane = tid % 64 + wave = tid // 64 + lds = fx.SharedAllocator().allocate(Smem).peek() + x8 = lds.x8.ptr + red = lds.red.ptr + tls = lds.tls.ptr + qs = lds.qs.ptr + + def stamp(k): + """``timeline``: s_memrealtime (100 MHz) of this CTA passing point k, kept + in LDS until the kernel ends.""" + # The compile-time and traced conditions cannot share one ``and``. + if const_expr(timeline): # noqa: SIM102 + if tid == 0: + fx.ptr_store(memrealtime(), tls + k) + + if const_expr(timeline): # noqa: SIM102 + if tid < TL_POINTS: # same wave as the stamping thread: ordered + fx.ptr_store(fx.Int64(0), tls + tid) + stamp(0) + emit_k1( + tid, bid, lane, wave, x8, red, qs, stamp, step, layer, eps, index_scale, + tokens, q_len, init_blocks, local_blocks, IndexHeads(ih_count, ih_own), + ( + ar, res, h_out, g_in, w_qkv, s_qkv, g_q, g_k, g_iq, g_ik, + cos_sin, positions, slot_mapping, k_cache, v_cache, k_scale, + v_scale, index_cache, q_out, + iq_out, block_table, seq_lens, iscore, scratch, + ), + bt_width, + ) # fmt: skip + if const_expr(timeline): # noqa: SIM102 + if tid < TL_POINTS: + fx.generic_store( + fx.inttoptr( + fx.PointerType.get(fx.Int64.ir_type, fx.AddressSpace.Global, 8), + tl + fx.Int64((bid * TL_POINTS + tid) * 8), + ), + fx.ptr_load(tls + tid), + ) + + @flyc.jit + def launch_pre_attn( + ar: Int64, + res: Int64, + h_out: Int64, + g_in: Int64, + w_qkv: Int64, + s_qkv: Int64, + g_q: Int64, + g_k: Int64, + g_iq: Int64, + g_ik: Int64, + cos_sin: Int64, + positions: Int64, + slot_mapping: Int64, + k_cache: Int64, + v_cache: Int64, + k_scale: Int64, + v_scale: Int64, + index_cache: Int64, + q_out: Int64, + iq_out: Int64, + block_table: Int64, + seq_lens: Int64, + bt_width: Int32, + q_len: Int32, + iscore: Int64, + scratch: Int64, + step: Int64, + layer: Int32, + tl: Int64, + stream: fx.Stream = _CURRENT_STREAM, + ): + pre_attn_kernel( + ar, + res, + h_out, + g_in, + w_qkv, + s_qkv, + g_q, + g_k, + g_iq, + g_ik, + cos_sin, + positions, + slot_mapping, + k_cache, + v_cache, + k_scale, + v_scale, + index_cache, + q_out, + iq_out, + block_table, + seq_lens, + bt_width, + q_len, + iscore, + scratch, + step, + layer, + tl, + ).launch(grid=(BLOCKS,), block=(THREADS,), stream=stream) + + return launch_pre_attn diff --git a/vllm/models/minimax_m3/amd/mono/layout.py b/vllm/models/minimax_m3/amd/mono/layout.py new file mode 100644 index 000000000000..38c2a0a0cd92 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/layout.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Where the mono layer kernel (K4) puts its tasks and data, per decode token count. + +Plain Python, evaluated at build time: which CTA runs each stage's tasks, the +scratch and peer-symmetric regions, and the LDS row pitches. Kept free of FlyDSL +so it is testable without a GPU. +""" + +from vllm.models.minimax_m3.amd.mono.config import ( + BLOCKS, + HEAD_DIM, + HIDDEN, + INTER, + LOCAL_Q_HEADS, + MAX_INDEX_BLOCKS, + MAX_SPARSE_KEYS, + MAX_TOKENS, + MOE_SLOTS, + N_ROUTED, + PAGE16, + SPARSE_BLOCK, + TOP_K, + TOPK_BLOCKS, + WAVES, +) + +H = LOCAL_Q_HEADS +SPLIT_KEYS = 256 # one gluon context partition per split task +N_SPLIT = MAX_SPARSE_KEYS // SPLIT_KEYS +O_K = H * HEAD_DIM +O_ROWS = 32 # hidden rows per o_proj / attention-reduce task +# from this many tokens the attention all-reduce is a reduce-scatter to the row +# group's owner rank (t % npes) plus an all-gather of the sums: half the bytes over +# each link for one more hop (smaller batches are latency-bound: one hop wins) +AG_RS_FROM = 12 +N_O = HIDDEN // O_ROWS +PAGES_PER_BLOCK = SPARSE_BLOCK // PAGE16 +N_ROUTER = N_ROUTED // WAVES # one expert per wave +UG_ROWS = 16 +UG_PER_SLOT = INTER // UG_ROWS +SH_ROWS = 8 # mid rows per shared-expert up / gate task (a half task) +SH_TASKS = INTER // SH_ROWS +DN_ROWS = 32 +N_DN = HIDDEN // DN_ROWS +# The router logits are written SCORE_COPIES times and CTA b polls copy +# b % SCORE_COPIES: 240 expert CTAs spinning on one 1 KB region serialize on its +# lines. K4 on 4 GPUs 37.3 -> 35.9 us; copies of xn / mid measured neutral and of +# the split outputs +0.3 us (the split CTAs pay the extra stores). +SCORE_COPIES = 8 + + +def _layout(items, align=256): + """Byte offsets of (value, tag) pair regions, 8 bytes per pair.""" + off, out = 0, {} + for name, pairs in items: + out[name] = off + off += (pairs * 8 + align - 1) // align * align + out["_bytes"] = off + return out + + +# plain words, not pairs: the selected pages then the context length, as the +# original selector emits them (read by the tests); one row per token +SPARSE_TABLE_WORDS = TOPK_BLOCKS * PAGES_PER_BLOCK + 2 +# a context past THREADS blocks, one-head build: each split task's TOPK_BLOCKS +# best (key, block) (indexer context parallelism: the peers' sel_cp) +SEL_LONG_PAIRS = N_SPLIT * TOPK_BLOCKS * 2 + +# every region holds MAX_TOKENS rows (token k at k * the one-token size) +SCRATCH = _layout( + [ + # indexer block scores, token k's block b at pair k * MAX_INDEX_BLOCKS + b + ("iscore", MAX_INDEX_BLOCKS * MAX_TOKENS), + ("sparse_table", SPARSE_TABLE_WORDS // 2 * MAX_TOKENS), + ("sp_o", N_SPLIT * H * HEAD_DIM // 2 * MAX_TOKENS), # packed bf16 pairs + ("sp_m", N_SPLIT * H * MAX_TOKENS), + ("sp_l", N_SPLIT * H * MAX_TOKENS), + ("attnq", O_K // 4 * MAX_TOKENS // 2), # merged attention, plain FP8 words + ("attnq_s", MAX_TOKENS), # and its per-token scale + # the finished attention all-reduce + residual, bf16 pairs (o tasks -> router) + ("a", HIDDEN // 2 * MAX_TOKENS), + ("xn8", HIDDEN // 4 * MAX_TOKENS), # MXFP8 xn, 4 fp8 per word + ("xsc", HIDDEN // 32 * MAX_TOKENS), # and its per-1x32 E8M0 scales + # copy c, token k at pairs [(c * MAX_TOKENS + k) * N_ROUTED, ..) + # per expert (routing key, sigmoid) pairs, see route_key + ("scores", 2 * N_ROUTED * MAX_TOKENS * SCORE_COPIES), + ("mid", MOE_SLOTS * INTER * MAX_TOKENS), + # routing record read by the tests: token k's expert ids, then weights + ("sel", 2 * MOE_SLOTS * MAX_TOKENS), + # token k's split task p's candidates: rank r at pairs 2 ((k N_SPLIT + p) + # TOPK_BLOCKS + r) + (0 key, 1 block) + ("sel_long", SEL_LONG_PAIRS * MAX_TOKENS), + # wide: the MXFP8 xn and scales as plain words (every token's row), and a + # done flag per router task + ("xn8p", HIDDEN // 4 * MAX_TOKENS // 2), + ("xscp", HIDDEN // 32 * MAX_TOKENS // 2), + ("xndone", N_ROUTER * MAX_TOKENS), + # and the mid rows as MXFP8 words, their E8M0 scales a word each (the down + # stage's LDS layout), a done flag per up / gate task pair + ("mid8p", MOE_SLOTS * INTER * MAX_TOKENS // 8), + ("midsp", MOE_SLOTS * INTER // 32 * MAX_TOKENS // 2), + ("ugdone", (TOP_K * MAX_TOKENS + 1) * UG_PER_SLOT // 2), + ] +) +SCRATCH_BYTES = SCRATCH["_bytes"] + + +def sym_layout(npes: int): + """Per-rank symmetric buffer: attn / ffn partial regions, [src rank][token][HIDDEN], + and the peers' index-selection candidates (``sel_cp``).""" + rows = npes * MAX_TOKENS * HIDDEN // 2 + # a_ag / ffn_ag: the finished attention / FFN all-reduce rows, each written by + # its owner rank (the reduce-scatter + all-gather form, S >= AG_RS_FROM) + ag = MAX_TOKENS * HIDDEN // 2 + return _layout( + [ + ("attn", rows), + ("ffn", rows), + ("a_ag", ag), + ("ffn_ag", ag), + # indexer context parallelism: token k's candidates for this rank's + # index head, source rank s's half h's r-th best at pairs + # 2 ((k N_SPLIT + 2 s + h) TOPK_BLOCKS + r) + (0 key, 1 block) + ("sel_cp", SEL_LONG_PAIRS * MAX_TOKENS), + ] + ) + + +def stage_bases(tokens: int): + """Task t of a stage runs on CTA (base + t) % BLOCKS; the merge (S > 1), o and + down tasks start past the tokens' split CTAs.""" + n_split = N_SPLIT * tokens + n_merge = tokens if tokens > 1 else 0 + router = (n_split + n_merge + N_O) % BLOCKS + return { + "split": 0, + "merge": n_split, + "o": n_split + n_merge, + "router": router, + "shared": _shared_base(tokens, router), + "ug": 0, + "down": n_split, + } + + +def ug_ctas_per_token(tokens: int): + """CTAs per token of the routed up / gate tasks, token k on CTAs k C .. (k + 1) + C - 1: a CTA stages one token's xn, not all (fused layer, S = 4: 57.9 -> 57.6 + us). None: round-robin -- S < 4 (S = 2: +0.4 us) or S not dividing the grid.""" + return BLOCKS // tokens if tokens >= 4 and BLOCKS % tokens == 0 else None + + +def ug_tasks_of(cta: int, tokens: int) -> int: + """Routed up / gate tasks CTA ``cta`` runs.""" + per_token = TOP_K * UG_PER_SLOT + cpt = ug_ctas_per_token(tokens) + if cpt is None: + return (per_token * tokens - cta + BLOCKS - 1) // BLOCKS + return (per_token - cta % cpt + cpt - 1) // cpt + + +def _shared_base(tokens: int, router: int) -> int: + """First of the SH_TASKS consecutive CTAs that run the shared expert's up / + gate before the routing lands. Ranked by the most routed up / gate tasks a + window CTA carries (they queue behind it); then, only if that is the most any + CTA carries, by its router CTAs (they start it late and would delay the last + routed tasks: S = 4, 4 us); then by its down CTAs (S = 2: 43.1 vs 43.7 us); + then by its router CTAs anyway.""" + n_router = N_ROUTER * tokens + n_split = N_SPLIT * tokens + most = max(ug_tasks_of(c, tokens) for c in range(BLOCKS)) + + def cost(base): + ctas = [(base + j) % BLOCKS for j in range(SH_TASKS)] + busiest = max(ug_tasks_of(c, tokens) for c in ctas) + in_router = sum((c - router) % BLOCKS < n_router for c in ctas) + in_down = sum(n_split <= c < n_split + N_DN for c in ctas) + return busiest, in_router if busiest == most else 0, in_down, in_router + + return min(range(BLOCKS), key=cost) + + +def shared_after_routing(tokens: int) -> bool: + """Every CTA carries the same number of routed up / gate tasks (S = 4): a + shared-expert CTA routes first and runs its shared half task while its first + routed task's weights are in flight, instead of holding its routing back + (S = 4 54.51 -> 54.09 us; S = 2, where the shared window sits on CTAs with one + task fewer, +0.1).""" + loads = [ug_tasks_of(c, tokens) for c in range(BLOCKS)] + return min(loads) == max(loads) + + +WIDE_FROM = 5 # token counts from here run the expert-grouped MoE + + +def wide_moe(tokens: int) -> bool: + """The MoE stage groups its work by expert (every distinct expert's weights + read once, a B column per token) instead of by (token, slot): past four tokens + the per-token tasks re-read the experts the tokens share, and their operands + no longer fit in LDS.""" + return tokens >= WIDE_FROM + + +# LDS rows of a token's MXFP8 xn and of its E8M0 scales, padded by 16 B: at the +# natural 1536 / 192-word rows the 16 tokens a MFMA B operand gathers sit in one +# bank group (16-way conflicts; ug spent ~3 us a task in them at S = 16) +XN8_ROW = HIDDEN // 4 + 4 +XSC_ROW = HIDDEN // 32 + 4 +# and the down stage's MXFP8 mid rows / their scales (16 lanes, 16 rows): 196 and +# 28 words put the rows' 16 B reads, and 4 B scale reads, in distinct banks +MID_ROW = INTER // 4 + 4 +MIDSC_ROW = INTER // 32 + 4 +# and the o stage's FP8 attention rows (512 words: the same 16-way conflict) +O_ROW = O_K // 4 + 4 + + +def pool_words(tokens: int) -> int: + """LDS words of the activation pool ``x``. Narrow: every token's bf16 xn (the + largest stage operand). Wide: the largest of the split's PV partials, the up / + gate stage's MXFP8 xn + scales + up partials + router sigmoids, and the down + stage's MXFP8 mids + scales + a zero row.""" + if not wide_moe(tokens): + return HIDDEN // 2 * tokens + rows = MOE_SLOTS * tokens + 1 + return max( + WAVES * O_K, + tokens * (XN8_ROW + XSC_ROW + N_ROUTED) + WAVES * 64 * 4, + rows * (MID_ROW + MIDSC_ROW), + ) diff --git a/vllm/models/minimax_m3/amd/mono/peer_buffer.py b/vllm/models/minimax_m3/amd/mono/peer_buffer.py new file mode 100644 index 000000000000..60766486a665 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/peer_buffer.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +"""Symmetric peer memory for the mono kernels' in-kernel all-reduces.""" + +import torch +import torch.distributed as dist +from aiter.ops.flydsl.quick_allreduce_int4_ipc import UncachedIpcHeap + + +class PeerBuffer: + """One uncached buffer per rank, every rank holding every peer's address. + + K4 pushes a layer's attention partials into every peer's ``attn`` region and + its FFN partials into the ``ffn`` region, then polls its own. Regions are + reused by every layer and step (the mailbox epoch tells a stale pair from a + fresh one) and no rank can overwrite a peer's unread pairs: to push layer L's + FFN partials a rank must have passed layer L's attention reduce, which needs + every peer's layer-L attention push, and a peer only pushes that after it + finished reading layer L-1's FFN region. So ranks are at most half a layer + apart and the two regions double-buffer each other. + """ + + def __init__(self, nbytes: int, group, rank: int, npes: int, device: torch.device): + self.local = UncachedIpcHeap.alloc_uncached(nbytes) + self._opened: list[int] = [] + addresses = [self.local] + if npes > 1: + handles = UncachedIpcHeap.gather_object_list_via_broadcast( + group, UncachedIpcHeap.get_mem_handle_bytes(self.local) + ) + addresses = [] + for peer, handle in enumerate(handles): + if peer == rank: + addresses.append(self.local) + else: + base = UncachedIpcHeap.open_mem_handle(handle) + self._opened.append(base) + addresses.append(base) + dist.barrier(group=group) + self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device) + + def close(self) -> None: + """Close the peer mappings and free the local buffer; idempotent.""" + opened, self._opened = self._opened, [] + for base in opened: + UncachedIpcHeap.close_mem_handle(base) + if self.local: + UncachedIpcHeap.free_device_mem(self.local) + self.local = 0 diff --git a/vllm/models/minimax_m3/amd/mono/runner.py b/vllm/models/minimax_m3/amd/mono/runner.py new file mode 100644 index 000000000000..048a4aaba6a8 --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/runner.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One decode step of 1..MAX_TOKENS tokens on the mono kernels. + +Every sparse MoE layer is one launch of the fused layer kernel; it takes +``(ar, res)`` -- the pair vLLM's input norm would receive (reduced MoE output, +residual) -- and leaves ``(ar, h_mid)`` for the next layer, alternating between +two buffers by layer parity. The embedding, the leading dense-FFN layers and the +final norm run the original modules. The K/V cache is read in vLLM's interleaved +page-16 layout under the fixed unit fp8 scale; the step's tables come from the +forward context. +""" + +from typing import Any, NamedTuple, cast + +import torch + +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, +) +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe +from vllm.models.minimax_m3.amd.mono import config as mono_config +from vllm.models.minimax_m3.amd.mono.config import ( + BLOCKS, + HEAD_DIM, + HIDDEN, + LOCAL_Q_HEADS, + MAX_TOKENS, + N_ROUTED, + PAGE16, + TOP_K, + TOPK_BLOCKS, + TP, + IndexHeads, + MonoUnsupported, +) +from vllm.models.minimax_m3.amd.mono.kernels.post_attn import build_post_attn_kernel +from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import K1_ARGS +from vllm.models.minimax_m3.amd.mono.kernels.pre_attn import ( + SCRATCH_BYTES as K1_SCRATCH_BYTES, +) +from vllm.models.minimax_m3.amd.mono.layout import SCRATCH_BYTES as K4_SCRATCH_BYTES +from vllm.models.minimax_m3.amd.mono.layout import sym_layout +from vllm.models.minimax_m3.amd.mono.peer_buffer import PeerBuffer +from vllm.models.minimax_m3.amd.mono.weights import SparseMoeLayer, build_sparse_layer +from vllm.platforms import current_platform + + +def _ptr(t: torch.Tensor) -> int: + return t.data_ptr() + + +class StepTables(NamedTuple): + """The per-step inputs every layer launch of a decode step shares.""" + + block_table: torch.Tensor # int32 [rows, width], one row per token + seq_lens: torch.Tensor # int32 [rows] + slot_mapping: torch.Tensor # int64, page-128 slots (K/V and index cache) + q_len: int # query tokens per request (speculative verify rows) + + +def token_rows( + block_table: torch.Tensor, seq_lens: torch.Tensor, q: int, n: int +) -> tuple[torch.Tensor, torch.Tensor]: + """One (block table row, seq_len) per token of the step: a request of q query + tokens becomes q rows, token j seeing ``seq_len - (q - 1 - j)`` keys.""" + if q == 1: + return block_table, seq_lens + reqs = n // q + back = torch.arange(q - 1, -1, -1, dtype=seq_lens.dtype, device=seq_lens.device) + return ( + block_table[:reqs].repeat_interleave(q, dim=0), + (seq_lens[:reqs, None] - back).reshape(-1), + ) + + +def dense_prefix(layers) -> int: + """Leading layers without an MoE FFN; every later layer must be sparse MoE.""" + n = 0 + while n < len(layers) and not layers[n].is_moe_layer: + n += 1 + if n == len(layers) or not all( + lyr.is_moe_layer and hasattr(lyr.self_attn, "indexer") for lyr in layers[n:] + ): + raise MonoUnsupported("layer mix: dense layers must precede the MoE layers") + return n + + +def _fused_shared_weight() -> float: + """The fused shared expert's slot weight, from aiter's routing metadata.""" + meta = rocm_aiter_moe.aiter_topK_meta_data + if meta is None: + raise MonoUnsupported("fused shared expert routing metadata") + total_w, total_ids = meta + if int(total_ids[0, TOP_K].item()) != N_ROUTED: + raise MonoUnsupported("fused shared expert id") + return float(total_w[0, TOP_K].item()) + + +class MonoDecodeRunner: + """Owns the per-layer weight views, scratch, peer buffers and compiled kernels.""" + + def __init__(self, model, config) -> None: + npes = get_tensor_model_parallel_world_size() + if npes != TP: + raise MonoUnsupported(f"TP {npes}") + self.model = model + self.rank = get_tensor_model_parallel_rank() + layers = list(model.layers) + # Enable indexer context parallelism: each rank computes all TP index + # heads, allowing local block scoring without cross-rank exchange. + heads = IndexHeads(count=TP, own=self.rank) + tp_group = get_tp_group().device_group + self.n_dense = n_dense = dense_prefix(layers) + self.sparse = [ + build_sparse_layer(layers[i], i, heads, rank=self.rank, tp_group=tp_group) + for i in range(n_dense, len(layers)) + ] + attn0 = layers[n_dense].self_attn + moe0 = layers[n_dense].block_sparse_moe + self.shared_fp8 = self.sparse[0].shared_fp8 + if any(lw.shared_fp8 != self.shared_fp8 for lw in self.sparse): + raise MonoUnsupported("shared expert fused in some layers only") + # a separate MXFP8 shared expert adds unscaled: routed_scaling_factor goes + # on the routed weights alone (route_scale), as the model's runner does + shared_weight = 1.0 if self.shared_fp8 else _fused_shared_weight() + self.sm_scale = attn0.scaling + if abs(self.sm_scale - HEAD_DIM**-0.5) > 1e-12: + raise MonoUnsupported(f"softmax scale {self.sm_scale}") + sel = attn0.indexer.impl + if sel.topk_blocks != TOPK_BLOCKS: + raise MonoUnsupported(f"indexer top-{sel.topk_blocks} blocks") + self.layer_name = attn0.layer_name + self.index_prefix = attn0.indexer.index_cache.prefix + # the caches this build viewed: vLLM re-binds them after its cudagraph + # memory-profiling capture, which must rebuild the runner + self._cache_ptrs = self.cache_ptrs(model, n_dense) + # the K/V page views (V offset to its side of the block) and how many page + # ids a block spans; the kernels are compiled for that geometry + self.kv_views = [lyr.self_attn.get_aiter_sparse_pa_kv_cache() for lyr in layers] + stride = attn0._aiter_sparse_pa_block_page_stride + mono_config.configure(kv_page_stride=stride, static_kv_scale=True) + self._build( + npes, heads, float(config.rms_norm_eps), float(moe0.routed_scaling_factor), + shared_weight, + float(getattr(moe0.experts, "routed_experts", moe0.experts).swiglu_limit), + sel.init_blocks, sel.local_blocks, + ) # fmt: skip + dev = self.h.device + n16 = self.kv_views[0][0].shape[0] + # the host's unit fp8 scale for every page-16 slot (K1 writes 1.0 back) + self.kv_scale = torch.ones(n16, 1, PAGE16, dtype=torch.float32, device=dev) + self.res0 = torch.zeros(MAX_TOKENS, HIDDEN, dtype=torch.bfloat16, device=dev) + + @staticmethod + def cache_ptrs(model, n_dense: int) -> tuple[int, int]: + """Addresses of the first MoE layer's KV cache and index cache.""" + attn = model.layers[n_dense].self_attn + return attn.kv_cache.data_ptr(), attn.indexer.index_cache.kv_cache.data_ptr() + + def stale(self) -> bool: + """Whether vLLM re-bound the caches this runner viewed (profiling capture).""" + return self.cache_ptrs(self.model, self.n_dense) != self._cache_ptrs + + def step_tables(self, n: int) -> StepTables: + """The current step's tables for ``n`` tokens, from the forward context.""" + fwd = get_forward_context() + decode = cast(Any, fwd.attn_metadata[self.layer_name]).decode + q = decode.decode_query_len + block_table, seq_lens = token_rows(decode.block_table, decode.seq_lens, q, n) + return StepTables(block_table, seq_lens, fwd.slot_mapping[self.layer_name], q) + + def start( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None, + aux_hidden_states: list[torch.Tensor] | None = None, + ) -> torch.Tensor: + """Embedding and the dense layers on the original modules; leaves the first + MoE layer's (ar, res) in ``ars[0][:n]`` and the returned residual (zeros + when layer 0 is already an MoE layer: (embeddings, 0)).""" + n = positions.numel() + hidden = ( + inputs_embeds + if inputs_embeds is not None + else self.model.embed_input_ids(input_ids) + ) + residual = None + if aux_hidden_states is not None: + self.model._maybe_add_hidden_state(aux_hidden_states, 0, hidden, residual) + for i in range(self.n_dense): + hidden, residual = self.model.layers[i](positions, hidden, residual) + if aux_hidden_states is not None: + self.model._maybe_add_hidden_state( + aux_hidden_states, i + 1, hidden, residual + ) + self.ars[0][:n].copy_(hidden.view(n, HIDDEN)) + return self.res0[:n] if residual is None else residual.view(n, HIDDEN) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: + """The final hidden states of one decode step (what the model returns).""" + aux_hidden_states: list[torch.Tensor] | None = ( + [] if self.model.aux_hidden_state_layers else None + ) + res = self.start(input_ids, positions, inputs_embeds, aux_hidden_states) + tables = self.step_tables(res.shape[0]) + for i, lw in enumerate(self.sparse): + res = self.run_sparse_layer(i, lw, tables, positions, res) + if aux_hidden_states is not None: + self.model._maybe_add_hidden_state( + aux_hidden_states, + lw.layer_id + 1, + self.ars[(i + 1) % 2][: res.shape[0]], + res, + ) + hidden = self.run_final_norm(res) + if aux_hidden_states is not None: + return hidden, aux_hidden_states + return hidden + + def run_sparse_layer( + self, + i: int, + lw: SparseMoeLayer, + tables: StepTables, + positions: torch.Tensor, + res: torch.Tensor, + ) -> torch.Tensor: + """Sparse layer ``i``: (ars[i % 2][:n], res) -> (ars[(i + 1) % 2][:n], the + returned residual).""" + n = res.shape[0] + block_table, seq_lens, slots, q_len = tables + k16, v16 = self.kv_views[lw.layer_id] + self.k4[n]( + _ptr(self.h), _ptr(self.q), _ptr(block_table), _ptr(seq_lens), + _ptr(k16), _ptr(v16), _ptr(self.kv_scale), _ptr(self.kv_scale), + _ptr(lw.w_o), _ptr(lw.s_o), _ptr(lw.g_post), + _ptr(lw.gate), _ptr(lw.bias), _ptr(lw.w13), _ptr(lw.s13), + _ptr(lw.w2), _ptr(lw.s2), + _ptr(lw.sh_args) if lw.shared_fp8 else 0, + _ptr(self.h_mids[(i + 1) % 2]), _ptr(self.ars[(i + 1) % 2]), + _ptr(self.scratch4), self.peers.local, _ptr(self.peers.addresses), + _ptr(self.step), self.rank, lw.layer_id, + block_table.shape[1], q_len, + 0, # timeline stamps (kernel debug build only) + _ptr(self.k1_args[i]), _ptr(positions), _ptr(slots), _ptr(res), + stream=torch.cuda.current_stream(), # flydsl rejects torch.Stream + ) # fmt: skip + return self.h_mids[(i + 1) % 2][:n] + + def _build( + self, npes, heads, eps, route_scale, shared_weight, swiglu_limit, + init_blocks, local_blocks, + ) -> None: # fmt: skip + """Compile the layer kernels and allocate scratch, peer buffers and the + per-step activation buffers; ``self.sparse`` and ``self.sm_scale`` are set.""" + device_index = torch.accelerator.current_device_index() + dev = torch.device("cuda", device_index) + cus = current_platform.get_device_properties(device_index).multi_processor_count + if cus != BLOCKS: + raise MonoUnsupported( + f"{cus} CUs: the kernels need all {BLOCKS} CTAs co-resident" + ) + # one layer kernel per decode batch size (every graph captures its own) + self.k4 = { + s: build_post_attn_kernel( + npes, + self.sm_scale, + eps, + route_scale, + shared_weight, + swiglu_limit, + init_blocks, + local_blocks, + s, + fuse_k1=True, + heads=heads, + shared_fp8=self.shared_fp8, + ) + for s in range(1, MAX_TOKENS + 1) + } + self.scratch1 = torch.zeros(K1_SCRATCH_BYTES, dtype=torch.uint8, device=dev) + self.scratch4 = torch.zeros(K4_SCRATCH_BYTES, dtype=torch.uint8, device=dev) + self.peers = PeerBuffer( + sym_layout(npes)["_bytes"], get_tp_group().cpu_group, self.rank, npes, dev + ) + self.step = torch.zeros(1, dtype=torch.int32, device=dev) + bf16 = torch.bfloat16 + # row k = token k of the step; sparse layer i reads ars[i % 2] (the + # previous layer's output) and writes ars[(i + 1) % 2], likewise h_mids + self.ars = [ + torch.empty(MAX_TOKENS, HIDDEN, dtype=bf16, device=dev) for _ in range(2) + ] + self.h_mids = [ + torch.empty(MAX_TOKENS, HIDDEN, dtype=bf16, device=dev) for _ in range(2) + ] + self.h = torch.empty(MAX_TOKENS, HIDDEN, dtype=bf16, device=dev) + self.q = torch.empty( + MAX_TOKENS, LOCAL_Q_HEADS * HEAD_DIM, dtype=bf16, device=dev + ) + self.iq = torch.empty(MAX_TOKENS, 1, HEAD_DIM, dtype=bf16, device=dev) + # each layer's K1 pointers that do not change per step (K1_ARGS order) + self.k1_args = [] + for i, lw in enumerate(self.sparse): + ptrs = { + "ar": _ptr(self.ars[i % 2]), "g_in": _ptr(lw.g_in), + "w_qkv": _ptr(lw.w_qkv), + "s_qkv": _ptr(lw.s_qkv), "g_q": _ptr(lw.g_q), "g_k": _ptr(lw.g_k), + "g_iq": _ptr(lw.g_iq), "g_ik": _ptr(lw.g_ik), + "cos_sin": _ptr(lw.cos_sin), + "index_cache": _ptr(lw.index_cache), "iq_out": _ptr(self.iq), + "scratch": _ptr(self.scratch1), + } # fmt: skip + self.k1_args.append( + torch.tensor([ptrs[a] for a in K1_ARGS], dtype=torch.int64, device=dev) + ) + + def close(self) -> None: + self.peers.close() + + def run_final_norm(self, res: torch.Tensor) -> torch.Tensor: + ar = self.ars[len(self.sparse) % 2] + hidden, _ = self.model.norm(ar[: res.shape[0]], res) + self.step.add_(1) + return hidden diff --git a/vllm/models/minimax_m3/amd/mono/weights.py b/vllm/models/minimax_m3/amd/mono/weights.py new file mode 100644 index 000000000000..ad7446e601cf --- /dev/null +++ b/vllm/models/minimax_m3/amd/mono/weights.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Zero-copy views (and two converted copies) of one vLLM ``MiniMaxM3DecoderLayer`` +for the mono kernels. + +The kernels read tensors through raw pointers in fixed layouts: a differently +laid-out weight would not fault, it would produce plausible garbage. So every +layout assumption is asserted once and a model that does not match is refused +(``MonoUnsupported``). The experts already have the kernels' layout (MXFP4, aiter +``shuffle_weight(16, 16)`` + ``e8m0_shuffle`` scales, fused shared expert in slot +128). The attention linears are MXFP8 in vLLM and the router gate is fp32, so +those are converted once when the runner is built: MXFP8 -> bf16 -> per-channel +fp8, preshuffled (16, 16); gate -> bf16. + +Without shared-expert fusion the routed table holds the 128 routed experts and +the shared expert stays the separate MXFP8 MLP; its weights are copied once into +the f8f6f4 MFMA's fp8 A-operand order (``mxfp8_a_tiles``, the same bytes) with +``e8m0_shuffle``-ordered scales, so the kernels run it in MXFP8. +""" + +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from aiter import dtypes +from aiter.ops.quant import pertoken_quant +from aiter.ops.shuffle import shuffle_weight + +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, +) +from vllm.models.minimax_m3.amd.mono.config import ( + HEAD_DIM, + HIDDEN, + INTER, + LOCAL_Q_HEADS, + N_ROUTED, + O_K, + ONE_INDEX_HEAD, + ROTARY_DIM, + TOP_K, + TOPK_BLOCKS, + TP, + IndexHeads, + MonoUnsupported, +) +from vllm.models.minimax_m3.amd.mono.kernels.post_attn import SH_ARGS + + +def _need(ok: bool, what: str) -> None: + if not ok: + raise MonoUnsupported(what) + + +def _bf16_vec(t: torch.Tensor, n: int, name: str) -> torch.Tensor: + _need( + t.dtype == torch.bfloat16 and t.numel() == n and t.is_contiguous(), + f"{name}: {t.dtype} {tuple(t.shape)}", + ) + return t.data + + +@dataclass(frozen=True) +class SparseMoeLayer: + """Everything the layer kernel reads for one sparse-attention MoE layer.""" + + layer_id: int + attn: object # MiniMaxM3SparseAttention: layer name, KV cache views, scales + index_cache: torch.Tensor # [blocks, block, head_dim] fp8, shared block ids + g_in: torch.Tensor + w_qkv: torch.Tensor + s_qkv: torch.Tensor + g_q: torch.Tensor + g_k: torch.Tensor + g_iq: torch.Tensor + g_ik: torch.Tensor + cos_sin: torch.Tensor + w_o: torch.Tensor + s_o: torch.Tensor + g_post: torch.Tensor + gate: torch.Tensor + bias: torch.Tensor + w13: torch.Tensor + s13: torch.Tensor + w2: torch.Tensor + s2: torch.Tensor + # the separate MXFP8 shared expert (``mxfp8_a_tiles``) and the int64 table of + # their addresses the kernels take (``SH_ARGS`` order); None when it is fused + sh_w13: torch.Tensor | None = None + sh_s13: torch.Tensor | None = None + sh_w2: torch.Tensor | None = None + sh_s2: torch.Tensor | None = None + sh_args: torch.Tensor | None = None + + @property + def shared_fp8(self) -> bool: + return self.sh_args is not None + + def shared_bytes(self) -> int: + """Bytes of the kernels' copy of the separate shared expert.""" + if not self.shared_fp8: + return 0 + nbytes = 0 + for t in (self.sh_w13, self.sh_s13, self.sh_w2, self.sh_s2): + assert t is not None + nbytes += t.numel() + return nbytes + + +def _ptpc_fp8_from_mxfp8( + linear, rows: int, cols: int, name: str +) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> per-channel fp8 + preshuffled (16, 16) + fp32 scale [rows], as the kernels' PTPC GEMV reads.""" + w = getattr(linear, "weight", None) + s = getattr(linear, "weight_scale", None) + _need( + w is not None and w.dtype == dtypes.fp8 and tuple(w.shape) == (rows, cols), + f"{name}: weight {None if w is None else (w.dtype, tuple(w.shape))}", + ) + _need( + s is not None and tuple(s.shape) == (rows, cols // 32), + f"{name}: scale {None if s is None else tuple(s.shape)}", + ) + assert w is not None and s is not None + bf = dequant_mxfp8_to_bf16(w.data, s.data.view(torch.uint8)) + q, scale = pertoken_quant(bf, quant_dtype=dtypes.fp8) + return ( + shuffle_weight(q, layout=(16, 16)).contiguous(), + scale.view(rows).to(torch.float32).contiguous(), + ) + + +def mxfp8_a_tiles( + linear, rows: int, cols: int, name: str +) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 linear (fp8 [rows, cols], e8m0 [rows, cols/32]) -> its bytes in the + f8f6f4 MFMA's fp8 A order and its scales in ``e8m0_shuffle`` order. + + A 16-row x 128-k tile is 2 KiB; lane g * 16 + r of a tile holds row r's k + 16 g .. + 16 then 64 + 16 g .. + 16 (the B operand's order), 32 bytes.""" + w = getattr(linear, "weight", None) + s = getattr(linear, "weight_scale", None) + _need( + w is not None and w.dtype == dtypes.fp8 and tuple(w.shape) == (rows, cols), + f"{name}: weight {None if w is None else (w.dtype, tuple(w.shape))}", + ) + _need( + s is not None and tuple(s.shape) == (rows, cols // 32), + f"{name}: scale {None if s is None else tuple(s.shape)}", + ) + assert w is not None and s is not None + # the row-major layout the native MXFP8 linear kernel keeps; a kernel that + # preshuffles or pads would make the tiles silently wrong + _need( + not getattr(w, "is_shuffled", False) and not getattr(s, "is_shuffled", False), + f"{name}: weight preshuffled", + ) + _need(rows % 32 == 0 and cols % 256 == 0, f"{name}: shape {rows} x {cols}") + tiles = ( + w.data.contiguous() + .view(torch.uint8) + .view(rows // 16, 16, cols // 128, 2, 4, 16) + .permute(0, 2, 4, 1, 3, 5) + .contiguous() + .view(-1) + ) + scales = ( + s.data.contiguous() + .view(torch.uint8) + .view(rows // 32, 2, 16, cols // 256, 2, 4) + .permute(0, 3, 5, 2, 4, 1) + .contiguous() + .view(-1) + ) + return tiles, scales + + +def _ptpc_fp8_from_bf16( + bf: torch.Tensor, rows: int +) -> tuple[torch.Tensor, torch.Tensor]: + """bf16 [rows, cols] -> per-channel fp8 preshuffled (16, 16) + fp32 scale [rows].""" + q, scale = pertoken_quant(bf, quant_dtype=dtypes.fp8) + return ( + shuffle_weight(q, layout=(16, 16)).contiguous(), + scale.view(rows).to(torch.float32).contiguous(), + ) + + +def _build_cp_qkv( + linear, rank: int, tp_group, layer_id: int +) -> tuple[torch.Tensor, torch.Tensor]: + """The wider fused projection of indexer context parallelism: every rank's + index q head, all-gathered once at build time (``IndexHeads(TP, rank)``). + + One rank's rows (2560): q(2048) | k(128) | v(128) | index_q(128) | index_k(128) + CP rows (2944): q(2048) | k(128) | v(128) | iq_rank0..3 | index_k(128) + + The gather's list order is the TP group's rank order, which ``IndexHeads.own`` + indexes; the rank's own slot is checked against its local rows.""" + one_rows = (LOCAL_Q_HEADS + 3 + 1) * HEAD_DIM # 2560 + w = getattr(linear, "weight", None) + s = getattr(linear, "weight_scale", None) + _need( + w is not None + and w.dtype == dtypes.fp8 + and tuple(w.shape) == (one_rows, HIDDEN), + f"layer {layer_id} qkv_proj CP: weight " + f"{None if w is None else (w.dtype, tuple(w.shape))}", + ) + _need( + s is not None and tuple(s.shape) == (one_rows, HIDDEN // 32), + f"layer {layer_id} qkv_proj CP: scale {None if s is None else tuple(s.shape)}", + ) + assert w is not None and s is not None + _need( + dist.get_rank(tp_group) == rank and dist.get_world_size(tp_group) == TP, + f"layer {layer_id} qkv_proj CP: TP group rank {dist.get_rank(tp_group)}" + f" of {dist.get_world_size(tp_group)}, not {rank} of {TP}", + ) + bf = dequant_mxfp8_to_bf16(w.data, s.data.view(torch.uint8)) + + iq_start = (LOCAL_Q_HEADS + 2) * HEAD_DIM # 2304 + iq_end = iq_start + HEAD_DIM # 2432 + my_iq = bf[iq_start:iq_end].contiguous() + + all_iq = [torch.empty_like(my_iq) for _ in range(TP)] + dist.all_gather(all_iq, my_iq, group=tp_group) + _need( + torch.equal(all_iq[rank], my_iq), + f"layer {layer_id} qkv_proj CP: gathered index q heads not in rank order", + ) + + cp_bf = torch.cat([bf[:iq_start], *all_iq, bf[iq_end:]], dim=0) + cp_rows = (LOCAL_Q_HEADS + 3 + TP) * HEAD_DIM # 2944 + assert cp_bf.shape[0] == cp_rows, f"CP QKV rows {cp_bf.shape[0]} != {cp_rows}" + return _ptpc_fp8_from_bf16(cp_bf, cp_rows) + + +def _cos_sin(attn, layer_id: int) -> torch.Tensor: + t = attn.rotary_emb.cos_sin_cache + _need( + t.dim() == 2 and t.shape[1] == ROTARY_DIM, + f"layer {layer_id}: cos/sin cache {tuple(t.shape)}", + ) + if t.dtype != torch.bfloat16: + t = t.to(torch.bfloat16) + return t.contiguous() + + +def build_sparse_layer( + layer, + layer_id: int, + heads: IndexHeads = ONE_INDEX_HEAD, + rank: int = 0, + tp_group=None, +) -> SparseMoeLayer: + attn = layer.self_attn + _need( + hasattr(attn, "indexer") and layer.is_moe_layer, + f"layer {layer_id}: not sparse MoE", + ) + # the model has one index head per rank; under indexer context parallelism + # the fused projection is built wider (every rank's), see _build_cp_qkv + _need( + (attn.num_heads, attn.num_kv_heads, attn.num_idx_heads, attn.head_dim) + == (LOCAL_Q_HEADS, 1, 1, HEAD_DIM), + f"layer {layer_id}: heads " + f"{attn.num_heads}/{attn.num_kv_heads}/{attn.num_idx_heads}", + ) + _need(attn.rotary_emb.rotary_dim == ROTARY_DIM, f"layer {layer_id}: rotary_dim") + sel = attn.indexer.impl # the selection config lives on the indexer impl + _need( + (sel.topk_blocks, sel.init_blocks, sel.local_blocks, attn.skip_index_topk) + == (TOPK_BLOCKS, 0, 1, False), + f"layer {layer_id}: sparse selection config", + ) + _need( + attn.use_aiter_sparse_pa and str(attn.kv_cache_dtype).startswith("fp8"), + f"layer {layer_id}: needs the AITER sparse PA path with an fp8 KV cache", + ) + index_cache = attn.indexer.index_cache.kv_cache + _need( + index_cache.numel() > 0 + and index_cache.dtype == dtypes.fp8 + and index_cache.dim() == 3 + and index_cache.shape[2] == HEAD_DIM + and index_cache.is_contiguous(), + f"layer {layer_id}: index cache {index_cache.dtype} {tuple(index_cache.shape)}", + ) + + if heads.count > 1: + _need(tp_group is not None, f"layer {layer_id}: CP requires tp_group") + w_qkv, s_qkv = _build_cp_qkv(attn.qkv_proj, rank, tp_group, layer_id) + else: + w_qkv, s_qkv = _ptpc_fp8_from_mxfp8( + attn.qkv_proj, heads.rows, HIDDEN, f"layer {layer_id} qkv_proj" + ) + + w_o, s_o = _ptpc_fp8_from_mxfp8( + attn.o_proj, HIDDEN, O_K, f"layer {layer_id} o_proj" + ) + + moe = layer.block_sparse_moe + # FusedMoEFactory returns a MoERunner wrapping the RoutedExperts module + experts = getattr(moe.experts, "routed_experts", moe.experts) + fused = moe.is_fused_shared_expert_enabled + E = N_ROUTED + 1 if fused else N_ROUTED + w13, w2 = experts.w13_weight, experts.w2_weight + _need( + w13.shape[0] == E, + f"layer {layer_id}: {'fused' if fused else 'separate'} shared expert " + f"(experts {w13.shape[0]})", + ) + shared = {} + if not fused: + mlp = moe.shared_experts + _need(mlp is not None, f"layer {layer_id}: no shared expert") + shared["sh_w13"], shared["sh_s13"] = mxfp8_a_tiles( + mlp.gate_up_proj, 2 * INTER, HIDDEN, f"layer {layer_id} shared gate_up" + ) + shared["sh_w2"], shared["sh_s2"] = mxfp8_a_tiles( + mlp.down_proj, HIDDEN, INTER, f"layer {layer_id} shared down" + ) + shared["sh_args"] = torch.tensor( + [shared[a].data_ptr() for a in SH_ARGS], + dtype=torch.int64, + device=shared["sh_w13"].device, + ) + _need(experts.top_k == TOP_K, f"layer {layer_id}: top_k {experts.top_k}") + _need( + getattr(w13, "is_shuffled", False) and getattr(w2, "is_shuffled", False), + f"layer {layer_id}: experts not in aiter shuffled layout", + ) + _need( + w13.numel() * w13.element_size() == E * 2 * INTER * HIDDEN // 2 + and w2.numel() * w2.element_size() == E * HIDDEN * INTER // 2, + f"layer {layer_id}: expert weight bytes", + ) + s13, s2 = experts.w13_weight_scale, experts.w2_weight_scale + _need( + s13.numel() == E * 2 * INTER * HIDDEN // 32 + and s2.numel() == E * HIDDEN * INTER // 32, + f"layer {layer_id}: expert scale bytes", + ) + gate = moe.gate.weight.data + _need(tuple(gate.shape) == (N_ROUTED, HIDDEN), f"layer {layer_id}: gate") + bias = moe.e_score_correction_bias + _need( + bias is not None and bias.dtype == torch.float32 and bias.numel() == N_ROUTED, + f"layer {layer_id}: bias", + ) + return SparseMoeLayer( + layer_id=layer_id, + attn=attn, + index_cache=index_cache, + g_in=_bf16_vec(layer.input_layernorm.weight, HIDDEN, "input_layernorm"), + w_qkv=w_qkv, + s_qkv=s_qkv, + g_q=_bf16_vec(attn.q_norm.weight, HEAD_DIM, "q_norm"), + g_k=_bf16_vec(attn.k_norm.weight, HEAD_DIM, "k_norm"), + g_iq=_bf16_vec(attn.index_q_norm.weight, HEAD_DIM, "index_q_norm"), + g_ik=_bf16_vec(attn.index_k_norm.weight, HEAD_DIM, "index_k_norm"), + cos_sin=_cos_sin(attn, layer_id), + w_o=w_o, + s_o=s_o, + g_post=_bf16_vec( + layer.post_attention_layernorm.weight, HIDDEN, "post_attention_layernorm" + ), + gate=gate.to(torch.bfloat16).contiguous(), + bias=bias.data, + w13=w13.data, + s13=s13.data, + w2=w2.data, + s2=s2.data, + **shared, + ) diff --git a/vllm/models/minimax_m3/amd/ops/flydsl_utils.py b/vllm/models/minimax_m3/amd/ops/flydsl_utils.py new file mode 100644 index 000000000000..890ab3714e74 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/flydsl_utils.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Helpers shared by the MiniMax-M3 FlyDSL kernels (gfx950).""" + +import functools + +import flydsl.expr as fx +import torch +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith, llvm +from flydsl.expr import rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +# Cache-policy bits of buffer loads and stores: sc1 is coherent across the XCDs' +# L2s, sc0 | sc1 across the system, nt streams. +SC0, NT, SC1 = 0x1, 0x2, 0x10 +LOG2E = 1.4426950409 # softmax runs in exp2 +# The first 16 dwords of kernel arguments arrive in SGPRs, so a kernel's first +# loads do not wait for s_load. +PRELOAD_HINTS = { + "llvm_options": { + "amdgpu-kernarg-preload": True, + "amdgpu-kernarg-preload-count": 16, + } +} + + +def run_compiled(exe, *args, waves_per_eu: int | None = None): + """Compiles and runs ``exe`` on the first call, then dispatches the cache.""" + import flydsl.compiler as flyc + from flydsl.compiler.kernel_function import CompilationContext + + cf = getattr(exe, "_cf", None) + if cf is None: + hints = {"waves_per_eu": waves_per_eu} if waves_per_eu else {} + with CompilationContext.compile_hints(hints): + exe._cf = flyc.compile(exe, *args) + else: + cf(*args) + + +@functools.cache +def cu_count(device: torch.device) -> int: + return torch.cuda.get_device_properties(device).multi_processor_count + + +@functools.cache +def unit_scale(device: torch.device) -> torch.Tensor: + return torch.ones(1, dtype=torch.float32, device=device) + + +def global_i32_ptr(addr_i64): + ptr_ty = fx.PointerType.get( + T.i32, address_space=fx.AddressSpace.Global, alignment=4 + ) + return fx.inttoptr(ptr_ty, fx.Int64(addr_i64)) + + +def bits(x): + return fx.Float32(x).bitcast(fx.Int32) + + +def as_f32(x): + return fx.Int32(x).bitcast(fx.Float32) + + +def permlane_swap(op, d_a, d_b): + """v_permlane{16,32}_swap; fx has no wrapper.""" + pair_ty = ir.Type.parse("!llvm.struct<(i32, i32)>") + res = op(pair_ty, fx.as_ir_value(d_a), fx.as_ir_value(d_b), False, False) + return ( + fx.Int32(llvm.extractvalue(T.i32, res, [0])), + fx.Int32(llvm.extractvalue(T.i32, res, [1])), + ) + + +def maxf_nn(a, b): + """v_max_f32 with nnan; without it LLVM canonicalizes both inputs first.""" + fm = ir.Attribute.parse("#arith.fastmath") + return fx.Float32( + arith.MaxNumFOp(fx.as_ir_value(a), fx.as_ir_value(b), fastmath=fm).result + ) + + +def xlane4(x, op): + """``op`` over lanes L, L^16, L^32, L^48, broadcast to all four.""" + a, b = permlane_swap(rocdl.permlane32_swap, bits(x), bits(x)) + m = op(as_f32(a), as_f32(b)) + a, b = permlane_swap(rocdl.permlane16_swap, bits(m), bits(m)) + return op(as_f32(a), as_f32(b)) + + +def perm(hi, lo, sel): + """v_perm_b32: byte i of the result is byte sel_i of (hi:lo).""" + return fx.Int32(rocdl.perm_b32(fx.Int32(hi), fx.Int32(lo), fx.Int32(sel))) + + +def exp2(x): + """v_exp_f32 without denormal range reduction.""" + return fx.Float32(fx.rocdl.exp2(fx.Float32.ir_type, fx.Float32(x).ir_value())) + + +def pk_bf16(a, b): + """Two fp32 -> one dword of two bf16 (RNE), ``a`` low.""" + return fx.Int32( + fx.rocdl.cvt_pk_bf16_f32( + fx.as_ir_value(fx.Float32(a)), fx.as_ir_value(fx.Float32(b)) + ) + ) + + +def pk_f16(a, b): + """Two fp32 -> one dword of two fp16 (RNE), ``a`` low.""" + return ( + Vec.from_elements([fx.Float32(a), fx.Float32(b)], fx.Float32) + .to(fx.Float16) + .bitcast(fx.Int32)[0] + ) + + +def cvt_fp4x2_f16(word, sel): + """The two e2m1 codes of byte ``sel`` of ``word`` as two fp16 (exact).""" + return Vec( + fx.rocdl.cvt_scalef32_pk_f16_fp4( + Vec.make_type(2, fx.Float16), + fx.as_ir_value(fx.Int32(word)), + fx.as_ir_value(fx.Float32(1.0)), + sel, + ) + ) + + +def cvt_fp8x2_f16(word, hi): + """E4M3 bytes 0, 1 (``hi``: 2, 3) of ``word`` as two fp16 (exact).""" + return Vec( + fx.rocdl.cvt_scalef32_pk_f16_fp8( + Vec.make_type(2, fx.Float16), + fx.as_ir_value(fx.Int32(word)), + fx.as_ir_value(fx.Float32(1.0)), + hi, + ) + ) + + +def mfma_16x16x32(a4, b4, acc, fp16: bool): + """v_mfma_f32_16x16x32_{bf16,f16} on 4-dword operands.""" + dtype = fx.Float16 if fp16 else fx.BFloat16 + op = fx.rocdl.mfma_f32_16x16x32_f16 if fp16 else fx.rocdl.mfma_f32_16x16x32_bf16 + a = Vec(a4).bitcast(dtype) + b = Vec(b4).bitcast(dtype) + return Vec(op(Vec.make_type(4, fx.Float32), [a, b, acc, 0, 0, 0])) + + +def mfma_fp8_16x16x128(a8, b8, acc4): + """v_mfma_f32_16x16x128_f8f6f4 on e4m3 operands; zero scales select the + unscaled instruction.""" + return Vec( + fx.rocdl.mfma_scale_f32_16x16x128_f8f6f4( + Vec.make_type(4, fx.Float32), [a8, b8, acc4, 0, 0, 0, 0, 0, 0] + ) + ) + + +def pack8(h0, h1): + """Two 4 x i32 halves -> the 8 x i32 operand of the fp8 16x16x128 MFMA.""" + a, b = Vec(h0), Vec(h1) + return Vec.from_elements( + [a[k] for k in range(4)] + [b[k] for k in range(4)], fx.Int32 + ) + + +# Byte selectors of the even and odd dims in key order. +PERM_EVEN, PERM_ODD = 0x06040200, 0x07050301 + + +def key_order4(d): + """4 dwords of e4m3 (16 dims) -> even dims, then odd dims per 8 (key order).""" + return Vec.from_elements( + [ + perm(d[1], d[0], PERM_EVEN), + perm(d[1], d[0], PERM_ODD), + perm(d[3], d[2], PERM_EVEN), + perm(d[3], d[2], PERM_ODD), + ], + fx.Int32, + ) + + +def dequant8(c, t_lo, t_hi): + """8 e2m1 codes (low nibble the even dim) and the group's e4m3 table -> + the e4m3 keys of the even dims, then of the odd dims.""" + ml = perm(t_hi, t_lo, c & fx.Int32(0x07070707)) + mh = perm(t_hi, t_lo, (c >> 4) & fx.Int32(0x07070707)) + sign = fx.Int32(0x80808080 - (1 << 32)) + return ml ^ ((c << 4) & sign), mh ^ (c & sign) diff --git a/vllm/models/minimax_m3/amd/ops/index_cp/__init__.py b/vllm/models/minimax_m3/amd/ops/index_cp/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/index_cp/__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/minimax_m3/amd/ops/index_cp/host.py b/vllm/models/minimax_m3/amd/ops/index_cp/host.py new file mode 100644 index 000000000000..91ecd0646f80 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/index_cp/host.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Buffers of the context-parallel decode indexer (TP4, gfx950). + +Rank ``me`` holds index head ``me``; ``kernel_cp_push`` scores a quarter of every +request's blocks for all four heads and stores head h's scores into rank h's +receive buffer. Buffers and flags are uncached memory mapped by the peers; the +call counts stay on the device, so the handshake replays in CUDA graphs. +""" + +import torch +import torch.distributed as dist + +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce +from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode import ( + kernel_cp_push as _kernel, +) + +WORLD = _kernel.WORLD + + +class _DevicePointer: + """A raw device allocation as a tensor (``__cuda_array_interface__``).""" + + def __init__(self, ptr: int, nbytes: int): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + + +class IndexCP: + def __init__( + self, rank: int, group, device: torch.device, max_rows: int, max_width: int + ): + self.rank, self.device = rank, device + self.max_rows, self.max_width = max_rows, max_width + score_bytes = max_rows * max_width * 4 + iq_bytes = WORLD * max_rows * _kernel.D + # Each rank's buffer, mapped by every peer (uncached on ROCm). + self.score_ptrs, self.iq_ptrs, self.flag_ptrs = ( + CustomAllreduce.create_shared_buffer(n, group=group) + for n in (score_bytes, iq_bytes, _kernel.FLAG_BYTES) + ) + self.score_rcv = torch.as_tensor( + _DevicePointer(self.score_ptrs[rank], score_bytes), device=device + ).view(torch.float32) + flags = torch.as_tensor( + _DevicePointer(self.flag_ptrs[rank], _kernel.FLAG_BYTES), device=device + ) + flags.zero_() + # Call counts (the flag values), equal on every rank. + self.counts = torch.zeros(_kernel.CNT_WORDS, dtype=torch.int32, device=device) + torch.accelerator.synchronize(device) + dist.barrier(group=group) + + def score_push( + self, + iq: torch.Tensor, + index_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + init_blocks: int, + local_blocks: int, + width: int, + grid: int = 0, + ) -> torch.Tensor: + """Score this rank's blocks of every head from its index-q rows ``iq``. + + Returns this rank's head over every block, [1, rows, width]: a view of + the receive buffer, valid until the next call. + """ + from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode.host import ( + minimax_m3_nvfp4_index_score_push_cp, + ) + + rows = iq.shape[0] + assert rows <= self.max_rows and width <= self.max_width + minimax_m3_nvfp4_index_score_push_cp( + iq, + index_cache, + self.score_ptrs, + width, + rows * width * 4, + block_table, + seq_lens, + init_blocks, + local_blocks, + self.rank, + self.counts, + self.flag_ptrs, + self.iq_ptrs, + self.max_rows, + grid=grid, + ) + return self.score_rcv[: rows * width].view(1, rows, width) + + +# None means the single-head decode path. +_CP: IndexCP | None = None + + +def init(device: torch.device, max_width: int) -> None: + """Create this rank's ``IndexCP``; a collective, called once outside capture.""" + global _CP + if _CP is not None: + return + from vllm.config import get_current_vllm_config + from vllm.distributed import get_tp_group + from vllm.platforms.rocm import on_gfx950 + + tp = get_tp_group() + # One set of buffers and flags: microbatches would share them. + if ( + tp.world_size != WORLD + or not on_gfx950() + or get_current_vllm_config().parallel_config.use_ubatching + ): + return + _CP = IndexCP(tp.rank_in_group, tp.cpu_group, device, _kernel.MAX_ROWS, max_width) + + +def get() -> IndexCP | None: + return _CP diff --git a/vllm/models/minimax_m3/amd/ops/index_topk_decode/__init__.py b/vllm/models/minimax_m3/amd/ops/index_topk_decode/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/index_topk_decode/__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/minimax_m3/amd/ops/index_topk_decode/host.py b/vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py new file mode 100644 index 000000000000..f3669b733c96 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/index_topk_decode/host.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch wrappers of the FlyDSL index top-k: the block ids AITER's +``pa_sparse_block_topk`` selects, in its order.""" + +import functools + +import torch + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + run_compiled, +) + +from . import prefill_kernel as _kernel + +MAX_SLOTS = 128 # AITER's TOPK_MAX_SLOTS: 8192 blocks + + +def _pow2_ceil(n: int) -> int: + return 1 << (n - 1).bit_length() if n >= 1 else 1 + + +@functools.lru_cache(maxsize=len(_kernel.PREFILL_STRIPS)) +def _get_prefill_kernel(strips: int): + return _kernel.compile_topk_prefill(strips) + + +@functools.lru_cache(maxsize=1) +def _get_decode_kernel(): + return _kernel.compile_topk_prefill(26, decode=True) + + +def supports_index_topk(topk_blocks: int, block_size: int) -> bool: + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() and topk_blocks == _kernel.TOPK and block_size == _kernel.BLK + + +def _check(score: torch.Tensor, topk_idx: torch.Tensor, max_seq_len: int) -> int: + """Validate the layouts; returns the score columns the kernel walks.""" + heads, rows, width = score.shape + assert score.dtype == torch.float32 and score.stride(2) == 1 + assert topk_idx.dtype == torch.int32 and topk_idx.stride(2) == 1 + assert topk_idx.shape == (heads, rows, _kernel.TOPK) + slots = _pow2_ceil((max_seq_len + _kernel.BLK * 64 - 1) // (_kernel.BLK * 64)) + assert slots <= MAX_SLOTS and width >= slots * 64 + return slots * 64 + + +def minimax_m3_index_topk_prefill_indices( + score: torch.Tensor, # [H, rows, S] fp32 + topk_idx: torch.Tensor, # [H, rows, 16] int32 + *, + max_seq_len: int, + num_valid_pages: torch.Tensor, # [>= rows] int32: row n's visible blocks +) -> None: + """Top-16 blocks of ragged prefill rows.""" + heads, rows, _ = score.shape + cols = _check(score, topk_idx, max_seq_len) + assert num_valid_pages.dtype == torch.int32 + assert num_valid_pages.is_contiguous() and num_valid_pages.numel() >= rows + if rows == 0: + return + tiles = _kernel.PREFILL_STRIPS + strips = next((s for s in tiles if max_seq_len <= s * _kernel.BLK * 64), tiles[-1]) + run_compiled( + _get_prefill_kernel(strips), + score.data_ptr(), + score.stride(0), + score.stride(1), + cols, + num_valid_pages.data_ptr(), + 1, + topk_idx.data_ptr(), + topk_idx.stride(0), + topk_idx.stride(1), + rows, + heads, + torch.cuda.current_stream(), + ) + + +def minimax_m3_index_topk_decode_indices( + score: torch.Tensor, # [H, num_reqs * query_len, S] fp32 + topk_idx: torch.Tensor, # [H, num_reqs * query_len, 16] int32 + seq_lens: torch.Tensor, # [>= num_reqs] int32 + *, + max_seq_len: int, + query_len: int, +) -> None: + """Top-16 blocks of decode rows, ``query_len`` rows per request.""" + heads, rows, _ = score.shape + cols = _check(score, topk_idx, max_seq_len) + assert query_len > 0 and rows % query_len == 0 + assert seq_lens.dtype == torch.int32 and seq_lens.is_contiguous() + assert seq_lens.numel() >= rows // query_len + if rows == 0: + return + run_compiled( + _get_decode_kernel(), + score.data_ptr(), + score.stride(0), + score.stride(1), + cols, + seq_lens.data_ptr(), + query_len, + topk_idx.data_ptr(), + topk_idx.stride(0), + topk_idx.stride(1), + rows, + heads, + torch.cuda.current_stream(), + ) diff --git a/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py b/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py new file mode 100644 index 000000000000..b18db79e055a --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/index_topk_decode/prefill_kernel.py @@ -0,0 +1,387 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Index top-k for MiniMax-M3 (FlyDSL, gfx950). + +One 64-lane workgroup per (row, head) returns the 16 highest-scoring block ids +in AITER's order, walking the row in windows of STRIPS x 64 blocks. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import PRELOAD_HINTS + +BLK = 128 # tokens per block +TOPK = 16 +PREFILL_STRIPS = (4, 8, 16, 20, 26) +CAP = 64 # candidates one rank pass holds (one per lane) +IDX_BITS = 13 # block index bits (slots * 64 <= 8192) +INT_MIN = -(1 << 31) +KEY_NEG_INF = (0xFF800000 ^ 0x7FFFFFFF) - (1 << 32) # signed key of -inf (NaN's) +KEY_POS_INF = 0x7F800000 + + +def _s32(v): + v &= 0xFFFFFFFF + return v - (1 << 32) if v >= 1 << 31 else v + + +def _key(bits): + """Signed key in AITER's order; NaNs land above +inf or below -inf.""" + return bits ^ ((bits >> 31) & 0x7FFFFFFF) + + +def _fold_nan(k): + return ((k > KEY_POS_INF) | (k < KEY_NEG_INF)).select(fx.Int32(KEY_NEG_INF), k) + + +def _ballot(c): + return fx.Int64(fx.rocdl.ballot(T.i64, c)) + + +def _popc(m): + """Set bits of a wave mask; two 32-bit counts keep the compare scalar.""" + lo = fx.math.ctpop(m.to(fx.Int32)).to(fx.Int32) + hi = fx.math.ctpop((m >> 32).to(fx.Int32)).to(fx.Int32) + return fx.Int32(lo) + fx.Int32(hi) + + +def _below(mask, base): + """The set bits of ``mask`` below this lane, plus ``base`` (v_mbcnt).""" + lo = mask.to(fx.Int32) + hi = (mask >> 32).to(fx.Int32) + r = fx.rocdl.mbcnt_lo(T.i32, lo.ir_value(), fx.Int32(base).ir_value()) + return fx.Int32(fx.rocdl.mbcnt_hi(T.i32, hi.ir_value(), r)) + + +def _imax(a, b): + return (a > b).select(a, b) + + +def _one(c): + return c.select(fx.Int32(1), fx.Int32(0)) + + +def _comp(k, i): + """(key, index) as one i64 whose signed order is AITER's.""" + return (k.to(fx.Int64) << 32) | (i.to(fx.Int64) & 0xFFFFFFFF) + + +def compile_topk_prefill(strips: int, *, decode: bool = False): + assert strips == 26 if decode else strips in PREFILL_STRIPS + STRIPS = strips + WIN = STRIPS * 64 + RANK_BLOCK = 8 + # Candidates, rank-loop padding, then one discard slot per lane. + TRASH = CAP + RANK_BLOCK + CAND = TRASH + 64 + name = f"m3_index_topk_{'decode' if decode else 'prefill'}_indices_w{STRIPS}" + + @fx.struct + class Shared: + cand: fx.Array[fx.Int32, 2 * CAND, 16] + sel: fx.Array[fx.Int32, 4 * TOPK, 16] + + @flyc.kernel(name=name, known_block_size=[64, 1, 1]) + def kernel( + arg_score: fx.Int64, + i64_s_head: fx.Int64, + i64_s_row: fx.Int64, + i32_width: fx.Int32, + arg_nvp: fx.Int64, + i32_query_len: fx.Int32, + arg_topk: fx.Int64, + i64_t_head: fx.Int64, + i64_t_row: fx.Int64, + ): + smem = fx.SharedAllocator().allocate(Shared).peek() + n, h = fx.block_idx.x, fx.block_idx.y + lane = fx.thread_idx.x + n64 = fx.Int32(n).to(fx.Int64) + h64 = fx.Int32(h).to(fx.Int64) + + def rsrc(base, nbytes): + """Buffer at an i64 base; per-row rebasing keeps offsets small.""" + return buffer_ops.create_buffer_resource_from_addr( + base, num_records_bytes=fx.Int64(nbytes) + ) + + # Reads past the row's width return 0. + score_row = arg_score + (h64 * i64_s_head + n64 * i64_s_row) * 4 + srs = rsrc(score_row, i32_width.to(fx.Int64) * 4) + + def load_window(srs, wbase): + return Vec.from_elements( + [ + fx.Int32( + buffer_ops.buffer_load( + srs, + wbase + (j * 64) + lane, + vec_width=1, + dtype=fx.Int32, + ) + ) + for j in range_constexpr(STRIPS) + ], + fx.Int32, + ) + + # The first window is loaded before the row's length is known. + bits = load_window(srs, 0) + + def scalar_at(base, idx64, nbytes): + return fx.Int32( + buffer_ops.buffer_load( + rsrc(base + idx64 * 4, nbytes), + fx.Int32(0), + vec_width=1, + is_scalar=True, + ) + ) + + if const_expr(decode): + req = (fx.Int32(n) // i32_query_len).to(fx.Int64) + seq = scalar_at(arg_nvp, req, fx.Int64(4)) + kv0 = seq - i32_query_len + fx.Int32(n) % i32_query_len + 1 + kv = (kv0 < 0).select(fx.Int32(0), kv0) + nblk = (kv + BLK - 1) // BLK + else: + nblk = scalar_at(arg_nvp, n64, fx.Int64(4)) + nlive = (nblk < i32_width).select(nblk, i32_width) + # Later windows read only the row's live blocks. + srs_live = rsrc(score_row, nlive.to(fx.Int64) * 4) + + def view(ptr, n_dw, width): + return fx.logical_divide( + fx.make_view(fx.recast_iter(fx.Int32, ptr), fx.make_layout(n_dw, 1)), + fx.make_layout(width, 1), + ) + + pairs = view(smem.cand.ptr, 2 * CAND, 2) + quads = view(smem.cand.ptr, 2 * CAND, 4) + sels = view(smem.sel.ptr, 4 * TOPK, 4) + atom64 = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32) + atom128 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + + def lds_put(atom, v, w, vals): + r = fx.make_rmem_tensor(fx.make_layout(w, 1), fx.Int32) + r.store(Vec.from_elements(vals, fx.Int32)) + fx.copy(atom, r, v) + + def lds_get(atom, v, w): + r = fx.make_rmem_tensor(fx.make_layout(w, 1), fx.Int32) + fx.copy(atom, v, r) + x = Vec(r.load()) + return [fx.Int32(x[e]) for e in range_constexpr(w)] + + def put_pair(pos, i, k): + lds_put(atom64, fx.slice(pairs, (None, pos)), 2, [i, k]) + + def wave_sum(c): + """Sum over the wave of a per-lane count in [0, 63].""" + t = fx.Int32(0) + for b in range_constexpr(6): + t = t + (_popc(_ballot((c & (1 << b)) != 0)) << b) + return t + + def compact(keys, wbase, base, pred): + """Append the keys ``pred`` keeps at ``base`` if they fit below CAP.""" + oks = [ + _one(pred(keys[j], wbase + (j * 64) + lane)) + for j in range_constexpr(STRIPS) + ] + cnt = oks[0] + for j in range_constexpr(1, STRIPS): + cnt = cnt + oks[j] + pos = fx.Int32(base) + end = fx.Int32(base) + for b in range_constexpr(6): # cnt <= 32 + bm = _ballot((cnt & (1 << b)) != 0) + pos = pos + (_below(bm, 0) << b) + end = end + (_popc(bm) << b) + # On overflow nothing is written; the caller rebuilds the candidates. + if end <= CAP: + for j in range_constexpr(STRIPS): + put_pair( + (oks[j] != 0).select(pos, TRASH + lane), + wbase + (j * 64) + lane, + keys[j], + ) + pos = pos + oks[j] + return end + + def exact_cut(keys, wbase): + """The window's 16th key v and the index cut ie among the keys tied + with it: (k > v) | (k == v & i >= ie) keeps exactly its top 16.""" + pe = fx.Int32(0) + b = fx.Int32(31) + while b >= 0: + t = (pe | (fx.Int32(1) << b)) ^ INT_MIN + c = fx.Int32(0) + for j in range_constexpr(STRIPS): + c = c + _one(keys[j] >= t) + pe = (wave_sum(c) >= TOPK).select(t ^ INT_MIN, pe) + b = b - 1 + v = pe ^ INT_MIN + c = fx.Int32(0) + for j in range_constexpr(STRIPS): + c = c + _one(keys[j] > v) + need = TOPK - wave_sum(c) + ie = fx.Int32(0) + b = fx.Int32(IDX_BITS - 1) + while b >= 0: + t = ie | (fx.Int32(1) << b) + c = fx.Int32(0) + for j in range_constexpr(STRIPS): + c = c + _one((keys[j] == v) & ((wbase + (j * 64) + lane) >= t)) + ie = (wave_sum(c) >= need).select(t, ie) + b = b - 1 + return v, ie + + def lane_max_bound(m): + p = fx.Int32(0) + for d in range_constexpr(12): + t = p | _s32(1 << (31 - d)) + count = _popc(_ballot(m >= (t ^ INT_MIN))) + p = (count >= TOPK).select(t, p) + return p ^ INT_MIN + + def window(bits, wbase, nrun, tprev): + """One window of keys merged into the running top 16 (sel).""" + keys = [] + for j in range_constexpr(STRIPS): + live = (wbase + (j * 64) + lane) < nlive + keys.append(live.select(_key(bits[j]), fx.Int32(INT_MIN))) + m = keys[0] + for j in range_constexpr(1, STRIPS): + m = _imax(m, keys[j]) + # Bound: the lane maxima in the first window, the running 16th later. + tw = fx.Int32(INT_MIN) + if nrun < TOPK: + tw = lane_max_bound(m) + ta = _imax(tw, tprev) + total = compact(keys, wbase, nrun, lambda k, i: k >= ta) + # Exact only if 16 keys above -inf bound it and no NaN is above +inf. + nan_hi = _ballot(m > KEY_POS_INF) != fx.Int64(0) + if (total > CAP) | (ta <= KEY_NEG_INF) | nan_hi: + wl = nlive - wbase + if wl <= CAP - nrun: + # Every live block of the window is a candidate. + put_pair( + (lane < wl).select(nrun + lane, TRASH + lane), + wbase + lane, + _fold_nan(keys[0]), + ) + total = nrun + wl + else: + # NaNs to -inf and re-bound; exact search if still too many. + fixed = [] + for j in range_constexpr(STRIPS): + live = (wbase + (j * 64) + lane) < nlive + fixed.append(live.select(_fold_nan(keys[j]), fx.Int32(INT_MIN))) + mf = fixed[0] + for j in range_constexpr(1, STRIPS): + mf = _imax(mf, fixed[j]) + tf = _imax(lane_max_bound(mf), tprev) + total = compact(fixed, wbase, nrun, lambda k, i: k >= tf) + if total > CAP: + v, ie = exact_cut(fixed, wbase) + total = compact( + fixed, + wbase, + nrun, + lambda k, i: (k > v) | ((k == v) & (i >= ie)), + ) + if lane < RANK_BLOCK: + put_pair(total + lane, fx.Int32(0), fx.Int32(INT_MIN)) + fx.gpu.barrier() + mi, mk = lds_get(atom64, fx.slice(pairs, (None, lane)), 2) + mw = _comp(mk, mi) + rank = fx.Int32(0) + for g in range_constexpr(CAP // RANK_BLOCK): + if total > RANK_BLOCK * g: + es = [ + lds_get( + atom128, fx.slice(quads, (None, RANK_BLOCK // 2 * g + u)), 4 + ) + for u in range_constexpr(RANK_BLOCK // 2) + ] + # All of a block's reads in flight before the compares. + fx.rocdl.sched_barrier(0) + for e in es: + rank = rank + _one(_comp(e[1], e[0]) > mw) + rank = rank + _one(_comp(e[3], e[2]) > mw) + if (lane < total) & (rank < TOPK): + lds_put(atom128, fx.slice(sels, (None, rank)), 4, [mi, mk, 0, mi]) + return (total < TOPK).select(total, fx.Int32(TOPK)) + + wbase = fx.Int32(0) + nrun = fx.Int32(0) + tprev = fx.Int32(INT_MIN + 1) # dead keys are INT_MIN + while (wbase == 0) | (wbase < nlive): + if wbase > 0: + bits = load_window(srs_live, wbase) + # The running top 16 back into cand; its 16th key bounds the window. + fx.gpu.barrier() + r = lds_get(atom128, fx.slice(sels, (None, lane & (TOPK - 1))), 4) + if lane < nrun: + put_pair(lane, r[0], r[1]) + k15 = fx.Int32(fx.rocdl.readlane(T.i32, r[1], fx.Int32(TOPK - 1))) + tprev = (nrun >= TOPK).select(k15, fx.Int32(INT_MIN + 1)) + fx.gpu.barrier() + nrun = window( + [fx.Int32(bits[j]) for j in range_constexpr(STRIPS)], + wbase, + nrun, + tprev, + ) + wbase = wbase + WIN + + fx.gpu.barrier() + r = lds_get(atom128, fx.slice(sels, (None, lane & (TOPK - 1))), 4) + blk = (lane < nrun).select(r[0], fx.Int32(-1)) + buffer_ops.buffer_store( + blk, + rsrc(arg_topk + (h64 * i64_t_head + n64 * i64_t_row) * 4, TOPK * 4), + lane, + mask=lane < TOPK, + ) + + @flyc.jit + def launch( + arg_score: fx.Int64, + i64_s_head: fx.Int64, + i64_s_row: fx.Int64, + i32_width: fx.Int32, + arg_nvp: fx.Int64, + i32_query_len: fx.Int32, + arg_topk: fx.Int64, + i64_t_head: fx.Int64, + i64_t_row: fx.Int64, + i32_rows: fx.Int32, + i32_heads: fx.Int32, + stream: fx.Stream, + ): + kernel( + arg_score, + i64_s_head, + i64_s_row, + i32_width, + arg_nvp, + i32_query_len, + arg_topk, + i64_t_head, + i64_t_row, + ).launch( + grid=(fx.Int64(i32_rows), fx.Int64(i32_heads), 1), + block=(64, 1, 1), + stream=stream, + ) + + launch.kernel_name = name + launch.compile_hints = PRELOAD_HINTS + return launch diff --git a/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/__init__.py b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/__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/minimax_m3/amd/ops/norm_rope_nvfp4_insert/host.py b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/host.py new file mode 100644 index 000000000000..d6590ee458a4 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/host.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch wrapper of the fused norm / RoPE / NVFP4 insert.""" + +import functools +import struct + +import torch + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + run_compiled, +) + +from . import kernel as _kernel + + +@functools.cache +def _get_kernel(nq: int, nkv: int, niq: int, rotary_dim: int, tp: int): + return _kernel.compile_norm_rope_insert(nq, nkv, niq, rotary_dim, tp) + + +def _token_pairs_per_wave(n: int) -> int: + return 1 if n <= 512 else 4 + + +def minimax_m3_norm_rope_nvfp4_insert( + qkv: torch.Tensor, # [n, (nq + 2 nkv + niq + 1) * 128] bf16 + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_q_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, # [max_pos, rotary_dim] bf16 + positions: torch.Tensor, # [n] int64 + slot_mapping: torch.Tensor, # [n] int64, -1 skips the token + index_slot_mapping: torch.Tensor, # [n] int64, -1 skips the token + kv_cache: torch.Tensor, # [blocks, 2 nkv, 128, 72] uint8 + index_cache: torch.Tensor, # [blocks, 128, 72] uint8 + k_scale: torch.Tensor, # fp32 global scales + v_scale: torch.Tensor, + q_out: torch.Tensor, # [n, nq * 128] bf16 + index_q_out: torch.Tensor, # [n, niq * 128] e4m3, row stride free + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + eps: float, +) -> None: + """Norm + RoPE of q / k / index-q / index-k and the NVFP4 K/V and index-K + inserts in one launch; k and index-k are not written back into ``qkv``.""" + n = qkv.shape[0] + if n == 0: + return + nq, nkv, niq = num_heads, num_kv_heads, num_index_heads + rotary_dim = cos_sin_cache.shape[1] + assert eps >= 2.0**-126 # the kernel's rsqrt skips ocml's denormal scaling + assert qkv.dtype == torch.bfloat16 and qkv.is_contiguous() + assert qkv.shape[1] == (nq + 2 * nkv + niq + 1) * _kernel.D + assert cos_sin_cache.dtype == torch.bfloat16 and cos_sin_cache.is_contiguous() + for w in (q_norm_weight, k_norm_weight, index_q_norm_weight, index_k_norm_weight): + assert w.dtype == torch.bfloat16 and w.is_contiguous() + assert q_out.dtype == torch.bfloat16 and q_out.is_contiguous() + assert index_q_out.dtype == torch.float8_e4m3fn and index_q_out.stride(1) == 1 + assert index_q_out.stride(0) % 4 == 0 + assert kv_cache.shape[1:] == (2 * nkv, _kernel.BLK, 72) + assert kv_cache.stride(1) == _kernel.BLK * 72 and kv_cache.stride(3) == 1 + assert index_cache.shape[1:] == (_kernel.BLK, 72) and index_cache.stride(2) == 1 + assert index_cache.stride(1) == 72 + run_compiled( + _get_kernel(nq, nkv, niq, rotary_dim, _token_pairs_per_wave(n)), + qkv.data_ptr(), + q_out.data_ptr(), + index_q_out.data_ptr(), + q_norm_weight.data_ptr(), + k_norm_weight.data_ptr(), + index_q_norm_weight.data_ptr(), + index_k_norm_weight.data_ptr(), + cos_sin_cache.data_ptr(), + positions.data_ptr(), + slot_mapping.data_ptr(), + index_slot_mapping.data_ptr(), + kv_cache.data_ptr(), + index_cache.data_ptr(), + k_scale.data_ptr(), + v_scale.data_ptr(), + n, + index_q_out.stride(0), + kv_cache.stride(0), + index_cache.stride(0), + cos_sin_cache.shape[0], + struct.unpack(" 4 fp32.""" + cos, sin = css_t + e = _bf16x4(xws_t) + ss = e[0] * e[0] + e[1] * e[1] + ss = ss + e[2] * e[2] + ss = ss + e[3] * e[3] + for m in (16, 8, 4, 2, 1): + ss = _xsum(ss, m) + r = _rsq(fx.math.fma(ss, fx.Float32(0.0078125), eps)) + nrm = [(e[i] * r) * (fx.Float32(1.0) + w[i]) for i in range_constexpr(4)] + out = [] + for i in range_constexpr(4): + pr = _xor_lane(bits(nrm[i]), half4, hi16).bitcast(fx.Float32) + ps = pr * sin[i] + roped = fx.math.fma(nrm[i], cos[i], first.select(-ps, ps)) + out.append(in_rope.select(roped, nrm[i])) + return out + + +def _maxnum(a, b): + return fx.Float32(_arith.MaxNumFOp(fx.as_ir_value(a), fx.as_ir_value(b)).result) + + +def _minnum(a, b): + return fx.Float32(_arith.MinNumFOp(fx.as_ir_value(a), fx.as_ir_value(b)).result) + + +def _dpp(src, ctrl, old=None, bank_mask=0xF): + """v_mov_b32_dpp; an undef ``old`` lets the move fold into its consumer.""" + o = _llvm.mlir_undef(T.i32) if old is None else fx.as_ir_value(old) + full = old is None and bank_mask == 0xF + return fx.Int32( + fx.rocdl.update_dpp(T.i32, o, fx.as_ir_value(src), ctrl, 0xF, bank_mask, full) + ) + + +def _swap16(v): + """permlane16_swap(v, v): (own, lane ^ 16) in rows 0 / 2, swapped in 1 / 3.""" + return permlane_swap(fx.rocdl.permlane16_swap, v, v) + + +def _xor_lane(v, m, hi16): + """Lane l ^ m of the 32-lane half; ``hi16`` is this lane's bit 4.""" + if m == 1: + return _dpp(v, 0xB1) # quad_perm [1, 0, 3, 2] + if m == 2: + return _dpp(v, 0x4E) # quad_perm [2, 3, 0, 1] + if m == 4: # row_shr:4 into banks 1, 3, then row_shl:4 into banks 0, 2 + return _dpp(v, 0x104, old=_dpp(v, 0x114, bank_mask=0xA), bank_mask=0x5) + if m == 8: + return _dpp(v, 0x128) # row_ror:8 + assert m == 16 + a, b = _swap16(v) + return hi16.select(a, b) + + +def _xsum(x, m): + """Adds lane l ^ m's value to ``x``.""" + if m == 16: + a, b = _swap16(bits(x)) + return a.bitcast(fx.Float32) + b.bitcast(fx.Float32) + return x + _xor_lane(bits(x), m, None).bitcast(fx.Float32) + + +def _xmax(x, m): + return _maxnum(x, _xor_lane(bits(x), m, None).bitcast(fx.Float32)) + + +def _rsq(x): + """v_rsq_f32; ocml's denormal scaling never applies as x >= eps >= 2^-126.""" + return fx.Float32(fx.rocdl.rsq(T.f32, fx.as_ir_value(x))) + + +def _e4m3_byte(x): + """RNE e4m3 byte of ``x``.""" + pk = fx.rocdl.cvt_scalef32_pk_fp8_f32( + T.vec(2, T.i16), + fx.as_ir_value(Vec.filled(2, 0, fx.Int16)), + fx.as_ir_value(x), + fx.as_ir_value(x), + fx.as_ir_value(fx.Float32(1.0)), + False, + ) + return Vec(pk).bitcast(fx.Int32)[0] & 0xFF + + +def _e4m3_value(byte): + pk = fx.rocdl.cvt_scalef32_pk_f32_fp8( + T.vec(2, T.f32), + fx.as_ir_value(fx.Int32(byte)), + fx.as_ir_value(fx.Float32(1.0)), + False, + ) + return fx.Float32(Vec(pk)[0]) + + +def _iq_fp8(a, b, old, hi): + """Two fp32 -> e4m3 bytes in the low (``hi``: high) half of ``old``.""" + return fx.Int32( + fx.rocdl.cvt_pk_fp8_f32( + T.i32, + fx.as_ir_value(fx.Float32(a)), + fx.as_ir_value(fx.Float32(b)), + fx.as_ir_value(fx.Int32(old)), + hi, + ) + ) + + +def _e2m1x2(a, b, old, sel): + """E2M1 codes of ``a``, ``b`` into byte ``sel`` of ``old``, ``a`` low.""" + return fx.Int32( + fx.rocdl.cvt_scalef32_pk_fp4_f32( + T.i32, + fx.as_ir_value(fx.Int32(old)), + fx.as_ir_value(a), + fx.as_ir_value(b), + fx.as_ir_value(fx.Float32(1.0)), + sel, + ) + ) + + +def _quant4(xs, gs, inv_gs): + """E2M1 codes of this lane's 4 values (16 bits) and the e4m3 scale byte of + their 16-value group (lanes l ^ 0..3).""" + a = [fx.math.absf(x) for x in xs] + m = _maxnum(_maxnum(a[0], a[1]), _maxnum(a[2], a[3])) + m = _xmax(_xmax(m, 1), 2) + sf = _minnum((m * as_f32(RCP6)) * inv_gs, as_f32(F448)) + byte = _e4m3_byte(sf) + sv = _e4m3_value(byte) + quant = (sv != 0.0).select(fx.Float32(1.0) / (gs * sv), fx.Float32(0.0)) + y = [x * quant for x in xs] + return _e2m1x2(y[2], y[3], _e2m1x2(y[0], y[1], fx.Int32(0), 0), 1), byte + + +def _v_row(state, row, r, n_run, p0, vs, inv_vs): + """Adds row ``r`` of a V run at octet positions p0.. to ``state``: new + nibbles and scale bytes, with their masks.""" + words, written = state[0:4], state[4] + sc, sc_mask = state[5:7], state[7:9] + nib, sf = _quant4(_bf16x4(row), vs, inv_vs) + in_run = n_run > r + p = p0 + r + sh = p * 4 + words = [ + wd | in_run.select(((nib >> (4 * i)) & 15) << sh, fx.Int32(0)) + for i, wd in enumerate(words) + ] + written = written | in_run.select(fx.Int32(15) << sh, fx.Int32(0)) + bsh = (p % 4) * 8 + for k in range_constexpr(2): + here = in_run & ((p // 4) == k) + sc[k] = sc[k] | here.select(sf << bsh, fx.Int32(0)) + sc_mask[k] = sc_mask[k] | here.select(fx.Int32(255) << bsh, fx.Int32(0)) + return [*words, written, *sc, *sc_mask] + + +def _gptr(addr, ty, align): + return fx.inttoptr( + fx.PointerType.get(ty, address_space=fx.AddressSpace.Global, alignment=align), + fx.Int64(addr), + ) + + +def _k_scale_off(tok, grp): + """Byte of (token, group) in a head's K scale region.""" + r = tok % 32 + n = (r // 8) * 4 + r % 4 + return ((((tok // 32) * 4 + grp // 2) * 16 + n) * 2 + (r % 8) // 4) * 2 + grp % 2 + + +def compile_norm_rope_insert(nq: int, nkv: int, niq: int, rotary_dim: int, tp: int): + """Build the kernel; ``tp`` token pairs per norm / RoPE wave, one per V wave.""" + assert rotary_dim in (8, 16, 32, 64, 128) + NR = nq + nkv + 1 + niq # norm + RoPE waves + S = NR + nkv * tp # waves per 2 * tp tokens + ROW = (nq + 2 * nkv + niq + 1) * D + HALF = rotary_dim // 2 + SIDE = nkv * BLK * 72 # bytes of one side of a main page + TPW = 2 * tp + name = f"m3_norm_rope_nvfp4_insert_q{nq}_kv{nkv}_iq{niq}_r{rotary_dim}_t{tp}" + + @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1]) + def kernel( + i32_n: fx.Int32, + i32_cs_rows: fx.Int32, + arg_qkv: fx.Int64, # [n, ROW] bf16 + arg_pos: fx.Int64, # positions [n] i64 + arg_slot: fx.Int64, # main slots [n] i64, -1 skips + arg_islot: fx.Int64, # index slots [n] i64, -1 skips + arg_cs: fx.Int64, # cos_sin_cache [i32_cs_rows, rotary_dim] bf16 + arg_qw: fx.Int64, # norm weights [128] bf16: q, k, index-q, index-k + arg_kw: fx.Int64, + arg_iqw: fx.Int64, + arg_ikw: fx.Int64, + arg_q: fx.Int64, # q_out [n, nq * 128] bf16 + arg_iq: fx.Int64, # index_q_out [n, niq * 128] e4m3, row stride i32_iq_st + arg_kv: fx.Int64, # [blocks, 2 * nkv, 128, 72] u8 + arg_idx: fx.Int64, # [blocks, 128, 72] u8 + arg_ks: fx.Int64, # fp32 global scales + arg_vs: fx.Int64, + i64_kv_st: fx.Int64, + i64_idx_st: fx.Int64, + i32_iq_st: fx.Int32, + i32_eps: fx.Int32, # fp32 bits + ): + tx, bx = fx.thread_idx.x, fx.block_idx.x + lane = tx % 64 + wave = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64)) + wid = bx * NW + wave + s = wid % S + hl = lane % 32 + hi16 = (hl & 16) != 0 + dim0 = hl * 4 + toks = [(wid // S) * TPW + t * 2 + lane // 32 for t in range_constexpr(tp)] + lives = [tk < i32_n for tk in toks] + + qkv = buffer_ops.create_buffer_resource_from_addr( + arg_qkv, num_records_bytes=fx.Int64(i32_n) * (ROW * 2) + ) + sl_r = buffer_ops.create_buffer_resource_from_addr( + arg_slot, num_records_bytes=fx.Int64(i32_n) * 8 + ) + + if s < NR: + is_q = s < nq + is_k = (s >= nq) & (s < nq + nkv) + is_ik = s == nq + nkv + is_kq = is_k | is_ik + # This wave's head column in the qkv row. + col = (s <= nq + nkv).select( + is_ik.select(fx.Int32(nq + 2 * nkv + niq), s), + s - (nq + nkv + 1) + (nq + 2 * nkv), + ) + w_addr = is_q.select( + arg_qw, is_k.select(arg_kw, is_ik.select(arg_ikw, arg_iqw)) + ) + wr = buffer_ops.create_buffer_resource_from_addr( + w_addr, num_records_bytes=D * 2 + ) + pos_r = buffer_ops.create_buffer_resource_from_addr( + arg_pos, num_records_bytes=fx.Int64(i32_n) * 8 + ) + isl_r = buffer_ops.create_buffer_resource_from_addr( + arg_islot, num_records_bytes=fx.Int64(i32_n) * 8 + ) + cs_r = buffer_ops.create_buffer_resource_from_addr( + arg_cs, num_records_bytes=fx.Int64(i32_cs_rows) * (rotary_dim * 2) + ) + xws, poss, slots = [], [], [] + for t in range_constexpr(tp): + xws.append( + Vec( + buffer_ops.buffer_load( + qkv, + toks[t] * (ROW // 2) + col * (D // 2) + hl * 2, + vec_width=2, + dtype=fx.Int32, + mask=lives[t], + ) + ) + ) + poss.append( + fx.Int32( + buffer_ops.buffer_load( + pos_r, + toks[t] * 2, + vec_width=1, + dtype=fx.Int32, + mask=lives[t], + ) + ) + ) + kslot = fx.Int32( + buffer_ops.buffer_load( + sl_r, + toks[t] * 2, + vec_width=1, + dtype=fx.Int32, + mask=lives[t] & is_k, + ) + ) + islot = fx.Int32( + buffer_ops.buffer_load( + isl_r, + toks[t] * 2, + vec_width=1, + dtype=fx.Int32, + mask=lives[t] & is_ik, + ) + ) + slots.append(is_k.select(kslot, islot)) + w = _bf16x4( + Vec(buffer_ops.buffer_load(wr, hl * 2, vec_width=2, dtype=fx.Int32)) + ) + first = dim0 < HALF + in_rope = dim0 < rotary_dim + ib = first.select(dim0, dim0 - HALF) + css = [] + for t in range_constexpr(tp): + cs_off = (poss[t] * rotary_dim + ib) // 2 + css.append( + ( + _bf16x4( + Vec( + buffer_ops.buffer_load( + cs_r, + cs_off, + vec_width=2, + dtype=fx.Int32, + mask=lives[t] & in_rope, + ) + ) + ), + _bf16x4( + Vec( + buffer_ops.buffer_load( + cs_r, + cs_off + HALF // 2, + vec_width=2, + dtype=fx.Int32, + mask=lives[t] & in_rope, + ) + ) + ), + ) + ) + eps = i32_eps.bitcast(fx.Float32) + qvs = [] + for t in range_constexpr(tp): + out = _norm_rope_out4( + xws[t], css[t], w, eps, hi16, first, in_rope, HALF // 4 + ) + lo, hi = pk_bf16(out[0], out[1]), pk_bf16(out[2], out[3]) + qvs.append(Vec.from_elements([lo, hi], fx.Int32)) + if is_q: + q_r = buffer_ops.create_buffer_resource_from_addr( + arg_q, num_records_bytes=fx.Int64(i32_n) * (nq * D * 2) + ) + for t in range_constexpr(tp): + buffer_ops.buffer_store( + qvs[t], + q_r, + toks[t] * (nq * D // 2) + s * (D // 2) + hl * 2, + mask=lives[t], + ) + if s > nq + nkv: + iq_r = buffer_ops.create_buffer_resource_from_addr( + arg_iq, num_records_bytes=fx.Int64(i32_n) * fx.Int64(i32_iq_st) + ) + iq_col = (s - (nq + nkv + 1)) * D + for t in range_constexpr(tp): + f = [ + _minnum(_maxnum(x, as_f32(F448N)), as_f32(F448)) + for x in _bf16x4(qvs[t]) + ] + word = _iq_fp8( + f[2], f[3], _iq_fp8(f[0], f[1], fx.Int32(0), False), True + ) + buffer_ops.buffer_store( + word, + iq_r, + toks[t] * i32_iq_st + iq_col + dim0, + mask=lives[t], + offset_is_bytes=True, + ) + if is_kq: + gs = is_k.select( + fx.Float32(_gptr(arg_ks, T.f32, 4)[0]), fx.Float32(1.0) + ) + inv_gs = fx.Float32(1.0) / gs + quants = [ + _quant4(_bf16x4(qvs[t]), gs, inv_gs) for t in range_constexpr(tp) + ] + h = s - nq + grp = hl // 4 + for t in range_constexpr(tp): + slot = slots[t] + tp_ = slot & (BLK - 1) + page = is_k.select( + arg_kv + fx.Int64(slot >> 7) * i64_kv_st, + arg_idx + fx.Int64(slot >> 7) * i64_idx_st, + ) + d_off = is_k.select(h * (BLK * 64), fx.Int32(0)) + tp_ * 64 + hl * 2 + s_off = is_k.select( + nkv * BLK * 64 + h * (BLK * 8) + _k_scale_off(tp_, grp), + BLK * 64 + tp_ * 8 + grp, + ) + d_ptr = _gptr(page + fx.Int64(d_off), T.i16, 2) + s_ptr = _gptr(page + fx.Int64(s_off), T.i8, 1) + packed = quants[t][0].to(fx.Int16) + sf8 = quants[t][1].to(fx.Int8) + if lives[t] & (slot >= 0): + d_ptr[0] = packed + if hl % 4 == 0: + s_ptr[0] = sf8 + else: + vh = (s - NR) // tp + tok = (wid // S) * TPW + ((s - NR) % tp) * 2 + lane // 32 + live = tok < i32_n + vcol = (nq + nkv + vh) * (D // 2) # V head column (dwords) + vs = fx.Float32(_gptr(arg_vs, T.f32, 4)[0]) + near_t = tok - 7 + hl + has_near = (hl < 16) & (near_t >= 0) & (near_t < i32_n) + near = has_near.select( + fx.Int32( + buffer_ops.buffer_load( + sl_r, near_t * 2, vec_width=1, dtype=fx.Int32, mask=has_near + ) + ), + fx.Int32(-1), + ) + slot = fx.Int32( + buffer_ops.buffer_load( + sl_r, tok * 2, vec_width=1, dtype=fx.Int32, mask=live + ) + ) + has_prev = live & (tok > 0) + prev = has_prev.select( + fx.Int32( + buffer_ops.buffer_load( + sl_r, tok * 2 - 2, vec_width=1, dtype=fx.Int32, mask=has_prev + ) + ), + fx.Int32(-1), + ) + # The first token of a run in one 8-token octet writes the run and + # keeps the octet's other nibbles and scales. Rows it can cover load + # before its slot is known. + rows = [ + Vec( + buffer_ops.buffer_load( + qkv, + (tok + r) * (ROW // 2) + vcol + hl * 2, + vec_width=2, + dtype=fx.Int32, + mask=tok + r < i32_n, + ) + ) + for r in range_constexpr(8) + ] + base = slot - slot % 8 + in_oct = (near >= base) & (near < base + 8) + votes = fx.Int64(fx.rocdl.ballot(T.i64, in_oct)) + mine = (votes >> fx.Int64((lane // 32) * 32)).to(fx.Int32) & 0xFFFF + n_run = fx.Int32(fx.math.ctpop(mine)) + run_first = (prev < 0) | ((prev >> 3) != (slot >> 3)) + tp_ = slot & (BLK - 1) + v_side = arg_kv + fx.Int64(slot >> 7) * i64_kv_st + SIDE + v_data = v_side + fx.Int64(vh * (BLK * 64) + (tp_ // 8) * (D * 4)) + v_sf = v_side + fx.Int64(nkv * BLK * 64 + vh * (BLK * 8) + tp_ - tp_ % 8) + dp = _gptr(v_data + fx.Int64(dim0 * 4), T.i32, 4) + sp = _gptr(v_sf + fx.Int64((hl // 4) * BLK), T.i32, 4) + go = live & (slot >= 0) & run_first + inv_vs = fx.Float32(1.0) / vs + # Runs of <= 4 tokens (decode) take the short tile. + for run_rows, take in ( + (rows[:4], go & (n_run <= 4)), + (rows, go & (n_run > 4)), + ): + if take: + old = [fx.Int32(dp[i]) for i in range_constexpr(4)] + old_sc = [fx.Int32(sp[i]) for i in range_constexpr(2)] + st = [fx.Int32(0)] * 9 + for r in range_constexpr(len(run_rows)): + st = _v_row(st, run_rows[r], r, n_run, tp_ % 8, vs, inv_vs) + keep = ~st[4] + for i in range_constexpr(4): + dp[i] = st[i] | (old[i] & keep) + sc0 = st[5] | (old_sc[0] & ~st[7]) + sc1 = st[6] | (old_sc[1] & ~st[8]) + if hl % 4 == 0: + sp[0] = sc0 + sp[1] = sc1 + + @flyc.jit + def launch( + arg_qkv: fx.Int64, + arg_q: fx.Int64, + arg_iq: fx.Int64, + arg_qw: fx.Int64, + arg_kw: fx.Int64, + arg_iqw: fx.Int64, + arg_ikw: fx.Int64, + arg_cs: fx.Int64, + arg_pos: fx.Int64, + arg_slot: fx.Int64, + arg_islot: fx.Int64, + arg_kv: fx.Int64, + arg_idx: fx.Int64, + arg_ks: fx.Int64, + arg_vs: fx.Int64, + i32_n: fx.Int32, + i32_iq_st: fx.Int32, + i64_kv_st: fx.Int64, + i64_idx_st: fx.Int64, + i32_cs_rows: fx.Int32, + i32_eps: fx.Int32, + stream: fx.Stream, + ): + waves = (i32_n + (TPW - 1)) // TPW * S + kernel( + i32_n, + i32_cs_rows, + arg_qkv, + arg_pos, + arg_slot, + arg_islot, + arg_cs, + arg_qw, + arg_kw, + arg_iqw, + arg_ikw, + arg_q, + arg_iq, + arg_kv, + arg_idx, + arg_ks, + arg_vs, + i64_kv_st, + i64_idx_st, + i32_iq_st, + i32_eps, + ).launch( + grid=(fx.Int64((waves + (NW - 1)) // NW), 1, 1), + block=(64 * NW, 1, 1), + stream=stream, + ) + + launch.kernel_name = name + launch.compile_hints = PRELOAD_HINTS + return launch diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/__init__.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/__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/minimax_m3/amd/ops/nvfp4_index_decode/host.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/host.py new file mode 100644 index 000000000000..37e260514bb0 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/host.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch wrappers of the FlyDSL NVFP4 decode index scorers.""" + +import functools + +import torch + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import run_compiled +from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.host import nvfp4_e4m3_lut +from vllm.utils.math_utils import cdiv + +from . import kernel as _kernel + +# 158 VGPRs = 3 waves per SIMD; more waves add a second round of workgroups. +TARGET_WAVES = 3072 +# Scores are buffered in LDS and stored by whole rows once the batch's score +# footprint reaches this; smaller footprints store directly. +BUFFER_MIN_BYTES = 8 << 20 +UNITS_PER_SLOT = 4 + + +def _end_bytes(t: torch.Tensor) -> int: + """Bytes from the data pointer to the end of the last element.""" + return ( + sum((n - 1) * s for n, s in zip(t.shape, t.stride())) + 1 + ) * t.element_size() + + +def _unit(max_blocks: int, slots: int, umax: int, footprint: int) -> int: + """Consecutive blocks per buffered unit (a power of two, 4..umax) giving + every wave slot >= UNITS_PER_SLOT units, or 0 for direct stores.""" + if footprint < BUFFER_MIN_BYTES: + return 0 + u = umax + while u > 4 and cdiv(max_blocks, u) < UNITS_PER_SLOT * slots: + u //= 2 + return u if cdiv(max_blocks, u) >= UNITS_PER_SLOT * slots else 0 + + +@functools.cache +def _get_kernel(rows: int, heads: int, buffered: bool): + return _kernel.compile_score_decode(rows, heads, buffered) + + +def _supported( + q_idx: torch.Tensor, + index_cache: torch.Tensor, + score: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + query_len: int, +) -> bool: + from vllm.platforms.rocm import on_gfx950 + + num_heads = q_idx.shape[1] + return ( + on_gfx950() + and num_heads in (1, 2, 4) + and 1 <= num_heads * query_len <= _kernel.MAX_ROWS + and q_idx.dtype == torch.float8_e4m3fn + and q_idx.shape[2] == _kernel.D + and q_idx.stride(2) == 1 + and q_idx.stride(0) % 16 == 0 + and q_idx.stride(1) % 16 == 0 + and q_idx.data_ptr() % 16 == 0 + and index_cache.dtype == torch.uint8 + and index_cache.stride(2) == 1 + and index_cache.stride(1) == _kernel.D // 2 + _kernel.D // 16 + and index_cache.stride(0) % 16 == 0 + and index_cache.data_ptr() % 16 == 0 + and score.dtype == torch.float32 + and score.stride(2) == 1 + and block_table.dtype == torch.int32 + and block_table.stride(1) == 1 + and seq_lens.dtype == torch.int32 + and seq_lens.is_contiguous() + ) + + +def minimax_m3_nvfp4_index_score_decode( + q_idx: torch.Tensor, # [num_reqs * query_len, H, 128] e4m3 + index_cache: torch.Tensor, # [num_blocks, 128, 72] uint8 + score: torch.Tensor, # [H, num_reqs * query_len, >= max_blocks] fp32 + block_table: torch.Tensor, # [num_reqs, >= max_blocks] int32 + seq_lens: torch.Tensor, # [num_reqs] int32 + init_blocks: int, + local_blocks: int, + query_len: int, + max_seq_len: int, +) -> None: + """Block scores of every decode row for blocks below cdiv(seq_len, 128).""" + total_q, num_heads, _ = q_idx.shape + num_reqs = total_q // query_len + if num_reqs == 0: + return + assert _supported(q_idx, index_cache, score, block_table, seq_lens, query_len) + rows = num_heads * query_len + max_blocks = max(1, cdiv(max_seq_len, _kernel.BLK)) + # A chunk of NW x (pages per wave) requests; with at most one page per + # wave, chunks of one request and each wave on its own positions. + pages_per_wave = cdiv(num_reqs * max_blocks, TARGET_WAVES) + s_bytes = _end_bytes(score) + assert s_bytes < 2**31 * 4, "score offsets must fit in int32 elements" + footprint = num_reqs * rows * max_blocks * 4 + chunk = 1 if pages_per_wave <= 1 else _kernel.NW * pages_per_wave + chunk_cap = _kernel.chunk_max(_kernel.padded_rows(rows)) + chunk = max(1, min(chunk, chunk_cap, num_reqs)) + nch = cdiv(num_reqs, chunk) + chunk = cdiv(num_reqs, nch) + per_wave = chunk < _kernel.NW + wanted = max(1, TARGET_WAVES // (_kernel.NW * nch)) + # Past 16 rows the buffered variant's flush costs more than it saves. + unit = ( + _unit( + max_blocks, + wanted * (_kernel.NW if per_wave else 1), + _kernel.UMAX, + footprint, + ) + if _kernel.padded_rows(rows) <= 16 + else 0 + ) + units = cdiv(max_blocks, max(1, unit)) + positions = cdiv(units, _kernel.NW) if per_wave else units + splits = max(1, min(positions, wanted)) + run_compiled( + _get_kernel(rows, num_heads, unit > 0), + q_idx.data_ptr(), + index_cache.data_ptr(), + nvfp4_e4m3_lut(q_idx.device).data_ptr(), + score.data_ptr(), + block_table.data_ptr(), + seq_lens.data_ptr(), + int(index_cache.stride(0)), + int(q_idx.stride(0)), + int(q_idx.stride(1)), + _end_bytes(q_idx), + int(score.stride(0)), + int(score.stride(1)), + s_bytes, + int(block_table.stride(0)), + _end_bytes(block_table), + int(num_reqs), + int(chunk), + int(nch), + int(splits), + int(init_blocks), + int(local_blocks), + int(per_wave), + int(unit), + torch.cuda.current_stream(), + # Keep the buffered variant at the plain one's occupancy. + waves_per_eu=3 if unit > 0 and _kernel.padded_rows(rows) <= 16 else None, + ) + + +@functools.cache +def _get_kernel_cp_push(): + from . import kernel_cp_push + + return kernel_cp_push.compile_score_push() + + +def minimax_m3_nvfp4_index_score_push_cp( + iq: torch.Tensor, # this rank's index head [num_reqs * 4, 128] e4m3 + index_cache: torch.Tensor, # [num_blocks, 128, 72] uint8 + dests: list[int], # head h's scores [rows, row_stride] fp32 on rank h + row_stride: int, + dest_bytes: int, + block_table: torch.Tensor, # [num_reqs, >= max_blocks] int32 + seq_lens: torch.Tensor, # [num_reqs] int32 + init_blocks: int, + local_blocks: int, + rank: int, + counts: torch.Tensor, # index_cp's device counters + flags: list[int], # rank p's flag words + iq_bufs: list[int], # rank p's index-q receive buffer [4, iq_rows, 128] e4m3 + iq_rows: int, + grid: int = 0, +) -> None: + """This rank's share of all four heads' decode scores, pushed to the heads' + ranks; returns once every rank's scores have arrived.""" + from . import kernel_cp_push + + rows = iq.shape[0] + assert rows % 4 == 0 and rows <= iq_rows + assert iq.dtype == torch.float8_e4m3fn and iq.shape[1] == kernel_cp_push.D + assert iq.stride(1) == 1 and iq.stride(0) == kernel_cp_push.D + num_reqs = rows // 4 + assert num_reqs <= kernel_cp_push.MAX_REQS + if not grid: + small = num_reqs <= kernel_cp_push.GRID_SMALL_REQS + grid = kernel_cp_push.GRID_SMALL if small else kernel_cp_push.GRID + assert 0 < grid <= kernel_cp_push.SC_MAX_WG + if num_reqs == 0: + return + assert dest_bytes < 2**31 * 4, "score offsets must fit in int32 elements" + run_compiled( + _get_kernel_cp_push(), + iq.data_ptr(), + index_cache.data_ptr(), + nvfp4_e4m3_lut(iq.device).data_ptr(), + *dests, + block_table.data_ptr(), + seq_lens.data_ptr(), + counts.data_ptr(), + *flags, + *iq_bufs, + int(iq_rows), + int(index_cache.stride(0)), + int(row_stride), + int(dest_bytes), + int(block_table.stride(0)), + _end_bytes(block_table), + int(num_reqs), + int(init_blocks), + int(local_blocks), + int(rank), + int(grid), + torch.cuda.current_stream(), + ) diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel.py new file mode 100644 index 000000000000..a80f10463a51 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel.py @@ -0,0 +1,564 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 decode index scorer for MiniMax-M3 (FlyDSL, gfx950). + +Per (head, query row, block): the max of q . k over the row's visible tokens, +with AITER's init / local sentinels, on e4m3_rn_satfinite(e2m1 x scale) keys. A +workgroup takes a chunk of up to 64 requests (lane r holds request r) and every +splits-th block position. The distinct pages at a position (a shared prefix page +is one page for the whole chunk) go round robin over the 4 waves; a wave +dequantizes its page into the A fragments of v_mfma_f32_16x16x128_f8f6f4, issuing +the next page's loads into the freed registers, and scores it against every +16-column group of Q rows holding the page. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + NT, + as_f32, + bits, + dequant8, + global_i32_ptr, + key_order4, + maxf_nn, + mfma_fp8_16x16x128, + pack8, + xlane4, +) +from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.kernel import with_sentinels + +NW = 4 +MAX_REQS = 64 # requests per chunk, one per lane +BLK = 128 # tokens per block +D = 128 # index head dim +PAGE_CODES = BLK * D // 2 # page: [128 x 64 B] e2m1 codes, [128 x 8 B] e4m3 scales +PAGE_BYTES = PAGE_CODES + BLK * 8 +MT = BLK // 16 # token tiles per page +UMAX = 8 # max consecutive blocks per LDS-buffered unit +MAX_ROWS = 32 # H * query_len + + +def padded_rows(rows: int) -> int: + """Columns a request takes in the 16-column groups: rows rounded up to a + power of two.""" + return 1 << (rows - 1).bit_length() + + +def chunk_max(rows: int) -> int: + """Requests per chunk: one per lane, and 256 Q rows (32 KB) in LDS.""" + return min(MAX_REQS, 256 // rows) + + +def block_sentinels(mx, visible, b, i32_init, i32_local): + """AITER's sentinels on block b of a row seeing ``visible`` tokens.""" + nb = ((visible < 0).select(fx.Int32(0), visible) + (BLK - 1)) // BLK + return with_sentinels(mx, b, nb, i32_init, i32_local) + + +def compile_score_decode(rows: int, heads: int, buffered: bool = False): + """Build the scorer. + + Args: + rows: H * query_len (1..32). + heads: index heads (1, 2 or 4). + buffered: gather a unit's scores in LDS and store whole rows (large + score footprints); otherwise store them directly. + + """ + assert 1 <= rows <= MAX_ROWS and heads in (1, 2, 4) and rows % heads == 0 + H = heads + QL = rows // H + R = padded_rows(rows) # a request's columns; column j = query row j // H + RPC = max(1, 16 // R) # requests per 16-column group + CPR = max(1, R // 16) # 16-column groups per request + GMASK = (1 << RPC) - 1 + QROWS = chunk_max(R) * R + QROW_BYTES = D + name = f"m3_nvfp4_index_score_decode_r{rows}h{H}" + ("_buf" if buffered else "") + SB_STRIDE = UMAX + 1 # padded against LDS bank conflicts + + @fx.struct + class Shared: + q: fx.Array[fx.Uint8, QROWS * QROW_BYTES, 16] + lut: fx.Array[fx.Int32, 512, 16] + seq: fx.Array[fx.Int32, MAX_REQS] + # The unit's scores as f32 bits, plus 64 dummy slots. + sbuf: fx.Array[fx.Int32, (QROWS * SB_STRIDE + 64) if buffered else 16, 16] + + @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1]) + def kernel( + arg_q: fx.Int64, # [num_reqs * QL, H, 128] e4m3 + arg_kv: fx.Int64, # [num_blocks, 128, 72] uint8 + arg_lut: fx.Int64, # nvfp4_e4m3_lut [256] int64 + arg_score: fx.Int64, # [H, num_reqs * QL, S] fp32 + arg_bt: fx.Int64, # [num_reqs, >= max_blocks] int32 + arg_seq: fx.Int64, # [num_reqs] int32 + i64_page_stride: fx.Int64, + i32_q_st: fx.Int32, # q strides (bytes) + i32_q_sh: fx.Int32, + i64_q_bytes: fx.Int64, + i32_s_sh: fx.Int32, # score strides (elements) + i32_s_st: fx.Int32, + i64_s_bytes: fx.Int64, + i32_bt_stride: fx.Int32, + i64_bt_bytes: fx.Int64, + i32_num_reqs: fx.Int32, + i32_chunk: fx.Int32, # requests per chunk + i32_nch: fx.Int32, + i32_splits: fx.Int32, # workgroups per chunk + i32_init: fx.Int32, + i32_local: fx.Int32, + i32_per_wave: fx.Int32, # 1: each wave walks its own positions + i32_unit: fx.Int32, # consecutive positions per unit, 0: none + ): + smem = fx.SharedAllocator().allocate(Shared).peek() + tx, pid = fx.thread_idx.x, fx.block_idx.x + lane = tx % 64 + wave = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64)) + col, kl = lane % 16, lane // 16 + ch = pid % i32_nch + split = pid // i32_nch + req0 = ch * i32_chunk + left = i32_num_reqs - req0 + nr = (left < i32_chunk).select(left, i32_chunk) + + # Lane r holds request req0 + r. + r_ok = lane < nr + seq_r = fx.Int32(global_i32_ptr(arg_seq)[r_ok.select(req0 + lane, fx.Int32(0))]) + seq_r = r_ok.select(seq_r, fx.Int32(0)) + nblk_r = (seq_r + (BLK - 1)) // BLK + first_row = seq_r - (QL - 1) + # Every row sees blocks before rbf_r whole: init / local ones there only + # need their sentinel. + rbf_r = ((first_row < 0).select(fx.Int32(0), first_row) + (BLK - 1)) // BLK + ll = nblk_r - i32_local + loc_last_r = (ll < 0).select(fx.Int32(0), ll) + fm_r = seq_r - QL # tokens past this are masked for some row + + btr = buffer_ops.create_buffer_resource_from_addr( + arg_bt, num_records_bytes=i64_bt_bytes + ) + sr = buffer_ops.create_buffer_resource_from_addr( + arg_score, num_records_bytes=i64_s_bytes + ) + qr = buffer_ops.create_buffer_resource_from_addr( + arg_q, num_records_bytes=i64_q_bytes + ) + + q16 = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.q.ptr), + fx.make_layout(QROWS * QROW_BYTES // 4, 1), + ), + fx.make_layout(4, 1), + ) + atom16 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + seq_l = smem.seq.ptr + scb = smem.sbuf.ptr + per_wave = i32_per_wave != fx.Int32(0) + if const_expr(buffered): + direct = per_wave | (i32_unit == fx.Int32(0)) + to_lds = (i32_per_wave == fx.Int32(0)) & (i32_unit != fx.Int32(0)) + usz = (i32_unit == fx.Int32(0)).select(fx.Int32(1), i32_unit) + else: + direct = i32_unit >= fx.Int32(0) + to_lds = i32_unit < fx.Int32(0) + usz = fx.Int32(1) + + def prologue(): + """Q rows, the e4m3 table and the request lengths -> LDS. + + Row i = request i // R, column j = i % R (query row j // H, head + j % H); 16 B column c8 sits at c8 ^ (i & 7), in the key's dim order. + """ + for k in range_constexpr(QROWS * 8 // (64 * NW)): + u = tx + k * 64 * NW + i = u // 8 + c8 = u % 8 + rr = i // R + j = i % R + # Missing requests' and padding rows stay unwritten; their + # scores are never stored. + if (rr < nr) & (j // H < QL): + tq = (req0 + rr) * QL + j // H + off = (tq * i32_q_st + (j % H) * i32_q_sh + c8 * 16) // 4 + d = Vec( + buffer_ops.buffer_load( + qr, off, vec_width=4, dtype=fx.Int32, mask=rr < nr + ) + ) + d = [fx.Int32(d[e]) for e in range_constexpr(4)] + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + r.store(key_order4(d)) + fx.copy(atom16, r, fx.slice(q16, (None, i * 8 + (c8 ^ (i & 7))))) + lut_g = global_i32_ptr(arg_lut) + lut_l = smem.lut.ptr + for e in range_constexpr(2): + lut_l[tx * 2 + e] = fx.Int32(lut_g[tx * 2 + e]) + seq_l[lane] = seq_r + + prologue() + fx.gpu.barrier() + + lut8 = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.lut.ptr), fx.make_layout(512, 1) + ), + fx.make_layout(2, 1), + ) + lut_atom = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32) + + def lut_load(sbyte): + r = fx.make_rmem_tensor(fx.make_layout(2, 1), fx.Int32) + fx.copy(lut_atom, fx.slice(lut8, (None, sbyte)), r) + v = Vec(r.load()) + return fx.Int32(v[0]), fx.Int32(v[1]) + + def q_frag(cg): + """B operand of column group cg (dims 32 kl .. +31 of Q row cg * 16 + + col).""" + i = cg * 16 + col + halves = [] + for h in range_constexpr(2): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + fx.copy( + atom16, fx.slice(q16, (None, i * 8 + ((2 * kl + h) ^ (i & 7)))), r + ) + halves.append(r.load()) + return pack8(halves[0], halves[1]) + + code_off = (col * 64 + kl * 16) // 4 # dwords + scale_off = col * 4 + kl # u16: groups 2 kl, 2 kl + 1 of token col + + def page_rsrc(page): + return buffer_ops.create_buffer_resource_from_addr( + arg_kv + fx.Int64(page) * i64_page_stride, num_records_bytes=PAGE_BYTES + ) + + def load_tile(kr, ok, t): + """Tile t of a page: this lane's 16 B of codes (token 16 t + col) and + its two group scales; no memory access when not ``ok``.""" + c = Vec( + buffer_ops.buffer_load( + kr, + code_off, + vec_width=4, + dtype=fx.Int32, + mask=ok, + soffset_bytes=t * 1024, + cache_modifier=NT, + ) + ) + sc = fx.Int32( + fx.Int16( + buffer_ops.buffer_load( + kr, + scale_off, + vec_width=1, + dtype=fx.Int16, + mask=ok, + soffset_bytes=PAGE_CODES + t * 128, + cache_modifier=NT, + ) + ).to(fx.Int32) + ) + return c, sc + + def dequant_tile(c, sc): + """The A fragment of a tile (e4m3 keys in the key's dim order).""" + t0 = lut_load(sc & 0xFF) + t1 = lut_load((sc >> 8) & 0xFF) + f = [] + for e in range_constexpr(4): + lo, hi = dequant8(fx.Int32(Vec(c)[e]), *(t0 if e < 2 else t1)) + f += [lo, hi] + return Vec.from_elements(f, fx.Int32) + + zero4 = fx.Vector.filled(4, 0.0, fx.Float32) + neg_inf = fx.Float32(float("-inf")) + + def frag_max(frags, qfs, masked, lims): + """Per column group: the max of its column's scores over the page. + + Tile t's results are reduced under tile t + 1's MFMAs. + """ + ng = len(qfs) + runs = [neg_inf] * ng + accs = [[None] * ng for _ in range(MT)] + + def reduce(t, runs): + out = [] + for g in range_constexpr(ng): + run = runs[g] + v = accs[t][g] + for e in range_constexpr(4): + x = v[e] + if const_expr(masked): + x = (lims[g] >= t * 16 + e).select(x, neg_inf) + run = maxf_nn(run, x) + out.append(run) + return out + + for t in range_constexpr(MT): + for g in range_constexpr(ng): + accs[t][g] = mfma_fp8_16x16x128(frags[t], qfs[g], zero4) + if const_expr(t >= 1): + runs = reduce(t - 1, runs) + runs = reduce(MT - 1, runs) + return [xlane4(r, maxf_nn) for r in runs] + + def row_value(mx, seq, q_off, b): + return block_sentinels(mx, seq - QL + q_off + 1, b, i32_init, i32_local) + + def put(i_row, jpos, v, cond, addr): + """Score v of Q row i_row at unit position jpos, to memory or to the + LDS buffer; lanes with nothing to buffer write a dummy slot.""" + if const_expr(not buffered): + buffer_ops.buffer_store(v, sr, addr, mask=cond) + else: + buffer_ops.buffer_store(v, sr, addr, mask=cond & direct) + idx = (cond & to_lds).select( + i_row * SB_STRIDE + jpos, QROWS * SB_STRIDE + lane + ) + (scb + idx).store(bits(v)) # `scb[idx] = v` would rebind scb + + def flush(ub): + """Store the unit at block ub: 4 consecutive blocks per thread and + pass.""" + NCH = UMAX // 4 + k = fx.Int32(0) + while k < fx.Int32(QROWS * NCH // (64 * NW)): + item = tx + k * (64 * NW) + i = item // NCH + c = item % NCH + rr = i // R + jc = i % R + qoff = jc // H + nb_c = (fx.Int32(seq_l[rr]) + (BLK - 1)) // BLK + b0c = ub + c * 4 + ok = (rr < nr) & (qoff < QL) & (c * 4 < i32_unit) + addr = (jc % H) * i32_s_sh + ((req0 + rr) * QL + qoff) * i32_s_st + b0c + for e in range_constexpr(4): + v = as_f32(fx.Int32(scb[i * SB_STRIDE + c * 4 + e])) + buffer_ops.buffer_store( + v, + sr, + addr + e, + mask=ok & (b0c + e < nb_c) & (c * 4 + e < i32_unit), + ) + k = k + 1 + + def score_cgs(frags, hm, masked_any, cgs, b, jpos): + """Scores of column groups cgs, masked where a request is at its last + blocks, with sentinels.""" + info = [] + for cg in cgs: + i = cg * 16 + col + crr = i // R + st_j = i % R + info.append((cg, crr, st_j // H, st_j % H)) + qfs = [q_frag(cg) for cg, _, _, _ in info] + seqs = [fx.Int32(seq_l[crr]) for _, crr, _, _ in info] + mxs = [neg_inf for _ in info] + if masked_any: + lims = [ + seq_c - QL + qoff - b * BLK - kl * 4 + for seq_c, (_, _, qoff, _) in zip(seqs, info) + ] + mxs = frag_max(frags, qfs, True, lims) + else: + mxs = frag_max(frags, qfs, False, None) + for mx, seq_c, (cg, crr, qoff, head) in zip(mxs, seqs, info): + hit = ((hm >> fx.Int64(crr)) & fx.Int64(1)) != fx.Int64(0) + ov = row_value(mx, seq_c, qoff, b) + otq = (req0 + crr) * QL + qoff + put( + cg * 16 + col, + jpos, + ov, + hit & (kl == 0) & (qoff < QL), + head * i32_s_sh + otq * i32_s_st + b, + ) + + def score_groups(frags, hm, mmask, b, jpos): + """Every column group holding the page (requests hm).""" + left_m = hm + while left_m != fx.Int64(0): + lead_req = fx.Int32(fx.math.cttz(left_m).to(fx.Int32)) + if const_expr(CPR == 1): + cg = lead_req // RPC + gbits = fx.Int64(GMASK) << fx.Int64(cg * RPC) + left_m = left_m & ~gbits + score_cgs(frags, hm, (mmask & gbits) != fx.Int64(0), [cg], b, jpos) + else: + gbits = fx.Int64(1) << fx.Int64(lead_req) + left_m = left_m & ~gbits + masked_any = (mmask & gbits) != fx.Int64(0) + score_cgs( + frags, + hm, + masked_any, + [lead_req * CPR + g for g in range(CPR)], + b, + jpos, + ) + + # The waves share positions in units of i32_unit blocks; page k of + # position `it` goes to wave (k + it) % NW. In per-wave mode (fewer + # requests than waves) wave w walks its own units instead. + u_step = per_wave.select(i32_splits * NW, i32_splits) + ub = per_wave.select(split * NW + wave, split) * usz + it = fx.Int32(0) + while fx.Int64(fx.rocdl.ballot(T.i64, ub < nblk_r)) != fx.Int64(0): + b = ub + jpos = fx.Int32(0) + while (jpos < usz) & ( + fx.Int64(fx.rocdl.ballot(T.i64, b < nblk_r)) != fx.Int64(0) + ): + in_r = b < nblk_r + skip = (b < rbf_r) & ((b < i32_init) | (b >= loc_last_r)) + need_r = in_r & ~skip + bid = fx.Int32( + buffer_ops.buffer_load( + btr, + need_r.select((req0 + lane) * i32_bt_stride + b, fx.Int32(0)), + vec_width=1, + dtype=fx.Int32, + ) + ) + bid = need_r.select(bid, fx.Int32(-1)) + # Sentinel-only blocks, stored by wave 0. + if (fx.Int64(fx.rocdl.ballot(T.i64, in_r & skip)) != fx.Int64(0)) & ( + (wave == 0) | per_wave + ): + # A runtime loop: unrolled, all store addresses stay live. + jr = fx.Int32(0) + while jr < fx.Int32(QL * H): + sv = row_value(neg_inf, seq_r, jr // H, b) + stq = (req0 + lane) * QL + jr // H + put( + lane * R + jr, + jpos, + sv, + in_r & skip, + (jr % H) * i32_s_sh + stq * i32_s_st + b, + ) + jr = jr + 1 + pend = fx.Int64(fx.rocdl.ballot(T.i64, need_r)) + mmask = fx.Int64( + fx.rocdl.ballot(T.i64, need_r & (b * BLK + (BLK - 1) > fm_r)) + ) + # This wave's pages, as the mask of their first lanes. + mine = fx.Int64(0) + k = it + while pend != fx.Int64(0): + lead = fx.Int32(fx.math.cttz(pend).to(fx.Int32)) + page = fx.Int32(fx.rocdl.readlane(T.i32, bid, lead)) + pend = pend & ~fx.Int64(fx.rocdl.ballot(T.i64, bid == page)) + mine = mine | (((k % NW) == wave) | per_wave).select( + fx.Int64(1) << fx.Int64(lead), fx.Int64(0) + ) + k = k + 1 + ok0 = mine != fx.Int64(0) + l0 = fx.Int32(fx.math.cttz(mine).to(fx.Int32)) + kr0 = page_rsrc( + fx.Int32(fx.rocdl.readlane(T.i32, bid, ok0.select(l0, fx.Int32(0)))) + ) + tiles = [load_tile(kr0, ok0, t) for t in range_constexpr(MT)] + while mine != fx.Int64(0): + lead = fx.Int32(fx.math.cttz(mine).to(fx.Int32)) + page = fx.Int32(fx.rocdl.readlane(T.i32, bid, lead)) + hm = fx.Int64(fx.rocdl.ballot(T.i64, bid == page)) + mine = mine & ~(fx.Int64(1) << fx.Int64(lead)) + okn = mine != fx.Int64(0) + ln = fx.Int32(fx.math.cttz(mine).to(fx.Int32)) + krn = page_rsrc( + fx.Int32( + fx.rocdl.readlane(T.i32, bid, okn.select(ln, fx.Int32(0))) + ) + ) + # Tile t of the next page loads into the registers tile t of + # this page frees; the barriers keep LLVM from hoisting the + # loads. + frags, nxt = [], [] + for t in range_constexpr(MT): + frags.append(dequant_tile(*tiles[t])) + nxt.append(load_tile(krn, okn, t)) + fx.rocdl.sched_barrier(0) + tiles = nxt + score_groups(frags, hm, mmask, b, jpos) + b = b + 1 + jpos = jpos + 1 + it = it + 1 + if const_expr(buffered): # noqa: SIM102 (the outer if is static) + if to_lds: + fx.gpu.barrier() + flush(ub) + fx.gpu.barrier() + ub = ub + u_step * usz + + @flyc.jit + def launch( + arg_q: fx.Int64, + arg_kv: fx.Int64, + arg_lut: fx.Int64, + arg_score: fx.Int64, + arg_bt: fx.Int64, + arg_seq: fx.Int64, + i64_page_stride: fx.Int64, + i32_q_st: fx.Int32, + i32_q_sh: fx.Int32, + i64_q_bytes: fx.Int64, + i32_s_sh: fx.Int32, + i32_s_st: fx.Int32, + i64_s_bytes: fx.Int64, + i32_bt_stride: fx.Int32, + i64_bt_bytes: fx.Int64, + i32_num_reqs: fx.Int32, + i32_chunk: fx.Int32, + i32_nch: fx.Int32, + i32_splits: fx.Int32, + i32_init: fx.Int32, + i32_local: fx.Int32, + i32_per_wave: fx.Int32, + i32_unit: fx.Int32, + stream: fx.Stream, + ): + kernel( + arg_q, + arg_kv, + arg_lut, + arg_score, + arg_bt, + arg_seq, + i64_page_stride, + i32_q_st, + i32_q_sh, + i64_q_bytes, + i32_s_sh, + i32_s_st, + i64_s_bytes, + i32_bt_stride, + i64_bt_bytes, + i32_num_reqs, + i32_chunk, + i32_nch, + i32_splits, + i32_init, + i32_local, + i32_per_wave, + i32_unit, + ).launch( + grid=(fx.Int64(i32_nch * i32_splits), 1, 1), + block=(64 * NW, 1, 1), + stream=stream, + ) + + launch.kernel_name = name + return launch diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel_cp_push.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel_cp_push.py new file mode 100644 index 000000000000..862144426914 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_decode/kernel_cp_push.py @@ -0,0 +1,611 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Context-parallel decode index scorer for MiniMax-M3 (FlyDSL, gfx950, TP4). + +Rank ``me`` scores blocks b with (b // CP_CH) % 4 == me for all four index heads +(bit for bit as ``kernel.py``) and stores head h's scores into rank h's buffer +(``index_cp.IndexCP``); it returns once every rank's scores have arrived. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops +from flydsl._mlir.dialects import llvm as _llvm +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + NT, + SC0, + SC1, + bits, + dequant8, + global_i32_ptr, + key_order4, + maxf_nn, + mfma_fp8_16x16x128, + pack8, + perm, + xlane4, +) +from vllm.models.minimax_m3.amd.ops.nvfp4_index_decode.kernel import block_sentinels + +WORLD = 4 # TP ranks = index heads, one per rank +BLK = 128 # tokens per block +D = 128 +CP_CH = 16 # blocks per run of one rank +PAGE_CODES = BLK * D // 2 # e2m1 codes, then [128 x 8 B] e4m3 group scales +PAGE_BYTES = PAGE_CODES + BLK * 8 +MT = BLK // 16 # token tiles per page +SYS = SC0 | SC1 +MAX_ROWS = 1024 # decode rows (max_num_seqs x query_len) the receive buffers hold +# Flag words per rank: [SC_FLAGS + p] source p's scores arrived, [SC_DONE + w] +# workgroup w done, [SC_IQ + p * SC_MAX_REQS + r] source p's index-q of request r. +SC_FLAGS = 64 +SC_DONE = 1024 +SC_MAX_WG = 1024 +SC_IQ = 2048 +SC_MAX_REQS = 256 +FLAG_BYTES = 12288 +# Device counters (int32); [CNT_SC] holds the call count. +CNT_SC = 128 +CNT_WORDS = 512 + + +def _flag_read(addr): + """An atomic add of 0 at system scope (never hoisted out of a spin loop).""" + return fx.Int32( + _llvm.AtomicRMWOp( + _llvm.AtomicBinOp.add, + global_i32_ptr(addr).llvm_ptr, + fx.Int32(0).ir_value(), + _llvm.AtomicOrdering.monotonic, + syncscope="", + alignment=4, + ).result + ) + + +def _sel4(p, a0, a1, a2, a3): + return (p == 0).select(a0, (p == 1).select(a1, (p == 2).select(a2, a3))) + + +GRID_SMALL, GRID = 512, 768 +GRID_SMALL_REQS = 32 +NW = 4 # waves per workgroup: quarter w of every run +QD = CP_CH // NW # blocks per quarter +WIN = 64 # pages per block-table window (one VGPR) +MAX_REQS = 64 * NW # one request per thread in the prefix sum +QL = 4 # DSpark rows per request +HEADS = 4 + + +def _uni(x): + return fx.Int32(fx.rocdl.readfirstlane(T.i32, fx.Int32(x))) + + +def _min(a, b): + return (a < b).select(a, b) + + +def _max(a, b): + return (a > b).select(a, b) + + +def _wave_scan(x, lane): + """Inclusive prefix sum over the wave's 64 lanes.""" + for d in range_constexpr(6): + y = fx.Int32(fx.rocdl.ds_bpermute(T.i32, (lane - (1 << d)) * 4, x)) + x = x + (lane >= (1 << d)).select(y, fx.Int32(0)) + return x + + +def compile_score_push(): + name = "m3_nvfp4_index_score_push_cp" + + @fx.struct + class Shared: + lut: fx.Array[fx.Int32, 512, 16] + stage: fx.Array[fx.Int32, 2 * 16 * NW * QD, 16] # [run % 2][col][quarter] + incl: fx.Array[fx.Int32, MAX_REQS] # inclusive prefix of the run counts + seq: fx.Array[fx.Int32, MAX_REQS] + wsum: fx.Array[fx.Int32, NW] + + @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1]) + def kernel( + arg_iq: fx.Int64, # this rank's head: [num_reqs * QL, 128] e4m3 + arg_kv: fx.Int64, # NVFP4 index pages [num_blocks, 128, 72] u8 + arg_lut: fx.Int64, # nvfp4_e4m3_lut [256] i64 + arg_d0: fx.Int64, # head h's scores [rows, S] fp32 (rank h's receive buffer) + arg_d1: fx.Int64, + arg_d2: fx.Int64, + arg_d3: fx.Int64, + arg_bt: fx.Int64, # block table [num_reqs, >= max_blocks] i32 + arg_seq: fx.Int64, # seq_lens [num_reqs] i32 + arg_cnt: fx.Int64, # device counters (CNT_*) + arg_f0: fx.Int64, # rank p's flag words + arg_f1: fx.Int64, + arg_f2: fx.Int64, + arg_f3: fx.Int64, + arg_q0: fx.Int64, # rank p's index-q receive buffer [4, rmax, 128] e4m3 + arg_q1: fx.Int64, + arg_q2: fx.Int64, + arg_q3: fx.Int64, + i32_rmax: fx.Int32, + i64_page_stride: fx.Int64, + i32_s: fx.Int32, # score row stride (elements) + i64_d_bytes: fx.Int64, + i32_bt_stride: fx.Int32, + i64_bt_bytes: fx.Int64, + i32_num_reqs: fx.Int32, + i32_init: fx.Int32, # leading blocks forced in (sentinel 1e30) + i32_local: fx.Int32, # trailing blocks forced in (sentinel 1e29) + i32_me: fx.Int32, + i32_nwg: fx.Int32, # the grid + ): + smem = fx.SharedAllocator().allocate(Shared).peek() + tx, pid = fx.thread_idx.x, fx.block_idx.x + lane = tx % 64 + wave = _uni(tx // 64) + col, kl = lane % 16, lane // 16 + st_qoff = col // HEADS # this lane's column: row st_qoff of head st_head + st_head = col % HEADS + incl_l, seq_l = smem.incl.ptr, smem.seq.ptr + + cr = buffer_ops.create_buffer_resource_from_addr( + arg_cnt, num_records_bytes=CNT_WORDS * 4 + ) + calls = ( + _uni( + buffer_ops.buffer_load( + cr, CNT_SC, vec_width=1, dtype=fx.Int32, cache_modifier=SC1 + ) + ) + + 1 + ) + btr = buffer_ops.create_buffer_resource_from_addr( + arg_bt, num_records_bytes=i64_bt_bytes + ) + qbytes = fx.Int64(i32_rmax) * (HEADS * D) + qs = (arg_q0, arg_q1, arg_q2, arg_q3) + fs = (arg_f0, arg_f1, arg_f2, arg_f3) + mine = _sel4(i32_me, arg_f0, arg_f1, arg_f2, arg_f3) + qr = buffer_ops.create_buffer_resource_from_addr( + _sel4(i32_me, arg_q0, arg_q1, arg_q2, arg_q3), num_records_bytes=qbytes + ) + drs = [ + buffer_ops.create_buffer_resource_from_addr( + d, num_records_bytes=i64_d_bytes + ) + for d in (arg_d0, arg_d1, arg_d2, arg_d3) + ] + + # LUT to LDS; per request this rank's run count, prefix-summed. + lut_g = global_i32_ptr(arg_lut) + lut_l = smem.lut.ptr + for i in range_constexpr(2): + lut_l[tx * 2 + i] = fx.Int32(lut_g[tx * 2 + i]) + r_ok = tx < i32_num_reqs + seq_t = fx.Int32(global_i32_ptr(arg_seq)[r_ok.select(tx, fx.Int32(0))]) + seq_t = r_ok.select(seq_t, fx.Int32(0)) + # This rank's index-q rows of requests pid, pid + nwg, ... to slot me of + # every rank; their flags follow below. + ir = buffer_ops.create_buffer_resource_from_addr( + arg_iq, num_records_bytes=fx.Int64(i32_num_reqs) * (QL * D) + ) + rs = pid + while (wave == 0) & (rs < i32_num_reqs): + v = buffer_ops.buffer_load( + ir, + (rs * QL * D + lane * 16) // 4, + vec_width=4, + dtype=fx.Int32, + mask=lane < QL * D // 16, + ) + for p in range_constexpr(WORLD): + buffer_ops.buffer_store( + v, + buffer_ops.create_buffer_resource_from_addr( + qs[p], num_records_bytes=qbytes + ), + ((i32_me * i32_rmax + rs * QL) * D + lane * 16) // 4, + mask=lane < QL * D // 16, + cache_modifier=SYS, + ) + rs = rs + i32_nwg + nb_t = (seq_t + (BLK - 1)) // BLK + runs_t = (nb_t // (WORLD * CP_CH)) + ( + nb_t % (WORLD * CP_CH) > i32_me * CP_CH + ).select(fx.Int32(1), fx.Int32(0)) + x = _wave_scan(runs_t, lane) + if lane == 63: + smem.wsum.ptr[wave] = x + seq_l[tx] = seq_t + fx.gpu.barrier() + for w in range_constexpr(NW - 1): + x = x + (wave > w).select(fx.Int32(smem.wsum.ptr[w]), fx.Int32(0)) + incl_l[tx] = x + fx.gpu.barrier() + + # Runs pid, pid + grid, ...: the busy workgroups stay one per CU. + total = _uni(incl_l[MAX_REQS - 1]) + nruns = (pid < total).select((total - pid + i32_nwg - 1) // i32_nwg, 0) + nslots = nruns * QD # pages of this wave + + lut8 = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.lut.ptr), fx.make_layout(512, 1) + ), + fx.make_layout(2, 1), + ) + lut_atom = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32) + stage16 = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.stage.ptr), + fx.make_layout(2 * 16 * NW * QD, 1), + ), + fx.make_layout(4, 1), + ) + atom16 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + + def lut_load(sbyte): + r = fx.make_rmem_tensor(fx.make_layout(2, 1), fx.Int32) + fx.copy(lut_atom, fx.slice(lut8, (None, sbyte)), r) + v = Vec(r.load()) + return fx.Int32(v[0]), fx.Int32(v[1]) + + code_off = (col * 64 + kl * 16) // 4 # dwords + # Scales 2 kl, 2 kl + 1 of token col as a dword: a u16 load's sign + # extension made the loop head wait for the whole next page. + scale_off = col * 2 + kl // 2 # dwords + # v_perm selectors: scale byte 2 (kl % 2) (+ 1) of the dword, zeros above + scale_sel = (kl % 2) * 2 + 0x0C0C0C00 + + def page_rsrc(page): + return buffer_ops.create_buffer_resource_from_addr( + arg_kv + fx.Int64(page) * i64_page_stride, num_records_bytes=PAGE_BYTES + ) + + def load_tile(kr, ok, t): + """Tile t: this lane's codes and scales; no memory access unless ok.""" + c = Vec( + buffer_ops.buffer_load( + kr, + code_off, + vec_width=4, + dtype=fx.Int32, + mask=ok, + soffset_bytes=t * 1024, + cache_modifier=NT, + ) + ) + sc = fx.Int32( + buffer_ops.buffer_load( + kr, + scale_off, + vec_width=1, + dtype=fx.Int32, + mask=ok, + soffset_bytes=PAGE_CODES + t * 128, + cache_modifier=NT, + ) + ) + return c, sc + + def dequant_tile(c, sc): + """The A fragment of a tile (e4m3 keys, even dims then odd dims per 8)""" + t0 = lut_load(perm(sc, sc, scale_sel)) + t1 = lut_load(perm(sc, sc, scale_sel + 1)) + f = [] + for e in range_constexpr(4): + lo, hi = dequant8(fx.Int32(Vec(c)[e]), *(t0 if e < 2 else t1)) + f += [lo, hi] + return Vec.from_elements(f, fx.Int32) + + def q_load(r, ok): + """This lane's 32 B of Q: row r * QL + st_qoff, head st_head.""" + off = ((st_head * i32_rmax + r * QL + st_qoff) * D + kl * 32) // 4 + a = Vec( + buffer_ops.buffer_load(qr, off, vec_width=4, dtype=fx.Int32, mask=ok) + ) + b = Vec( + buffer_ops.buffer_load( + qr, off + 4, vec_width=4, dtype=fx.Int32, mask=ok + ) + ) + return a, b + + def q_frag(a, b): + """B operand: the 32 B in the key's dim order.""" + halves = [] + for v in (a, b): + d = [fx.Int32(Vec(v)[e]) for e in range_constexpr(4)] + halves.append(key_order4(d)) + return pack8(halves[0], halves[1]) + + zero4 = fx.Vector.filled(4, 0.0, fx.Float32) + neg_inf = fx.Float32(float("-inf")) + + def reduce(v, t, run, masked, lim): + for e in range_constexpr(4): + x = v[e] + if const_expr(masked): + x = (lim >= t * 16 + e).select(x, neg_inf) + run = maxf_nn(run, x) + return run + + def page_pass(tiles, qf, krn, okn, masked, lim): + """Max of this lane's column over the page; tile t of the next page + loads into tile t's registers (sched barriers stop LLVM hoisting).""" + accs = [] + nxt = [] + for t in range_constexpr(MT): + f = dequant_tile(*tiles[t]) + nxt.append(load_tile(krn, okn, t)) + accs.append(mfma_fp8_16x16x128(f, qf, zero4)) + fx.rocdl.sched_barrier(0) + mx = neg_inf + if masked: + run_m = neg_inf + for t in range_constexpr(MT): + run_m = reduce(accs[t], t, run_m, True, lim) + mx = xlane4(run_m, maxf_nn) + else: + run_u = neg_inf + for t in range_constexpr(MT): + run_u = reduce(accs[t], t, run_u, False, None) + mx = xlane4(run_u, maxf_nn) + return mx, nxt + + def row_value(mx, seq, q_off, b): + return block_sentinels(mx, seq - QL + q_off + 1, b, i32_init, i32_local) + + fx.rocdl.s_waitcnt(vmcnt=0) + rs = pid + while (wave == 0) & (rs < i32_num_reqs): + for p in range_constexpr(WORLD): + buffer_ops.buffer_store( + calls, + buffer_ops.create_buffer_resource_from_addr( + fs[p], num_records_bytes=FLAG_BYTES + ), + SC_IQ + i32_me * SC_MAX_REQS + rs, + mask=lane == 0, + cache_modifier=SYS, + ) + rs = rs + i32_nwg + + w0 = fx.Int32(0) + while w0 < nslots: + # One page per lane: page s is block 4 w + s % 4 of run pid + (s // 4) + # * grid, of request rj (binary search in the prefix). + s = w0 + lane + k = pid + (s // QD) * i32_nwg + rj = fx.Int32(0) + for bs in range_constexpr(8): + step = MAX_REQS >> (bs + 1) + rj = (fx.Int32(incl_l[rj + step - 1]) <= k).select(rj + step, rj) + valid = s < nslots + rj = valid.select(rj, fx.Int32(0)) + excl = (rj > 0).select(fx.Int32(incl_l[_max(rj - 1, 0)]), fx.Int32(0)) + seqj = fx.Int32(seq_l[rj]) + nblkj = (seqj + (BLK - 1)) // BLK + bj = ((k - excl) * WORLD + i32_me) * CP_CH + wave * QD + s % QD + frow = seqj - (QL - 1) + rbf = ((frow < 0).select(fx.Int32(0), frow) + (BLK - 1)) // BLK + ll = nblkj - i32_local + loc_last = (ll < 0).select(fx.Int32(0), ll) + # Init / local blocks every row sees whole only need their sentinel. + skip = (bj < rbf) & ((bj < i32_init) | (bj >= loc_last)) + need = valid & (bj < nblkj) & ~skip + pgj = fx.Int32( + buffer_ops.buffer_load( + btr, + need.select(rj * i32_bt_stride + bj, fx.Int32(0)), + vec_width=1, + dtype=fx.Int32, + ) + ) + pgj = need.select(pgj, fx.Int32(-1)) + n = _min(nslots - w0, WIN) + + p_0 = fx.Int32(fx.rocdl.readlane(T.i32, pgj, 0)) + r_0 = fx.Int32(fx.rocdl.readlane(T.i32, rj, 0)) + ok0 = p_0 >= 0 + kr0 = page_rsrc(ok0.select(p_0, fx.Int32(0))) + fx.rocdl.sched_barrier(0) + tiles = [] + for t in range_constexpr(MT): + tiles.append(load_tile(kr0, ok0, t)) + fx.rocdl.sched_barrier(0) + # Wait for the index-q rows; lane i polls source i % 4 of its request. + word = mine + fx.Int64((SC_IQ + (lane % WORLD) * SC_MAX_REQS + rj) * 4) + want = lane < n + f = calls + if want: + f = _flag_read(word) + while fx.Int64(fx.rocdl.ballot(T.i64, want & (f < calls))) != 0: + if want & (f < calls): + f = _flag_read(word) + qa, qb = q_load(r_0, n > 0) + # Nothing pending where the loop head merges with the back edge. + fx.rocdl.s_waitcnt(vmcnt=0) + qf = q_frag(qa, qb) + qna, qnb = qa, qb + sw = n < 0 # false: Q of the next page already loaded + vq = [neg_inf] * QD + j = fx.Int32(0) + while j < n: + r = fx.Int32(fx.rocdl.readlane(T.i32, rj, j)) + b = fx.Int32(fx.rocdl.readlane(T.i32, bj, j)) + seq = fx.Int32(fx.rocdl.readlane(T.i32, seqj, j)) + if sw: + qf = q_frag(qna, qnb) + jn = j + 1 + okj = jn < n + jc = okj.select(jn, fx.Int32(0)) + pgn = fx.Int32(fx.rocdl.readlane(T.i32, pgj, jc)) + rn = fx.Int32(fx.rocdl.readlane(T.i32, rj, jc)) + okn = okj & (pgn >= 0) + krn = page_rsrc(okn.select(pgn, fx.Int32(0))) + sw = okj & (rn != r) + qna, qnb = q_load(rn, sw) + fx.rocdl.sched_barrier(0) + mx, tiles = page_pass( + tiles, + qf, + krn, + okn, + b * BLK + (BLK - 1) > seq - QL, + seq - QL + st_qoff - b * BLK - kl * 4, + ) + ov = row_value(mx, seq, st_qoff, b) + e = (w0 + j) % QD + vq = [(e == i).select(ov, v) for i, v in enumerate(vq)] + if e == QD - 1: + # The run's quarters meet in LDS; wave m % NW stores it. + m = (w0 + j) // QD + slot = m % 2 + val = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + val.store(Vec.from_elements([bits(v) for v in vq], fx.Int32)) + if kl == 0: + fx.copy( + atom16, + val, + fx.slice(stage16, (None, (slot * 16 + col) * NW + wave)), + ) + fx.gpu.barrier() + if wave == m % NW: + got = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + fx.copy( + atom16, + fx.slice(stage16, (None, (slot * 16 + col) * NW + kl)), + got, + ) + off = (r * QL + st_qoff) * i32_s + (b - b % CP_CH) + kl * QD + for h in range_constexpr(HEADS): + buffer_ops.buffer_store( + got.load(), + drs[h], + off, + mask=st_head == h, + cache_modifier=SYS, + ) + j = jn + w0 = w0 + WIN + + # sc0 sc1 stores are system-visible once vmcnt drains; a release fence + # per workgroup would write back L2. The last workgroup (dispatched last, + # so its spin cannot starve the others) collects the done words, signals + # every rank and waits for all four. + fx.rocdl.s_waitcnt(vmcnt=0) + fx.gpu.barrier() + own = buffer_ops.create_buffer_resource_from_addr( + mine, num_records_bytes=FLAG_BYTES + ) + buffer_ops.buffer_store( + calls, own, SC_DONE + pid, mask=tx == 0, cache_modifier=SYS + ) + if pid == i32_nwg - 1: + # Thread t polls done words t, t + 256, ... + def done_min(): + lo = calls + for i in range_constexpr(SC_MAX_WG // (64 * NW)): + w = tx + i * 64 * NW + f = _flag_read( + mine + fx.Int64((SC_DONE + _min(w, SC_MAX_WG - 1)) * 4) + ) + lo = _min(lo, (w < i32_nwg).select(f, calls)) + return lo + + lo = done_min() + while fx.Int64(fx.rocdl.ballot(T.i64, lo < calls)) != 0: + lo = done_min() + fx.gpu.barrier() + _llvm.FenceOp(_llvm.AtomicOrdering.release, syncscope="") + for p in range_constexpr(WORLD): + fr = buffer_ops.create_buffer_resource_from_addr( + fs[p], num_records_bytes=FLAG_BYTES + ) + buffer_ops.buffer_store( + calls, fr, SC_FLAGS + i32_me, mask=tx == 0, cache_modifier=SYS + ) + buffer_ops.buffer_store(calls, cr, CNT_SC, mask=tx == 0, cache_modifier=SYS) + if tx < WORLD: + f = _flag_read(mine + fx.Int64((SC_FLAGS + tx) * 4)) + while f < calls: + f = _flag_read(mine + fx.Int64((SC_FLAGS + tx) * 4)) + fx.gpu.barrier() + _llvm.FenceOp(_llvm.AtomicOrdering.acquire, syncscope="") + + @flyc.jit + def launch( + arg_iq: fx.Int64, + arg_kv: fx.Int64, + arg_lut: fx.Int64, + arg_d0: fx.Int64, + arg_d1: fx.Int64, + arg_d2: fx.Int64, + arg_d3: fx.Int64, + arg_bt: fx.Int64, + arg_seq: fx.Int64, + arg_cnt: fx.Int64, + arg_f0: fx.Int64, + arg_f1: fx.Int64, + arg_f2: fx.Int64, + arg_f3: fx.Int64, + arg_q0: fx.Int64, + arg_q1: fx.Int64, + arg_q2: fx.Int64, + arg_q3: fx.Int64, + i32_rmax: fx.Int32, + i64_page_stride: fx.Int64, + i32_s: fx.Int32, + i64_d_bytes: fx.Int64, + i32_bt_stride: fx.Int32, + i64_bt_bytes: fx.Int64, + i32_num_reqs: fx.Int32, + i32_init: fx.Int32, + i32_local: fx.Int32, + i32_me: fx.Int32, + i32_nwg: fx.Int32, + stream: fx.Stream, + ): + kernel( + arg_iq, + arg_kv, + arg_lut, + arg_d0, + arg_d1, + arg_d2, + arg_d3, + arg_bt, + arg_seq, + arg_cnt, + arg_f0, + arg_f1, + arg_f2, + arg_f3, + arg_q0, + arg_q1, + arg_q2, + arg_q3, + i32_rmax, + i64_page_stride, + i32_s, + i64_d_bytes, + i32_bt_stride, + i64_bt_bytes, + i32_num_reqs, + i32_init, + i32_local, + i32_me, + i32_nwg, + ).launch(grid=(fx.Int64(i32_nwg), 1, 1), block=(64 * NW, 1, 1), stream=stream) + + launch.kernel_name = name + return launch + + +assert CNT_SC < CNT_WORDS and SC_DONE + SC_MAX_WG <= FLAG_BYTES // 4 diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/__init__.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/__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/minimax_m3/amd/ops/nvfp4_index_prefill/host.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/host.py new file mode 100644 index 000000000000..2530d9f87eb5 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/host.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch wrapper of the FlyDSL NVFP4 prefill index scorers.""" + +import functools +import math + +import torch + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import cu_count, run_compiled +from vllm.utils.math_utils import cdiv + +from . import kernel as _kernel + +# Workgroups to aim for: more, shorter segments balance the tail of the causal +# windows; a segment keeps >= GROUP blocks. +TARGET_WORKGROUPS = 2048 +# Segments are added until the last round of workgroups is at least this full. +MIN_ROUND_FILL = 0.9 +WORKGROUPS_PER_CU = 2 +MAX_SEGMENTS = 1024 +# Small work (<= 128 query rows, or below SMALL_WORK) takes ``kernel_shared`` +# with 128-row tiles, aiming at SMALL_TARGET workgroups below MID_WORK. +SMALL_WORK = 4096 +MID_WORK = 16384 +SMALL_TARGET = 512 +# Tests: pin the kernel (True: kernel_shared) or its 32-row tiles per wave. +PIN_SHARED: bool | None = None +PIN_MQ: int | None = None + +_NVFP4_E4M3_LUT: dict[torch.device, torch.Tensor] = {} + + +def nvfp4_e4m3_lut(device: torch.device) -> torch.Tensor: + """[256] int64: byte j of entry s is e4m3_rn_satfinite(m_j * e4m3(s)) for the + E2M1 magnitudes m = 0, 0.5, 1, 1.5, 2, 3, 4, 6.""" + lut = _NVFP4_E4M3_LUT.get(device) + if lut is None: + # Device ops only, so a first call under graph capture is fine. + scales = torch.arange(256, dtype=torch.uint8, device=device) + scales = scales.view(torch.float8_e4m3fn).float() + code = torch.arange(8, device=device) + exp, man = code // 2, (code % 2).float() + mags = torch.where( + exp == 0, man * 0.5, torch.exp2((exp - 1).float()) * (1.0 + man * 0.5) + ) + # e2m1 * e4m3 is exact in fp32; saturate, then round once to E4M3. + values = (scales[:, None] * mags[None, :]).clamp(-448.0, 448.0) + lut = values.to(torch.float8_e4m3fn).view(torch.int64).reshape(256) + _NVFP4_E4M3_LUT[device] = lut + return lut + + +@functools.cache +def _get_kernel(heads: int): + return _kernel.compile_score_prefill(heads) + + +def _kernel_shared(): + # Imported late: the decode host imports this module. + from . import kernel_shared + + return kernel_shared + + +@functools.cache +def _get_kernel_shared(heads: int, mq: int): + return _kernel_shared().compile_score_prefill_shared(heads, mq) + + +def _supported(q_idx, index_cache, score, *int32_tensors) -> bool: + from vllm.platforms.rocm import on_gfx950 + + heads = q_idx.shape[1] + return ( + on_gfx950() + and q_idx.dtype == torch.float8_e4m3fn + and heads in (1, 2, 4) + and q_idx.shape[2] == _kernel.D + and q_idx.stride(2) == 1 + and q_idx.stride(1) == _kernel.D + and q_idx.stride(0) == heads * _kernel.D + and score.shape[0] == heads + and score.stride(2) == 1 + and index_cache.dtype == torch.uint8 + and index_cache.stride(2) == 1 + and index_cache.stride(1) == _kernel.D // 2 + _kernel.D // 16 + and index_cache.stride(0) % 16 == 0 + and index_cache.data_ptr() % 16 == 0 + and all(t.dtype == torch.int32 and t.is_contiguous() for t in int32_tensors) + ) + + +def minimax_m3_nvfp4_index_score_prefill( + q_idx: torch.Tensor, # [total_q, H, 128] e4m3 + index_cache: torch.Tensor, # [num_blocks, 128, 72] uint8 + score: torch.Tensor, # [H, total_q, >= cdiv(max_seq_len, 128)] fp32 + block_table: torch.Tensor, # [batch, >= max_blocks] int32 + cu_seqlens_q: torch.Tensor, # [batch + 1] int32 + seq_lens: torch.Tensor, # [batch] int32 + init_blocks: int, + local_blocks: int, + max_query_len: int, + max_seq_len: int, +) -> None: + """Block scores of every prefill row for 1, 2 or 4 index heads.""" + batch = cu_seqlens_q.shape[0] - 1 + total_q, heads, _ = q_idx.shape + if total_q == 0 or block_table.shape[0] == 0: + return + assert _supported(q_idx, index_cache, score, block_table, cu_seqlens_q, seq_lens) + max_blocks = cdiv(max_seq_len, _kernel.BLK) + assert max_blocks <= score.shape[2] + work = max_blocks * cdiv(max_query_len, 256) * batch * heads + small = max_query_len * heads <= 128 or work < SMALL_WORK + shared = small if PIN_SHARED is None else PIN_SHARED + mq = PIN_MQ or (1 if small else 2) + target = SMALL_TARGET if shared and work < MID_WORK else TARGET_WORKGROUPS + tile_q = _kernel_shared().tile_q(mq) if shared else _kernel.TILE_Q + qt = cdiv(max_query_len * heads, tile_q) + min_seg = 1 if shared else _kernel.GROUP + max_nseg = max(1, min(MAX_SEGMENTS, max_blocks // min_seg)) + nseg = max(1, min(max_nseg, target // max(1, qt * batch))) + slots = WORKGROUPS_PER_CU * cu_count(q_idx.device) + while nseg < max_nseg: + rounds = qt * batch * nseg / slots + if rounds / math.ceil(rounds) >= MIN_ROUND_FILL: + break + nseg += 1 + lut = () if shared else (nvfp4_e4m3_lut(q_idx.device).data_ptr(),) + run_compiled( + _get_kernel_shared(heads, mq) if shared else _get_kernel(heads), + q_idx.data_ptr(), + index_cache.data_ptr(), + *lut, + score.data_ptr(), + block_table.data_ptr(), + cu_seqlens_q.data_ptr(), + seq_lens.data_ptr(), + int(index_cache.stride(0)), + int(block_table.stride(0)), + int(score.stride(1)), + int(nseg), + int(qt), + int(total_q), + int(batch), + int(init_blocks), + int(local_blocks), + int(score.stride(0)), + int(cdiv(qt * nseg * batch, 8) * 8), + torch.cuda.current_stream(), + ) diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel.py new file mode 100644 index 000000000000..248fc990f954 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel.py @@ -0,0 +1,614 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 prefill index scorer for MiniMax-M3 (FlyDSL, gfx950). + +Per (query row, block): the max of q . k over the row's causal tokens, with +AITER's init / local sentinels, on e4m3_rn_satfinite(e2m1 x scale) keys. A +workgroup takes 512 query rows (token * H + head) x a segment of the request's +blocks; Q^T stays in VGPRs and each 16-token tile is one +v_mfma_f32_16x16x128_f8f6f4 per 16-row N-tile. Raw NVFP4 pages arrive in LDS by +DMA one block ahead, and block b + 1 is dequantized into the other key slot +under block b's MFMAs. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr import rocdl as _rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + as_f32, + bits, + dequant8, + global_i32_ptr, + key_order4, + maxf_nn, + mfma_fp8_16x16x128, + pack8, + permlane_swap, +) + +NW = 4 +ROWS_PER_WAVE = 128 +TILE_Q = NW * ROWS_PER_WAVE +BLK = 128 # tokens per block +D = 128 # index head dim +ROW_BYTES = D # one e4m3 key row in LDS +BLOCK_BYTES = BLK * ROW_BYTES +RS_T = ROW_BYTES // 16 # 16 B tiles per LDS row +SLOT_T = BLK * RS_T # 16 B tiles per key slot +N_SLOTS = 2 +LDS_BYTES = N_SLOTS * BLOCK_BYTES +PAGE_CODES = BLK * D // 2 # page: [128 x 64 B] e2m1 codes, [128 x 8 B] e4m3 scales +PAGE_BYTES = PAGE_CODES + BLK * 8 +GROUP = 4 # blocks per score store (16 B per row) +# SPAN groups of a row are stored back to back (the first SPAN - 1 wait in LDS): +# half-written 128 B score lines get evicted from L2 and cost extra write bytes. +SPAN = 4 +STAGE_T = (SPAN - 1) * NW * 2 * 64 # 16 B LDS tiles: (group, wave, row half, lane) +NI = ROWS_PER_WAVE // 16 # N-tiles per wave +MT = BLK // 16 # M-tiles per block +SENTINEL_INIT = 1e30 # AITER's forced-selection scores +SENTINEL_LOCAL = 1e29 +STAGE_W = 2048 + 256 # raw bytes per wave and block: 32 tokens' codes + scales +RAW_SLOTS = 2 +# M-tile -> code dwords dequantized under its MFMAs. +DQ_AT = {1: (0,), 2: (1,), 3: (2,), 4: (3,), 5: (4,), 6: (5,), 7: (6, 7)} + + +def _ir(v): + return v.ir_value() if hasattr(v, "ir_value") else v + + +def _dma(rsrc, lds_addr, voff, size): + """buffer_load ... lds: ``size`` bytes per lane into LDS at ``lds_addr``.""" + lds = llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<3>"), _ir(fx.Int32(lds_addr))) + fx.rocdl.raw_ptr_buffer_load_lds( + rsrc, + lds.result, + fx.Int32(size), + fx.Int32(voff), + fx.Int32(0), + fx.Int32(0), + aux=0, + ) + + +def _xlane_max4_pair(x, y): + """Max over lanes L, L^16, L^32, L^48 of two values at once, broadcast.""" + a, b = permlane_swap(_rocdl.permlane32_swap, bits(x), bits(y)) + m = bits(maxf_nn(as_f32(a), as_f32(b))) + a, b = permlane_swap(_rocdl.permlane16_swap, m, m) + big = bits(maxf_nn(as_f32(a), as_f32(b))) + a, b = permlane_swap(_rocdl.permlane32_swap, big, big) + return as_f32(a), as_f32(b) + + +def tile_rows(pid, i32_qt, i32_nseg, i32_nb, arg_cu, arg_seq, heads, tile_q): + """Workgroup pid's request b, query tile and context segment. + + Returns (live, b, seg, seq_start, seq_len, prefix, q_rows, row0); padded + workgroups read request 0 and are not live. + """ + n_wg = i32_qt * i32_nseg * i32_nb + q_tile = pid % i32_qt + seg = (pid // i32_qt) % i32_nseg + b = pid // (i32_qt * i32_nseg) + bc = (pid < n_wg).select(b, fx.Int32(0)) + cu = global_i32_ptr(arg_cu) + seq_start = fx.Int32(cu[bc]) + q_len = fx.Int32(cu[bc + 1]) - seq_start + seq_len = fx.Int32(global_i32_ptr(arg_seq)[bc]) + prefix = seq_len - q_len + q_rows = q_len * heads + row0 = q_tile * tile_q + live = (pid < n_wg) & (row0 < q_rows) + return live, b, seg, seq_start, seq_len, prefix, q_rows, row0 + + +def segment_blocks(prefix, row0, seq_len, seg, i32_nseg, heads, tile_q, align): + """The tile's nblk causal blocks and its segment [blk0, blk1) of them.""" + hi_pos = prefix + (row0 + tile_q) // heads + hi = (hi_pos < seq_len).select(hi_pos, seq_len) + nblk = (hi + (BLK - 1)) // BLK + seg_len = (nblk + i32_nseg - 1) // i32_nseg + if align > 1: + seg_len = (seg_len + (align - 1)) & ~(align - 1) + blk0 = seg * seg_len + blk1e = blk0 + seg_len + blk1 = (blk1e < nblk).select(blk1e, nblk) + return nblk, blk0, blk1 + + +def score_resources(arg_q, arg_score, i32_total_q, i32_hs, i32_S, heads): + """Buffer resources of the Q rows and the scores.""" + qr = buffer_ops.create_buffer_resource_from_addr( + arg_q, num_records_bytes=fx.Int64(i32_total_q) * (heads * D) + ) + sr = buffer_ops.create_buffer_resource_from_addr( + arg_score, + num_records_bytes=( + fx.Int64(i32_hs) * (heads - 1) + fx.Int64(i32_total_q) * fx.Int64(i32_S) + ) + * 4, + ) + return qr, sr + + +def page_lookup(arg_bt, bt_row, end): + """Page of a block, clamped to the last block before ``end``. + + A constant-space s_load: before an s_buffer_load LLVM would wait for the + in-flight LDS DMAs. + """ + bt = llvm.IntToPtrOp(ir.Type.parse("!llvm.ptr<4>"), _ir(fx.Int64(arg_bt))).result + + def page_of(blk): + bi = (blk < end).select(blk, end - 1) + p = buffer_ops.get_element_ptr( + bt, byte_offset=_ir((bt_row + bi) * 4), elem_type=T.i8 + ) + return fx.Int32(llvm.LoadOp(T.i32, p, alignment=4, invariant=True).result) + + return page_of + + +def with_sentinels(v, blk, nb, i32_init, i32_local): + """AITER's init / local sentinels on block blk of a row of nb causal blocks.""" + is_local = (blk >= nb - i32_local) & (blk < nb) + is_init = (blk < i32_init) & (blk < nb) + return is_local.select( + fx.Float32(SENTINEL_LOCAL), is_init.select(fx.Float32(SENTINEL_INIT), v) + ) + + +def compile_score_prefill(heads: int = 1): + """Build the scorer for 1, 2 or 4 index heads (sharing one index-K head).""" + assert heads in (1, 2, 4), heads + H = heads + name = f"m3_nvfp4_index_score_prefill_tq{TILE_Q}_h{H}" + + @fx.struct + class Shared: + a: fx.Array[fx.Uint8, LDS_BYTES, 16] # e4m3 key slots + lut: fx.Array[fx.Int32, 512, 16] + raw: fx.Array[fx.Uint8, RAW_SLOTS * NW * STAGE_W, 16] # DMA'd pages + stage: fx.Array[fx.Int32, STAGE_T * 4, 16] # staged score groups + + @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1]) + def kernel( + arg_q: fx.Int64, # [total_q, H, 128] e4m3 + arg_kv: fx.Int64, # [num_blocks, 128, 72] uint8 + arg_lut: fx.Int64, # nvfp4_e4m3_lut [256] int64 + arg_score: fx.Int64, # [H, total_q, S] fp32 + arg_bt: fx.Int64, # [batch, >= max_blocks] int32 + arg_cu: fx.Int64, # cu_seqlens_q [batch + 1] int32 + arg_seq: fx.Int64, # seq_lens [batch] int32 + i64_page_stride: fx.Int64, + i32_bt_stride: fx.Int32, + i32_S: fx.Int32, + i32_nseg: fx.Int32, + i32_qt: fx.Int32, # query tiles per request + i32_total_q: fx.Int32, + i32_nb: fx.Int32, + i32_init: fx.Int32, + i32_local: fx.Int32, + i32_hs: fx.Int32, # score head stride + ): + smem = fx.SharedAllocator().allocate(Shared).peek() + tx, pid = fx.thread_idx.x, fx.block_idx.x + lane = tx % 64 + wave = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64)) + l16, q16 = lane % 16, lane // 16 + live, b, seg, seq_start, seq_len, prefix, q_rows, row0 = tile_rows( + pid, i32_qt, i32_nseg, i32_nb, arg_cu, arg_seq, H, TILE_Q + ) + if live: + nblk, blk0, blk1 = segment_blocks( + prefix, row0, seq_len, seg, i32_nseg, H, TILE_Q, GROUP + ) + if blk0 < blk1: + qr, sr = score_resources( + arg_q, arg_score, i32_total_q, i32_hs, i32_S, H + ) + page_of = page_lookup(arg_bt, b * i32_bt_stride, nblk) + + # Lane l of wave w owns token 32 w + l / 2, dims 64 (l % 2) .. + # +63. The DMA writes LDS lane-linearly, so each lane reads back + # its own code and scale bytes. + raw_base_s = fx.Int32( + fx.rocdl.readfirstlane(T.i32, fx.Int32(fx.ptrtoint(smem.raw.ptr))) + ) + code_voff = [wave * 2048 + lane * 32 + h * 16 for h in range(2)] + scale_voff = PAGE_CODES + wave * 256 + lane * 4 + + def dma_raw(page, st): + kr = buffer_ops.create_buffer_resource_from_addr( + arg_kv + fx.Int64(page) * i64_page_stride, + num_records_bytes=PAGE_BYTES, + ) + base = raw_base_s + (st * NW + wave) * STAGE_W + for h in range_constexpr(2): + _dma(kr, base + h * 1024, code_voff[h], 16) + _dma(kr, base + 2048, scale_voff, 4) + + raw_t = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.raw.ptr), + fx.make_layout(RAW_SLOTS * NW * STAGE_W // 4, 1), + ), + fx.make_layout(4, 1), + ) + raw_d = smem.raw.ptr + + def read_raw(st): + """This lane's 8 code dwords and its scale dword.""" + b4 = (st * NW + wave) * (STAGE_W // 4) + out = [] + for h in range_constexpr(2): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + fx.copy( + fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32), + fx.slice(raw_t, (None, b4 // 4 + h * 64 + lane)), + r, + ) + v = Vec(r.load()) + out += [fx.Int32(v[i]) for i in range_constexpr(4)] + sc = fx.Int32( + fx.make_view( + fx.recast_iter(fx.Int32, raw_d), + fx.make_layout(RAW_SLOTS * NW * STAGE_W // 4, 1), + )[b4 + 512 + lane] + ) + return out + [sc] + + pg = [page_of(blk0 + i) for i in range_constexpr(2)] + for i in range_constexpr(2): + dma_raw(pg[i], i) + pg_next = page_of(blk0 + 2) + + lut_g = global_i32_ptr(arg_lut) + lut_l = smem.lut.ptr + for i in range_constexpr(2): + lut_l[tx * 2 + i] = fx.Int32(lut_g[tx * 2 + i]) + + # Rows past q_len * H read the next request's Q and are masked at + # the store. + rows = [row0 + wave * ROWS_PER_WAVE + ni * 16 + l16 for ni in range(NI)] + qpos = [prefix + r // H for r in rows] + # Q^T in the key's dim order (even dims, then odd dims per 8). + qf = [] + for ni in range_constexpr(NI): + qrow = (seq_start * H + rows[ni]) * ROW_BYTES + halves = [] + for h in range_constexpr(2): + d = Vec( + buffer_ops.buffer_load( + qr, + (qrow + h * 64 + q16 * 16) // 4, + vec_width=4, + dtype=fx.Int32, + ) + ) + d = [fx.Int32(d[i]) for i in range_constexpr(4)] + halves.append(key_order4(d)) + qf.append(pack8(halves[0], halves[1])) + + lds16 = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.a.ptr), + fx.make_layout(LDS_BYTES // 4, 1), + ), + fx.make_layout(4, 1), + ) + lds_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + lut8 = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.lut.ptr), fx.make_layout(512, 1) + ), + fx.make_layout(2, 1), + ) + lut_atom = fx.make_copy_atom(fx.UniversalCopy64b(), fx.Int32) + + def lds_load16(tile): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + fx.copy(lds_atom, fx.slice(lds16, (None, tile)), r) + return r.load() + + def lds_store16(tile, vals): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + r.store(Vec.from_elements(vals, fx.Int32)) + fx.copy(lds_atom, r, fx.slice(lds16, (None, tile))) + + def lut_load(sbyte): + r = fx.make_rmem_tensor(fx.make_layout(2, 1), fx.Int32) + fx.copy(lut_atom, fx.slice(lut8, (None, sbyte)), r) + v = Vec(r.load()) + return fx.Int32(v[0]), fx.Int32(v[1]) + + # 16 B column c of key row t sits at c ^ ((t >> 1) & 7). + t_row = wave * 32 + (lane >> 1) + hh = lane & 1 + wr_base = t_row * RS_T + wr_sw = (t_row >> 1) & 7 + + def dq_tables(raw): + sc = raw[8] + return [ + lut_load((sc >> (8 * g)) & 0xFF) for g in range_constexpr(4) + ] + + def dq_step(raw, tabs, dslot, k, pend): + """Code dword k -> 8 keys; every second one stores 16 B.""" + lo, hi = dequant8(raw[k], *tabs[k // 2]) + if const_expr(k % 2 == 0): + return [lo, hi] + col = (hh * 4 + k // 2) ^ wr_sw + lds_store16(dslot * SLOT_T + wr_base + col, pend + [lo, hi]) + return [] + + zero4 = fx.Vector.filled(4, 0.0, fx.Float32) + neg_inf = fx.Float32(float("-inf")) + rd_row = l16 * RS_T + rd_col = [q16 ^ (l16 >> 1), (4 + q16) ^ (l16 >> 1)] + + def compute_block(slot, blk, masked, st, dslot): + """This wave's scores of ``blk``; dequantizes the next block's + raw page ``st`` into key slot ``dslot`` on the way.""" + if const_expr(masked): + vis = [ + qpos[ni] - (blk * BLK + q16 * 4) + for ni in range_constexpr(NI) + ] + run = [neg_inf for _ in range(NI)] + raw = read_raw(st) + tabs = dq_tables(raw) + + def read_a(mt): + base = slot * SLOT_T + mt * 16 * RS_T + rd_row + return pack8( + lds_load16(base + rd_col[0]), lds_load16(base + rd_col[1]) + ) + + accs = [[None] * NI, [None] * NI] + + def issue(mt, a): + acc = accs[mt % 2] + for ni in range_constexpr(NI): + acc[ni] = mfma_fp8_16x16x128(a, qf[ni], zero4) + + def reduce(mt): + acc = accs[mt % 2] + for ni in range_constexpr(NI): + v = acc[ni] + for i in range_constexpr(4): + x = v[i] + if const_expr(masked): + x = (vis[ni] >= mt * 16 + i).select(x, neg_inf) + run[ni] = maxf_nn(run[ni], x) + + pend: list = [] + a_next = read_a(0) + for mt in range_constexpr(MT): + a = a_next + if const_expr(mt + 1 < MT): + a_next = read_a(mt + 1) + issue(mt, a) + if const_expr(mt >= 1): + reduce(mt - 1) + ks = DQ_AT.get(mt, ()) + for k in ks: + pend = dq_step(raw, tabs, dslot, k, pend) + n_valu = (16 if not masked else 48) + 10 * len(ks) + n_dsw = sum(k % 2 for k in ks) + for i in range_constexpr(NI): + fx.rocdl.sched_group_barrier(0x008, 1, 0) + fx.rocdl.sched_group_barrier( + 0x002, (n_valu + NI - 1) // NI, 0 + ) + if const_expr(i < 2): + fx.rocdl.sched_group_barrier(0x100, 1, 0) + if const_expr(i == 4 and n_dsw): + fx.rocdl.sched_group_barrier(0x200, n_dsw, 0) + fx.rocdl.sched_barrier(0) + reduce(MT - 1) + out = [] + for p in range_constexpr(NI // 2): + x, y = _xlane_max4_pair(run[2 * p], run[2 * p + 1]) + out += [x, y] + return out + + stage16 = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.stage.ptr), + fx.make_layout(STAGE_T * 4, 1), + ), + fx.make_layout(4, 1), + ) + + def stage_tile(k, h): + return ((k * NW + wave) * 2 + h) * 64 + lane + + def store_group(g0, sc): + """Scores sc[j][ni] of blocks g0 .. g0 + 3, sentinels applied. + + Lane q16 = k stores rows 2k and 2k + 1. A group waits in LDS + unless it ends its SPAN piece or the segment; that one stores + the piece's staged groups and itself back to back. + """ + full = g0 + GROUP <= blk1 + pos = (g0 >> 2) & (SPAN - 1) + piece0 = g0 & ~(SPAN * GROUP - 1) + first = (piece0 < blk0).select( + (blk0 >> 2) & (SPAN - 1), fx.Int32(0) + ) + flush = (pos == SPAN - 1) | (g0 + GROUP >= blk1) + + def store_now(vals, off, valid): + """16 B per row, single dwords for a partial tail.""" + if full: + buffer_ops.buffer_store( + fx.Vector.from_elements(vals, fx.Float32), + sr, + off, + mask=valid, + ) + else: + for j in range_constexpr(GROUP): + buffer_ops.buffer_store( + vals[j], sr, off + j, mask=valid & (g0 + j < blk1) + ) + + def flush_piece(h, vals, off, valid): + for k in range_constexpr(SPAN - 1): + t = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + fx.copy( + lds_atom, fx.slice(stage16, (None, stage_tile(k, h))), t + ) + buffer_ops.buffer_store( + Vec(t.load()), + sr, + off - (pos - k) * GROUP, + mask=valid & (first <= k) & (pos > k), + ) + store_now(vals, off, valid) + + def stage_group(h, vals): + t = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + t.store(Vec.from_elements([bits(x) for x in vals], fx.Int32)) + fx.copy( + lds_atom, t, fx.slice(stage16, (None, stage_tile(pos, h))) + ) + + for h in range_constexpr(2): + row = rows[h] + vals = [sc[j][h] for j in range_constexpr(GROUP)] + for k in range_constexpr(1, NW): + is_k = q16 == k + row = is_k.select(rows[2 * k + h], row) + vals = [ + is_k.select(sc[j][2 * k + h], vals[j]) + for j in range_constexpr(GROUP) + ] + valid = row < q_rows + tok = row // H + nb_row = (prefix + tok) // BLK + 1 + for j in range_constexpr(GROUP): + blk = g0 + j + vals[j] = with_sentinels( + vals[j], blk, nb_row, i32_init, i32_local + ) + off = (row % H) * i32_hs + (seq_start + tok) * i32_S + g0 + if flush: + flush_piece(h, vals, off, valid) + else: + stage_group(h, vals) + + n_groups = (blk1 - blk0 + (GROUP - 1)) // GROUP + # Groups every row of the tile sees whole need no causal mask. + v = prefix + row0 // H - (BLK - 1) + last_full = (v < 0).select(fx.Int32(-1), v // BLK) + d = last_full - (GROUP - 1) - blk0 + nf = (d < 0).select(fx.Int32(0), d // GROUP + 1) + n_full = (nf < n_groups).select(nf, n_groups) + + # Prologue: block blk0 dequantized into slot 0. + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0) + fx.gpu.barrier() + raw0 = read_raw(0) + tabs0 = dq_tables(raw0) + pend: list = [] + for k in range_constexpr(8): + pend = dq_step(raw0, tabs0, 0, k, pend) + fx.gpu.barrier() + + def group(g0, pg_next, masked): + """Blocks g0 .. g0 + 3: per block b, DMA block b + 2's page, + score b while dequantizing b + 1, then wait and barrier.""" + sc = [] + for j in range_constexpr(GROUP): + blk = g0 + j + dma_raw(pg_next, (j + 2) % RAW_SLOTS) + pg_next = page_of(blk + 3) + sc.append( + compute_block( + j % N_SLOTS, + blk, + masked, + (j + 1) % RAW_SLOTS, + (j + 1) % N_SLOTS, + ) + ) + fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0) + fx.gpu.barrier() + store_group(g0, sc) + return pg_next + + def one(st): + return fx.Int32(st[0] if isinstance(st, (list, tuple)) else st) + + for g, st in range( # type: ignore[call-overload] + fx.Index(0), fx.Index(n_full), fx.Index(1), init=[_ir(pg_next)] + ): + pgn = group(blk0 + fx.Int32(g) * GROUP, one(st), False) + res = yield [_ir(pgn)] + for g, st in range( # type: ignore[call-overload] + fx.Index(n_full), + fx.Index(n_groups), + fx.Index(1), + init=[_ir(one(res))], + ): + pgn = group(blk0 + fx.Int32(g) * GROUP, one(st), True) + res = yield [_ir(pgn)] + + @flyc.jit + def launch( + arg_q: fx.Int64, + arg_kv: fx.Int64, + arg_lut: fx.Int64, + arg_score: fx.Int64, + arg_bt: fx.Int64, + arg_cu: fx.Int64, + arg_seq: fx.Int64, + i64_page_stride: fx.Int64, + i32_bt_stride: fx.Int32, + i32_S: fx.Int32, + i32_nseg: fx.Int32, + i32_qt: fx.Int32, + i32_total_q: fx.Int32, + i32_nb: fx.Int32, + i32_init: fx.Int32, + i32_local: fx.Int32, + i32_hs: fx.Int32, + i32_grid: fx.Int32, + stream: fx.Stream, + ): + kernel( + arg_q, + arg_kv, + arg_lut, + arg_score, + arg_bt, + arg_cu, + arg_seq, + i64_page_stride, + i32_bt_stride, + i32_S, + i32_nseg, + i32_qt, + i32_total_q, + i32_nb, + i32_init, + i32_local, + i32_hs, + # Without the hint LLVM picks 264 VGPRs, one wave per SIMD. + value_attrs={"rocdl.waves_per_eu": 2}, + ).launch(grid=(fx.Int64(i32_grid), 1, 1), block=(64 * NW, 1, 1), stream=stream) + + launch.kernel_name = name + return launch diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py new file mode 100644 index 000000000000..390dbc075c0f --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 prefill index scorer for small work (FlyDSL, gfx950). + +The scores of ``kernel.py`` with one dequantized block per workgroup: a +workgroup is 4 waves x 32 MQ query rows x a segment of the causal window. Per +block, wave w dequantizes token tile w into e4m3 A fragments in LDS (double +buffered, one barrier per block) while the MFMAs of the previous block run; +every wave multiplies all four token tiles against its query rows with +v_mfma_f32_32x32x64_f8f6f4. Blocks before the wave's first query position skip +the causal mask in a loop of their own. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops +from flydsl._mlir import ir as _mlir_ir +from flydsl._mlir.dialects import arith as _arith +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr import rocdl as _rocdl +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + as_f32, + bits, + cvt_fp4x2_f16, + cvt_fp8x2_f16, + maxf_nn, + permlane_swap, +) +from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.kernel import ( + BLK, + PAGE_BYTES, + PAGE_CODES, + D, + page_lookup, + score_resources, + segment_blocks, + tile_rows, + with_sentinels, +) + +NW = 4 +MQ_CHOICES = (1, 2) # 32-row query tiles per wave +TILES = BLK // 32 # 32-token tiles per block +KS = 2 # MFMA k-steps per key row +FL = 32 # fragment bytes per lane +FRAG_BYTES = TILES * KS * 64 * FL # one block's A fragments +RING = 4 # fragment reads in flight ahead of the MFMA + + +def _sat448_fp8x2(x): + """Clamp 2 x f16 to +-448; the hardware fp8 convert does not saturate.""" + fm = _mlir_ir.Attribute.parse("#arith.fastmath") + hi = fx.as_ir_value(fx.Vector.filled(2, 448.0, fx.Float16)) + lo = fx.as_ir_value(fx.Vector.filled(2, -448.0, fx.Float16)) + v = _arith.MaxNumFOp(fx.as_ir_value(x), lo, fastmath=fm).result + return _arith.MinNumFOp(v, hi, fastmath=fm).result + + +def _fp8_pair(w, sv): + """Code dword w x splat scale sv -> 2 dwords of e4m3.""" + out = [] + for h in range_constexpr(2): + p = fx.as_ir_value(fx.Vector.filled(2, 0, fx.Int16)) + for k in range_constexpr(2): + x = _sat448_fp8x2(cvt_fp4x2_f16(w, 2 * h + k) * sv) + p = fx.rocdl.cvt_scalef32_pk_fp8_f16( + Vec.make_type(2, fx.Int16), + p, + x, + fx.as_ir_value(fx.Float32(1.0)), + k == 1, + ) + out.append(fx.Int32(Vec(p).bitcast(fx.Int32)[0])) + return out + + +def _mfma_fp8_32x32x64(a8, b8, acc16): + """v_mfma_f32_32x32x64_f8f6f4 on e4m3; lane (r, h) gets + D[8 (i // 4) + 4 h + i % 4][r].""" + return Vec( + fx.rocdl.mfma_scale_f32_32x32x64_f8f6f4( + Vec.make_type(16, fx.Float32), [a8, b8, acc16, 0, 0, 0, 0, 0, 0] + ) + ) + + +def _xlane_max32(x): + """max(x[L], x[L ^ 32]).""" + a, b = permlane_swap(_rocdl.permlane32_swap, bits(x), bits(x)) + return maxf_nn(as_f32(a), as_f32(b)) + + +def tile_q(mq: int) -> int: + """Query rows per workgroup.""" + return NW * 32 * mq + + +def compile_score_prefill_shared(heads: int = 1, mq: int = 2): + """Build the scorer for 1, 2 or 4 index heads and ``mq`` 32-row query tiles + per wave.""" + assert heads in (1, 2, 4) and mq in MQ_CHOICES, (heads, mq) + H, MQ = heads, mq + ROWS_PER_WAVE = 32 * MQ + TILE_Q = tile_q(MQ) + name = f"m3_nvfp4_index_score_prefill_fp8_shared_tq{TILE_Q}_h{H}" + + @fx.struct + class Shared: + frag: fx.Array[fx.Uint8, 2 * FRAG_BYTES, 16] # A fragments of 2 blocks + + @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1]) + def kernel( + arg_q: fx.Int64, # [total_q, H, 128] e4m3 + arg_kv: fx.Int64, # [num_blocks, 128, 72] uint8 + arg_score: fx.Int64, # [H, total_q, S] fp32 + arg_bt: fx.Int64, # [batch, >= max_blocks] int32 + arg_cu: fx.Int64, # cu_seqlens_q [batch + 1] int32 + arg_seq: fx.Int64, # seq_lens [batch] int32 + i64_page_stride: fx.Int64, + i32_bt_stride: fx.Int32, + i32_S: fx.Int32, + i32_nseg: fx.Int32, + i32_qt: fx.Int32, # query tiles per request + i32_total_q: fx.Int32, + i32_nb: fx.Int32, + i32_init: fx.Int32, + i32_local: fx.Int32, + i32_hs: fx.Int32, # score head stride + ): + smem = fx.SharedAllocator().allocate(Shared).peek() + tx, pid = fx.thread_idx.x, fx.block_idx.x + lane = tx % 64 + wave = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64)) + r32, h = lane % 32, lane // 32 + live, b, seg, seq_start, seq_len, prefix, q_rows, row0 = tile_rows( + pid, i32_qt, i32_nseg, i32_nb, arg_cu, arg_seq, H, TILE_Q + ) + if live: + nblk, blk0, blk1 = segment_blocks( + prefix, row0, seq_len, seg, i32_nseg, H, TILE_Q, 1 + ) + if blk0 < blk1: + qr, sr = score_resources( + arg_q, arg_score, i32_total_q, i32_hs, i32_S, H + ) + page_of = page_lookup(arg_bt, b * i32_bt_stride, blk1) + + tok = wave * 32 + r32 + + def load_raw(blk): + """Token tok's codes and scales of dims 64 h .. +63; zeros past + the segment.""" + ok = blk < blk1 + kr = buffer_ops.create_buffer_resource_from_addr( + arg_kv + fx.Int64(page_of(blk)) * i64_page_stride, + num_records_bytes=PAGE_BYTES, + ) + w = [] + for i in range_constexpr(2): + v = Vec( + buffer_ops.buffer_load( + kr, + (tok * 64 + h * 32 + i * 16) // 4, + vec_width=4, + dtype=fx.Int32, + mask=ok, + ) + ) + w += [fx.Int32(v[e]) for e in range_constexpr(4)] + sc = fx.Int32( + buffer_ops.buffer_load( + kr, + (PAGE_CODES + tok * 8 + h * 4) // 4, + vec_width=1, + dtype=fx.Int32, + mask=ok, + ) + ) + return w + [sc] + + frag128 = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.frag.ptr), + fx.make_layout(2 * FRAG_BYTES // 4, 1), + ), + fx.make_layout(4, 1), + ) + atom16 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + + def frag_slot(buf, n, s, half=0): + return ( + buf * (FRAG_BYTES // 16) + + ((n * KS + s) * 64 + lane) * (FL // 16) + + half + ) + + def store16(words, slot): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + r.store(Vec.from_elements(words, fx.Int32)) + fx.copy(atom16, r, fx.slice(frag128, (None, slot))) + + def load16(slot): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + fx.copy(atom16, fx.slice(frag128, (None, slot)), r) + return [fx.Int32(x) for x in Vec(r.load())] + + def dequant_step(raw, buf, s): + """K-step s (dims 64 h + 32 s .. +31) of this wave's tile.""" + scl = cvt_fp8x2_f16(raw[8], s == 1) + sv = [ + Vec.from_elements([scl[g], scl[g]], fx.Float16) for g in (0, 1) + ] + k = [] + for e in range_constexpr(4): + k += _fp8_pair(raw[4 * s + e], sv[e // 2]) + store16(k[:4], frag_slot(buf, wave, s, 0)) + store16(k[4:], frag_slot(buf, wave, s, 1)) + + def read_frag(buf, n, s): + return Vec.from_elements( + load16(frag_slot(buf, n, s, 0)) + + load16(frag_slot(buf, n, s, 1)), + fx.Int32, + ) + + # Rows past q_len * H are zero and never stored. + rows = [row0 + wave * ROWS_PER_WAVE + 32 * m + r32 for m in range(MQ)] + qpos = [prefix + rw // H for rw in rows] + qf = [] + for m in range_constexpr(MQ): + qrow = (seq_start * H + rows[m]) * D + d = [] + for i in range_constexpr(4): + v = Vec( + buffer_ops.buffer_load( + qr, + (qrow + h * 64 + i * 16) // 4, + vec_width=4, + dtype=fx.Int32, + mask=rows[m] < q_rows, + ) + ) + d += [fx.Int32(v[e]) for e in range_constexpr(4)] + qf.append( + [ + Vec.from_elements(d[8 * s : 8 * s + 8], fx.Int32) + for s in range(KS) + ] + ) + + zero16 = fx.Vector.filled(16, 0.0, fx.Float32) + neg_inf = fx.Float32(float("-inf")) + st_tok = [rw // H for rw in rows] + st_out = [ + (rw % H) * i32_hs + (seq_start + t) * i32_S + for rw, t in zip(rows, st_tok) + ] + st_nb = [qp // BLK + 1 for qp in qpos] + qpos_min = prefix + (row0 + wave * ROWS_PER_WAVE) // H + + def max16(acc): + v = acc[0] + for i in range_constexpr(1, 16): + v = maxf_nn(v, acc[i]) + return v + + def max16_masked(acc, lim): + v = neg_inf + for i in range_constexpr(16): + x = (lim >= 8 * (i // 4) + i % 4).select(acc[i], neg_inf) + v = maxf_nn(v, x) + return v + + def block(blk, buf, raw_next, masked): + """Scores of block blk for this wave's rows; dequantizes + raw_next into the other buffer meanwhile.""" + runs = [neg_inf] * MQ + jobs = [(n, s) for n in range(TILES) for s in range(KS)] + ring = [None] * len(jobs) + for j in range_constexpr(RING - 1): + ring[j] = read_frag(buf, *jobs[j]) + for n in range_constexpr(TILES): + accs = [zero16] * MQ + for s in range_constexpr(KS): + j = n * KS + s + if const_expr(j + RING - 1 < len(jobs)): + ring[j + RING - 1] = read_frag(buf, *jobs[j + RING - 1]) + if const_expr(j % TILES == 0): + dequant_step(raw_next, 1 - buf, j // TILES) + for m in range_constexpr(MQ): + accs[m] = _mfma_fp8_32x32x64(ring[j], qf[m][s], accs[m]) + base = blk * BLK + 32 * n + 4 * h + for m in range_constexpr(MQ): + if const_expr(masked): + mx = max16_masked(accs[m], qpos[m] - base) + else: + mx = max16(accs[m]) + runs[m] = maxf_nn(runs[m], mx) + for m in range_constexpr(MQ): + v = _xlane_max32(runs[m]) + v = with_sentinels(v, blk, st_nb[m], i32_init, i32_local) + buffer_ops.buffer_store( + v, + sr, + st_out[m] + blk, + mask=(rows[m] < q_rows) & (h == 0), + ) + + raw0 = load_raw(blk0) + for s in range_constexpr(KS): + dequant_step(raw0, 0, s) + raw = load_raw(blk0 + 1) + fx.gpu.barrier() + + # Blocks below dsplit need no causal mask. Two loops rather than a + # branch per block keep the wave within 256 registers; waves split + # at different blocks but pass the same barriers. + dfirst = (qpos_min + 1) // BLK + dsplit = (dfirst < blk0).select(blk0, dfirst) + dsplit = (dsplit > blk1).select(blk1, dsplit) + blk = blk0 + buf = fx.Int32(0) + while blk < dsplit: + nxt = load_raw(blk + 2) + block(blk, buf, raw, False) + raw = nxt + fx.gpu.barrier() + blk = blk + 1 + buf = 1 - buf + while blk < blk1: + nxt = load_raw(blk + 2) + block(blk, buf, raw, True) + raw = nxt + fx.gpu.barrier() + blk = blk + 1 + buf = 1 - buf + + @flyc.jit + def launch( + arg_q: fx.Int64, + arg_kv: fx.Int64, + arg_score: fx.Int64, + arg_bt: fx.Int64, + arg_cu: fx.Int64, + arg_seq: fx.Int64, + i64_page_stride: fx.Int64, + i32_bt_stride: fx.Int32, + i32_S: fx.Int32, + i32_nseg: fx.Int32, + i32_qt: fx.Int32, + i32_total_q: fx.Int32, + i32_nb: fx.Int32, + i32_init: fx.Int32, + i32_local: fx.Int32, + i32_hs: fx.Int32, + i32_grid: fx.Int32, + stream: fx.Stream, + ): + kernel( + arg_q, + arg_kv, + arg_score, + arg_bt, + arg_cu, + arg_seq, + i64_page_stride, + i32_bt_stride, + i32_S, + i32_nseg, + i32_qt, + i32_total_q, + i32_nb, + i32_init, + i32_local, + i32_hs, + ).launch(grid=(fx.Int64(i32_grid), 1, 1), block=(64 * NW, 1, 1), stream=stream) + + launch.kernel_name = name + return launch diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/__init__.py b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/__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/minimax_m3/amd/ops/nvfp4_prefill/host.py b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/host.py new file mode 100644 index 000000000000..27052b662673 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/host.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch wrapper of the FlyDSL NVFP4 sparse prefill attention.""" + +import functools + +import numpy as np +import torch + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + LOG2E, + cu_count, + run_compiled, + unit_scale, +) + +from . import kernel as _kernel +from .kernel import BLK, HG, TOPK, D + + +@functools.cache +def _get_kernel(rounds: int): + return _kernel.compile_nvfp4_prefill_fp16(rounds) + + +def supports_nvfp4_sparse_prefill( + num_heads: int, num_kv_heads: int, head_dim: int, topk_blocks: int, block_size: int +) -> bool: + from vllm.platforms.rocm import on_gfx950 + + return ( + on_gfx950() + and head_dim == D + and num_heads == num_kv_heads * HG + and topk_blocks == TOPK + and block_size == BLK + ) + + +def minimax_m3_nvfp4_sparse_prefill( + q: torch.Tensor, # [T, H, 128] bf16 + kv_cache: torch.Tensor, # [pages, 2 * Hkv, 128, 72] uint8 + topk: torch.Tensor, # [Hkv, T, 16] int32 logical blocks, -1 = none + block_table: torch.Tensor, # [reqs, max_blocks] int32 + query_req_id: torch.Tensor, # [T] int32 + query_abs_pos: torch.Tensor, # [T] int32 + num_kv_heads: int, + sm_scale: float, + out: torch.Tensor, # [T, H, 128] bf16 + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, +) -> None: + """Token t attends to the causal tokens of its ``topk`` blocks.""" + n_tok, n_heads, head_dim = q.shape + if n_tok == 0: + return + assert head_dim == D and n_heads == num_kv_heads * HG + assert q.dtype == torch.bfloat16 and out.dtype == torch.bfloat16 + assert out.shape == q.shape + assert q.is_contiguous() and out.is_contiguous() + assert kv_cache.dtype == torch.uint8 and tuple(kv_cache.shape[1:]) == ( + 2 * num_kv_heads, + BLK, + 72, + ) + assert kv_cache.stride(1) == BLK * 72 and kv_cache.stride(3) == 1 + assert topk.dtype == torch.int32 and tuple(topk.shape) == ( + num_kv_heads, + n_tok, + TOPK, + ) + assert topk.stride(2) == 1 and topk.stride(1) == TOPK + for t in (block_table, query_req_id, query_abs_pos): + assert t.dtype == torch.int32 and t.stride(-1) == 1 + assert kv_cache.stride(0) < 2**31 + unit = unit_scale(q.device) + ks = unit if k_scale is None else k_scale + vs = unit if v_scale is None else v_scale + c = float(np.float32(sm_scale) * np.float32(LOG2E)) + # Up to one workgroup per CU the one-round combine is ~1 us faster. + n_wg = (n_tok + 3) // 4 * num_kv_heads + rounds = 2 if n_wg > cu_count(q.device) else 1 + run_compiled( + _get_kernel(rounds), + q.data_ptr(), + kv_cache.data_ptr(), + topk.data_ptr(), + block_table.data_ptr(), + query_req_id.data_ptr(), + query_abs_pos.data_ptr(), + out.data_ptr(), + ks.data_ptr(), + vs.data_ptr(), + n_tok, + n_heads, + num_kv_heads, + topk.stride(0), + block_table.stride(0), + kv_cache.stride(0), + c, + torch.cuda.current_stream(q.device), + ) diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/kernel.py b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/kernel.py new file mode 100644 index 000000000000..bedf654f12c4 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_prefill/kernel.py @@ -0,0 +1,507 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 sparse prefill attention for MiniMax-M3 (FlyDSL, gfx950). + +A workgroup takes 4 query tokens x the 16 query heads of one KV head and walks +the union of the pages they selected. K and V are the exact e2m1 x e4m3 +products in fp16; Q and P are fp16; scores, softmax and accumulation are fp32; +the output is bf16. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + cvt_fp4x2_f16, + cvt_fp8x2_f16, + exp2, + global_i32_ptr, + maxf_nn, + mfma_16x16x32, + pk_bf16, + pk_f16, + xlane4, +) + +BLK = 128 # tokens per page +D = 128 +HG = 16 # query heads per KV head (the MFMA N dimension) +NW = 4 +KPW = BLK // NW # keys per wave +TOK_PER_WG = 4 +TOPK = 16 +MAXU = TOK_PER_WG * TOPK # union capacity: one lane per (token, slot) +DATA_B = BLK * D // 2 # e2m1 bytes per head per side +SCALE_B = BLK * D // 16 # e4m3 scale bytes per head per side +NDT = D // 16 # d-tiles of O^T +NKS = D // 32 # k-steps of S^T +RAW = 2 * 4 + 1 + 8 + 2 # NVFP4 dwords a lane loads per page +Q_ROW = D * 2 // 16 + 1 # 16 B LDS tiles per Q^T row, padded by one + + +def _bf16_to_f16(v): + return Vec(v).bitcast(fx.BFloat16).to(fx.Float32).to(fx.Float16).bitcast(fx.Int32) + + +def kv_lane_loader(arg_cache, page_bytes, kvh, hkv, w, n, g): + """The loader of this lane's NVFP4 K/V dwords of a page. + + S^T tile st row r is key 32 w + 8 (r // 4) + 4 st + r % 4, so lane + (n, g)'s P^T / V^T key slots are keys 32 w + 8 g + j. + """ + side_b = hkv * (DATA_B + SCALE_B) + k_data = kvh * DATA_B + k_scl = hkv * DATA_B + kvh * SCALE_B + v_data = side_b + kvh * DATA_B + v_scl = side_b + hkv * DATA_B + kvh * SCALE_B + k_rows = [w * KPW + 8 * (n // 4) + 4 * st + n % 4 for st in range_constexpr(2)] + k_load = [(k_data + r * (D // 2) + g * 16) // 4 for r in k_rows] + k_scale_load = (k_scl + ((w * 4 + g) * 16 + n) * 4) // 4 + v_load = (v_data + ((w * KPW + 8 * g) // 8 * 16 + n) * 32) // 4 + v_scale_load = (v_scl + (n // 2) * BLK + w * KPW + 8 * g) // 4 + + def load(page): + """This lane's RAW NVFP4 dwords of ``page``.""" + pr = buffer_ops.create_buffer_resource_from_addr( + arg_cache + fx.Int64(page) * fx.Int64(page_bytes), + num_records_bytes=page_bytes, + ) + raw = [] + for st in range_constexpr(2): + kd = fx.Vector( + buffer_ops.buffer_load(pr, k_load[st], vec_width=4, dtype=fx.Int32) + ) + raw += [kd[i] for i in range_constexpr(4)] + raw.append( + fx.Int32( + buffer_ops.buffer_load(pr, k_scale_load, vec_width=1, dtype=fx.Int32) + ) + ) + for h in range_constexpr(2): + vd = fx.Vector( + buffer_ops.buffer_load(pr, v_load + 4 * h, vec_width=4, dtype=fx.Int32) + ) + raw += [vd[i] for i in range_constexpr(4)] + vs = fx.Vector( + buffer_ops.buffer_load(pr, v_scale_load, vec_width=2, dtype=fx.Int32) + ) + raw += [vs[i] for i in range_constexpr(2)] + return raw + + return load + + +def compile_nvfp4_prefill_fp16(rounds: int = 2): + """Build the prefill kernel. + + Args: + rounds: rounds of the cross-wave combine. 2 halves its LDS so two + workgroups share a CU; 1 keeps Q^T in registers and combines all + tokens at once (grids of at most one workgroup per CU). + + """ + assert rounds in (1, 2) + tpr = TOK_PER_WG // rounds # tokens per combine round + wpe = rounds # waves per SIMD + name = f"m3_nvfp4_sparse_prefill_fp16matched_tq{TOK_PER_WG}_ksplit" + if rounds == 1: + name += "_r1" + + @fx.struct + class Shared: + # Q^T during the page loop, then the partial O per (wave, token, d-tile, + # lane) in 16 B tiles; (max, sum) per (wave, token, lane). + o: fx.Array[fx.Int32, NW * tpr * NDT * 64 * 4, 16] + ml: fx.Array[fx.Float32, NW * TOK_PER_WG * 2 * 64] + + @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1]) + def kernel( + arg_q: fx.Int64, # [T, H, 128] bf16 + arg_cache: fx.Int64, # [pages, 2 * Hkv, 128, 72] uint8 + arg_topk: fx.Int64, # [Hkv, T, 16] int32 logical blocks, -1 = none + arg_bt: fx.Int64, # [reqs, stride] int32 + arg_req: fx.Int64, # [T] int32 + arg_pos: fx.Int64, # [T] int32 + arg_out: fx.Int64, # [T, H, 128] bf16 + arg_ks: fx.Int64, # fp32 global K scale + arg_vs: fx.Int64, # fp32 global V scale + i32_T: fx.Int32, + i32_H: fx.Int32, + i32_hkv: fx.Int32, + i32_topk_hs: fx.Int32, # topk stride over KV heads + i32_bt_s: fx.Int32, + i32_page_byte: fx.Int32, + f32_c: fx.Float32, # sm_scale * log2(e) + ): + smem = fx.SharedAllocator().allocate(Shared).peek() + tx = fx.thread_idx.x + bidx, kvh = fx.block_idx.x, fx.block_idx.y + lane = tx % 64 + w = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64)) + n, g = lane % 16, lane // 16 + topk = global_i32_ptr(arg_topk) + bt = global_i32_ptr(arg_bt) + req_p = global_i32_ptr(arg_req) + pos_p = global_i32_ptr(arg_pos) + f32p = fx.PointerType.get( + T.f32, address_space=fx.AddressSpace.Global, alignment=4 + ) + c = f32_c * fx.Float32(fx.inttoptr(f32p, fx.Int64(arg_ks))[0]) + v_scale = fx.Float32(fx.inttoptr(f32p, fx.Int64(arg_vs))[0]) + neg_inf = fx.Float32(float("-inf")) + topk_base = kvh * i32_topk_hs + + lds = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.o.ptr), + fx.make_layout(NW * tpr * NDT * 64 * 4, 1), + ), + fx.make_layout(4, 1), + ) + lds_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + qr = buffer_ops.create_buffer_resource_from_addr( + arg_q, num_records_bytes=fx.Int64(i32_T) * fx.Int64(i32_H) * (D * 2) + ) + if const_expr(rounds == 2): + # Q^T in LDS; thread (token, head, quarter) stages 64 B. + tq, hq, qq = tx // 64, (tx // 4) % HG, tx % 4 + tok_q = bidx * TOK_PER_WG + tq + tok_q = (tok_q < i32_T).select(tok_q, fx.Int32(0)) + q_src = (tok_q * i32_H + kvh * HG + hq) * (D // 2) + qq * 16 + for i in range_constexpr(4): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + q4 = buffer_ops.buffer_load( + qr, q_src + 4 * i, vec_width=4, dtype=fx.Int32 + ) + r.store(_bf16_to_f16(q4)) + q_tile = (tq * HG + hq) * Q_ROW + qq * 4 + i + fx.copy(lds_atom, r, fx.slice(lds, (None, q_tile))) + + def q_frags(t): + """Token t's Q^T, one B operand per k-step.""" + if const_expr(rounds == 1): + return q_regs[t] + out = [] + for ks in range_constexpr(NKS): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + q_tile = (t * HG + n) * Q_ROW + g * 4 + ks + fx.copy(lds_atom, fx.slice(lds, (None, q_tile)), r) + out.append(Vec(r.load())) + return out + + # Per token: the page of its own block and its causal limit. + diag, lim, q_regs = [], [], [] + for t in range_constexpr(TOK_PER_WG): + tok = bidx * TOK_PER_WG + t + ok = tok < i32_T + tc = ok.select(tok, fx.Int32(0)) + pos = fx.Int32(fx.rocdl.readfirstlane(T.i32, fx.Int32(pos_p[tc]))) + req = fx.Int32(fx.rocdl.readfirstlane(T.i32, fx.Int32(req_p[tc]))) + diag.append( + fx.Int32( + fx.rocdl.readfirstlane( + T.i32, fx.Int32(bt[req * i32_bt_s + pos // BLK]) + ) + ) + ) + lim.append(pos % BLK - w * KPW) + if const_expr(rounds == 1): + q_row = (tc * i32_H + kvh * HG + n) * (D // 2) + g * 16 + q_regs.append( + [ + _bf16_to_f16( + buffer_ops.buffer_load( + qr, q_row + ks * 4, vec_width=4, dtype=fx.Int32 + ) + ) + for ks in range_constexpr(NKS) + ] + ) + + load_page = kv_lane_loader(arg_cache, i32_page_byte, kvh, i32_hkv, w, n, g) + + # Lane (token lane // 16, slot lane % 16): the page it selected, -1 none. + lane_tok = bidx * TOK_PER_WG + lane // TOPK + lane_tok_ok = lane_tok < i32_T + lane_tok_clamped = lane_tok_ok.select(lane_tok, fx.Int32(0)) + blk = fx.Int32(topk[topk_base + lane_tok_clamped * TOPK + lane % TOPK]) + lane_blk_ok = lane_tok_ok & (blk >= 0) + lane_page = fx.Int32( + bt[ + fx.Int32(req_p[lane_tok_clamped]) * i32_bt_s + + lane_blk_ok.select(blk, fx.Int32(0)) + ] + ) + lane_page = lane_blk_ok.select(lane_page, fx.Int32(-1)) + # The union's first page is the lowest valid lane's: load it before the + # dedup runs. + valid = fx.Int64(fx.rocdl.ballot(T.i64, lane_page >= 0)) + first = (valid != fx.Int64(0)).select( + fx.Int32(fx.math.cttz(valid).to(fx.Int32)), fx.Int32(0) + ) + + def page_at(mask): + """Page of the lowest lane left in ``mask``, the first page when empty.""" + idx = (mask != fx.Int64(0)).select( + fx.Int32(fx.math.cttz(mask).to(fx.Int32)), first + ) + pg = fx.Int32(fx.rocdl.readlane(T.i32, lane_page, idx)) + return (pg >= 0).select(pg, fx.Int32(0)) + + raw0 = load_page(page_at(valid)) + fx.rocdl.sched_barrier(0) + # Union: the first lane of each distinct page. + n_dup = fx.Int32(0) + for j in range_constexpr(MAXU - 1): + pj = fx.Int32(fx.rocdl.readlane(T.i32, lane_page, j)) + n_dup = n_dup + ((lane_page == pj) & (lane > j)).select( + fx.Int32(1), fx.Int32(0) + ) + umask = fx.Int64(fx.rocdl.ballot(T.i64, (lane_page >= 0) & (n_dup == 0))) + n_u = fx.Int32(fx.math.ctpop(umask).to(fx.Int32)) + + def dequant(raw): + """K frags kf[st][ks] and V^T frags vf[dt] of a page in fp16.""" + kf = [] + for st in range_constexpr(2): + ksc = cvt_fp8x2_f16(raw[8], st == 1) + row = [] + for ks in range_constexpr(NKS): + sv = Vec.from_elements([ksc[ks // 2], ksc[ks // 2]], fx.Float16) + row.append( + Vec.from_elements( + [ + (cvt_fp4x2_f16(raw[st * 4 + ks], i) * sv).bitcast( + fx.Int32 + )[0] + for i in range_constexpr(4) + ], + fx.Int32, + ) + ) + kf.append(row) + vs = [] + for h in range_constexpr(2): + vs += [ + cvt_fp8x2_f16(raw[17 + h], False), + cvt_fp8x2_f16(raw[17 + h], True), + ] + # A V byte holds two consecutive keys: already the operand's fp16 pair. + vf = [ + Vec.from_elements( + [ + (cvt_fp4x2_f16(raw[9 + dt], i) * vs[i]).bitcast(fx.Int32)[0] + for i in range_constexpr(4) + ], + fx.Int32, + ) + for dt in range_constexpr(NDT) + ] + return kf, vf + + zero4 = Vec.filled(4, 0.0, fx.Float32) + key_off = [ + 8 * g + 4 * st + e for st in range_constexpr(2) for e in range_constexpr(4) + ] + if const_expr(rounds == 2): + fx.gpu.barrier() # Q^T staged + # Iteration u dequantizes page u and loads page u + 1. + init = [umask] + init += [neg_inf for _ in range_constexpr(TOK_PER_WG)] + [ + fx.Float32(0.0) for _ in range_constexpr(TOK_PER_WG) + ] + init += [zero4 for _ in range_constexpr(TOK_PER_WG * NDT)] + init += raw0 + n_st = 1 + 2 * TOK_PER_WG + TOK_PER_WG * NDT + results = init + for u, st_ in range( # type: ignore[call-overload] + fx.Index(0), + fx.Index(n_u), + fx.Index(1), + init=[fx.as_ir_value(x) for x in init], + ): + rem = fx.Int64(st_[0]) + m = [fx.Float32(st_[1 + t]) for t in range_constexpr(TOK_PER_WG)] + l_run = [ + fx.Float32(st_[1 + TOK_PER_WG + t]) for t in range_constexpr(TOK_PER_WG) + ] + o = [ + [ + Vec(st_[1 + 2 * TOK_PER_WG + t * NDT + dt]) + for dt in range_constexpr(NDT) + ] + for t in range_constexpr(TOK_PER_WG) + ] + raw = [fx.Int32(st_[n_st + i]) for i in range_constexpr(RAW)] + page = page_at(rem) + rem = rem & (rem - fx.Int64(1)) + raw_next = load_page(page_at(rem)) + kf, vf = dequant(raw) + hits = fx.Int64(fx.rocdl.ballot(T.i64, lane_page == page)) + for t in range_constexpr(TOK_PER_WG): + m_t, l_t, o_t = m[t], l_run[t], o[t] + if (hits.shrui(fx.Int64(TOPK * t)) & fx.Int64(0xFFFF)) != fx.Int64(0): + qf = q_frags(t) + a = [] + for st in range_constexpr(2): + acc = zero4 + for ks in range_constexpr(NKS): + acc = mfma_16x16x32(kf[st][ks], qf[ks], acc, True) + a += [acc[e] for e in range_constexpr(4)] + if page == diag[t]: + a = [ + (key_off[i] <= lim[t]).select(a[i], neg_inf) + for i in range_constexpr(8) + ] + mx = a[0] + for i in range_constexpr(1, 8): + mx = maxf_nn(mx, a[i]) + # The running max is kept in scaled units (c > 0). + m_new = maxf_nn(m_t, xlane4(mx, maxf_nn) * c) + m_use = (m_new == neg_inf).select(fx.Float32(0.0), m_new) + alpha = exp2(m_t - m_use) + p = [exp2(fx.math.fma(x, c, -m_use)) for x in a] + ps = p[0] + for i in range_constexpr(1, 8): + ps = ps + p[i] + l_t = l_t * alpha + ps + pf = Vec.from_elements( + [pk_f16(p[2 * j], p[2 * j + 1]) for j in range_constexpr(4)], + fx.Int32, + ) + # alpha is exactly 1 unless some row's max grew. + if fx.Int64(fx.rocdl.ballot(T.i64, m_new != m_t)) != fx.Int64(0): + av = Vec.filled(4, alpha, fx.Float32) + o_t = [x * av for x in o_t] + o_t = [ + mfma_16x16x32(vf[dt], pf, o_t[dt], True) + for dt in range_constexpr(NDT) + ] + m_t = m_new + m[t], l_run[t], o[t] = m_t, l_t, o_t + state = ( + [rem] + m + l_run + [x for token_out in o for x in token_out] + raw_next + ) + results = yield [fx.as_ir_value(x) for x in state] + + # Combine the waves' partials, tpr tokens per round; wave w finalizes + # token w in round w // tpr. + ml = smem.ml.ptr + if const_expr(rounds == 2): + fx.gpu.barrier() + for t in range_constexpr(TOK_PER_WG): + ml[((w * TOK_PER_WG + t) * 2) * 64 + lane] = fx.Float32(results[1 + t]) + ml[((w * TOK_PER_WG + t) * 2 + 1) * 64 + lane] = fx.Float32( + results[1 + TOK_PER_WG + t] + ) + for rnd in range_constexpr(rounds): + for tt in range_constexpr(tpr): + for dt in range_constexpr(NDT): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + ri = 1 + 2 * TOK_PER_WG + (tpr * rnd + tt) * NDT + dt + r.store(Vec(results[ri]).bitcast(fx.Int32)) + fx.copy( + lds_atom, + r, + fx.slice(lds, (None, ((w * tpr + tt) * NDT + dt) * 64 + lane)), + ) + fx.gpu.barrier() + if w // tpr == rnd: + t = w + ms = [ + fx.Float32(ml[((ww * TOK_PER_WG + t) * 2) * 64 + lane]) + for ww in range_constexpr(NW) + ] + big = ms[0] + for ww in range_constexpr(1, NW): + big = maxf_nn(big, ms[ww]) + big = (big == neg_inf).select(fx.Float32(0.0), big) + wt = [ + fx.Float32(fx.math.exp2(ms[ww] - big)) for ww in range_constexpr(NW) + ] + l_tot = fx.Float32(0.0) + for ww in range_constexpr(NW): + l_w = ml[((ww * TOK_PER_WG + t) * 2 + 1) * 64 + lane] + l_tot = l_tot + fx.Float32(l_w) * wt[ww] + l_tot = xlane4(l_tot, lambda a, b: a + b) + inv = (l_tot > fx.Float32(0.0)).select(v_scale / l_tot, fx.Float32(0.0)) + acc = [zero4 for _ in range_constexpr(NDT)] + for ww in range_constexpr(NW): + wv = Vec.filled(4, wt[ww], fx.Float32) + for dt in range_constexpr(NDT): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + o_tile = ((ww * tpr + t - tpr * rnd) * NDT + dt) * 64 + lane + fx.copy(lds_atom, fx.slice(lds, (None, o_tile)), r) + acc[dt] = acc[dt] + Vec(r.load()).bitcast(fx.Float32) * wv + # Lane (n, g) holds dims [32 g, 32 g + 32) of head n in order. + tok = bidx * TOK_PER_WG + t + words = [ + [ + pk_bf16(acc[2 * j][e] * inv, acc[2 * j + 1][e] * inv) + for j in range_constexpr(4) + ] + for e in range_constexpr(4) + ] + outr = buffer_ops.create_buffer_resource_from_addr( + arg_out, + num_records_bytes=fx.Int64(i32_T) * fx.Int64(i32_H) * (D * 2), + ) + o_row = (tok * i32_H + kvh * HG + n) * (D // 2) + g * 16 + for e in range_constexpr(4): + val = Vec.from_elements(words[e], fx.Int32) + buffer_ops.buffer_store(val, outr, o_row + e * 4, mask=tok < i32_T) + fx.gpu.barrier() + + @flyc.jit + def launch( + arg_q: fx.Int64, + arg_cache: fx.Int64, + arg_topk: fx.Int64, + arg_bt: fx.Int64, + arg_req: fx.Int64, + arg_pos: fx.Int64, + arg_out: fx.Int64, + arg_ks: fx.Int64, + arg_vs: fx.Int64, + i32_T: fx.Int32, + i32_H: fx.Int32, + i32_hkv: fx.Int32, + i32_topk_hs: fx.Int32, + i32_bt_s: fx.Int32, + i32_page_b: fx.Int32, + f32_c: fx.Float32, + stream: fx.Stream, + ): + kernel( + arg_q, + arg_cache, + arg_topk, + arg_bt, + arg_req, + arg_pos, + arg_out, + arg_ks, + arg_vs, + i32_T, + i32_H, + i32_hkv, + i32_topk_hs, + i32_bt_s, + i32_page_b, + f32_c, + value_attrs={ + "rocdl.waves_per_eu": wpe, + "rocdl.flat_work_group_size": "256,256", + }, + ).launch( + grid=((i32_T + TOK_PER_WG - 1) // TOK_PER_WG, i32_hkv, 1), + block=(64 * NW, 1, 1), + stream=stream, + ) + + launch.kernel_name = name + return launch diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/__init__.py b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/__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/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py new file mode 100644 index 000000000000..d266c4e9d00c --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Launch wrapper of the FlyDSL NVFP4 sparse decode.""" + +import ctypes +import functools + +import numpy as np +import torch + +from vllm.logger import init_logger +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + LOG2E, + cu_count, + run_compiled, + unit_scale, +) +from vllm.models.minimax_m3.amd.ops.nvfp4_index_prefill.host import nvfp4_e4m3_lut + +from . import kernel as _kernel + +logger = init_logger(__name__) + +_HIP_DEVICE_ATTRIBUTE_NUMBER_OF_XCCS = 10018 +_SUPPORTED_XCD_COUNT = 8 + +MAX_SPLITS = 8 +# Launches with more (row, KV head) workgroup rows than this do not split. +MAX_SPLIT_ROWS = 1024 + + +@functools.cache +def _get_kernel(splits: int, topk: int, xcds: int): + return _kernel.compile_nvfp4_sparse_decode(splits, topk, xcds) + + +def _query_xcd_count(device_index: int) -> int: + """The device's XCD count, or zero when HIP cannot report it.""" + try: + get_attribute = ctypes.CDLL("libamdhip64.so").hipDeviceGetAttribute + get_attribute.argtypes = [ + ctypes.POINTER(ctypes.c_int), + ctypes.c_int, + ctypes.c_int, + ] + get_attribute.restype = ctypes.c_int + value = ctypes.c_int() + status = get_attribute( + ctypes.byref(value), _HIP_DEVICE_ATTRIBUTE_NUMBER_OF_XCCS, device_index + ) + return value.value if status == 0 and value.value > 0 else 0 + except (AttributeError, OSError, TypeError, ctypes.ArgumentError): + return 0 + + +@functools.cache +def _xcd_count(device_index: int) -> int: + """8 on gfx950 (the topology the XCD grid was tuned on), else 0.""" + from vllm.platforms.rocm import on_gfx950 + + queried = _query_xcd_count(device_index) + count = queried if on_gfx950() and queried == _SUPPORTED_XCD_COUNT else 0 + logger.info( + "MiniMax-M3 NVFP4 sparse decode: XCD-aware grid %s on device %d " + "(reported XCDs=%d)", + "on" if count else "off", + device_index, + queried, + ) + return count + + +@functools.cache +def _split_buffers(device: torch.device): + """Split partials and tickets; allocated in the warm-up, before capture.""" + part = torch.empty( + MAX_SPLIT_ROWS * MAX_SPLITS * 64 * _kernel.PART_F, + dtype=torch.float32, + device=device, + ) + return part, torch.zeros(MAX_SPLIT_ROWS, dtype=torch.int32, device=device) + + +def decode_splits(num_rows: int, num_blocks: int, cus: int) -> int: + """Splits of a row's blocks. + + The kernel is VALU-bound, so the cost model is rounds of workgroups per CU + times blocks per split plus about one block of fixed cost; the fewest splits + among the minima, each split at least two blocks. + """ + best_cost, best = None, 1 + for s in range(1, MAX_SPLITS + 1): + per = -(-num_blocks // s) + if s > 1 and per < 2: + break + cost = -(-(num_rows * s) // cus) * (per + 1) + if best_cost is None or cost < best_cost: + best_cost, best = cost, s + return best + + +def minimax_m3_nvfp4_sparse_decode_fly( + q: torch.Tensor, # [rows, H, 128] bf16 + kv_cache: torch.Tensor, # [pages, 2 * Hkv, 128, 72] uint8 + topk_idx: torch.Tensor, # [Hkv, rows, topk] int32 logical blocks, -1 = none + block_table: torch.Tensor, # [reqs, max_blocks] int32 + seq_lens: torch.Tensor, # [reqs] int32 + sm_scale: float, + output: torch.Tensor, # [rows, H, 128] bf16 + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + decode_query_len: int = 1, + splits: int | None = None, + xcd_count: int | None = None, +) -> None: + """Sparse decode attention over the NVFP4 K/V pages. + + Row i is query row i % decode_query_len of request i // decode_query_len + and attends to the tokens of its ``topk_idx`` blocks up to its position. + """ + rows, num_heads, head_dim = q.shape + if rows == 0: + return + num_kv_heads = kv_cache.shape[1] // 2 + topk = topk_idx.shape[-1] + assert rows == seq_lens.shape[0] * decode_query_len + assert head_dim == _kernel.D and num_heads == num_kv_heads * _kernel.HG + assert q.dtype == torch.bfloat16 and output.dtype == torch.bfloat16 + assert q.is_contiguous() and output.is_contiguous() + assert kv_cache.dtype == torch.uint8 and kv_cache.stride(1) == _kernel.BLK * 72 + assert kv_cache.stride(3) == 1 and kv_cache.stride(0) < 2**31 + assert topk_idx.dtype == torch.int32 and topk_idx.shape[0] == num_kv_heads + assert topk_idx.shape[1] >= rows + assert topk_idx.stride(2) == 1 and topk_idx.stride(1) == topk + for t in (block_table, seq_lens): + assert t.dtype == torch.int32 and t.stride(-1) == 1 + if xcd_count is None: + xcd_count = _xcd_count(q.device.index) + xcds = xcd_count if xcd_count and xcd_count > 1 else 0 + if xcds: + group = xcds * decode_query_len + grid_x = -(-rows // group) * group + else: + grid_x = rows + if splits is None: + wg_rows = grid_x * num_kv_heads + splits = ( + 1 + if wg_rows > MAX_SPLIT_ROWS + else decode_splits(rows * num_kv_heads, topk, cu_count(q.device)) + ) + assert splits == 1 or grid_x * num_kv_heads <= MAX_SPLIT_ROWS + part, cnt = _split_buffers(q.device) + unit = unit_scale(q.device) + ks = unit if k_scale is None else k_scale + vs = unit if v_scale is None else v_scale + c = float(np.float32(sm_scale) * np.float32(LOG2E)) + run_compiled( + _get_kernel(splits, topk, xcds), + q.data_ptr(), + kv_cache.data_ptr(), + topk_idx.data_ptr(), + block_table.data_ptr(), + seq_lens.data_ptr(), + output.data_ptr(), + ks.data_ptr(), + vs.data_ptr(), + nvfp4_e4m3_lut(q.device).data_ptr(), + part.data_ptr(), + cnt.data_ptr(), + part.numel() * 4, + rows, + decode_query_len, + num_heads, + num_kv_heads, + topk_idx.stride(0), + block_table.stride(0), + kv_cache.stride(0), + c, + grid_x, + torch.cuda.current_stream(q.device), + ) diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/kernel.py b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/kernel.py new file mode 100644 index 000000000000..96188426f25c --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/kernel.py @@ -0,0 +1,576 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 sparse decode attention for MiniMax-M3 (FlyDSL, gfx950). + +One workgroup = one query row x the 16 query heads of one KV head x one split of +the row's top-k blocks. K and V are dequantized in registers to +e4m3_rn_satfinite(e2m1 x block scale), the values AITER's FP8 sparse PA reads, +and widened to bf16; scores and softmax are fp32, P is bf16. Wave w owns keys +[32 w, 32 w + 32) of every block: + + S^T = K Q^T v_mfma_f32_16x16x32_bf16, A = K, B = Q^T + O^T += V^T P^T A = V^T, B = P^T straight from the S^T accumulators + +The waves combine through LDS. With ``splits`` > 1 every split stores its +unnormalized partial and the last split of a row to arrive combines them in +split order, so the result does not depend on arrival order. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops +from flydsl._mlir import ir as _ir +from flydsl._mlir.dialects import arith as _arith +from flydsl._mlir.dialects import llvm as _llvm +from flydsl.expr import const_expr, range_constexpr +from flydsl.expr.typing import T +from flydsl.expr.typing import Vector as Vec + +from vllm.models.minimax_m3.amd.ops.flydsl_utils import ( + PRELOAD_HINTS, + SC1, + as_f32, + bits, + cvt_fp4x2_f16, + cvt_fp8x2_f16, + exp2, + global_i32_ptr, + maxf_nn, + mfma_16x16x32, + perm, + pk_bf16, + xlane4, +) +from vllm.models.minimax_m3.amd.ops.nvfp4_prefill.kernel import ( + BLK, + HG, + KPW, + NDT, + NKS, + NW, + RAW, + D, + kv_lane_loader, +) + +PART_F = 36 # fp32 per (split, row, lane) partial: O (NDT x 4), max, sum, pad +SIGN4 = 0x80808080 - (1 << 32) +MAG4 = 0x07070707 + + +def _nnan(): + return _ir.Attribute.parse("#arith.fastmath") + + +def _bf16x2_e4m3(word, hi): + """E4M3 bytes 0, 1 (``hi``: 2, 3) of ``word`` -> one dword of two bf16.""" + return Vec( + fx.rocdl.cvt_scalef32_pk_bf16_fp8( + Vec.make_type(2, fx.BFloat16), + fx.as_ir_value(fx.Int32(word)), + fx.as_ir_value(fx.Float32(1.0)), + hi, + ) + ).bitcast(fx.Int32)[0] + + +def _k_bf16x8(codes, tab_lo, tab_hi): + """Eight K codes under one block scale -> four dwords of bf16 pairs. + + ``tab_lo`` / ``tab_hi`` hold the e4m3 keys of magnitudes 0-3 / 4-7 for the + block scale (``nvfp4_e4m3_lut``). + """ + even = perm(tab_hi, tab_lo, codes & MAG4) ^ ((codes << 4) & SIGN4) + odd = perm(tab_hi, tab_lo, (codes >> 4) & MAG4) ^ (codes & SIGN4) + lo = perm(odd, even, 0x05010400) + hi = perm(odd, even, 0x07030602) + return [ + _bf16x2_e4m3(lo, False), + _bf16x2_e4m3(lo, True), + _bf16x2_e4m3(hi, False), + _bf16x2_e4m3(hi, True), + ] + + +def _v_bf16x8(codes, scale_pairs): + """One V dword (eight keys of one dim) -> four dwords of bf16 pairs. + + The fp16 products are exact; they are clamped to +-448 and rounded once to + e4m3. + """ + big = Vec.filled(2, 448.0, fx.Float16) + small = Vec.filled(2, -448.0, fx.Float16) + f8 = [None, None] + for j in range_constexpr(4): + v = cvt_fp4x2_f16(codes, j) * scale_pairs[j] + v = _arith.MinNumFOp( + fx.as_ir_value(v), fx.as_ir_value(big), fastmath=_nnan() + ).result + v = _arith.MaxNumFOp(v, fx.as_ir_value(small), fastmath=_nnan()).result + old = Vec.filled(2, 0, fx.Int16) if j % 2 == 0 else f8[j // 2] + f8[j // 2] = Vec( + fx.rocdl.cvt_scalef32_pk_fp8_f16( + Vec.make_type(2, fx.Int16), + fx.as_ir_value(old), + v, + fx.as_ir_value(fx.Float32(1.0)), + j % 2 == 1, + ) + ) + out = [] + for h in range_constexpr(2): + word = f8[h].bitcast(fx.Int32)[0] + out += [_bf16x2_e4m3(word, False), _bf16x2_e4m3(word, True)] + return out + + +def _global_atomic_add(addr_i64, value): + """Relaxed agent-scope atomic add on a global i32; returns the old value.""" + return fx.Int32( + _llvm.AtomicRMWOp( + _llvm.AtomicBinOp.add, + global_i32_ptr(addr_i64).llvm_ptr, + fx.Int32(value).ir_value(), + _llvm.AtomicOrdering.monotonic, + syncscope="agent", + alignment=4, + ).result + ) + + +def compile_nvfp4_sparse_decode(splits: int = 1, topk: int = 16, xcds: int = 0): + """Build the decode kernel. + + Args: + splits: workgroups per (row, KV head), each over a contiguous share of + the row's blocks. + topk: top-k width (one lane per slot). + xcds: if > 0, workgroup w runs on XCD w % xcds and a request's rows and + a row's splits are placed on one XCD. + + """ + assert 1 <= topk <= 64 and splits >= 1 + SPLITS, TOPK, XCDS = splits, topk, xcds + # Split partials need sc1 unless all splits of a row run on one XCD. + PART_SCOPE = 0 if XCDS else SC1 + name = f"m3_nvfp4_sparse_decode_k{TOPK}_s{SPLITS}" + (f"_x{XCDS}" if XCDS else "") + + @fx.struct + class Shared: + o: fx.Array[fx.Int32, NW * NDT * 64 * 4, 16] + ml: fx.Array[fx.Float32, NW * 2 * 64] + lut: fx.Array[fx.Int32, 512] + ticket: fx.Array[fx.Int32, 4] + + @flyc.kernel(name=name, known_block_size=[64 * NW, 1, 1]) + def kernel( + i32_rows: fx.Int32, + i32_ql: fx.Int32, # rows per request + i32_topk_hs: fx.Int32, # topk stride over KV heads + i32_bt_s: fx.Int32, + i32_H: fx.Int32, + i32_hkv: fx.Int32, + arg_seq: fx.Int64, # [reqs] int32 + arg_topk: fx.Int64, # [Hkv, >= rows, TOPK] int32 blocks, -1 = none + arg_bt: fx.Int64, # [reqs, stride] int32 + arg_lut: fx.Int64, # [256] int64 K magnitude tables + arg_q: fx.Int64, # [rows, H, 128] bf16 + arg_cache: fx.Int64, # [pages, 2 * Hkv, 128, 72] uint8 + arg_ks: fx.Int64, # fp32 global K scale + arg_vs: fx.Int64, # fp32 global V scale + arg_out: fx.Int64, # [rows, H, 128] bf16 + arg_part: fx.Int64, # [grid rows * Hkv, S, 64, PART_F] fp32 + arg_cnt: fx.Int64, # [grid rows * Hkv] int32 tickets, zero + i64_part_bytes: fx.Int64, + i32_page_byte: fx.Int32, + f32_c: fx.Float32, # sm_scale * log2(e) + ): + smem = fx.SharedAllocator().allocate(Shared).peek() + tx = fx.thread_idx.x + bx, kvh = fx.block_idx.x, fx.block_idx.y + lane = tx % 64 + w = fx.Int32(fx.rocdl.readfirstlane(T.i32, tx // 64)) + n, g = lane % 16, lane // 16 + topk = global_i32_ptr(arg_topk) + bt = global_i32_ptr(arg_bt) + seq_p = global_i32_ptr(arg_seq) + f32p = fx.PointerType.get( + T.f32, address_space=fx.AddressSpace.Global, alignment=4 + ) + c = f32_c * fx.Float32(fx.inttoptr(f32p, fx.Int64(arg_ks))[0]) + v_scale = fx.Float32(fx.inttoptr(f32p, fx.Int64(arg_vs))[0]) + neg_inf = fx.Float32(float("-inf")) + lut_g = global_i32_ptr(arg_lut) + smem.lut.ptr[tx] = fx.Int32(lut_g[tx]) + smem.lut.ptr[tx + 256] = fx.Int32(lut_g[tx + 256]) + + if const_expr(XCDS > 0): + # grid.x is padded to XCDS x ql so a request's rows share an XCD. + xcd = bx % XCDS + slot = bx // XCDS + row = ((slot // i32_ql) * XCDS + xcd) * i32_ql + slot % i32_ql + else: + row = bx + row_ok = row < i32_rows + rowc = row_ok.select(row, fx.Int32(0)) + req = rowc // i32_ql + pos = fx.Int32(seq_p[req]) - i32_ql + rowc % i32_ql + # Padding rows (seq_len 0 or past i32_rows) see nothing. + pos = (row_ok & (pos >= 0)).select(pos, fx.Int32(-1)) + pos = fx.Int32(fx.rocdl.readfirstlane(T.i32, pos)) + req = fx.Int32(fx.rocdl.readfirstlane(T.i32, req)) + self_blk = (pos >= 0).select(pos // BLK, fx.Int32(-1)) + + # Lane = top-k slot. + slot_ok = lane < TOPK + blk = fx.Int32( + topk[kvh * i32_topk_hs + rowc * TOPK + slot_ok.select(lane, fx.Int32(0))] + ) + ok = slot_ok & (blk >= 0) & (blk <= self_blk) + page = fx.Int32(bt[req * i32_bt_s + ok.select(blk, fx.Int32(0))]) + page = ok.select(page, fx.Int32(-1)) + blk = ok.select(blk, fx.Int32(-1)) + vmask = fx.Int64(fx.rocdl.ballot(T.i64, ok)) + n_v = fx.Int32(fx.math.ctpop(vmask).to(fx.Int32)) + if const_expr(SPLITS > 1): + below = (fx.Int64(1) << fx.Int64(lane)) - fx.Int64(1) + rank = fx.Int32(fx.math.ctpop(vmask & below).to(fx.Int32)) + chunk = (n_v + (SPLITS - 1)) // SPLITS + lo = fx.block_idx.z * chunk + vmask = fx.Int64( + fx.rocdl.ballot(T.i64, ok & (rank >= lo) & (rank < lo + chunk)) + ) + n_v = fx.Int32(fx.math.ctpop(vmask).to(fx.Int32)) + + lim = pos % BLK - w * KPW + qr = buffer_ops.create_buffer_resource_from_addr( + arg_q, num_records_bytes=fx.Int64(i32_rows) * fx.Int64(i32_H) * (D * 2) + ) + q_row = (rowc * i32_H + kvh * HG + n) * (D // 2) + g * 16 + qf = [ + fx.Vector( + buffer_ops.buffer_load(qr, q_row + ks * 4, vec_width=4, dtype=fx.Int32) + ) + for ks in range_constexpr(NKS) + ] + + load_block = kv_lane_loader(arg_cache, i32_page_byte, kvh, i32_hkv, w, n, g) + + def dequant_k(raw): + kf = [] + for st in range_constexpr(2): + tabs = [] + for grp in range_constexpr(2): + sb = raw[8].shrui(fx.Int32(8 * (2 * st + grp))) & fx.Int32(0xFF) + tabs.append( + ( + fx.Int32(smem.lut.ptr[sb * 2]), + fx.Int32(smem.lut.ptr[sb * 2 + 1]), + ) + ) + kf.append( + [ + Vec.from_elements( + _k_bf16x8(raw[st * 4 + ks], *tabs[ks // 2]), fx.Int32 + ) + for ks in range_constexpr(NKS) + ] + ) + return kf + + def dequant_v(raw): + sp = [ + cvt_fp8x2_f16(raw[17 + j // 2], j % 2 == 1) for j in range_constexpr(4) + ] + return [ + Vec.from_elements(_v_bf16x8(raw[9 + dt], sp), fx.Int32) + for dt in range_constexpr(NDT) + ] + + fx.gpu.barrier() # LUT in LDS + + zero4 = Vec.filled(4, 0.0, fx.Float32) + key_off = [ + 8 * g + 4 * st + e for st in range_constexpr(2) for e in range_constexpr(4) + ] + first = (n_v > 0).select( + fx.Int32(fx.math.cttz(vmask).to(fx.Int32)), fx.Int32(0) + ) + + def slot_at(mask): + """The lowest slot left in ``mask``, the first one when empty.""" + return (mask != fx.Int64(0)).select( + fx.Int32(fx.math.cttz(mask).to(fx.Int32)), first + ) + + def page_of(slot): + pg = fx.Int32(fx.rocdl.readlane(T.i32, page, slot)) + return (pg >= 0).select(pg, fx.Int32(0)) + + def attend(kf, vf, cur_blk, live, m_t, l_t, o_t): + """Fold one block into (max, sum, O); a dead block is fully masked.""" + a = [] + for st in range_constexpr(2): + acc = zero4 + for ks in range_constexpr(NKS): + acc = mfma_16x16x32(kf[st][ks], qf[ks], acc, False) + a += [acc[e] for e in range_constexpr(4)] + diag = cur_blk == self_blk + dead = live == fx.Int32(0) + a = [ + (dead | (diag & (key_off[i] > lim))).select(neg_inf, a[i]) + for i in range_constexpr(8) + ] + mx = a[0] + for i in range_constexpr(1, 8): + mx = maxf_nn(mx, a[i]) + # The running max is kept in scaled units (c > 0). + m_new = maxf_nn(m_t, xlane4(mx, maxf_nn) * c) + m_use = (m_new == neg_inf).select(fx.Float32(0.0), m_new) + alpha = exp2(m_t - m_use) + p = [exp2(fx.math.fma(x, c, -m_use)) for x in a] + ps = p[0] + for i in range_constexpr(1, 8): + ps = ps + p[i] + pf = Vec.from_elements( + [pk_bf16(p[2 * j], p[2 * j + 1]) for j in range_constexpr(4)], + fx.Int32, + ) + av = Vec.filled(4, alpha, fx.Float32) + o_t = [ + mfma_16x16x32(vf[dt], pf, o_t[dt] * av, False) + for dt in range_constexpr(NDT) + ] + return m_new, l_t * alpha + ps, o_t + + # Two blocks per iteration, the next pair's loads in flight; an odd last + # block is masked. + init = [vmask, neg_inf, fx.Float32(0.0)] + init += [zero4 for _ in range_constexpr(NDT)] + init += load_block(page_of(slot_at(vmask))) + n_st = 3 + NDT + results = init + for u, st_ in range( # type: ignore[call-overload] + fx.Index(0), + fx.Index((n_v + 1) // 2), + fx.Index(1), + init=[fx.as_ir_value(x) for x in init], + ): + rem = fx.Int64(st_[0]) + m_t = fx.Float32(st_[1]) + l_t = fx.Float32(st_[2]) + o_t = [Vec(st_[3 + dt]) for dt in range_constexpr(NDT)] + raw_a = [fx.Int32(st_[n_st + i]) for i in range_constexpr(RAW)] + blk_a = fx.Int32(fx.rocdl.readlane(T.i32, blk, slot_at(rem))) + rem_b = rem & (rem - fx.Int64(1)) + raw_b = load_block(page_of(slot_at(rem_b))) + fx.rocdl.sched_barrier(0) + kf_a = dequant_k(raw_a) + vf_a = dequant_v(raw_a) + rem_n = rem_b & (rem_b - fx.Int64(1)) + fx.rocdl.sched_barrier(0) + raw_n = load_block(page_of(slot_at(rem_n))) + fx.rocdl.sched_barrier(0) + res = attend(kf_a, vf_a, blk_a, fx.Int32(1), m_t, l_t, o_t) + m_t, l_t, o_t = res[0], res[1], res[2] + blk_b = fx.Int32(fx.rocdl.readlane(T.i32, blk, slot_at(rem_b))) + live_b = (rem_b != fx.Int64(0)).select(fx.Int32(1), fx.Int32(0)) + kf_b = dequant_k(raw_b) + vf_b = dequant_v(raw_b) + res = attend(kf_b, vf_b, blk_b, live_b, m_t, l_t, o_t) + m_t, l_t, o_t = res[0], res[1], res[2] + state = [rem_n, m_t, l_t] + o_t + raw_n + results = yield [fx.as_ir_value(x) for x in state] + + def store_out(acc, inv): + outr = buffer_ops.create_buffer_resource_from_addr( + arg_out, + num_records_bytes=fx.Int64(i32_rows) * fx.Int64(i32_H) * (D * 2), + ) + o_row = (row * i32_H + kvh * HG + n) * (D // 2) + g * 16 + for e in range_constexpr(4): + val = Vec.from_elements( + [ + pk_bf16(acc[2 * j][e] * inv, acc[2 * j + 1][e] * inv) + for j in range_constexpr(4) + ], + fx.Int32, + ) + buffer_ops.buffer_store(val, outr, o_row + e * 4, mask=row < i32_rows) + + # Combine the waves' partials in wave 0. + o_lds = fx.logical_divide( + fx.make_view( + fx.recast_iter(fx.Int32, smem.o.ptr), + fx.make_layout(NW * NDT * 64 * 4, 1), + ), + fx.make_layout(4, 1), + ) + lds_atom = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32) + ml = smem.ml.ptr + ml[(w * 2) * 64 + lane] = fx.Float32(results[1]) + ml[(w * 2 + 1) * 64 + lane] = fx.Float32(results[2]) + for dt in range_constexpr(NDT): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + r.store(Vec(results[3 + dt]).bitcast(fx.Int32)) + fx.copy(lds_atom, r, fx.slice(o_lds, (None, (w * NDT + dt) * 64 + lane))) + fx.gpu.barrier() + if w == 0: + ms = [fx.Float32(ml[(ww * 2) * 64 + lane]) for ww in range_constexpr(NW)] + big = ms[0] + for ww in range_constexpr(1, NW): + big = maxf_nn(big, ms[ww]) + big_raw = big + big = (big == neg_inf).select(fx.Float32(0.0), big) + wt = [fx.Float32(fx.math.exp2(ms[ww] - big)) for ww in range_constexpr(NW)] + l_tot = fx.Float32(0.0) + for ww in range_constexpr(NW): + l_tot = l_tot + fx.Float32(ml[(ww * 2 + 1) * 64 + lane]) * wt[ww] + l_tot = xlane4(l_tot, lambda x, y: x + y) + acc = [zero4 for _ in range_constexpr(NDT)] + for ww in range_constexpr(NW): + wv = Vec.filled(4, wt[ww], fx.Float32) + for dt in range_constexpr(NDT): + r = fx.make_rmem_tensor(fx.make_layout(4, 1), fx.Int32) + fx.copy( + lds_atom, + fx.slice(o_lds, (None, (ww * NDT + dt) * 64 + lane)), + r, + ) + acc[dt] = acc[dt] + Vec(r.load()).bitcast(fx.Float32) * wv + if const_expr(SPLITS > 1): + tid = row * i32_hkv + kvh + pr = buffer_ops.create_buffer_resource_from_addr( + arg_part, num_records_bytes=i64_part_bytes + ) + row_f = (tid * SPLITS + fx.block_idx.z) * 64 + lane + for dt in range_constexpr(NDT): + buffer_ops.buffer_store( + acc[dt].bitcast(fx.Int32), + pr, + row_f * PART_F + dt * 4, + cache_modifier=PART_SCOPE, + ) + buffer_ops.buffer_store( + Vec.from_elements([bits(big_raw), bits(l_tot)], fx.Int32), + pr, + row_f * PART_F + NDT * 4, + cache_modifier=PART_SCOPE, + ) + # The partial is stored before the ticket is taken. + fx.rocdl.s_waitcnt(vmcnt=0) + if lane == 0: + smem.ticket.ptr[0] = _global_atomic_add( + arg_cnt + fx.Int64(tid) * 4, 1 + ) + fx.rocdl.s_waitcnt(lgkmcnt=0) + if ( + fx.Int32( + fx.rocdl.readfirstlane(T.i32, fx.Int32(smem.ticket.ptr[0])) + ) + == SPLITS - 1 + ): + rows_f = [ + (tid * SPLITS + sp) * 64 + lane + for sp in range_constexpr(SPLITS) + ] + mls = [ + Vec( + buffer_ops.buffer_load( + pr, + rf * PART_F + NDT * 4, + vec_width=2, + dtype=fx.Int32, + cache_modifier=PART_SCOPE, + ) + ) + for rf in rows_f + ] + bigs = [as_f32(x[0]) for x in mls] + top = bigs[0] + for sp in range_constexpr(1, SPLITS): + top = maxf_nn(top, bigs[sp]) + top = (top == neg_inf).select(fx.Float32(0.0), top) + ws = [fx.Float32(fx.math.exp2(b - top)) for b in bigs] + l_all = fx.Float32(0.0) + for sp in range_constexpr(SPLITS): + l_all = l_all + as_f32(mls[sp][1]) * ws[sp] + inv_all = (l_all > fx.Float32(0.0)).select( + v_scale / l_all, fx.Float32(0.0) + ) + acc_all = [zero4 for _ in range_constexpr(NDT)] + for sp in range_constexpr(SPLITS): + wv = Vec.filled(4, ws[sp], fx.Float32) + for dt in range_constexpr(NDT): + part = Vec( + buffer_ops.buffer_load( + pr, + rows_f[sp] * PART_F + dt * 4, + vec_width=4, + dtype=fx.Int32, + cache_modifier=PART_SCOPE, + ) + ).bitcast(fx.Float32) + acc_all[dt] = acc_all[dt] + part * wv + store_out(acc_all, inv_all) + if lane == 0: + global_i32_ptr(arg_cnt + fx.Int64(tid) * 4)[0] = fx.Int32(0) + else: + inv = (l_tot > fx.Float32(0.0)).select(v_scale / l_tot, fx.Float32(0.0)) + store_out(acc, inv) + + @flyc.jit + def launch( + arg_q: fx.Int64, + arg_cache: fx.Int64, + arg_topk: fx.Int64, + arg_bt: fx.Int64, + arg_seq: fx.Int64, + arg_out: fx.Int64, + arg_ks: fx.Int64, + arg_vs: fx.Int64, + arg_lut: fx.Int64, + arg_part: fx.Int64, + arg_cnt: fx.Int64, + i64_part_bytes: fx.Int64, + i32_rows: fx.Int32, + i32_ql: fx.Int32, + i32_H: fx.Int32, + i32_hkv: fx.Int32, + i32_topk_hs: fx.Int32, + i32_bt_s: fx.Int32, + i32_page_b: fx.Int32, + f32_c: fx.Float32, + i32_grid_x: fx.Int32, + stream: fx.Stream, + ): + kernel( + i32_rows, + i32_ql, + i32_topk_hs, + i32_bt_s, + i32_H, + i32_hkv, + arg_seq, + arg_topk, + arg_bt, + arg_lut, + arg_q, + arg_cache, + arg_ks, + arg_vs, + arg_out, + arg_part, + arg_cnt, + i64_part_bytes, + i32_page_b, + f32_c, + ).launch( + grid=(i32_grid_x, i32_hkv, SPLITS), + block=(64 * NW, 1, 1), + stream=stream, + ) + + launch.kernel_name = name + launch.compile_hints = PRELOAD_HINTS + return launch diff --git a/vllm/models/minimax_m3/amd/sparse_attention_msa.py b/vllm/models/minimax_m3/amd/sparse_attention_msa.py index a2c49d13abf0..44ad045db0d9 100644 --- a/vllm/models/minimax_m3/amd/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/amd/sparse_attention_msa.py @@ -126,7 +126,6 @@ def forward( kv_cache: torch.Tensor, output: torch.Tensor, *, - query_fp8: torch.Tensor | None = None, decode_sparse_table: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: from vllm.models.minimax_m3.amd.ops.sparse_pa import ( diff --git a/vllm/models/minimax_m3/amd/sparse_attention_nvfp4.py b/vllm/models/minimax_m3/amd/sparse_attention_nvfp4.py new file mode 100644 index 000000000000..72f816956438 --- /dev/null +++ b/vllm/models/minimax_m3/amd/sparse_attention_nvfp4.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax M3 block-sparse attention over an NVFP4 K/V cache on ROCm (gfx950). + +The cache spec is NVIDIA's, but the bytes inside a page are laid out for the +FlyDSL kernels, so pages are not interchangeable with NVIDIA's. +""" + +from dataclasses import replace +from typing import ClassVar + +import torch + +from vllm.config.cache import CacheDType +from vllm.forward_context import get_forward_context +from vllm.models.minimax_m3.common.sparse_attention import ( + MiniMaxM3SparseBackend, + MiniMaxM3SparseImpl, + MiniMaxM3SparseMetadata, + MiniMaxM3SparseMetadataBuilder, + minimax_m3_query_token_positions, +) +from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim +from vllm.v1.attention.backend import AttentionLayer, CommonAttentionMetadata +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout + + +class MiniMaxM3SparseRocmNvfp4MetadataBuilder(MiniMaxM3SparseMetadataBuilder): + """Adds each prefill token's request and absolute position.""" + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> MiniMaxM3SparseMetadata: + metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) + p = metadata.prefill + if p is not None and p.query_req_id is None: + p.query_req_id, p.query_abs_pos = minimax_m3_query_token_positions( + p.cu_seqlens_q, p.context_lens, metadata.num_prefill_tokens + ) + return metadata + + +class MiniMaxM3SparseRocmNvfp4Backend(MiniMaxM3SparseBackend): + """MiniMax M3 sparse attend over NVFP4 K/V pages (ROCm, gfx950).""" + + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["nvfp4"] + + @staticmethod + def get_builder_cls() -> type["MiniMaxM3SparseRocmNvfp4MetadataBuilder"]: + return MiniMaxM3SparseRocmNvfp4MetadataBuilder + + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + """Store each head's K and V as two slots of packed data + block scales.""" + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4: + return spec + return replace( + spec, + num_head_slots=2 * spec.num_kv_heads, + state_content_bytes=nvfp4_kv_cache_full_dim(spec.head_size) + * get_dtype_size(spec.dtype), + ) + + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: + # The kernels read the page stride from stride(0) and address pages in + # 64 bits, so block-outer packing works too (opt in with BLHNC). + return (KVCacheLayout.LBHNC, KVCacheLayout.BLHNC) + + +def rocm_nvfp4_unsupported_reason( + num_heads: int, num_kv_heads: int, head_dim: int, topk_blocks: int, block_size: int +) -> str | None: + """Why the ROCm NVFP4 attend cannot serve this layer, None when it can.""" + from vllm.models.minimax_m3.amd.ops.nvfp4_prefill.host import ( + supports_nvfp4_sparse_prefill, + ) + + if not supports_nvfp4_sparse_prefill( + num_heads, num_kv_heads, head_dim, topk_blocks, block_size + ): + return ( + "needs gfx950, head_dim 128, 16 query heads per KV head and 16 " + f"selected 128-token blocks (got num_heads={num_heads}, " + f"num_kv_heads={num_kv_heads}, head_dim={head_dim}, " + f"topk_blocks={topk_blocks}, block_size={block_size})" + ) + return None + + +class MiniMaxM3SparseRocmNvfp4Impl(MiniMaxM3SparseImpl): + """FlyDSL NVFP4 sparse decode (bf16 compute) and prefill (fp16 compute).""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + reason = rocm_nvfp4_unsupported_reason( + self.num_heads, + self.num_kv_heads, + self.head_size, + self.topk_blocks, + self.block_size, + ) + if reason is not None: + raise ValueError(f"MiniMax M3 NVFP4 KV cache on ROCm {reason}") + + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + kv_cache: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + from vllm.models.minimax_m3.amd.ops.nvfp4_prefill.host import ( + minimax_m3_nvfp4_sparse_prefill, + ) + from vllm.models.minimax_m3.amd.ops.nvfp4_sparse_decode.host import ( + minimax_m3_nvfp4_sparse_decode_fly, + ) + + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return output # profiling run; caches unbound + md = attn_metadata[layer.layer_name] # type: ignore[attr-defined] + assert isinstance(md, MiniMaxM3SparseMetadata) + + nd = md.num_decode_tokens + num_tokens = md.num_actual_tokens + topk = layer.topk_indices_buffer # type: ignore[attr-defined] + assert topk is not None + q = query[:num_tokens].view(-1, self.num_heads, self.head_size) + out = output[:num_tokens].view(-1, self.num_heads, self.head_size) + k_scale = layer._k_scale # type: ignore[attr-defined] + v_scale = layer._v_scale # type: ignore[attr-defined] + + if md.num_decodes > 0: + d = md.decode + assert d is not None + minimax_m3_nvfp4_sparse_decode_fly( + q[:nd].contiguous(), + kv_cache, + topk[:, :nd, :], + d.block_table, + d.seq_lens, + self.scale, + out[:nd], + k_scale=k_scale, + v_scale=v_scale, + decode_query_len=d.decode_query_len, + ) + if md.num_prefills > 0: + p = md.prefill + assert p is not None + assert p.query_req_id is not None and p.query_abs_pos is not None + minimax_m3_nvfp4_sparse_prefill( + q[nd:], + kv_cache, + topk[:, nd:num_tokens, :], + p.block_table, + p.query_req_id, + p.query_abs_pos, + self.num_kv_heads, + self.scale, + out[nd:], + k_scale=k_scale, + v_scale=v_scale, + ) + return output diff --git a/vllm/models/minimax_m3/common/dspark.py b/vllm/models/minimax_m3/common/dspark.py new file mode 100644 index 000000000000..58fecf2bc353 --- /dev/null +++ b/vllm/models/minimax_m3/common/dspark.py @@ -0,0 +1,413 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax-M3.1 DSpark draft model. + +The Qwen3 DSpark block drafter (context-KV precompute, non-causal query block, +Markov and confidence heads) built from MiniMax-M3 layers: Gemma RMSNorm +(``x * (1 + w)``) everywhere, partial RoPE, and a dense SwiGLU-OAI MLP. The +checkpoint ships no token embedding or LM head; both are shared with the target. +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm import _custom_ops as ops +from vllm.config import CacheConfig, VllmConfig +from vllm.model_executor.layers.activation import SiluAndMulWithClamp +from vllm.model_executor.layers.attention.attention import get_attention_context +from vllm.model_executor.layers.fusion.quant_activation import get_input_quant_key +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic +from vllm.model_executor.models.qwen3_dflash import ( + DFlashQwen3Attention, + DFlashQwen3DecoderLayer, + _resolve_layer_attention, +) +from vllm.model_executor.models.qwen3_dspark import ( + Qwen3DSparkForCausalLM, + Qwen3DSparkModel, +) +from vllm.model_executor.models.utils import WeightsMapper +from vllm.models.minimax_m3.common.ops.dspark_context_kv import ( + DSparkContextKVWriter, + rotary_dims_supported, +) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3TextConfig +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) + +_CHECKPOINT_PREFIX = "language_model.model.dspark." + + +def _shared_slot_mapping( + slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, +) -> torch.Tensor | None: + """The context slot mapping when every draft layer shares it, else None + (dummy runs, and drafts whose layers span several KV cache groups).""" + if slot_mapping is None or isinstance(slot_mapping, torch.Tensor): + return slot_mapping + first = slot_mapping[0] + if first is None or any( + s is None or s.data_ptr() != first.data_ptr() for s in slot_mapping + ): + return None + return first + + +def _gemma_rms_norm(hidden_size: int, eps: float) -> nn.Module: + """The target's Gemma RMSNorm (fused FlashInfer kernels on CUDA). + + vLLM's GemmaRMSNorm hands the IR op an fp32 ``1 + w`` weight, which the + vllm_c kernel rejects for bf16 inputs, so every call ran the unfused native + path (about ten kernels). + """ + if current_platform.is_rocm(): + from vllm.models.minimax_m3.amd.model import MiniMAXGemmaRMSNorm + else: + from vllm.models.minimax_m3.nvidia.model import MiniMAXGemmaRMSNorm + return MiniMAXGemmaRMSNorm(hidden_size, eps=eps) + + +def dspark_qknorm_rope_kv_insert( + qkv: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + num_heads: int, + num_kv_heads: int, + rotary_dim: int, + eps: float, + q_fp8_out: torch.Tensor | None, + layer_name: LayerNameType, +) -> None: + """The fused QK-norm/RoPE on ``qkv``, also storing the draft tokens' k/v. + + A custom op so the slot mapping is read when it runs: read in the compiled + forward, it is fixed at trace time, when the draft layers have none. + """ + _, layer, kv_cache, slot_mapping = get_attention_context( + _resolve_layer_name(layer_name) + ) + insert = slot_mapping is not None + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + q_norm_weight, + k_norm_weight, + cos_sin_cache, + positions, + num_heads, + num_kv_heads, + rotary_dim, + eps, + slot_mapping=slot_mapping, + kv_cache=kv_cache if insert else None, + block_size=kv_cache.size(2) if insert else 0, + kv_cache_dtype="fp8" if insert and kv_cache.dtype == torch.uint8 else "auto", + q_fp8_out=q_fp8_out, + q_fp8_scale=layer._q_scale_float, + ) + + +def dspark_qknorm_rope_kv_insert_fake( + qkv: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + num_heads: int, + num_kv_heads: int, + rotary_dim: int, + eps: float, + q_fp8_out: torch.Tensor | None, + layer_name: LayerNameType, +) -> None: + return + + +direct_register_custom_op( + op_name="minimax_m3_dspark_qknorm_rope_kv_insert", + op_func=dspark_qknorm_rope_kv_insert, + mutates_args=["qkv", "q_fp8_out"], + fake_impl=dspark_qknorm_rope_kv_insert_fake, +) + + +class MiniMaxM3DSparkAttention(DFlashQwen3Attention): + def __init__( + self, + config: MiniMaxM3TextConfig, + sliding_window: int | None, + causal: bool, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str, + query_pre_quantized: bool = False, + ) -> None: + super().__init__( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + rope_parameters={ + "rope_theta": config.rope_theta, + "partial_rotary_factor": config.partial_rotary_factor, + }, + max_position=config.max_position_embeddings, + head_dim=config.head_dim, + rms_norm_eps=config.rms_norm_eps, + sliding_window=sliding_window, + causal=causal, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + query_pre_quantized=query_pre_quantized, + ) + self.q_norm = _gemma_rms_norm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = _gemma_rms_norm(self.head_dim, eps=config.rms_norm_eps) + # With an FP8 KV cache the backend wants an FP8 query, and the fused + # kernel below can write one. Producing it there instead of letting the + # attention layer quantize afterwards saves a pass over q. + self._emit_fp8_query = query_pre_quantized + # FLASHINFER_TS writes an FP8 query's output as FP16 (FA4 only bf16); + # an MXFP8 o_proj quantizes FP16 as readily and writes the model dtype. + self._attn_output_dtype = ( + torch.float16 + if query_pre_quantized + and self.attn.attn_backend.get_name() == "FLASHINFER_TS" + and get_input_quant_key(self.o_proj) == kMxfp8Dynamic + else None + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + # The target's fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, + # in place: one launch instead of the q-norm, k-norm and RoPE kernels. + # q_fp8_out additionally stores the normed/roped q as E4M3, which is + # what an FP8-KV backend consumes. + q_fp8 = ( + torch.empty( + (qkv.shape[0], self.q_size), + dtype=current_platform.fp8_dtype(), + device=qkv.device, + ) + if self._emit_fp8_query + else None + ) + # Also stores the draft tokens' k/v, so the attention layer gets q alone. + torch.ops.vllm.minimax_m3_dspark_qknorm_rope_kv_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + self.rotary_emb.cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + self.rotary_emb.rotary_dim, + self.q_norm.variance_epsilon, + q_fp8, + _encode_layer_name(self.attn.layer_name), + ) + q = qkv[:, : self.q_size] if q_fp8 is None else q_fp8 + # output_dtype must be stated: it otherwise follows the query, which + # may be FP8. + attn_output = self.attn( + q, None, None, output_dtype=self._attn_output_dtype or hidden_states.dtype + ) + output, _ = self.o_proj(attn_output) + return output + + +class MiniMaxM3DSparkMLP(nn.Module): + """Dense SwiGLU-OAI MLP: gate * sigmoid(alpha * gate) * (up + beta).""" + + def __init__( + self, + config: MiniMaxM3TextConfig, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + if config.hidden_act != "swigluoai": + raise ValueError( + f"Unsupported MiniMax DSpark activation {config.hidden_act!r}; " + "only swigluoai is supported." + ) + intermediate_size = config.dspark_ffn_hidden_size + self.gate_up_proj = MergedColumnParallelLinear( + config.hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = SiluAndMulWithClamp( + swiglu_limit=config.swiglu_limit, + alpha=config.swiglu_alpha, + beta=config.swiglu_beta, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x, _ = self.down_proj(self.act_fn(gate_up)) + return x + + +class MiniMaxM3DSparkDecoderLayer(DFlashQwen3DecoderLayer): + def __init__( + self, + vllm_config: VllmConfig, + *, + config: MiniMaxM3TextConfig, + layer_idx: int, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + if not getattr(config, "dspark_use_dense_ffn", False): + raise NotImplementedError("MiniMax DSpark supports dense FFN drafts only.") + self.hidden_size = config.hidden_size + sliding_window, causal = _resolve_layer_attention(config, layer_idx) + # An FP8 draft cache means the backend reads an FP8 query, which the + # fused QK-norm/RoPE kernel emits directly. + kv_cache_dtype = getattr(cache_config, "cache_dtype", "auto") + self.self_attn = MiniMaxM3DSparkAttention( + config, + sliding_window=sliding_window, + causal=causal, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + query_pre_quantized=str(kv_cache_dtype).startswith("fp8"), + ) + self.mlp = MiniMaxM3DSparkMLP(config, quant_config, prefix=f"{prefix}.mlp") + hidden_size, eps = config.hidden_size, config.rms_norm_eps + self.input_layernorm = _gemma_rms_norm(hidden_size, eps=eps) + self.post_attention_layernorm = _gemma_rms_norm(hidden_size, eps=eps) + + +class MiniMaxM3DSparkModel(Qwen3DSparkModel): + decoder_layer_cls = MiniMaxM3DSparkDecoderLayer + quant_config_mapper = WeightsMapper( + orig_to_new_prefix={_CHECKPOINT_PREFIX: "model."} + ) + + def __init__( + self, + *, + vllm_config: VllmConfig, + start_layer_id: int = 0, + prefix: str = "", + ) -> None: + super().__init__( + vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix + ) + hidden_size, eps = self.config.hidden_size, self.config.rms_norm_eps + self.hidden_norm = _gemma_rms_norm(hidden_size, eps=eps) + self.norm = _gemma_rms_norm(hidden_size, eps=eps) + + def _normalize_context_states(self, context_states: torch.Tensor) -> torch.Tensor: + return self.hidden_norm(context_states) + + def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor: + return torch.stack( + [layer.self_attn.k_norm(k) for layer, k in zip(self.layers, all_k)] + ) + + def _context_kv_writer(self) -> DSparkContextKVWriter | None: + """The fused context K/V writer, or None where the base path runs.""" + if not hasattr(self, "_ctx_kv_writer"): + if any(a.kv_cache.numel() == 0 for a in self._attn_layers): + return None # KV caches not bound yet: decide once they are + rotary_dim = self._rope_cos_sin_cache.shape[-1] + supported = ( + self._rope_is_neox + and self.layers[0].self_attn.v_scale is None + and rotary_dims_supported(rotary_dim, self._head_dim) + and self._rope_cos_sin_cache.dtype == torch.bfloat16 + and DSparkContextKVWriter.supported(self._attn_layers) + ) + self._ctx_kv_writer = ( + DSparkContextKVWriter(self._attn_layers) if supported else None + ) + return self._ctx_kv_writer + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + """One kernel norms, ropes and stores the context K/V of every layer, + straight from the projection output.""" + slot_mapping = _shared_slot_mapping(context_slot_mapping) + writer = ( + self._context_kv_writer() + if slot_mapping is not None + and hasattr(self, "_num_attn_layers") + and context_states.dtype == torch.bfloat16 + else None + ) + if writer is None: + super().precompute_and_store_context_kv( + context_states, context_positions, context_slot_mapping + ) + return + all_kv = self._project_context_kv_token_major( + context_states, + context_states.shape[0], + self._num_attn_layers, + self._num_kv_heads, + self._head_dim, + ) + writer( + all_kv, + self._k_norm_weights, + self._rms_norm_eps, + context_positions, + self._rope_cos_sin_cache, + self._rope_cos_sin_cache.shape[-1], + slot_mapping, + ) + + +class MiniMaxM3DSparkForCausalLM(Qwen3DSparkForCausalLM): + model_cls = MiniMaxM3DSparkModel + + @staticmethod + def _rename_checkpoint_weights( + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, tensor in weights: + name = name.removeprefix(_CHECKPOINT_PREFIX) + name = name.replace(".decoder_layer.", ".") + if name.startswith("final_norm."): + name = "norm." + name.removeprefix("final_norm.") + # MiniMax stores the MXFP8 E8M0 block scales as weight_scale_inv. + name = name.replace(".weight_scale_inv", ".weight_scale") + yield name, tensor + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + return super().load_weights(self._rename_checkpoint_weights(weights)) diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index 0d7ee1d2f9ba..c86d1fcd7b33 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -115,20 +115,30 @@ def __init__( backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend, ) -> None: super().__init__() + # Per-token storage width in elements of ``dtype``. + self.head_size = head_dim if indexer_kv_dtype in ("fp8", "fp8_e4m3"): - cache_dtype = torch.float8_e4m3fn + cache_dtype = query_dtype = torch.float8_e4m3fn elif indexer_kv_dtype == "bf16": - cache_dtype = torch.bfloat16 + cache_dtype = query_dtype = torch.bfloat16 + elif indexer_kv_dtype == "nvfp4": + # Packed NVFP4 pages scored by the MSA Q8KV4 indexer with e4m3 Q. + cache_dtype, query_dtype = torch.uint8, torch.float8_e4m3fn + self.head_size = head_dim // 2 + head_dim // 16 else: raise NotImplementedError( f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the " - "MiniMax M3 indexer cache (only 'bf16' or 'fp8'/'fp8_e4m3')." + "MiniMax M3 indexer cache (only 'bf16', 'fp8'/'fp8_e4m3' or " + "'nvfp4')." ) self.kv_cache = torch.tensor([]) self.head_dim = head_dim self.indexer_kv_dtype = indexer_kv_dtype - # Side-cache storage dtype: bf16, or e4m3 for the fp8 score path. + # Side-cache storage dtype: bf16, e4m3 for the fp8 score path, or uint8 + # for packed NVFP4. self.dtype = cache_dtype + # Index-query dtype the fused qk-norm/rope kernel emits. + self.query_dtype = query_dtype self.prefix = prefix self.cache_config = cache_config # Impl-chosen backend -> each impl gets its own builder (get_attn_backend). @@ -147,7 +157,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=1, - head_size=self.head_dim, + head_size=self.head_size, dtype=self.dtype, cache_role=SparseCacheRole.INDEXER, ) @@ -392,6 +402,11 @@ def __init__( backend_cls=type(self).indexer_backend_cls, ) + def reserve_workspace(self) -> None: + """Reserve the workspace-manager scratch ``forward`` needs; called from + the memory-profiling run, which skips ``forward``.""" + return + def forward( self, index_query: torch.Tensor, @@ -501,8 +516,9 @@ def select_indexer_impl_cls( On Blackwell (SM100) with ``topk_blocks == 16`` (the only width fmha_sm100's ``sparse_topk_select`` kernel supports), the fmha_sm100 score + top-k path is - used for both bf16 and fp8 index caches. Everything else falls back to the - Triton indexer, with fp8 restricted to CUDA platforms that advertise fp8 + used for both bf16 and fp8 index caches. An nvfp4 index cache uses the MSA + Q8KV4/Q8KV8 paged indexers (SM100/SM103/SM107 only). Everything else falls back to + the Triton indexer, with fp8 restricted to CUDA platforms that advertise fp8 support. """ if get_current_vllm_config().attention_config.minimax_m3_indexer_cp: @@ -530,10 +546,9 @@ def select_indexer_impl_cls( tp_size, ) return MiniMaxM3IndexerTritonCPImpl - if indexer_kv_dtype in ("mxfp4", "nvfp4"): + if indexer_kv_dtype == "mxfp4": raise NotImplementedError( - f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) " - "CuteDSL indexer impl." + "indexer_kv_dtype='mxfp4' is not supported by the MiniMax M3 indexer." ) is_cuda = current_platform.is_cuda() is_fp8 = indexer_kv_dtype in ("fp8", "fp8_e4m3") @@ -542,6 +557,28 @@ def select_indexer_impl_cls( f"indexer_kv_dtype={indexer_kv_dtype!r} requires CUDA fp8 support." ) is_sm100 = is_cuda and current_platform.is_device_capability_family(100) + if indexer_kv_dtype == "nvfp4": + if not ( + is_cuda + and ( + current_platform.is_device_capability(100) + or current_platform.is_device_capability(103) + or current_platform.is_device_capability(107) + ) + ): + raise NotImplementedError( + "indexer_kv_dtype='nvfp4' requires an SM100, SM103 or SM107 GPU." + ) + from vllm.models.minimax_m3.nvidia.indexer_nvfp4 import ( + MiniMaxM3IndexerNvfp4Impl, + ) + + logger.info_once( + "MiniMax M3 indexer: selected NVFP4 (MSA Q8KV4 decode, Q8KV8 prefill) " + "[topk_blocks=%d]", + topk_blocks, + ) + return MiniMaxM3IndexerNvfp4Impl use_msa = ( is_sm100 and topk_blocks == 16 @@ -628,6 +665,9 @@ def index_cache(self) -> MiniMaxM3IndexerCache: def num_index_heads(self) -> int: return self.impl.num_index_heads + def reserve_workspace(self) -> None: + self.impl.reserve_workspace() + def forward( self, index_query: torch.Tensor, diff --git a/vllm/models/minimax_m3/common/mm_preprocess.py b/vllm/models/minimax_m3/common/mm_preprocess.py index 958e4f189b78..602aee361f46 100644 --- a/vllm/models/minimax_m3/common/mm_preprocess.py +++ b/vllm/models/minimax_m3/common/mm_preprocess.py @@ -40,24 +40,28 @@ ) from vllm.transformers_utils.configs.minimax_m3 import MiniMaxM3Config from vllm.transformers_utils.processors.minimax_m3 import ( - MIN_SHORT_SIDE_PIXEL, MiniMaxM3VLImageProcessor, MiniMaxM3VLVideoProcessor, MiniMaxVLProcessor, - smart_resize, + compute_aligned_size, + resolve_max_long_side_pixel, ) +# Worst-case long-side caps in pixels, matching the processor's detail=high +# tiers. Profiling budgets are sized against these; request-level +# ``detail``/``max_long_side_pixel`` kwargs cannot exceed them. +_IMAGE_MAX_LONG_SIDE_PIXEL = 3584 +_VIDEO_MAX_LONG_SIDE_PIXEL = 2188 + # Upper bound on the number of frames used to build the dummy video during # memory profiling. Sized to the worst-case video the processor accepts: -# ``max_total_pixels // max_pixels_per_frame`` = 301,056,000 // 602,112 = 500 -# frames, each at the video processor's per-frame ``max_pixels`` (768 * 28 * 28 -# = 602,112). This reaches the true worst-case ~192,000 vision tokens, but only -# because the dummy video is sized via ``get_video_size_with_most_features()`` -# (the video ``max_pixels`` bound), not the smaller image bound. Without a cap, +# ``max_total_pixels // max_pixels_per_frame`` = 301,056,000 // (2184 * 2184) +# = 63 frames, each a square at the video long-side cap (2188 px floored to +# the patch*merge factor, 28). Without a cap, # ``_get_max_video_frames(seq_len)`` with M3's large ``max_model_len`` yields # ~1400 frames, producing a multi-GB dummy tensor that overflows the # multimodal encoder cache. -_MAX_FRAMES_PER_VIDEO = 500 +_MAX_FRAMES_PER_VIDEO = 63 class MiniMaxM3VLProcessingInfo(BaseProcessingInfo): @@ -101,34 +105,46 @@ def _get_vision_info( image_height: int, num_frames: int, image_processor, + mm_kwargs: Mapping[str, object], + modality: str, ) -> tuple[ImageSize, int]: """Compute resized image size and number of vision tokens. - Mirrors the processor's Qwen-style ``smart_resize`` (area bound by - ``max_pixels``) so token counts match the actual processor output. + Mirrors the processor's ``compute_aligned_size`` so token counts + match the actual processor output. """ patch_size: int = image_processor.patch_size merge_size: int = image_processor.merge_size temporal_patch_size: int = image_processor.temporal_patch_size factor = patch_size * merge_size - max_pixels: int = image_processor.max_pixels - # Long-side resize spec (opt-in). ``image_processor`` is the *video* - # processor when counting video tokens, so read the bounds off it. - max_long_side_pixel = getattr(image_processor, "max_long_side_pixel", None) - min_short_side_pixel = getattr( - image_processor, "min_short_side_pixel", MIN_SHORT_SIDE_PIXEL + + # ``image_processor`` is the *video* processor when counting video + # tokens, so modality-scoped request kwargs (flat keys overlaid with + # ``images_kwargs``/``videos_kwargs`` via ``modality=``) and processor + # defaults are read per modality. Per-key fallback mirrors the + # processor, whose own defaults feed ``resolve_max_long_side_pixel`` + # the same way. + merged = self.ctx.get_merged_mm_kwargs(mm_kwargs) + # Kwargs scoped to this modality (``images_kwargs`` / ``videos_kwargs``) + # override the flat ones. + scoped = merged.get("images_kwargs" if modality == "image" else "videos_kwargs") + if isinstance(scoped, Mapping): + merged = {**merged, **scoped} + detail = merged.get("detail", getattr(image_processor, "detail", None)) + max_long_side_pixel = merged.get( + "max_long_side_pixel", getattr(image_processor, "max_long_side_pixel", None) + ) + resolved_long_side = resolve_max_long_side_pixel( + detail, max_long_side_pixel, modality=modality ) - new_h, new_w = smart_resize( + # No volumetric cap here: token counting must not raise; the video + # processor enforces its volume cap on the real inputs. + new_h, new_w = compute_aligned_size( image_height, image_width, - factor=factor, - max_pixels=max_pixels, - max_long_side_pixel=max_long_side_pixel, - min_short_side_pixel=min_short_side_pixel, - # Token counting must not raise; the volumetric/area cap is enforced - # in the processor's _preprocess on the real inputs. - max_total_pixels=None, + factor, + max_long_side_pixel=resolved_long_side, ) grid_h = new_h // patch_size grid_w = new_w // patch_size @@ -153,6 +169,8 @@ def get_num_image_tokens( image_height=image_height, num_frames=1, image_processor=image_processor, + mm_kwargs=mm_kwargs, + modality="image", ) return n @@ -170,48 +188,27 @@ def get_num_video_tokens( image_height=image_height, num_frames=num_frames, image_processor=image_processor, + mm_kwargs=mm_kwargs, + modality="video", ) return n def get_image_size_with_most_features(self) -> ImageSize: - # Largest square (a multiple of patch_size*merge_size) whose area is - # within the image processor's bound — this yields the most vision - # tokens for one image. With the long-side spec the square side is - # capped by ``max_long_side_pixel`` (and the fixed ``max_total_pixels``); - # otherwise it is bound by the ``max_pixels`` area. + # Worst case is a detail=high image: a square whose side is the max + # long side (3584 px, a multiple of patch_size*merge_size). image_processor = self.get_image_processor() factor = image_processor.patch_size * image_processor.merge_size - max_long_side_pixel = getattr(image_processor, "max_long_side_pixel", None) - if max_long_side_pixel is not None: - side_px = min( - max_long_side_pixel, - math.isqrt(image_processor.max_total_pixels), - ) - else: - side_px = math.isqrt(image_processor.max_pixels) - side = max(factor, (side_px // factor) * factor) + side = (_IMAGE_MAX_LONG_SIDE_PIXEL // factor) * factor return ImageSize(width=side, height=side) def get_video_size_with_most_features(self) -> ImageSize: - # Per-frame size that yields the most vision tokens, bound by the - # *video* processor's ``max_pixels`` (which differs from the image - # bound). Token count depends only on area, so maximize the area - # achievable with both sides a multiple of patch_size*merge_size rather - # than picking the largest square — a square (e.g. 756x756 for M3's - # 602,112 bound) leaves area on the table, undercounting frames. + # Worst single frame is a detail=high square capped by the video + # long-side max (2188 px), floored to a patch_size*merge_size + # multiple (2184). video_processor = self.get_video_processor() factor = video_processor.patch_size * video_processor.merge_size - per_frame_pixels = video_processor.max_pixels - max_long_side_pixel = getattr(video_processor, "max_long_side_pixel", None) - if max_long_side_pixel is not None: - # Long-side spec: a frame's worst case is a square capped by - # ``max_long_side_pixel`` (per-frame area, not the volumetric cap). - per_frame_pixels = min(per_frame_pixels, max_long_side_pixel**2) - units = per_frame_pixels // (factor * factor) # h_u * w_u - h_u = math.isqrt(units) - while units % h_u: - h_u -= 1 - return ImageSize(width=(units // h_u) * factor, height=h_u * factor) + side = (_VIDEO_MAX_LONG_SIDE_PIXEL // factor) * factor + return ImageSize(width=side, height=side) def get_max_image_tokens(self) -> int: image_processor = self.get_image_processor() @@ -220,7 +217,9 @@ def get_max_image_tokens(self) -> int: image_width=size.width, image_height=size.height, image_processor=image_processor, - mm_kwargs={}, + # Budgets must cover the worst case (detail=high); the default + # detail=auto cap would undercount. + mm_kwargs={"max_long_side_pixel": _IMAGE_MAX_LONG_SIDE_PIXEL}, ) def _get_max_video_frames(self, max_tokens: int) -> int: @@ -233,7 +232,8 @@ def _get_max_video_frames(self, max_tokens: int) -> int: image_height=size.height, num_frames=num_frames + 1, image_processor=video_processor, - mm_kwargs={}, + # Worst case (detail=high); see get_max_image_tokens. + mm_kwargs={"max_long_side_pixel": _VIDEO_MAX_LONG_SIDE_PIXEL}, ) if next_n > max_tokens: break @@ -265,7 +265,8 @@ def get_max_video_tokens( image_height=size.height, num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts), image_processor=video_processor, - mm_kwargs={}, + # Worst case (detail=high); see get_max_image_tokens. + mm_kwargs={"max_long_side_pixel": _VIDEO_MAX_LONG_SIDE_PIXEL}, ) @@ -487,6 +488,11 @@ def load_bytes( backend: VideoDecoderBackend = "opencv", **kwargs, ) -> tuple[npt.NDArray, dict[str, Any]]: + # fps=-1 (unspecified) falls back to the default 1 fps. + if fps == -1: + fps = 1 + if not 0.2 <= fps <= 5: + raise ValueError(f"video fps must be in [0.2, 5], got {fps}") return super().load_bytes( data, num_frames=num_frames, @@ -526,6 +532,15 @@ def compute_frames_index_to_sample( break indices.append(target_frame) prev_kept_ts = target_frame / video_fps + + # Always keep the last frame when it is far enough (>= 1/fps, up to + # eps) from the previously kept frame, matching the SFT frame + # sampler's constant-mode behavior. + last_frame_idx = total_frames - 1 + last_ts = last_frame_idx / video_fps + if indices and indices[-1] != last_frame_idx and last_ts - prev_kept_ts > eps: + indices.append(last_frame_idx) + if not indices: indices = [0] return indices diff --git a/vllm/models/minimax_m3/common/ops/dspark_context_kv.py b/vllm/models/minimax_m3/common/ops/dspark_context_kv.py new file mode 100644 index 000000000000..514f052051d6 --- /dev/null +++ b/vllm/models/minimax_m3/common/ops/dspark_context_kv.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused DSpark context-KV write for all draft layers in one CUDA kernel. + +Replaces the permute copy, per-layer K-norm launches + stack, positions.repeat, +rotary_embedding and per-layer reshape_and_cache_flash calls in +``DFlashQwen3Model.precompute_and_store_context_kv``. K/V are read straight +from the fused projection output ``all_kv`` ``[num_ctx, L, 2, nkv, 128]`` and +written into each layer's paged cache (FLASH_ATTN, FLASHINFER_TS) +``[num_blocks, nkv, block_size, 2 * 128]`` (K at ``[..., :128]``, V at +``[..., 128:]``), matching ``reshape_and_cache_flash``: fp8 e4m3 of +``x / scale`` for kv_cache_dtype fp8, raw bf16 otherwise; slots < 0 are padding +and are not written. + +K gets the Gemma RMSNorm of its layer (``x * rsqrt(mean(x^2) + eps) * (1 + w)`` +in fp32, materialized to bf16 like the unfused flashinfer path) followed by +partial NeoX RoPE with the token's position (see +``fused_minimax_m3_dspark_context_kv``). Each layer's cache is a separate +allocation, addressed as layer 0's cache plus a per-layer element offset, and +each fp8 scale as layer 0's k scale plus a float offset; the offset tables are +built once and cached (keyed by the data_ptrs) so the call stays CUDA-graph +safe. +""" + +import torch + +from vllm import _custom_ops as ops +from vllm.platforms import current_platform + +_SUPPORTED_KV_DTYPES = ("auto", "bfloat16", "fp8", "fp8_e4m3") +_HEAD_DIM = 128 + + +def rotary_dims_supported(rotary_dim: int, head_dim: int) -> bool: + """The kernel rotates the leading 16, 32 or 64 dims of a 128-dim head.""" + return head_dim == _HEAD_DIM and rotary_dim in (16, 32, 64) + + +class DSparkContextKVWriter: + """Writes precomputed context K/V into every draft layer's KV cache. + + Built from the draft's ``Attention`` layers; the cache- and scale-offset + tables are cached, keyed by their data_ptrs, so calls captured in a CUDA + graph keep addressing the same memory. + """ + + def __init__(self, attn_layers: list) -> None: + layers = list(attn_layers) + error = self._validate(layers) + if error is not None: + raise ValueError(f"unsupported DSpark context-KV layout: {error}") + self._layers = layers + cache = layers[0].kv_cache + self._num_layers = len(layers) + self._num_kv_heads = cache.shape[1] + self._kv_cache_dtype = "fp8" if cache.dtype == torch.uint8 else "auto" + self._table_key: tuple | None = None + self._cache_offsets: torch.Tensor | None = None + self._scale_offsets: torch.Tensor | None = None + + @staticmethod + def _validate(layers: list) -> str | None: + if not layers: + return "no layers" + ref = None + for a in layers: + kv_dtype = getattr(a, "kv_cache_dtype", None) + if kv_dtype not in _SUPPORTED_KV_DTYPES: + return f"kv_cache_dtype {kv_dtype!r}" + cache = getattr(a, "kv_cache", None) + if not isinstance(cache, torch.Tensor) or cache.numel() == 0: + return "KV cache not bound" + if cache.dim() != 4 or cache.shape[3] != 2 * _HEAD_DIM: + return f"unexpected cache shape {tuple(cache.shape)}" + block_size = cache.shape[2] + if block_size & (block_size - 1): + return f"block size {block_size} is not a power of two" + if not cache.is_contiguous(): + return "cache is not contiguous" + if kv_dtype in ("fp8", "fp8_e4m3"): + if cache.dtype != torch.uint8: + return f"fp8 cache has dtype {cache.dtype}" + elif cache.dtype != torch.bfloat16: + return f"cache has dtype {cache.dtype}" + for name in ("_k_scale", "_v_scale"): + scale = getattr(a, name, None) + if ( + not isinstance(scale, torch.Tensor) + or scale.dtype != torch.float32 + or not scale.is_cuda + ): + return f"{name} is not a CUDA fp32 tensor" + key = (cache.shape, cache.stride(), cache.dtype, kv_dtype) + if ref is None: + ref = key + elif key != ref: + return "layers do not share cache shape/stride/dtype" + # Caches are addressed from layer 0's as 16-element aligned offsets. + delta = cache.data_ptr() - layers[0].kv_cache.data_ptr() + if delta % (16 * cache.element_size()): + return "layer caches are not 16-element aligned to each other" + return None + + @staticmethod + def supported(attn_layers: list) -> bool: + # The fp8 store is OCP E4M3; ROCm fnuz caches keep the base path. + if not current_platform.is_cuda() or not hasattr( + torch.ops._C, "fused_minimax_m3_dspark_context_kv" + ): + return False + return DSparkContextKVWriter._validate(list(attn_layers)) is None + + def _ensure_tables(self) -> None: + layers = self._layers + key = tuple( + (a.kv_cache.data_ptr(), a._k_scale.data_ptr(), a._v_scale.data_ptr()) + for a in layers + ) + if key == self._table_key: + return + cache0, scale0 = layers[0].kv_cache, layers[0]._k_scale + + def table(offsets: list[int]) -> torch.Tensor: + return torch.tensor(offsets, dtype=torch.int64, device=cache0.device) + + self._cache_offsets = table( + [ + (a.kv_cache.data_ptr() - cache0.data_ptr()) // cache0.element_size() + for a in layers + ] + ) + # [2, L]: k scales, then v scales, as float offsets from layer 0's k. + self._scale_offsets = table( + [ + (getattr(a, name).data_ptr() - scale0.data_ptr()) // 4 + for name in ("_k_scale", "_v_scale") + for a in layers + ] + ) + self._table_key = key + + def __call__( + self, + all_kv: torch.Tensor, + k_norm_weight: torch.Tensor, + eps: float, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + rotary_dim: int, + slot_mapping: torch.Tensor, + ) -> None: + num_ctx, L, _, nkv, _ = all_kv.shape + assert self._num_layers == L and self._num_kv_heads == nkv + if num_ctx == 0: + return + self._ensure_tables() + ops.fused_minimax_m3_dspark_context_kv( + all_kv, + k_norm_weight, + cos_sin_cache, + positions, + slot_mapping, + self._layers[0].kv_cache, + self._cache_offsets, + self._layers[0]._k_scale, + self._scale_offsets, + rotary_dim, + eps, + self._kv_cache_dtype, + ) diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 73ab067eb94f..4cc28f0ffd1e 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -264,7 +264,7 @@ def __init__( ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) - # Stable context-length buffer for decode cudagraph replays. + # Stable context-length buffer for the prefill metadata. self.context_len_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, dtype=torch.int32, @@ -306,14 +306,13 @@ def build( assert num_decodes + num_prefills == num_reqs assert num_decode_tokens + num_prefill_tokens == num_tokens - # Decode-first batch: context lengths into the stable cudagraph buffer. - context_lens = self.context_len_buffer[:num_reqs] - context_lens.copy_( - common_attn_metadata.compute_num_computed_tokens(), non_blocking=True - ) - prefill_metadata: MiniMaxM3SparsePrefillMetadata | None = None if num_prefills > 0: + # Only prefill reads the context lengths; decode-only steps skip it. + context_lens = self.context_len_buffer[:num_reqs] + context_lens.copy_( + common_attn_metadata.compute_num_computed_tokens(), non_blocking=True + ) seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound assert seq_lens_cpu is not None prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:] @@ -439,11 +438,12 @@ def forward( query: torch.Tensor, kv_cache: torch.Tensor, output: torch.Tensor, - *, - query_fp8: torch.Tensor | None = None, ) -> torch.Tensor: """Attend the queries to the indexer-selected blocks. Per kernel. + ``query`` is E4M3 for impls with ``use_cutlass_decode``, else the model + dtype. + The indexer has already written the top-k block ids into ``layer.topk_indices_buffer`` (decode at ``[:, :nd]``, prefill at ``[:, nd:num_tokens]``); the attend reads them from there. @@ -463,8 +463,6 @@ def forward( query: torch.Tensor, kv_cache: torch.Tensor, output: torch.Tensor, - *, - query_fp8: torch.Tensor | None = None, ) -> torch.Tensor: attn_metadata = get_forward_context().attn_metadata if not isinstance(attn_metadata, dict): @@ -563,6 +561,17 @@ def select_main_backend_and_impl_cls( "MiniMax M3 sparse attention supports only the 'nvfp4' NVFP4 KV " f"cache dtype, got {kv_cache_dtype}" ) + if use_nvfp4 and current_platform.is_rocm(): + from vllm.models.minimax_m3.amd.sparse_attention_nvfp4 import ( + MiniMaxM3SparseRocmNvfp4Backend, + MiniMaxM3SparseRocmNvfp4Impl, + ) + + logger.info_once( + "MiniMax M3 sparse attention selected ROCm NVFP4 (topk_blocks=%s)", + topk_blocks, + ) + return MiniMaxM3SparseRocmNvfp4Backend, MiniMaxM3SparseRocmNvfp4Impl if use_nvfp4 and (use_aiter_sparse_pa or not use_msa): raise ValueError( "MiniMax M3 sparse attention supports an NVFP4 KV cache only with " @@ -590,10 +599,11 @@ def select_main_backend_and_impl_cls( MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl, MiniMaxM3SparseMSANvfp4Backend, + MiniMaxM3SparseMSANvfp4Impl, ) if use_nvfp4: - return MiniMaxM3SparseMSANvfp4Backend, MiniMaxM3SparseMSAImpl + return MiniMaxM3SparseMSANvfp4Backend, MiniMaxM3SparseMSANvfp4Impl return MiniMaxM3SparseMSABackend, MiniMaxM3SparseMSAImpl return MiniMaxM3SparseBackend, MiniMaxM3SparseTritonImpl diff --git a/vllm/models/minimax_m3/common/vision_tower.py b/vllm/models/minimax_m3/common/vision_tower.py index c7aba484f2a3..b730b07093bb 100644 --- a/vllm/models/minimax_m3/common/vision_tower.py +++ b/vllm/models/minimax_m3/common/vision_tower.py @@ -59,11 +59,12 @@ def __init__(self, config: PreTrainedConfig) -> None: stride=(temporal_patch_size, patch_size, patch_size), bias=False, ) + # The checkpoint stores this weight in fp32 (the vision tower's only + # fp32 tensor); keep it that way so the conv runs at full precision. + self.patch_embedding.to(torch.float32) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: # pixel_values: (N, C * temporal_patch_size * patch_size²) - if self.patch_embedding.weight.dtype != pixel_values.dtype: - self.patch_embedding = self.patch_embedding.to(pixel_values.dtype) x = pixel_values.reshape( pixel_values.shape[0], self.num_channels, @@ -71,7 +72,8 @@ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: self.patch_size, self.patch_size, ) - return self.patch_embedding(x).reshape(x.shape[0], -1) + out = self.patch_embedding(x.to(self.patch_embedding.weight.dtype)) + return out.reshape(x.shape[0], -1).to(pixel_values.dtype) class MiniMaxVLAttention(nn.Module): @@ -657,7 +659,9 @@ def __init__( prefix=maybe_prefix(prefix, "patch_merge_mlp"), ) - self.dtype = self.vision_model.embeddings.patch_embedding.weight.dtype + # Tower compute dtype; the patch embedding weight stays fp32 (see + # MiniMaxVLPatchEmbed), so take it from the projector instead. + self.dtype = self.multi_modal_projector.linear_1.weight.dtype self.out_hidden_size = text_hidden_size def forward( diff --git a/vllm/models/minimax_m3/nvidia/indexer_nvfp4.py b/vllm/models/minimax_m3/nvidia/indexer_nvfp4.py new file mode 100644 index 000000000000..501d81f772a7 --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/indexer_nvfp4.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 index-cache indexer impl for MiniMax M3 (SM100/SM103/SM107). + +The index-K side cache holds packed NVFP4 pages ``[num_blocks, 128, 72]`` (E2M1 +codes, then E4M3 block scales, no global scale) written by the fused +qk-norm/rope insert, and index-Q is E4M3. Both sides use the MSA paged indexers +of ``fmha_sm100`` with fused top-k, which write 15 score-ranked history pages +plus the query's local page last (``init_blocks=0, local_blocks=1``) as +logical block ids: + +* Decode scores the NVFP4 pages with the Q8KV4 decode indexer, which reads the + 1..8 queries of each request in place (they take the last slots of its + 8-query tile) and writes their top-k straight into the shared buffer. + Uniform decode batches are captured in full CUDA graphs, so each padded + batch size and query length gets its own CUDA-graph wrapper whose buffers + never move; the builder re-plans it outside capture every step. +* Prefill has no NVFP4 kernel. Each layer dequantizes the history pages of the + prefill requests into compact FP8 workspace pages (what the decode kernel + scores) and runs the Q8KV8 prefill indexer on them, in request chunks whose + pages fit the gather budget. + +All index heads of the rank share the index keys and select their own pages. +Prefill and decode score every head of the rank in one launch. Prefill plans +are rebuilt every step and only run eagerly (mixed batches are never captured +whole). +""" + +from dataclasses import dataclass, field +from typing import Any, ClassVar + +import torch + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.forward_context import get_forward_context +from vllm.models.minimax_m3.common.indexer import ( + MiniMaxM3IndexerBackend, + MiniMaxM3IndexerImpl, + MiniMaxM3IndexerMetadata, + MiniMaxM3IndexerMetadataBuilder, +) +from vllm.models.minimax_m3.nvidia.indexer_nvfp4_decode_cp import ( + Nvfp4DecodeCPLayer, + Nvfp4DecodeCPPlan, + get_nvfp4_indexer_decode_cp, +) +from vllm.models.minimax_m3.nvidia.ops.gather_nvfp4_index_k import ( + minimax_m3_gather_nvfp4_index_k_fp8, +) +from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backend import ( + AttentionBackend, + CommonAttentionMetadata, + max_decode_query_len, +) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.worker.workspace import current_workspace_manager + +PAGE_SIZE = 128 +HEAD_DIM = 128 +TOPK_BLOCKS = 16 +# Queries per request the Q8KV4 decode indexer scores. +DECODE_QUERY_LEN = 8 +# Longest block table the MSA indexers accept. +MAX_PAGES = 8192 +# Index heads the Q8KV4 decode indexer scores in one launch. +DECODE_NUM_HEADS = (1, 2, 4) + + +def _msa_indexers(): + """Import the MSA indexer wrappers. + + Through ``fmha_sm100.sparse``, which binds the package's csrc JIT loader. + """ + from vllm.third_party.fmha_sm100.sparse import ( + BatchDecodeIndexerQ8KV4Wrapper, + BatchPrefillIndexerQ8KV8Wrapper, + ) + + return BatchDecodeIndexerQ8KV4Wrapper, BatchPrefillIndexerQ8KV8Wrapper + + +def max_gather_pages(vllm_config: VllmConfig) -> int: + """FP8 pages one prefill chunk may gather; always fits one full request.""" + tokens = max( + vllm_config.model_config.max_model_len, + 4 * vllm_config.scheduler_config.max_num_batched_tokens, + ) + return cdiv(tokens, PAGE_SIZE) + + +class MiniMaxM3IndexerNvfp4Backend(MiniMaxM3IndexerBackend): + """Indexer side-cache backend selecting the NVFP4 builder.""" + + @staticmethod + def get_builder_cls() -> type["MiniMaxM3IndexerNvfp4MetadataBuilder"]: + return MiniMaxM3IndexerNvfp4MetadataBuilder + + +@dataclass +class MiniMaxM3IndexerNvfp4DecodeMetadata: + wrapper: Any # planned CUDA-graph wrapper, None when decode CP is on + num_reqs: int # includes zero-length full-CUDA-graph padding requests + query_len: int # uniform, <= DECODE_QUERY_LEN + # Set instead of the wrapper when decode CP is on. + cp_plan: Nvfp4DecodeCPPlan | None = None + + +@dataclass +class MiniMaxM3IndexerNvfp4PrefillChunk: + wrapper: Any # planned BatchPrefillIndexerQ8KV8Wrapper + # Token range within the prefill region (offsets from num_decode_tokens). + token_start: int + token_end: int + block_table: torch.Tensor # [n, max_blocks] physical pages + num_pages: torch.Tensor # [n] int32 history pages to gather + page_offsets: torch.Tensor # [n] int32 first compact page of each request + total_pages: int + max_pages: int + + +@dataclass +class MiniMaxM3IndexerNvfp4Metadata(MiniMaxM3IndexerMetadata): + """The base ``prefill``/``decode`` fields are unused on this path.""" + + decode_q8kv4: MiniMaxM3IndexerNvfp4DecodeMetadata | None = None + prefill_chunks: list[MiniMaxM3IndexerNvfp4PrefillChunk] = field( + default_factory=list + ) + + +class MiniMaxM3IndexerNvfp4MetadataBuilder(MiniMaxM3IndexerMetadataBuilder): + """Plans the MSA decode and prefill indexers once per step.""" + + 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) + if self.max_decode_query_len > DECODE_QUERY_LEN: + raise NotImplementedError( + "The MiniMax M3 NVFP4 indexer scores at most " + f"{DECODE_QUERY_LEN} decode queries per request, got " + f"{self.max_decode_query_len} (use num_speculative_tokens <= " + f"{DECODE_QUERY_LEN - 1})." + ) + if self.num_index_heads not in DECODE_NUM_HEADS: + raise NotImplementedError( + "The MiniMax M3 NVFP4 indexer serves " + f"{DECODE_NUM_HEADS} index heads per rank, got " + f"{self.num_index_heads}." + ) + self.max_gather_pages = max_gather_pages(vllm_config) + # Decode wrapper per padded decode batch size and query length. + self.decode_wrappers: dict[tuple[int, int], Any] = {} + # Decode CP (opt-in); static, so all ranks and graphs agree. + self.decode_cp = get_nvfp4_indexer_decode_cp( + num_index_heads=self.num_index_heads, + max_decode_query_len=self.max_decode_query_len, + ) + # Pages a request can span; decode scores at most this many. + self.max_decode_pages = min( + cdiv(vllm_config.model_config.max_model_len, PAGE_SIZE), MAX_PAGES + ) + # One page-score scratch for every decode wrapper: they never run at + # once, and a buffer per batch size would hold 256 KiB per request at a + # 1M context for every batch size the engine sees. Decode batches, + # padded to a CUDA graph size, never exceed max_decode_reqs. + self.decode_score_buffer: torch.Tensor | None = None + if self.decode_cp is None: + max_decode_reqs = max( + vllm_config.scheduler_config.max_num_seqs, + vllm_config.compilation_config.max_cudagraph_capture_size or 0, + ) + decode_wrapper_cls, _ = _msa_indexers() + self.decode_score_buffer = torch.empty( + decode_wrapper_cls.score_buffer_size( + max_decode_reqs, + self.max_decode_pages, + num_heads=self.num_index_heads, + ), + dtype=torch.float32, + device=device, + ) + + def _decode_wrapper( + self, block_table: torch.Tensor, seq_lens: torch.Tensor, query_len: int + ) -> Any: + num_reqs = seq_lens.shape[0] + wrapper = self.decode_wrappers.get((num_reqs, query_len)) + if wrapper is None: + # The runner's block table and seq_lens are persistent buffers, so + # the wrapper reads them in place and plan() copies nothing. A + # strided (width-capped) block table gets its own buffer. + if not block_table.is_contiguous(): + block_table = torch.empty( + block_table.shape, dtype=torch.int32, device=self.device + ) + decode_wrapper_cls, _ = _msa_indexers() + wrapper = self.decode_wrappers[num_reqs, query_len] = decode_wrapper_cls( + num_heads=self.num_index_heads, + use_cuda_graph=True, + block_table_buffer=block_table, + seq_lens_buffer=seq_lens, + score_buffer=self.decode_score_buffer, + ) + return wrapper + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> MiniMaxM3IndexerNvfp4Metadata: + num_reqs = common_attn_metadata.num_reqs + num_tokens = common_attn_metadata.num_actual_tokens + seq_lens = common_attn_metadata.seq_lens + block_table = common_attn_metadata.block_table_tensor + qsl_cpu = common_attn_metadata.query_start_loc_cpu + + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=self.reorder_batch_threshold, + require_uniform=True, + ) + ) + assert num_decodes + num_prefills == num_reqs + assert num_decode_tokens + num_prefill_tokens == num_tokens + + decode: MiniMaxM3IndexerNvfp4DecodeMetadata | None = None + if num_decode_tokens > 0: + # num_decodes counts the zero-length requests that pad a full CUDA + # graph batch, so it is the batch size the graph was captured with. + query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes] + query_len = int(query_lens_cpu[0]) + assert torch.all((query_lens_cpu == query_len) | (query_lens_cpu == 0)) + num_pages = min(block_table.shape[1], self.max_decode_pages) + if self.decode_cp is not None: + decode = MiniMaxM3IndexerNvfp4DecodeMetadata( + wrapper=None, + num_reqs=num_decodes, + query_len=query_len, + cp_plan=self.decode_cp.plan( + block_table[:num_decodes, :num_pages], + seq_lens[:num_decodes], + ), + ) + else: + decode_block_table = block_table[:num_decodes, :num_pages] + decode_seq_lens = seq_lens[:num_decodes] + wrapper = self._decode_wrapper( + decode_block_table, decode_seq_lens, query_len + ) + wrapper.plan(decode_block_table, decode_seq_lens, query_len=query_len) + decode = MiniMaxM3IndexerNvfp4DecodeMetadata( + wrapper=wrapper, + num_reqs=num_decodes, + query_len=query_len, + ) + + prefill_chunks: list[MiniMaxM3IndexerNvfp4PrefillChunk] = [] + if num_prefills > 0: + prefill_chunks = self._build_prefill_chunks( + common_attn_metadata, num_decodes + ) + + return MiniMaxM3IndexerNvfp4Metadata( + seq_lens=seq_lens, + max_seq_len=common_attn_metadata.max_seq_len, + slot_mapping=common_attn_metadata.slot_mapping, + num_actual_tokens=num_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + decode_q8kv4=decode, + prefill_chunks=prefill_chunks, + ) + + def _build_prefill_chunks( + self, common_attn_metadata: CommonAttentionMetadata, num_decodes: int + ) -> list[MiniMaxM3IndexerNvfp4PrefillChunk]: + _, prefill_wrapper_cls = _msa_indexers() + num_reqs = common_attn_metadata.num_reqs + seq_lens = common_attn_metadata.seq_lens + block_table = common_attn_metadata.block_table_tensor + qsl_cpu = common_attn_metadata.query_start_loc_cpu[num_decodes : num_reqs + 1] + # Host-side lengths keep planning free of device syncs. An upper bound + # (async spec decode) only over-sizes the plan and the gather: the + # kernels mask by the device seq_lens and never score the extra pages. + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound + assert seq_lens_cpu is not None + seq_lens_cpu = seq_lens_cpu[num_decodes:num_reqs].to(torch.int32) + query_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32) + # Only pages before the last query's local page are ever scored. + history_pages = (seq_lens_cpu - 1) // PAGE_SIZE + device = seq_lens.device + + chunks = [] + start, num_prefills = 0, seq_lens_cpu.shape[0] + while start < num_prefills: + end, total_pages = start, 0 + while end < num_prefills and ( + end == start + or total_pages + int(history_pages[end]) <= self.max_gather_pages + ): + total_pages += int(history_pages[end]) + end += 1 + assert total_pages <= self.max_gather_pages + num_chunk_reqs = end - start + max_seq_len = int(seq_lens_cpu[start:end].max()) + width = cdiv(max_seq_len, PAGE_SIZE) + # [cu_seqlens_q | num_pages | page_offsets | compact block table] + # in one pinned buffer, so a single H2D copy is truly asynchronous + # (a pageable one may synchronize with the stream). + staging = torch.empty( + 3 * num_chunk_reqs + 1 + num_chunk_reqs * width, + dtype=torch.int32, + pin_memory=True, + ) + cu_seqlens_q, pages, offsets, table = staging.split( + [ + num_chunk_reqs + 1, + num_chunk_reqs, + num_chunk_reqs, + num_chunk_reqs * width, + ] + ) + torch.sub(qsl_cpu[start : end + 1], qsl_cpu[start], out=cu_seqlens_q) + pages.copy_(history_pages[start:end]) + torch.cumsum(pages, 0, dtype=torch.int32, out=offsets) + offsets.sub_(pages) + cols = torch.arange(width, dtype=torch.int32) + table.view(num_chunk_reqs, width).copy_( + torch.where(cols < pages[:, None], offsets[:, None] + cols, 0) + ) + staging_gpu = staging.to(device, non_blocking=True) + cu_seqlens_q_gpu, pages_gpu, offsets_gpu, table_gpu = staging_gpu.split( + [ + num_chunk_reqs + 1, + num_chunk_reqs, + num_chunk_reqs, + num_chunk_reqs * width, + ] + ) + lo, hi = num_decodes + start, num_decodes + end + token_start = int(qsl_cpu[start] - qsl_cpu[0]) + token_end = int(qsl_cpu[end] - qsl_cpu[0]) + wrapper = prefill_wrapper_cls(num_heads=self.num_index_heads) + wrapper.plan( + cu_seqlens_q_gpu, + seq_lens[lo:hi], + table_gpu.view(num_chunk_reqs, width), + total_q=token_end - token_start, + max_seqlen_q=int(query_lens_cpu[start:end].max()), + max_seqlen_k=max_seq_len, + ) + chunks.append( + MiniMaxM3IndexerNvfp4PrefillChunk( + wrapper=wrapper, + token_start=token_start, + token_end=token_end, + block_table=block_table[lo:hi], + num_pages=pages_gpu, + page_offsets=offsets_gpu, + total_pages=total_pages, + max_pages=int(pages.max()), + ) + ) + start = end + return chunks + + +class MiniMaxM3IndexerNvfp4Impl(MiniMaxM3IndexerImpl): + """Decode: Q8KV4 on NVFP4 pages. Prefill: FP8 gather + Q8KV8.""" + + indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerNvfp4Backend + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + config = ( + self.topk_blocks, + self.init_blocks, + self.local_blocks, + self.block_size, + self.index_head_dim, + self.score_type, + ) + if config != (TOPK_BLOCKS, 0, 1, PAGE_SIZE, HEAD_DIM, "max"): + raise NotImplementedError( + "The MiniMax M3 NVFP4 indexer supports topk_blocks=16, " + "init_blocks=0, local_blocks=1, 128-token blocks, index head " + "dim 128 and max scoring; got (topk_blocks, init_blocks, " + "local_blocks, block_size, index_head_dim, score_type)=" + f"{config}." + ) + vllm_config = get_current_vllm_config() + self.max_gather_pages = max_gather_pages(vllm_config) + # Per-layer decode-CP buffers (collective: same layer order on all ranks). + self.decode_cp = get_nvfp4_indexer_decode_cp( + num_index_heads=self.num_index_heads, + max_decode_query_len=max_decode_query_len(vllm_config), + ) + self.decode_cp_layer: Nvfp4DecodeCPLayer | None = ( + None if self.decode_cp is None else self.decode_cp.new_layer() + ) + + def reserve_workspace(self) -> None: + current_workspace_manager().get_simultaneous( + ((self.max_gather_pages, PAGE_SIZE, HEAD_DIM), torch.float8_e4m3fn) + ) + + def forward( + self, + index_query: torch.Tensor, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return None, None # profiling run; caches unbound + md = attn_metadata[self.index_cache.prefix] + assert isinstance(md, MiniMaxM3IndexerNvfp4Metadata) + + nd = md.num_decode_tokens + # Not sliced to the actual tokens: a padded full-CUDA-graph decode + # batch also scores its padding rows, and every captured tensor must + # keep its address and shape. + index_q = index_query.view(-1, self.num_index_heads, self.index_head_dim) + kv = self.index_cache.kv_cache + # The kernels write logical block ids straight into the shared + # [T, H, 16] buffer the attend reads. + buf = self.topk_indices_buffer + assert buf is not None + + if md.decode_q8kv4 is not None: + d = md.decode_q8kv4 + rows = d.num_reqs * d.query_len + if d.cp_plan is not None: + assert self.decode_cp is not None and self.decode_cp_layer is not None + self.decode_cp.forward( + self.decode_cp_layer, + d.cp_plan, + index_q[:rows], + kv, + buf[:rows], + d.query_len, + ) + else: + d.wrapper.run(index_q[:rows], kv, out=buf[:rows]) + for chunk in md.prefill_chunks: + self._prefill( + index_q[nd + chunk.token_start : nd + chunk.token_end], + kv, + chunk, + buf[nd + chunk.token_start : nd + chunk.token_end], + ) + return None, None + + def _prefill( + self, + index_q: torch.Tensor, + kv: torch.Tensor, + chunk: MiniMaxM3IndexerNvfp4PrefillChunk, + out: torch.Tensor, + ) -> None: + (k_pages,) = current_workspace_manager().get_simultaneous( + ((max(chunk.total_pages, 1), PAGE_SIZE, HEAD_DIM), torch.float8_e4m3fn) + ) + minimax_m3_gather_nvfp4_index_k_fp8( + kv, + chunk.block_table, + chunk.num_pages, + chunk.page_offsets, + chunk.max_pages, + k_pages, + ) + chunk.wrapper.run(index_q, k_pages, out=out) diff --git a/vllm/models/minimax_m3/nvidia/indexer_nvfp4_decode_cp.py b/vllm/models/minimax_m3/nvidia/indexer_nvfp4_decode_cp.py new file mode 100644 index 000000000000..012515855628 --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/indexer_nvfp4_decode_cp.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Context-parallel decode for the MiniMax M3 NVFP4 indexer (opt-in, TP4). + +At TP4 each rank owns one of the 4 index heads, but the index-K cache is +replicated, so every rank reads the whole cache per sparse layer. With +``--attention-config '{"indexer_decode_cp": true}'``, rank ``r`` +instead: + +1. exchanges index-q heads so it holds all 4 (``_q_exchange_kernel``); +2. scores all 4 heads over pages ``j * CP + r`` with the stock Q8KV4 kernel + on a compact block table, reading 1/CP of the index-K; +3. takes the exact local top 16 per (query, head) and sends it to the head's + owner (``_local_topk_push_kernel``); +4. merges the owner's ``16 * CP`` candidates into the indexer output: 15 + pages by (score desc, page asc), then the own page (``_merge_kernel``). + +The merge is exact: the global top 16 lies in the union of per-rank top 16s. +The stock decode ranks 16-bit quantized scores, so the two may differ at ties. + +The chunking is written against one contract, the MSA Q8KV4 decode scorer's: +eight query slots per request, ``(seq_len - 1) / 128`` pages scored, and only +the pages below a query's own page stored. The three places that depend on it +are ``_compact_block_table_kernel`` (the length it is handed), +``_local_topk_push_kernel`` (the cut that length drops) and ``_set_width`` +(room for the page that length invents). + +Both exchanges are NVLink stores into symmetric buffers whose payload is the +flag (``0xFF`` E4M3 NaN for queries, ``-1`` for keys); consumers restore it. +Buffers are per layer, shapes are fixed per batch size, and nothing syncs with +the host, so the path is CUDA-graph safe. It needs symmetric memory across the +tensor-parallel ranks and falls back to the stock decode without it. + +The three Triton kernels are PDL-launched and wait for their producer, so +each starts its prologue in the tail of the kernel before it. Only the push +releases its dependent early, right after its own wait, which is what lets +the merge overlap it; the exchange must not, because the score kernel after +it reads the whole tile. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + get_tp_group, +) +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv +from vllm.v1.attention.ops.cp_common import direct_cp_enabled + +logger = init_logger(__name__) + +PAGE_SIZE = 128 +HEAD_DIM = 128 +TOPK_BLOCKS = 16 +DECODE_QUERY_LEN = 8 # query slots per request in a Q8KV4 decode tile +MAX_PAGES = 8192 +NUM_INDEX_HEADS = 4 +SUPPORTED_TP_SIZES = (4,) +# Rows with at most this many ranked pages take a single top-k. +TOPK_CHUNK = 256 +# Groups of the two-level top-k over longer rows. +MAX_TOPK_GROUPS = 128 + +_TOPK = tl.constexpr(TOPK_BLOCKS) +_RANKED = tl.constexpr(TOPK_BLOCKS - 1) +_PAGE = tl.constexpr(PAGE_SIZE) +_SLOTS = tl.constexpr(DECODE_QUERY_LEN) +_WORDS = tl.constexpr(HEAD_DIM // 4) # int32 words per index-q head + + +@triton.jit +def _peer_ptr(ptrs, rank, dtype: tl.constexpr): + return tl.load(ptrs + rank).to(tl.pointer_type(dtype)) + + +@triton.jit +def _compact_block_table_kernel( + bt_ptr, # [B, P] int32 + seq_lens_ptr, # [B] int32 + cbt_ptr, # [B, W] int32 out: compact page j -> global page j * CP + RANK + cseq_ptr, # [B] int32 out: lengths the score kernel sees + seq_out_ptr, # [B] int32 out: real lengths (stable buffer) + P, + W, + stride_bt, + stride_cbt, + CP: tl.constexpr, + RANK: tl.constexpr, + BLOCK: tl.constexpr, +): + """Compact page j = global page ``j * CP + RANK``, and the length the score + kernel is handed for it. + + That length is a lie: the scorer stores only the pages below a query's own + page, so ``n_local * 128 + 8`` puts all eight query slots one page past + the chunk and stores every owned page, uncausally. The real cut comes back + in ``_local_topk_push_kernel``. ``W`` keeps the invented page addressable, + since a scorer that clamps the own page to its last score column would + drop the chunk's last page.""" + b = tl.program_id(0) + seq = tl.load(seq_lens_ptr + b) + tl.store(seq_out_ptr + b, seq) + hist = tl.minimum(tl.maximum((seq - 1) // _PAGE, 0), P) + n_local = tl.minimum((hist + (CP - 1 - RANK)) // CP, W - 1) + for start in range(0, W, BLOCK): + j = start + tl.arange(0, BLOCK) + src = tl.minimum(j * CP + RANK, P - 1) + v = tl.load(bt_ptr + b.to(tl.int64) * stride_bt + src, mask=j < W) + tl.store(cbt_ptr + b.to(tl.int64) * stride_cbt + j, v, mask=j < W) + tl.store(cseq_ptr + b, n_local * _PAGE + _SLOTS) + + +@triton.jit +def _has_ff_byte(v): + u = v.to(tl.uint32, bitcast=True) + ff = ( + ((u & 0xFF) == 0xFF) + | (((u >> 8) & 0xFF) == 0xFF) + | (((u >> 16) & 0xFF) == 0xFF) + | ((u >> 24) == 0xFF) + ) + return tl.max(ff.to(tl.int32), axis=0) > 0 + + +@triton.jit +def _q_exchange_kernel( + q_ptr, # [rows, 32] int32 view of this rank's e4m3 index-q head + tile_ptrs, # [CP] int64: q tile [B, 8, H, 128] bytes of every rank + H: tl.constexpr, + CP: tl.constexpr, + RANK: tl.constexpr, + HEAD: tl.constexpr, + QLEN: tl.constexpr, +): + """Store this rank's index-q head into every rank's tile, then wait for + the other heads. The PDL wait is for index-q's producer.""" + tl.extra.cuda.gdc_wait() + t = tl.program_id(0) + b = t // QLEN + slot = (_SLOTS - QLEN) + t % QLEN + w = tl.arange(0, _WORDS) + row = tl.load(q_ptr + t * _WORDS + w) + dst_off = ((b * _SLOTS + slot) * H + HEAD) * _WORDS + w + for dst in tl.static_range(CP): + tl.store(_peer_ptr(tile_ptrs, dst, tl.int32) + dst_off, row) + local = _peer_ptr(tile_ptrs, RANK, tl.int32) + hw = tl.arange(0, H * _WORDS) + off = (b * _SLOTS + slot) * H * _WORDS + hw + v = tl.load(local + off, volatile=True) + pending = _has_ff_byte(v) + while pending: + v = tl.load(local + off, volatile=True) + pending = _has_ff_byte(v) + + +@triton.jit +def _page_keys(row_ptr, j, n, CP: tl.constexpr, RANK: tl.constexpr): + """Key ``1 << 48 | ordered(score) << 16 | (0xFFFF - page)`` for ``j < n``, + else 0. Unique; orders by score desc, then page asc.""" + valid = j < n + page = j * CP + RANK + score = tl.load(row_ptr + j, mask=valid, other=0.0) + bits = score.to(tl.uint32, bitcast=True) + ordered = bits ^ tl.where((bits >> 31) != 0, 0xFFFFFFFF, 0x80000000) + key = (1 << 48) | (ordered.to(tl.int64) << 16) | (0xFFFF - page).to(tl.int64) + return tl.where(valid, key, 0) + + +@triton.jit +def _local_topk_push_kernel( + score_ptr, # [B, 8 * H, W] fp32 compact scores (row = slot * H + head) + seq_lens_ptr, # [B] int32 + recv_ptrs, # [CP] int64: candidates [CP, rows_max, 16] int64 of every rank + rows_max, + stride_sb, + stride_sr, + H: tl.constexpr, + CP: tl.constexpr, + RANK: tl.constexpr, + QLEN: tl.constexpr, + WIDTH: tl.constexpr, + CHUNK: tl.constexpr, + NUM_GROUPS: tl.constexpr, + GROUP: tl.constexpr, +): + """Per (query, head): exact top 16 of the owned pages, sent to the head's + owner. Long rows use a two-level select (top 16 groups, then top 16). + + PDL-launched: the wait is for the scores, and releasing the merge only + after it keeps the merge's own wait a guarantee that the score kernel has + read the q tile.""" + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + t = tl.program_id(0) + h = tl.program_id(1) + b = t // QLEN + i = t % QLEN + slot = (_SLOTS - QLEN) + i + own = (tl.load(seq_lens_ptr + b) - QLEN + i) // _PAGE + if own + 1 > _TOPK: + # The causal cut the compact length dropped: the rank's pages below + # the query's own page, bounded by the width like the scored pages (a + # length past the page table ranks what it holds, as the stock TopK). + n_rank = tl.minimum((own + (CP - 1 - RANK)) // CP, WIDTH - 1) + row_ptr = score_ptr + b.to(tl.int64) * stride_sb + (slot * H + h) * stride_sr + if n_rank <= CHUNK: + j = tl.arange(0, CHUNK) + best = tl.topk(_page_keys(row_ptr, j, n_rank, CP, RANK), _TOPK) + else: + g = tl.arange(0, NUM_GROUPS) + s = tl.arange(0, GROUP) + keys = _page_keys( + row_ptr, g[:, None] * GROUP + s[None, :], n_rank, CP, RANK + ) + top_g = tl.topk(tl.max(keys, axis=1), _TOPK) + page = (0xFFFF - (top_g & 0xFFFF)).to(tl.int32) + grp = tl.where(top_g > 0, ((page - RANK) // CP) // GROUP, NUM_GROUPS) + cand = _page_keys( + row_ptr, grp[:, None] * GROUP + s[None, :], n_rank, CP, RANK + ) + best = tl.topk(tl.reshape(cand, [_TOPK * GROUP]), _TOPK) + k = tl.arange(0, _TOPK) + dst = _peer_ptr(recv_ptrs, h, tl.int64) + tl.store(dst + (RANK * rows_max + t) * _TOPK + k, best) + + +@triton.jit +def _merge_kernel( + seq_lens_ptr, # [B] int32 + recv_ptrs, # [CP] int64 + tile_ptrs, # [CP] int64 + out_ptr, # [rows, 1, 16] int32 + rows_max, + out_stride, + H: tl.constexpr, + CP: tl.constexpr, + RANK: tl.constexpr, + QLEN: tl.constexpr, +): + """Merge all ranks' candidates for this rank's head and restore sentinels. + + PDL-launched after every push program has started, so polling here never + blocks a push.""" + t = tl.program_id(0) + b = t // QLEN + i = t % QLEN + own = (tl.load(seq_lens_ptr + b) - QLEN + i) // _PAGE + k = tl.arange(0, _TOPK) + out = out_ptr + t.to(tl.int64) * out_stride + k + if own + 1 <= _TOPK: + tl.store(out, tl.where(k < own + 1, k, -1)) + # No push for this row: wait for the score before resetting the q tile. + tl.extra.cuda.gdc_wait() + else: + recv = _peer_ptr(recv_ptrs, RANK, tl.int64) + p = tl.arange(0, CP) + off = (p[:, None] * rows_max + t) * _TOPK + k[None, :] + keys = tl.load(recv + off, volatile=True) + pending = tl.min(tl.min(keys, axis=1), axis=0) == -1 + while pending: + keys = tl.load(recv + off, volatile=True) + pending = tl.min(tl.min(keys, axis=1), axis=0) == -1 + best = tl.topk(tl.reshape(keys, [CP * _TOPK]), _TOPK) + page = (0xFFFF - (best & 0xFFFF)).to(tl.int32) + tl.store(out, tl.where(k == _RANKED, own, page)) + tl.store(recv + off, tl.full([CP, _TOPK], -1, tl.int64)) + # The score kernel has read this q tile row; re-arm it. + slot = (_SLOTS - QLEN) + i + tile = _peer_ptr(tile_ptrs, RANK, tl.int32) + hw = tl.arange(0, H * _WORDS) + tl.store( + tile + (b * _SLOTS + slot) * H * _WORDS + hw, + tl.full(hw.shape, -1, tl.int32), + ) + + +@dataclass +class Nvfp4DecodeCPPlan: + """Per padded batch size: the 4-head score wrapper and its buffers.""" + + # The wrapper owns the compact [B, 8 * H, width] fp32 scores. + wrapper: Any # BatchDecodeIndexerQ8KV4Wrapper(num_heads=4), CUDA-graph mode + block_table: torch.Tensor # [B, width] compact pages + score_seq_lens: torch.Tensor # [B] lengths the score kernel sees + seq_lens: torch.Tensor # [B] real lengths + num_reqs: int + + +@dataclass +class Nvfp4DecodeCPLayer: + """One sparse layer's symmetric buffers.""" + + # Keep both alive: freeing one unmaps it while peers still write to it. + tile: torch.Tensor # local q tile, int32 words + recv: torch.Tensor # local candidates [CP, rows_max, 16] int64 + tile_ptrs: torch.Tensor # [CP] int64 + recv_ptrs: torch.Tensor # [CP] int64 + + +class Nvfp4IndexerDecodeCP: + """Per-process state shared by the builder and every sparse layer.""" + + def __init__(self, *, max_reqs: int, device: torch.device) -> None: + from vllm.third_party.fmha_sm100.sparse import BatchDecodeIndexerQ8KV4Wrapper + + self._wrapper_cls = BatchDecodeIndexerQ8KV4Wrapper + self.cp = get_tensor_model_parallel_world_size() + self.rank = get_tensor_model_parallel_rank() + self.head = self.rank # one index head per rank at TP4 + # Compact page width, fixed by the first plan() (see _set_width). + self.width = 0 + self.num_groups = 0 + self.group = 0 + self.max_reqs = max_reqs + self.rows_max = max_reqs * DECODE_QUERY_LEN + self.device = device + self._plans: dict[int, Nvfp4DecodeCPPlan] = {} + self._tile_words = max_reqs * DECODE_QUERY_LEN * NUM_INDEX_HEADS * HEAD_DIM // 4 + self._recv_words = self.cp * self.rows_max * TOPK_BLOCKS + + def new_layer(self) -> Nvfp4DecodeCPLayer: + """Collective: every TP rank builds its sparse layers in the same order.""" + import torch.distributed._symmetric_memory as symm_mem + + group_name = get_tp_group().device_group.group_name + tile = symm_mem.empty(self._tile_words, dtype=torch.int32, device=self.device) + tile.fill_(-1) + recv = symm_mem.empty(self._recv_words, dtype=torch.int64, device=self.device) + recv.fill_(-1) + tile_ptrs = symm_mem.rendezvous(tile, group_name).buffer_ptrs + recv_ptrs = symm_mem.rendezvous(recv, group_name).buffer_ptrs + return Nvfp4DecodeCPLayer( + tile=tile, + recv=recv, + tile_ptrs=torch.tensor(tile_ptrs, dtype=torch.int64, device=self.device), + recv_ptrs=torch.tensor(recv_ptrs, dtype=torch.int64, device=self.device), + ) + + def _set_width(self, max_pages: int) -> None: + """Fix the compact page width and the two-level top-k shape. + + One column past the pages a rank can own, so the page the compact + length invents stays inside the score buffer and the chunk's last page + is stored (see ``_compact_block_table_kernel``). Both levels of the + top-k select over power-of-two tiles that cover the width. + """ + width = cdiv(min(max_pages, MAX_PAGES), self.cp) + 1 + if self.width == 0: + padded = 1 << (width - 1).bit_length() + self.width = width + self.group = max(1, padded // MAX_TOPK_GROUPS) + self.num_groups = padded // self.group + assert width <= self.width, "the block table grew after the first plan" + + def plan( + self, block_table: torch.Tensor, seq_lens: torch.Tensor + ) -> Nvfp4DecodeCPPlan: + """Outside CUDA graph capture, once per step.""" + self._set_width(block_table.shape[1]) + num_reqs = seq_lens.shape[0] + p = self._plans.get(num_reqs) + if p is None: + assert num_reqs <= self.max_reqs + i32 = dict(dtype=torch.int32, device=self.device) + bt = torch.empty(num_reqs, self.width, **i32) + score_seq_lens = torch.empty(num_reqs, **i32) + p = self._plans[num_reqs] = Nvfp4DecodeCPPlan( + wrapper=self._wrapper_cls( + num_heads=NUM_INDEX_HEADS, + use_cuda_graph=True, + block_table_buffer=bt, + seq_lens_buffer=score_seq_lens, + ), + block_table=bt, + score_seq_lens=score_seq_lens, + seq_lens=torch.empty(num_reqs, **i32), + num_reqs=num_reqs, + ) + _compact_block_table_kernel[(num_reqs,)]( + block_table, + seq_lens, + p.block_table, + p.score_seq_lens, + p.seq_lens, + min(block_table.shape[1], MAX_PAGES), + self.width, + block_table.stride(0), + p.block_table.stride(0), + CP=self.cp, + RANK=self.rank, + BLOCK=1024, + ) + p.wrapper.plan(p.block_table, p.score_seq_lens) + return p + + def forward( + self, + layer: Nvfp4DecodeCPLayer, + plan: Nvfp4DecodeCPPlan, + index_q: torch.Tensor, # [rows, 1, 128] e4m3, this rank's head + k_cache: torch.Tensor, + out: torch.Tensor, # [rows, 1, 16] int32 + query_len: int, + ) -> None: + rows = index_q.shape[0] + assert rows == plan.num_reqs * query_len and out.shape[0] == rows + _q_exchange_kernel[(rows,)]( + index_q.view(torch.int32).view(rows, HEAD_DIM // 4), + layer.tile_ptrs, + H=NUM_INDEX_HEADS, + CP=self.cp, + RANK=self.rank, + HEAD=self.head, + QLEN=query_len, + num_warps=1, + launch_pdl=True, + ) + q_tile = layer.tile[ + : plan.num_reqs * DECODE_QUERY_LEN * NUM_INDEX_HEADS * HEAD_DIM // 4 + ] + scores = plan.wrapper._run_scores( + q_tile.view(torch.float8_e4m3fn).view(-1, NUM_INDEX_HEADS, HEAD_DIM), + k_cache, + ) + _local_topk_push_kernel[(rows, NUM_INDEX_HEADS)]( + scores, + plan.seq_lens, + layer.recv_ptrs, + self.rows_max, + scores.stride(0), + scores.stride(1), + H=NUM_INDEX_HEADS, + CP=self.cp, + RANK=self.rank, + QLEN=query_len, + WIDTH=self.width, + CHUNK=TOPK_CHUNK, + NUM_GROUPS=self.num_groups, + GROUP=self.group, + num_warps=4, + launch_pdl=True, + ) + _merge_kernel[(rows,)]( + plan.seq_lens, + layer.recv_ptrs, + layer.tile_ptrs, + out, + self.rows_max, + out.stride(0), + H=NUM_INDEX_HEADS, + CP=self.cp, + RANK=self.rank, + QLEN=query_len, + num_warps=1, + launch_pdl=True, + ) + + +_STATE: Nvfp4IndexerDecodeCP | None = None +_RESOLVED = False + + +def get_nvfp4_indexer_decode_cp( + *, num_index_heads: int, max_decode_query_len: int +) -> Nvfp4IndexerDecodeCP | None: + """The process's decode-CP state, or None (stock decode). Built once.""" + global _STATE, _RESOLVED + if _RESOLVED: + return _STATE + _RESOLVED = True + vllm_config = get_current_vllm_config() + if not vllm_config.attention_config.indexer_decode_cp: + return None + tp = get_tensor_model_parallel_world_size() + text_config = vllm_config.model_config.hf_text_config + total_heads = text_config.sparse_attention_config["sparse_num_index_heads"] + why = None + if tp not in SUPPORTED_TP_SIZES: + why = f"tensor parallel size {tp} (needs one of {SUPPORTED_TP_SIZES})" + elif total_heads != NUM_INDEX_HEADS or num_index_heads != 1: + why = f"{total_heads} index heads, {num_index_heads} per rank" + elif vllm_config.parallel_config.decode_context_parallel_size > 1: + why = "decode context parallelism is enabled" + elif max_decode_query_len > DECODE_QUERY_LEN: + why = f"decode query length {max_decode_query_len} > {DECODE_QUERY_LEN}" + elif not direct_cp_enabled(get_tp_group(), torch.int32, use_direct=None): + why = "the tensor-parallel ranks have no symmetric memory between them" + if why is not None: + logger.warning( + "indexer_decode_cp ignored: %s; using the stock NVFP4 indexer decode.", + why, + ) + return None + sched = vllm_config.scheduler_config + _STATE = Nvfp4IndexerDecodeCP( + max_reqs=sched.max_num_seqs, + device=torch.device("cuda", torch.accelerator.current_device_index()), + ) + logger.info_once( + "MiniMax M3 NVFP4 indexer: decode context parallelism over %d TP ranks " + "(each scores 1/%d of the index-K pages for all %d index heads).", + tp, + tp, + NUM_INDEX_HEADS, + ) + return _STATE diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index d67e9ea19d90..6afe4defdb9d 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -18,11 +18,18 @@ from torch import nn from transformers import PreTrainedConfig +import vllm.envs as envs from vllm import _custom_ops as ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.forward_context import get_forward_context +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMulWithClamp from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention.attention import set_default_quant_scales @@ -35,6 +42,28 @@ GateLinear, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + MEGA_MOE_NVFP4, + DeepGemmMegaMoEExperts, + MegaMoEActivation, +) +from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( + TrtLlmNvFp4ExpertsModular, + TrtLlmNvFp4ExpertsMonolithic, +) +from vllm.model_executor.layers.fused_moe.moe_output import ( + MoEOutput, + UnfinalizedMoEOutput, +) +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import _unpack +from vllm.model_executor.layers.fusion.quant_activation import ( + QuantizedActivation, + as_quantized_activation, + get_input_quant_key, +) from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, MinimaxM3QKVParallelLinearWithIndexer, @@ -43,6 +72,13 @@ ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp8Dynamic, + kNvfp4DynamicToken, +) from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -63,6 +99,7 @@ AutoWeightsLoader, PPMissingLayer, WeightsMapper, + extract_layer_index, init_vllm_registered_model, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -71,6 +108,11 @@ spec_decode_needs_target_embed, ) from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model +from vllm.models.common.ops.sequence_parallel import ( + sp_all_gather, + sp_padding_mask, + sp_shard, +) from vllm.models.minimax_m3.common.encoder_cudagraph import ( MiniMaxM3EncoderCudaGraphMixin, ) @@ -87,7 +129,9 @@ ) from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.utils.flashinfer import is_flashinfer_autotuning from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -95,6 +139,44 @@ get_kv_quant_mode, ) +logger = init_logger(__name__) + +# Largest batch whose MoE top-k finalize runs inside the next layer's fused +# all-reduce + input norm + MXFP8 quantization. Larger batches finalize in the +# MoE and use the regular fused all-reduce + norm. +FUSED_MOE_FINALIZE_MAX_TOKENS = 32 + +_QKV_FUSED_EPILOGUE_MIN_TOKENS = 512 + + +def _supports_qkv_fused_epilogue(attn: "MiniMaxM3SparseAttention") -> bool: + """Whether minimax_m3_qkv_index_fused_epilogue can replace this layer's + qkv_proj + norm/RoPE/cache-insert chain.""" + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, + ) + + if not current_platform.is_device_capability_family(100): + return False + # Both keep the weight K-major and its scales F8_128x4-swizzled. + kernel = getattr(attn.qkv_proj.quant_method, "kernel", None) + if not isinstance( + kernel, (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel) + ): + return False + nkv = attn.num_kv_heads + return ( + attn.use_nvfp4_kv + and attn.indexer_kv_dtype == "nvfp4" + and nkv in (1, 2, 4) + and attn.num_heads == 16 * nkv + and attn.num_idx_heads == nkv + and attn.head_dim == attn.idx_head_dim == 128 + and attn.rotary_emb.rotary_dim == 64 + and attn.hidden_size % 128 == 0 + ) + def _sparse_attention_layer_ids(config: PreTrainedConfig) -> set[int]: """Layer ids whose attention runs the extra sparse "index" branch.""" @@ -115,6 +197,73 @@ def _is_moe_layer(config: PreTrainedConfig, layer_id: int) -> bool: return moe_layer_freq[layer_id] != 0 +def _use_mega_moe(vllm_config: VllmConfig) -> bool: + """Whether MoE layers run the DeepGEMM NVFP4 MegaMoE kernel. + + Opt in with ``--kernel-config '{"moe_backend": "deep_gemm_mega_moe"}'``. Only + checkpoints with NVFP4 routed experts (MiniMax-M3.1) are supported. + """ + if vllm_config.kernel_config.moe_backend != "deep_gemm_mega_moe": + return False + moe_quant_algo = getattr(vllm_config.quant_config, "moe_quant_algo", None) + if moe_quant_algo != "NVFP4": + raise NotImplementedError( + "MiniMax M3 MegaMoE requires a checkpoint with NVFP4 routed experts " + f"(moe_quant_algo: NVFP4), got moe_quant_algo={moe_quant_algo!r}. " + "Pick a different moe backend for this checkpoint." + ) + if not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "MiniMax M3 MegaMoE requires expert parallel. Enable it with " + "--enable-expert-parallel." + ) + return True + + +# MiniMax-M3.1 QAT quantizes the FC2 input with a fixed global scale of 16; +# checkpoints carry a placeholder ``w2.input_global_scale`` of 1.0. +MEGA_MOE_FC2_INPUT_SCALE = 1.0 / 16.0 + + +def _mega_moe_expert_params_mapping( + num_experts: int, +) -> list[tuple[str, str, int, str]]: + """Map ``experts.{e}.w{1,2,3}.*`` to the MegaMoE ``w13_*``/``w2_*`` params.""" + return [ + ( + "experts.w2_" if shard_id == "w2" else "experts.w13_", + f"experts.{expert_id}.{shard_id}.", + expert_id, + shard_id, + ) + for expert_id in range(num_experts) + for shard_id in ("w1", "w2", "w3") + ] + + +def _mega_moe_expert_weight( + name: str, weight: torch.Tensor +) -> tuple[str, torch.Tensor] | None: + """Convert a compressed-tensors NVFP4 routed-expert tensor to the ModelOpt + convention ``DeepGemmMegaMoEExperts`` loads, or None to drop it. + + The global scales become dequantization scales. The FC1 input is quantized + per token, so its input scale is dropped, and the FC2 input scale is the + fixed QAT scale whatever the checkpoint says. + """ + prefix, _, suffix = name.rpartition(".") + if suffix == "weight_packed": + return f"{prefix}.weight", weight + if suffix == "weight_global_scale": + return f"{prefix}.weight_scale_2", 1.0 / weight.float() + if suffix == "input_global_scale": + if not prefix.endswith(".w2"): + return None + scale = torch.full_like(weight, MEGA_MOE_FC2_INPUT_SCALE, dtype=torch.float32) + return f"{prefix}.input_scale", scale + return name, weight + + class MiniMAXGemmaRMSNorm(nn.Module): """Gemma-style RMS normalization backed by FlashInfer kernels. @@ -155,6 +304,7 @@ def __init__( intermediate_size: int, quant_config: QuantizationConfig | None = None, reduce_results: bool = True, + disable_tp: bool = False, prefix: str = "", ) -> None: super().__init__() @@ -163,6 +313,7 @@ def __init__( [intermediate_size] * 2, bias=False, quant_config=quant_config, + disable_tp=disable_tp, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( @@ -171,6 +322,7 @@ def __init__( bias=False, quant_config=quant_config, reduce_results=reduce_results, + disable_tp=disable_tp, prefix=f"{prefix}.down_proj", ) if config.hidden_act != "swigluoai": @@ -184,10 +336,35 @@ def __init__( alpha=config.swiglu_alpha, beta=config.swiglu_beta, ) + from vllm.models.minimax_m3.nvidia.ops import dense_gemm1_mxfp8 + + # At or below the shared-experts stream threshold the persistent kernel + # would contend with the routed experts on the aux stream. + self.fused_gemm1_min_tokens: int | None = None + if dense_gemm1_mxfp8.can_fuse_gemm1( + self.gate_up_proj, self.down_proj, self.act_fn + ): + self.fused_gemm1_min_tokens = max( + dense_gemm1_mxfp8.MIN_TOKENS, + envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD, + ) def forward(self, x: torch.Tensor) -> torch.Tensor: - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) + num_tokens = x.shape[0] + # Run unfused under FlashInfer autotuning so gate_up_proj is tuned too + if ( + self.fused_gemm1_min_tokens is not None + and num_tokens > self.fused_gemm1_min_tokens + and not is_flashinfer_autotuning() + ): + from vllm.models.minimax_m3.nvidia.ops.dense_gemm1_mxfp8 import ( + fused_gemm1, + ) + + x = fused_gemm1(x, self.gate_up_proj, self.act_fn) + else: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) x, _ = self.down_proj(x) return x @@ -203,6 +380,7 @@ def __init__( quant_config: QuantizationConfig | None = None, reduce_results: bool = True, prefix: str = "", + use_mega_moe: bool = False, ) -> None: super().__init__() self.tp_size = get_tensor_model_parallel_world_size() @@ -212,6 +390,10 @@ def __init__( f"the number of experts {config.num_local_experts}." ) + self.use_mega_moe = use_mega_moe + self.reduce_results = reduce_results + self.scoring_func = config.scoring_func + self.top_k = config.num_experts_per_tok self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) self.n_shared_experts = getattr(config, "n_shared_experts", None) @@ -239,6 +421,10 @@ def __init__( ) self.shared_experts: MiniMaxM3MLP | None = None + if self.use_mega_moe: + self._init_mega_moe_experts(config, quant_config, prefix) + return + if self.n_shared_experts: self.shared_experts = MiniMaxM3MLP( config=config, @@ -265,20 +451,125 @@ def __init__( swiglu_beta=config.swiglu_beta, routed_scaling_factor=self.routed_scaling_factor, router_logits_dtype=self.gate.out_dtype, + # The NVFP4 routed experts are QAT'd with a per-token FC1 input + # scale; their ``input_global_scale`` is a 1.0 placeholder that + # the quantization config cannot tell from a calibrated scale. + qat_activation_key=kNvfp4DynamicToken, shared_experts=self.shared_experts, quant_config=quant_config, reduce_results=reduce_results, prefix=f"{prefix}.experts", ) + def _init_mega_moe_experts( + self, + config: PreTrainedConfig, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + if config.hidden_act != "swigluoai": + raise ValueError( + f"MiniMax M3 MegaMoE requires swigluoai, got {config.hidden_act}." + ) + # Hidden states are replicated across TP ranks, so each rank routes + # only its own token shard through the EP-wide kernel. + self.use_sequence_parallel = self.tp_size > 1 + + def replicated_shared_experts() -> MiniMaxM3MLP: + # Replicated, so the kernel can fuse it. + return MiniMaxM3MLP( + config=config, + intermediate_size=config.intermediate_size * self.n_shared_experts, + quant_config=quant_config, + reduce_results=False, + disable_tp=True, + prefix=f"{prefix}.shared_experts", + ) + + shared_experts = None + if self.n_shared_experts: + # Only describes the weights: a fused layer allocates its own. + with torch.device("meta"): + shared_experts = replicated_shared_experts() + ep_group = get_ep_group() + num_experts = config.num_local_experts + if num_experts % ep_group.world_size != 0: + raise ValueError( + f"MiniMax M3 MegaMoE needs num_experts={num_experts} divisible " + f"by the EP size {ep_group.world_size}." + ) + num_local_experts = num_experts // ep_group.world_size + self.experts = DeepGemmMegaMoEExperts( + get_current_vllm_config(), + num_experts=num_experts, + num_local_experts=num_local_experts, + experts_start_idx=ep_group.rank_in_group * num_local_experts, + top_k=self.top_k, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + activation=MegaMoEActivation( + "swiglu", + clamp=config.swiglu_limit, + alpha=config.swiglu_alpha, + beta=config.swiglu_beta, + ), + weight_format=MEGA_MOE_NVFP4, + shared_experts=shared_experts, + prefix=f"{prefix}.experts", + ) + if shared_experts is not None and not self.experts.has_fused_shared_experts: + self.shared_experts = replicated_shared_experts() + @staticmethod def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None: assert param.size() == loaded_weight.size() param.data.copy_(loaded_weight.to(torch.float32)) + @property + def ffn_all_reduce_deferred(self) -> bool: + if self.use_mega_moe: + return not self.reduce_results + return self.experts.moe_config.skip_final_all_reduce + + def defer_finalize(self, max_num_tokens: int) -> bool: + """Leave the routed top-k reduction of calls up to ``max_num_tokens`` + to the consumer; return whether the experts accepted. + + Only the TRTLLM-Gen NVFP4 experts can stop after GEMM2. + """ + if self.use_mega_moe: + return False + experts_cls = getattr( + self.experts.routed_experts.quant_method, "experts_cls", None + ) + if experts_cls not in (TrtLlmNvFp4ExpertsMonolithic, TrtLlmNvFp4ExpertsModular): + return False + moe_config = self.experts.moe_config + moe_config.defer_moe_finalize(max_num_tokens) + return moe_config.use_deferred_moe_finalize + + def defers_finalize(self, num_tokens: int) -> bool: + return ( + not self.use_mega_moe + and self.experts.moe_config.should_defer_moe_finalize(num_tokens) + ) + + def forward_unfinalized(self, hidden_states: torch.Tensor) -> MoEOutput: + """``forward`` with the routed top-k reduction and all-reduce left open.""" + router_logits, _ = self.gate(hidden_states) + # The runner's custom op returns tensors only, so run its body directly. + shared_output, routed = _unpack( + self.experts._forward_impl(hidden_states, router_logits, hidden_states) + ) + assert shared_output is not None + assert isinstance(routed, UnfinalizedMoEOutput) + return MoEOutput(routed=routed, shared_output=shared_output) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) + if self.use_mega_moe: + return self._forward_mega_moe(hidden_states) # router_logits: (num_tokens, n_experts); GateLinear casts to fp32. router_logits, _ = self.gate(hidden_states) @@ -288,6 +579,54 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return final_hidden_states.view(num_tokens, hidden_dim) + def _forward_mega_moe(self, hidden_states: torch.Tensor) -> torch.Tensor: + if not self.use_sequence_parallel: + return self._forward_mega_moe_local(hidden_states) + + num_tokens = hidden_states.shape[0] + # Routing and the experts read the padding mask from the forward + # context; hand them this rank's shard of it. + forward_context = None + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + forward_context = get_forward_context() + is_padding = forward_context.is_padding + forward_context.is_padding = sp_padding_mask( + is_padding[:num_tokens] if is_padding is not None else None, + hidden_states, + ) + try: + local_output = self._forward_mega_moe_local(sp_shard(hidden_states)) + finally: + if forward_context is not None: + forward_context.is_padding = is_padding + if self.reduce_results: + return sp_all_gather(local_output)[:num_tokens] + + # Leave the other ranks' rows zero, so the deferred TP all-reduce (fused + # into the next RMSNorm) gathers them. + shard_size, hidden_dim = local_output.shape + tp_rank = get_tensor_model_parallel_rank() + output = local_output.new_zeros(shard_size * self.tp_size, hidden_dim) + output[tp_rank * shard_size : (tp_rank + 1) * shard_size] = local_output + return output[:num_tokens] + + def _forward_mega_moe_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + scoring_func=self.scoring_func, + e_score_correction_bias=self.e_score_correction_bias.data + if self.e_score_correction_bias is not None + else None, + topk=self.top_k, + renormalize=True, + routed_scaling_factor=self.routed_scaling_factor, + ) + return self.experts( + hidden_states, topk_weights, topk_ids, shared_experts=self.shared_experts + ) + class MiniMaxM3Attention(nn.Module): """Dense attention with per-head QK norm and partial RoPE.""" @@ -556,6 +895,26 @@ def __init__( compilation_config.static_forward_context[self.layer_name] = self self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache + self.use_qkv_fused_epilogue = _supports_qkv_fused_epilogue(self) + if self.use_qkv_fused_epilogue: + from vllm.models.minimax_m3.nvidia.ops.qkv_index_fused_epilogue import ( + register_qkv_index_fused_epilogue_warmup, + ) + + logger.info_once("MiniMax-M3: QKV+indexer GEMM with fused epilogue enabled") + block_size = vllm_config.cache_config.block_size + # The query is E4M3-only when CUTLASS decode serves the layer. + q_is_fp8 = getattr(self.impl, "use_cutlass_decode", False) + if vllm_config.kernel_config.enable_jit_warmup and block_size is not None: + register_qkv_index_fused_epilogue_warmup( + has_q=not q_is_fp8, + has_qfp8=q_is_fp8, + nkv=self.num_kv_heads, + block_size=block_size, + weight_bytes=self.qkv_proj.output_size_per_partition + * self.hidden_size, + ) + def get_attn_backend(self) -> type[MiniMaxM3SparseBackend]: return self.attn_backend @@ -570,37 +929,74 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) - def _allocate_query_fp8(self, qkv: torch.Tensor) -> torch.Tensor | None: - if not getattr(self.impl, "use_cutlass_decode", False): - return None + def _qkv_proj_fused_epilogue( + self, + hidden_states: torch.Tensor | QuantizedActivation, + positions: torch.Tensor, + q: torch.Tensor, + index_q: torch.Tensor, + slot_mapping: torch.Tensor, + index_slot_mapping: torch.Tensor, + ) -> None: + """qkv_proj + QK-norm/RoPE + KV/index-cache insert as one MXFP8 kernel, + writing q (bf16 or E4M3, as allocated) and index_q.""" + from vllm.models.minimax_m3.nvidia.ops.qkv_index_fused_epilogue import ( + minimax_m3_qkv_index_fused_epilogue, + ) + + weight = self.qkv_proj.weight + if weight.shape[1] != self.hidden_size: + weight = weight.t() # the CuTe-DSL kernel keeps a [K, N] view + qa = as_quantized_activation(hidden_states, kMxfp8Dynamic) + if qa is not None: + xq, xsf = qa.data, qa.scale + else: + assert isinstance(hidden_states, torch.Tensor) + xq, xsf = mxfp8_e4m3_quantize(hidden_states, is_sf_swizzled_layout=True) + q_is_fp8 = q.dtype == torch.float8_e4m3fn + minimax_m3_qkv_index_fused_epilogue( + xq, + xsf, + weight, + self.qkv_proj.weight_scale, + positions, + self.rotary_emb.cos_sin_cache, + self.q_norm.weight, + self.k_norm.weight, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.q_norm.variance_epsilon, + slot_mapping=slot_mapping, + index_slot_mapping=index_slot_mapping, + kv_cache=self.kv_cache, + index_cache=self.indexer.index_cache.kv_cache, + q_out=None if q_is_fp8 else q, + index_q_out=index_q, + q_fp8_out=q if q_is_fp8 else None, + q_fp8_scale=self._q_scale_float, + k_scale=self._k_scale_float, + v_scale=self._v_scale_float, + ) + + def _allocate_query( + self, num_tokens: int, dtype: torch.dtype, device: torch.device + ) -> torch.Tensor: + # CUTLASS decode reads an E4M3 query, so the insert emits only that. + fp8 = getattr(self.impl, "use_cutlass_decode", False) return torch.empty( - (qkv.shape[0], self.q_size), - dtype=torch.float8_e4m3fn, - device=qkv.device, + (num_tokens, self.q_size), + dtype=torch.float8_e4m3fn if fp8 else dtype, + device=device, ) def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, ) -> torch.Tensor: - # Single fused projection emitting [q | k | v | index_q | index_k]. - qkv, _ = self.qkv_proj(hidden_states) - - # Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on the - # main (q/k) and index (index_q/index_k) branches, all read straight out - # of the single fused ``qkv`` tensor (the "5 results"). Once the paged - # caches are bound the kernel also inserts k/v and the index key into - # them; the initial memory-profiling run (caches unbound, no slot_mapping) - # short-circuits to zeros below. k/v and index_k are rewritten in place - # inside qkv (and scatter-inserted into the caches); q and index_q are - # de-interleaved - # straight into the dedicated contiguous ``q``/``index_q`` buffers below. - cos_sin_cache = self.rotary_emb.cos_sin_cache rotary_dim = self.rotary_emb.rotary_dim eps = self.q_norm.variance_epsilon - num_tokens = qkv.shape[0] fwd_slot_mapping = get_forward_context().slot_mapping if ( @@ -608,49 +1004,83 @@ def forward( or self.layer_name not in fwd_slot_mapping ): # Memory-profiling run: caches not yet bound, slot_mapping is empty. - return qkv.new_zeros((num_tokens, self.hidden_size)) + qkv, _ = self.qkv_proj(hidden_states) + self.indexer.reserve_workspace() + return qkv.new_zeros((qkv.shape[0], self.hidden_size)) main_slot_mapping = fwd_slot_mapping[self.layer_name] index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix] - q = qkv.new_empty((num_tokens, self.q_size)) - # With query_fp8, q is emitted only in fp8; the attend dequantizes it - # into ``q`` for the tokens that run on Triton decode. - query_fp8 = self._allocate_query_fp8(qkv) - # index_q matches the index-K cache dtype (e4m3 for the fp8 score path); - # the fused kernel emits fp8 directly when this buffer is e4m3. - index_q = qkv.new_empty( - (num_tokens, self.index_q_size), - dtype=self.indexer.index_cache.dtype, + # FlashInfer autotunes the GEMMs its dummy runs call: take the unfused + # path there so qkv_proj is tuned for the M < 512 fallback. + if isinstance(hidden_states, QuantizedActivation): + num_tokens = hidden_states.orig_shape[0] + dtype, device = hidden_states.orig_dtype, hidden_states.data.device + else: + num_tokens = hidden_states.shape[0] + dtype, device = hidden_states.dtype, hidden_states.device + fused = ( + self.use_qkv_fused_epilogue + and num_tokens >= _QKV_FUSED_EPILOGUE_MIN_TOKENS + and not is_flashinfer_autotuning() ) - ops.fused_minimax_m3_qknorm_rope_kv_insert( - qkv, - self.q_norm.weight, - self.k_norm.weight, - cos_sin_cache, - positions, - self.num_heads, - self.num_kv_heads, - rotary_dim, - eps, - self.index_q_norm.weight, - self.index_k_norm.weight, - self.num_idx_heads, - main_slot_mapping, - index_slot_mapping, - self.kv_cache, - self.indexer.index_cache.kv_cache, - self.kv_cache.size(2), # paged-cache block size - q if query_fp8 is None else None, - index_q, - self.kv_cache_dtype, - q_fp8_out=query_fp8, - q_fp8_scale=self._q_scale_float, - kv_k_scale=self._k_scale if self.use_nvfp4_kv else None, - kv_v_scale=self._v_scale if self.use_nvfp4_kv else None, + if not fused: + # Single fused projection emitting [q | k | v | index_q | index_k]. + qkv, _ = self.qkv_proj(hidden_states) + q = self._allocate_query(num_tokens, dtype, device) + q_is_fp8 = q.dtype == torch.float8_e4m3fn + # index_q is e4m3 for the fp8 and nvfp4 index caches, else bf16; the + # fused kernel emits fp8 directly when this buffer is e4m3. + index_q = torch.empty( + (num_tokens, self.index_q_size), + dtype=self.indexer.index_cache.query_dtype, + device=device, ) + if fused: + self._qkv_proj_fused_epilogue( + hidden_states, + positions, + q, + index_q, + main_slot_mapping, + index_slot_mapping, + ) + else: + # Horizontally-fused per-head Gemma QK-norm + partial NeoX RoPE on + # the main (q/k) and index (index_q/index_k) branches, all read + # straight out of qkv. k/v and index_k are scatter-inserted into the + # caches; q and index_q are de-interleaved into the dedicated + # contiguous buffers above. With an E4M3 q nothing reads qkv + # afterwards, so k and index_k are not rewritten in place. + ops.fused_minimax_m3_qknorm_rope_kv_insert( + qkv, + self.q_norm.weight, + self.k_norm.weight, + cos_sin_cache, + positions, + self.num_heads, + self.num_kv_heads, + rotary_dim, + eps, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.num_idx_heads, + main_slot_mapping, + index_slot_mapping, + self.kv_cache, + self.indexer.index_cache.kv_cache, + self.kv_cache.size(2), # paged-cache block size + None if q_is_fp8 else q, + index_q, + self.kv_cache_dtype, + q_fp8_out=q if q_is_fp8 else None, + q_fp8_scale=self._q_scale_float, + kv_k_scale=self._k_scale if self.use_nvfp4_kv else None, + kv_v_scale=self._v_scale if self.use_nvfp4_kv else None, + write_qkv=not q_is_fp8, + ) - output = torch.empty_like(q) - attn_output = self._run_attention(q, query_fp8, index_q, output) + output = torch.empty((num_tokens, self.q_size), dtype=dtype, device=device) + attn_output = self._run_attention(q, index_q, output) output, _ = self.o_proj(attn_output) return output @@ -658,7 +1088,6 @@ def forward( def _run_attention( self, query: torch.Tensor, - query_fp8: torch.Tensor | None, index_query: torch.Tensor, output: torch.Tensor, ) -> torch.Tensor: @@ -666,13 +1095,7 @@ def _run_attention( # metadata and can't be captured into a cudagraph. The indexer writes its # top-k into the shared ``topk_indices_buffer``; the attend reads it back. self.indexer(index_query) - return self.impl.forward( - self, - query, - self.kv_cache, - output, - query_fp8=query_fp8, - ) + return self.impl.forward(self, query, self.kv_cache, output) class MiniMaxM3DecoderLayer(nn.Module): @@ -737,12 +1160,18 @@ def __init__( ) self.is_moe_layer = force_moe or _is_moe_layer(config, layer_id) if self.is_moe_layer: + use_mega_moe = _use_mega_moe(vllm_config) + if use_mega_moe and is_mtp_block: + raise NotImplementedError( + "MiniMax M3 MegaMoE does not support MTP draft layers." + ) self.block_sparse_moe = MiniMaxM3MoE( config=config, layer_id=layer_id, quant_config=quant_config, reduce_results=reduce_results, prefix=f"{prefix}.block_sparse_moe", + use_mega_moe=use_mega_moe, ) else: self.mlp = MiniMaxM3MLP( @@ -764,9 +1193,29 @@ def __init__( def forward( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | MoEOutput, residual: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor | MoEOutput, torch.Tensor]: + hidden_states, residual = self.forward_input_norm(hidden_states, residual) + return self.forward_after_input_norm(positions, hidden_states, residual) + + def forward_input_norm( + self, hidden_states: torch.Tensor | MoEOutput, residual: torch.Tensor | None + ) -> tuple[torch.Tensor | QuantizedActivation, torch.Tensor]: + if isinstance(hidden_states, MoEOutput): + # The preceding MoE left a small batch unfinalized: finalize, + # all-reduce, norm and MXFP8-quantize the qkv_proj input at once. + assert residual is not None + from vllm.models.minimax_m3.nvidia.ops.finalize_allreduce_norm import ( + finalize_allreduce_gemma_rms_norm_mxfp8, + ) + + return finalize_allreduce_gemma_rms_norm_mxfp8( + hidden_states, + residual, + self.input_layernorm.weight, + self.input_layernorm.variance_epsilon, + ) if self.fuse_input_allreduce and residual is not None: hidden_states, residual = fused_allreduce_gemma_rms_norm( hidden_states, residual, self.input_layernorm @@ -777,6 +1226,14 @@ def forward( hidden_states = self.input_layernorm(hidden_states) else: hidden_states, residual = self.input_layernorm(hidden_states, residual) + return hidden_states, residual + + def forward_after_input_norm( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor | QuantizedActivation, + residual: torch.Tensor, + ) -> tuple[torch.Tensor | MoEOutput, torch.Tensor]: hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, @@ -785,6 +1242,11 @@ def forward( hidden_states, residual = fused_allreduce_gemma_rms_norm( hidden_states, residual, self.post_attention_layernorm ) + if self.is_moe_layer and self.block_sparse_moe.defers_finalize( + hidden_states.shape[0] + ): + # The next layer's input norm finalizes it with the all-reduce. + return self.block_sparse_moe.forward_unfinalized(hidden_states), residual ffn = self.block_sparse_moe if self.is_moe_layer else self.mlp hidden_states = ffn(hidden_states) return hidden_states, residual @@ -794,7 +1256,7 @@ def ffn_all_reduce_deferred(self) -> bool: """This layer's FFN output is left un-reduced; the caller fuses the all-reduce into the next RMSNorm.""" if self.is_moe_layer: - return self.block_sparse_moe.experts.moe_config.skip_final_all_reduce + return self.block_sparse_moe.ffn_all_reduce_deferred return not self.mlp.down_proj.reduce_results @@ -808,6 +1270,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_text_config quant_config = vllm_config.quant_config self.config = config + self.use_mega_moe = _use_mega_moe(vllm_config) self.vocab_size = config.vocab_size @@ -871,6 +1334,54 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): layer.fuse_input_allreduce = idx > 0 and prev_defers prev_defers = layer.ffn_all_reduce_deferred self.fuse_final_norm_allreduce = prev_defers + if current_platform.is_cuda(): + self._init_fused_moe_finalize(vllm_config) + + def _init_fused_moe_finalize(self, vllm_config: VllmConfig) -> None: + """Let small batches leave a MoE's top-k finalize to the next layer's + input norm, which already fuses the TP all-reduce, and have it quantize + the normed rows to the MXFP8 its qkv_proj reads. + + Each MoE whose successor can take the output gets the hint here, and + ``defers_finalize`` decides per call. The final norm feeds the LM head + in BF16, so the last layer keeps its own finalize. Builds a collective + kernel: every TP rank decides alike. + """ + from vllm.models.minimax_m3.nvidia.ops.finalize_allreduce_norm import ( + init_fused_moe_finalize, + supports_fused_moe_finalize, + ) + + if not supports_fused_moe_finalize(vllm_config): + return + layers = self.layers[self.start_layer : self.end_layer] + deferring = [ + layer.block_sparse_moe + for layer, successor in zip(layers, layers[1:]) + if layer.is_moe_layer + and successor.fuse_input_allreduce + and get_input_quant_key(successor.self_attn.qkv_proj) == kMxfp8Dynamic + and layer.block_sparse_moe.defer_finalize(FUSED_MOE_FINALIZE_MAX_TOKENS) + ] + if not deferring: + return + # The modular experts return the router's FP32 routing weights; the + # monolithic kernel routes itself in BF16. + experts_cls = getattr( + deferring[0].experts.routed_experts.quant_method, "experts_cls", None + ) + init_fused_moe_finalize( + hidden_size=self.config.hidden_size, + max_num_tokens=FUSED_MOE_FINALIZE_MAX_TOKENS, + top_k=self.config.num_experts_per_tok, + fp32_weights=experts_cls is TrtLlmNvFp4ExpertsModular, + ) + logger.info_once( + "MiniMax M3: MoE top-k finalize fused into the next all-reduce + " + "input norm + MXFP8 quantization for up to %d tokens (%d layers).", + FUSED_MOE_FINALIZE_MAX_TOKENS, + len(deferring), + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -894,23 +1405,30 @@ def forward( residual = intermediate_tensors["residual"] remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) + + # Aux hidden state k is the residual stream entering layer k. Under TP + # it is only complete after layer k's input norm (or the final norm) + # finishes the previous FFN's deferred all-reduce, so it is read there. + # Cloned, as the norms may update `residual` in place. aux_hidden_states: list[torch.Tensor] = [] - if get_pp_group().is_first_rank: - self._maybe_add_hidden_state( - aux_hidden_states, self.start_layer, hidden_states, residual - ) - # Preserve the entry tap before in-place residual updates. - if aux_hidden_states: - aux_hidden_states[0] = aux_hidden_states[0].clone() for idx, layer in enumerate( self.layers[self.start_layer : self.end_layer], start=self.start_layer ): - hidden_states, residual = layer(positions, hidden_states, residual) - self._maybe_add_hidden_state( - aux_hidden_states, idx + 1, hidden_states, residual + hidden_states, residual = layer.forward_input_norm(hidden_states, residual) + # A later PP stage's start boundary is captured by the stage before. + if idx in self.aux_hidden_state_layers and ( + idx > self.start_layer or get_pp_group().is_first_rank + ): + aux_hidden_states.append(residual.clone()) + hidden_states, residual = layer.forward_after_input_norm( + positions, hidden_states, residual ) if not get_pp_group().is_last_rank: + # Under PP every FFN reduces its own output: no all-reduce is pending. + self._maybe_add_hidden_state( + aux_hidden_states, self.end_layer, hidden_states, residual + ) return IntermediateTensors( { "hidden_states": hidden_states, @@ -920,11 +1438,13 @@ def forward( ) if self.fuse_final_norm_allreduce: - hidden_states, _ = fused_allreduce_gemma_rms_norm( + hidden_states, residual = fused_allreduce_gemma_rms_norm( hidden_states, residual, self.norm ) else: - hidden_states, _ = self.norm(hidden_states, residual) + hidden_states, residual = self.norm(hidden_states, residual) + if self.end_layer in self.aux_hidden_state_layers: + aux_hidden_states.append(residual.clone()) aux_hidden_states = remote_aux + aux_hidden_states if len(aux_hidden_states) > 0: @@ -932,6 +1452,8 @@ def forward( return hidden_states def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + if self.use_mega_moe: + return _mega_moe_expert_params_mapping(self.config.num_local_experts) # Checkpoint experts use w1=gate, w2=down, w3=up. return fused_moe_make_expert_params_mapping( self, @@ -975,6 +1497,23 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if "weight_scale_inv" in name: name = name.replace("weight_scale_inv", "weight_scale") + if self.use_mega_moe and ".block_sparse_moe.experts." in name: + converted = _mega_moe_expert_weight(name, loaded_weight) + if converted is None: + continue + name, loaded_weight = converted + if self.use_mega_moe and ".shared_experts." in name: + if is_pp_missing_parameter(name, self): + continue + moe = self.layers[extract_layer_index(name)].block_sparse_moe + if moe.experts.has_fused_shared_experts: + name = name.replace(".gate_proj.", ".w1.") + name = name.replace(".up_proj.", ".w3.") + loaded_params.add( + moe.experts.load_shared_expert_weight(name, loaded_weight) + ) + continue + for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue diff --git a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py index 970b31c70233..f12232c0ee99 100644 --- a/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py +++ b/vllm/models/minimax_m3/nvidia/msa_cutlass_sparse_decode.py @@ -22,11 +22,22 @@ # fmha_sm100 plans one row per query head. Keep every cached plan within the # fixed planner allocation used by the MSA decode kernel. _MAX_QUERY_HEAD_ROWS = 65536 +# The Q8KV4 kernel (every NVFP4 cache) plans without that allocation. Without +# split KV it launches sequences x query tokens in gridDim.z, which also keeps +# the sequence count under its 16-bit work-item field. +_MAX_Q8KV4_QUERY_TOKENS = 65535 _MAX_DECODE_QUERY_LEN = 32 _DEFAULT_MIN_CUTLASS_BATCH_SIZE = 16 _B300_EAGLE3_QUERY_HEAD_ROWS = 1024 _B300_EAGLE3_MAX_QUERY_LEN = 8 _B300_EAGLE3_HEAD_GEOMETRIES = {(64, 4), (32, 2), (16, 1)} +# Q heads per KV head the Q8KV4 NVFP4 decode kernel serves. +_Q8KV4_GQA_RATIOS = (8, 16) +# Q8KV4 block-scale staging of the NVFP4 KV cache. MiniMax-M3.1 writes it with +# unit global scales, so block scales are e4m3(amax / 6) and every +# code * block_scale already fits E4M3: read it unstaged. Shift 3 is for block +# scales that use the full E4M3 range next to a calibrated global scale. +NVFP4_BLOCK_SCALE_SHIFT = 0 def is_nvfp4_kv_cache(kv_cache_dtype: str) -> bool: @@ -61,13 +72,37 @@ class MSACutlassDecodeMetadata: page_table: torch.Tensor +def uses_q8kv4_decode( + *, + kv_cache_dtype: str, + num_q_heads: int, + num_kv_heads: int, +) -> bool: + """Return whether sparse decode runs on the MSA Q8KV4 kernel: always for an + NVFP4 cache, which is never read with a BF16 query.""" + if not is_nvfp4_kv_cache(kv_cache_dtype): + return False + if ( + num_q_heads % num_kv_heads != 0 + or num_q_heads // num_kv_heads not in _Q8KV4_GQA_RATIOS + ): + raise ValueError( + "MiniMax M3 NVFP4 KV cache decodes only on the MSA Q8KV4 kernel, " + f"which serves {_Q8KV4_GQA_RATIOS} Q heads per KV head; got " + f"{num_q_heads} Q heads over {num_kv_heads} KV heads" + ) + return True + + @triton.jit def _update_runtime_metadata_kernel( seq_lens_ptr, kv_segment_lens_ptr, qo_offset_ptr, + request_seq_lens_ptr, num_rows: tl.constexpr, decode_query_len: tl.constexpr, + HAS_REQUEST_SEQ_LENS: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) @@ -81,6 +116,10 @@ def _update_runtime_metadata_kernel( seq_len - decode_query_len + local_query, mask=mask, ) + if HAS_REQUEST_SEQ_LENS: + tl.store( + request_seq_lens_ptr + request, seq_len, mask=mask & (local_query == 0) + ) @dataclass @@ -101,11 +140,14 @@ def _build_plan( num_kv_heads: int, page_size: int, topk_blocks: int, + q8kv4: bool, ) -> Any: from vllm.third_party.fmha_sm100.api import fmha_sm100_plan qo_lens_cpu = torch.full((batch,), decode_query_len, dtype=torch.int32) - kv_lens_cpu = initial_seq_lens_cpu + # Zero-length full-CUDA-graph padding requests would keep the Q8KV4 + # kernel from planning the batch; lengths are refreshed every step. + kv_lens_cpu = initial_seq_lens_cpu.clamp(min=1) plan = fmha_sm100_plan( qo_lens_cpu, kv_lens_cpu, @@ -120,12 +162,20 @@ def _build_plan( use_fp8_kvcache=True, split_prefill_decode=False, device=device, - # The Q8KV4 route keeps plan-time lengths and its own page offsets, - # which this cached plan does not refresh. - decode_backend="kv_mode3", + decode_backend="q8kv4" if q8kv4 else "kv_mode3", + block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT, ) plan_info = plan[3] + # The Q8KV4 plan addresses pages per request, through the same flat + # block table as the per-token rows below. + q8kv4_plan = plan_info.get("q8kv4") + if q8kv4_plan is not None: + q8kv4_plan["kv_indptr"].copy_( + torch.arange(batch + 1, dtype=torch.int32, device=device).mul_( + page_table_stride + ) + ) row_starts = ( torch.arange(batch, dtype=torch.int32, device=device) .mul_(page_table_stride) @@ -155,6 +205,7 @@ def prepare( num_kv_heads: int, page_size: int, topk_blocks: int, + q8kv4: bool = False, ) -> MSACutlassDecodeMetadata: batch = int(seq_lens.shape[0]) if ( @@ -196,6 +247,7 @@ def prepare( num_kv_heads, page_size, topk_blocks, + q8kv4, ) plan = self.plans.get(key) if plan is None: @@ -209,17 +261,21 @@ def prepare( num_kv_heads=num_kv_heads, page_size=page_size, topk_blocks=topk_blocks, + q8kv4=q8kv4, ) self.plans[key] = plan plan_info = plan[3] + q8kv4_plan = plan_info.get("q8kv4") num_rows = batch * decode_query_len _update_runtime_metadata_kernel[(triton.cdiv(num_rows, 128),)]( seq_lens, plan_info["kv_segment_lens"], plan_info["qo_offset"], + q8kv4_plan["seq_lens"] if q8kv4_plan is not None else seq_lens, num_rows=num_rows, decode_query_len=decode_query_len, + HAS_REQUEST_SEQ_LENS=q8kv4_plan is not None, BLOCK_SIZE=128, ) return MSACutlassDecodeMetadata( @@ -288,6 +344,7 @@ def should_prepare_decode_metadata( ) -> bool: """Return whether a graph shape can use the CUTLASS decode path.""" total_q = batch_size * decode_query_len + nvfp4 = is_nvfp4_kv_cache(kv_cache_dtype) if min_batch_size is None: min_batch_size = _min_cutlass_batch_size( decode_query_len, num_q_heads, num_kv_heads @@ -302,8 +359,12 @@ def should_prepare_decode_metadata( topk_blocks=topk_blocks, ) and 1 <= decode_query_len <= _MAX_DECODE_QUERY_LEN - and (batch_size >= min_batch_size or is_nvfp4_kv_cache(kv_cache_dtype)) - and total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS + and (batch_size >= min_batch_size or nvfp4) + and ( + total_q <= _MAX_Q8KV4_QUERY_TOKENS + if nvfp4 + else total_q * num_q_heads <= _MAX_QUERY_HEAD_ROWS + ) ) @@ -318,9 +379,14 @@ def prepare_decode_metadata( num_kv_heads: int, page_size: int, topk_blocks: int, + q8kv4: bool = False, plan_cache: MSACutlassDecodePlanCache | None = None, ) -> MSACutlassDecodeMetadata: - """Prepare graph-stable runtime metadata for one sparse decode step.""" + """Prepare graph-stable runtime metadata for one sparse decode step. + + ``q8kv4`` plans the Q8KV4 kernel, which serves every NVFP4 cache (see + ``uses_q8kv4_decode``). + """ cache = plan_cache or MSACutlassDecodePlanCache() return cache.prepare( block_table, @@ -331,6 +397,7 @@ def prepare_decode_metadata( num_kv_heads=num_kv_heads, page_size=page_size, topk_blocks=topk_blocks, + q8kv4=q8kv4, ) diff --git a/vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py b/vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py new file mode 100644 index 000000000000..f18819269172 --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/ops/dense_gemm1_mxfp8.py @@ -0,0 +1,664 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP8 dense GEMM1 (gate_up) with fused clamped-SwiGLU + MXFP8 quant epilogue. + +One persistent SM100 tcgen05 MXFP8 GEMM tuned for prefill shapes (BM=128, +BN=128 per gate/up half, BK=128, cta_group=2; 6-warp CTA: warps 0-3 +epilogue, warp 4 MMA, warp 5 TMA) computes ``gate = X @ Wg^T`` and +``up = X @ Wu^T`` and applies clamped SwiGLU and MXFP8 quantization. No +BF16 intermediate. + +The tmem accumulator is double-buffered with a 32-col overlap (even tiles +[0,256), odd [224,480), SF fixed [500,512)): the epilogue reads the shared +32-col group first, so the next tile's MMA waits on a 32-col read instead +of 256. The single epilogue warpgroup processes a tile in 32-col chunks +(load -> act+quant -> store), overlap chunk first. + +Tile order: flat (n-fastest) when the B operand fits in half of L2, else +grouped m-fastest (GROUP_M=32) so B panels stay L2-resident (dense TP1: +B=151MB > 126MB L2). +""" + +from functools import cache + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import ( + Float8E4M3FN, + Float8E8M0FNU, + Float32, + Int32, + Int64, + Uint8, + Uint16, + Uint32, + Uint64, + cute, +) +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + +from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem +from vllm.cute_utils.cvt import fp32x4_to_fp8x4 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic +from vllm.platforms import current_platform + + +class Sm100DenseGemm1Mxfp8: + def __init__(self, GROUP_M: int, use_pdl: bool = True) -> None: + BM, BN, BK = 128, 128, 128 + self.cta_tile = (BM, BN, BK) + self.GROUP_M = GROUP_M + self.use_pdl = use_pdl + + smem_bytes = cutlass.memory.SmemAllocator.capacity_in_bytes() + # A + Bg/2 + Bu/2 tiles + SFA + SFBg + SFBu (2-CTA: B halves, SF full) + self.stage_size = (BM + BN) * BK + (BM + 256) * (BK // 32) + self.num_stages = smem_bytes // self.stage_size + + @cute.jit + def prepare_AB(self, A: cute.Tensor, BM: int, BK: int) -> cpasync.TmaInfo: + tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=tcgen05.CtaGroup.TWO) + swizzle_128B = cute.make_swizzle(3, 4, 3) + s_layout = cute.make_layout((BM, BK, self.num_stages), stride=(BK, 1, BM * BK)) + s_layout = cute.make_composed_layout(swizzle_128B, 0, s_layout) + return cpasync.make_tiled_tma_atom(tma_op, A, s_layout, (BM, BK)) + + @cute.jit + def prepare_SF( + self, + SF: cute.Tensor, + R: Int32, + Ksf: Int32, + BR: int, + BKsf: int, + tma_op: cpasync.CopyG2STileBaseOp, + ) -> cpasync.TmaInfo: + g_layout = cute.make_layout((Ksf * 128, cute.ceil_div(R, 128))) + SF_i64 = cute.recast_tensor(cute.make_tensor(SF.iterator, g_layout), Int64) + tiler = (128 * BKsf // 8, BR // 128) + s_layout = cute.make_layout((*tiler, self.num_stages)) + return cpasync.make_tiled_tma_atom(tma_op, SF_i64, s_layout, tiler) + + @cute.jit + def __call__( + self, + A: cute.Tensor, + B: cute.Tensor, + SFA: cute.Tensor, + SFB: cute.Tensor, + Q: cute.Tensor, + SFQ: cute.Tensor, + limit: Float32, + alpha: Float32, + beta: Float32, + grid_size: Int32, + stream: CUstream, + ) -> None: + M, K = A.shape + N2, _ = B.shape # 2 * N (gate|up) + BM, BN, BK = self.cta_tile + + A_tma = self.prepare_AB(A, BM, BK) + B_tma = self.prepare_AB(B, BN, BK) + SFA_tma = self.prepare_SF( + SFA, + M, + K // 32, + BM, + BK // 32, + cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO), + ) + SFB_tma = self.prepare_SF( + SFB, + N2, + K // 32, + BN, + BK // 32, + cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO), + ) + + self.kernel(A_tma, B_tma, SFA_tma, SFB_tma, Q, SFQ, limit, alpha, beta).launch( + grid=(grid_size, 1, 1), + block=(6 * 32, 1, 1), + cluster=(2, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + A_tma: cpasync.TmaInfo, + B_tma: cpasync.TmaInfo, + SFA_tma: cpasync.TmaInfo, + SFB_tma: cpasync.TmaInfo, + Q_tensor: cute.Tensor, + SFQ_tensor: cute.Tensor, + limit: Float32, + alpha: Float32, + beta: Float32, + ) -> None: + tid, _, _ = cute.arch.thread_idx() + raw_bid, _, _ = cute.arch.block_idx() + num_bids, _, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(tid // 32) + cta_rank = raw_bid % 2 + + BM, BN, BK = self.cta_tile + GROUP_M = self.GROUP_M + num_stages = self.num_stages + + smem = cutlass.memory.SmemAllocator() + sA = smem.allocate_tensor( + Float8E4M3FN, + A_tma.smem_layout.outer, + byte_alignment=128, + swizzle=A_tma.smem_layout.inner, + ) + sB = smem.allocate_tensor( + Float8E4M3FN, + B_tma.smem_layout.outer, + byte_alignment=128, + swizzle=B_tma.smem_layout.inner, + ) + sSFA = smem.allocate_tensor(Int64, SFA_tma.smem_layout, byte_alignment=128) + # SFB slots: rank r's multicast copy fills slot r of BOTH CTAs; + # slice shape/stride per slot matches the atom's (64,1):(1,64) tile + sSFB = smem.allocate_tensor( + Int64, + cute.make_layout((64, 1, 2, num_stages), stride=(1, 64, 64, 128)), + byte_alignment=128, + ) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + tmem_full_mbar = smem.allocate_array(Int64, 2) + tmem_partial_mbar = smem.allocate_array(Int64, 1) + tmem_empty_mbar = smem.allocate_array(Int64, 2) + taddr = smem.allocate(Int32, 4) + + # allocate named barrier + BAR_EPI = 1 + BAR_TMEM_ALLOC = 2 + + M, K = A_tma.tma_tensor.shape + N2, _ = B_tma.tma_tensor.shape + N = N2 // 2 + grid_m = cute.ceil_div(M, BM * 2) * 2 + grid_n = N // BN + total_tiles = grid_m * grid_n + # output SF swizzle: NUM_K_TILES = ceil(N/32/4) tiles of 512B per 128 rows + NUM_KT = cute.ceil_div(N // 32, 4) + + def compute_bid(bid, GROUP_M, grid_m, grid_n): + if cutlass.const_expr(GROUP_M == 0): + bid_m = (bid // (grid_n * 2)) * 2 + bid % 2 + bid_n = (bid // 2) % grid_n + else: + group_size = GROUP_M * grid_n + group_id = bid // group_size + group_m = min(grid_m - group_id * GROUP_M, GROUP_M) + local_bid = bid - group_id * group_size + bid_m = group_id * GROUP_M + local_bid % group_m + bid_n = local_bid // group_m + return bid_m, bid_n + + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 2) + cute.arch.mbarrier_init(tma_empty_mbar + i, 1) + for i in cutlass.range_constexpr(2): + cute.arch.mbarrier_init(tmem_full_mbar + i, 1) + cute.arch.mbarrier_init(tmem_empty_mbar + i, 256) + cute.arch.mbarrier_init(tmem_partial_mbar, 256) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(A_tma.atom) + cpasync.prefetch_descriptor(B_tma.atom) + cpasync.prefetch_descriptor(SFA_tma.atom) + cpasync.prefetch_descriptor(SFB_tma.atom) + + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + if warp_id == 5: + # TMA warp + tma_stage = 0 + parity = 1 + tma_full_mbar_ = to_cta0_smem(tma_full_mbar) + + gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK)) + gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN, BK)) + gSFA_tiles = cute.zipped_divide( + SFA_tma.tma_tensor, (128 * (BK // 32) // 8, BM // 128) + ) + gSFB_tiles = cute.zipped_divide( + SFB_tma.tma_tensor, (128 * (BK // 32) // 8, BN // 128) + ) + + for work_id in range(raw_bid, total_tiles, num_bids): + bid_m, bid_n = compute_bid(work_id, GROUP_M, grid_m, grid_n) + + for iter_k in range(cute.ceil_div(K, BK)): + mbar = tma_full_mbar_ + tma_stage + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, parity) + with cute.arch.elect_one(): + mbarrier.arrive_expect_tx(mbar, self.stage_size, "cluster") + simple_tma_copy( + A_tma.atom, + gA_tiles[None, (bid_m, iter_k)], + sA[None, None, tma_stage], + mbar, + ) + simple_tma_copy( + B_tma.atom, + gB_tiles[None, (cta_rank * grid_n + bid_n, iter_k)], + sB[None, None, tma_stage], + mbar, + ) + simple_tma_copy( + SFA_tma.atom, + gSFA_tiles[None, (iter_k, bid_m)], + sSFA[None, None, tma_stage], + mbar, + ) + # TMA multicast. rank0 loads gate, rank1 loads up + simple_tma_copy( + SFB_tma.atom, + gSFB_tiles[None, (iter_k, bid_n + cta_rank * grid_n)], + sSFB[None, None, cta_rank, tma_stage], + mbar, + mcast_mask=Uint16(3), + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + parity ^= 1 + + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + elif warp_id == 4: + # MMA warp (2-CTA collective alloc; only cta_rank 0 issues) + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=2 * 32) + + tma_stage = 0 + tma_full_parity = 0 + + sdesc = _tcgen05.make_sdesc_128B_swizzle(0) + # SF smem descriptor: no swizzle, SBO = 8 * 16 bytes. + sdesc_sf = Uint64((8 << 32) | (1 << 46)) + idesc = _tcgen05.make_mxfp8_idesc(BM * 2, BN * 2) + multicast_mask = Uint16(3) + + # each [128,128] MXFP8 tile needs 4 tmem columns for SF + sfa_tmem = 500 + sfbg_tmem = sfa_tmem + 4 + sfbu_tmem = sfbg_tmem + 4 + + if cta_rank == 0: + tmem_stage = 0 + tmem_parity = 1 + for _ in range(raw_bid, total_tiles, num_bids): + acc_base = 224 * tmem_stage + cute.arch.mbarrier_wait(tmem_empty_mbar + tmem_stage, tmem_parity) + # wait for the previous tile's overlap-region read; the + # single partial barrier phases once per tile + cute.arch.mbarrier_wait(tmem_partial_mbar, tmem_stage ^ 1) + _tcgen05.fence_after_thread_sync() + + for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1): + a_desc = sdesc | ( + sA[None, None, tma_stage].iterator.toint() >> 4 + ) + b_desc = sdesc | ( + sB[None, None, tma_stage].iterator.toint() >> 4 + ) + sfa_desc = sdesc_sf | ( + sSFA[None, None, tma_stage].iterator.toint() >> 4 + ) + sfbg_desc = sdesc_sf | ( + sSFB[None, None, 0, tma_stage].iterator.toint() >> 4 + ) + sfbu_desc = sdesc_sf | ( + sSFB[None, None, 1, tma_stage].iterator.toint() >> 4 + ) + + cute.arch.mbarrier_wait( + tma_full_mbar + tma_stage, tma_full_parity + ) + _tcgen05.fence_after_thread_sync() + + _tcgen05.cp(sfa_tmem, sfa_desc, "32x128b", "warpx4", 2) + _tcgen05.cp(sfbg_tmem, sfbg_desc, "32x128b", "warpx4", 2) + _tcgen05.cp(sfbu_tmem, sfbu_desc, "32x128b", "warpx4", 2) + + for k in cutlass.range_constexpr(BK // 32): + enable_input_d = iter_k > 0 or k > 0 + # select this k-block's SF byte inside the staged + # 4-col tmem atoms (idesc SF-id fields) + idesc_k = idesc + ((k << 4) | (k << 29)) + _tcgen05.mma_mxfp8( + acc_base, + a_desc, + b_desc, + idesc_k, + sfa_tmem, + sfbg_tmem, + enable_input_d, + 2, + ) + a_desc += 32 >> 4 + b_desc += 32 >> 4 + _tcgen05.commit(tma_empty_mbar + tma_stage, multicast_mask, 2) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_full_parity ^= 1 + + _tcgen05.commit(tmem_full_mbar + tmem_stage, multicast_mask, 2) + tmem_stage ^= 1 + if tmem_stage == 0: + tmem_parity ^= 1 + + else: + # epilogue warps (one warpgroup; tid < 128 indexes the tmem lane) + if warp_id == 0: + _tcgen05.alloc(taddr, cta_group=2) + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=2 * 32) + + exp_scale = -1.4426950216293334961 * alpha + + # ((32,4), (M, N/BN)) + Q_ = cute.zipped_divide(Q_tensor, (1, cute.make_layout((32, 4)))) + Q_ = Q_[(0, None), None] + + f8x32_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + Float8E4M3FN, + num_bits_per_copy=256, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, + ) + + partial_ = to_cta0_smem(tmem_partial_mbar) + empty_ = to_cta0_smem(tmem_empty_mbar) + + tmem_stage = 0 + tmem_parity = 0 + for work_id in range(raw_bid, total_tiles, num_bids): + bid_m, bid_n = compute_bid(work_id, GROUP_M, grid_m, grid_n) + m = bid_m * BM + tid + acc_base = 224 * tmem_stage + + if warp_id == 0: + cute.arch.mbarrier_wait(tmem_full_mbar + tmem_stage, tmem_parity) + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=4 * 32) + _tcgen05.fence_after_thread_sync() + + rGate = cute.make_rmem_tensor(32, Float32) + rUp = cute.make_rmem_tensor(32, Float32) + + # cc=0 is the overlap chunk: cols [224,256) (up[96:128] on + # even tiles, gate[0:32] on odd) plus the matching 32 cols of + # the other operand. cc=1..3: the remaining chunk-pairs in + # order. Each chunk: load -> act+quant -> store. + for cc in cutlass.range_constexpr(4): + if cutlass.const_expr(cc == 0): + c = 3 - 3 * tmem_stage # even tiles: 3; odd: 0 + gate_col = 96 + 128 * tmem_stage + up_col = 224 + 128 * tmem_stage + else: + c = cc - 1 + tmem_stage # even: 0,1,2; odd: 1,2,3 + gate_col = acc_base + c * 32 + up_col = acc_base + 128 + c * 32 + rGate.store(_tcgen05.ld(tid // 32 * 32, gate_col, "32x32b", 32)) + rUp.store(_tcgen05.ld(tid // 32 * 32, up_col, "32x32b", 32)) + _tcgen05.wait_ld() + if cutlass.const_expr(cc == 0): + _tcgen05.fence_before_thread_sync() + mbarrier.arrive(partial_, "cluster") + if cutlass.const_expr(cc == 3): + _tcgen05.fence_before_thread_sync() + mbarrier.arrive(empty_ + tmem_stage, "cluster") + + q_rmem = cute.make_rmem_tensor(32, Float8E4M3FN) + q_i32 = cute.recast_tensor(q_rmem, Uint32) + + # activation + for j in cutlass.range_constexpr(32): + rGate[j] = cute.arch.fmin(rGate[j], limit) + rUp[j] = cute.arch.fmin(cute.arch.fmax(rUp[j], -limit), limit) + + af = cute.make_rmem_tensor(32, Float32) + for j in cutlass.range(32, vectorize=True): + e = cute.math.exp2(rGate[j] * exp_scale, fastmath=True) + gate = rGate[j] * cute.arch.rcp_approx(1.0 + e) + af[j] = gate * (rUp[j] + beta) + + # quantization + amax0 = cute.math.abs(af[0]) + amax1 = cute.math.abs(af[1]) + for j in cutlass.range_constexpr(15): + amax0 = cute.arch.fmax(amax0, cute.math.abs(af[j * 2 + 2])) + amax1 = cute.arch.fmax(amax1, cute.math.abs(af[j * 2 + 3])) + amax = cute.arch.fmax(amax0, amax1) + + # Scale byte via exact integer math: + # ceil(log2(amax/448)) = E - 8 + (amax > 1.75 * 2^E), + # E = exponent(amax); amax=0 (all-zero block) clamps to 0. + bits = amax.bitcast(Int32) + exp_bits = bits & 0x7F800000 + thr = (exp_bits | 0x600000).bitcast(Float32) # 1.75 * 2^E + sb = (exp_bits >> 23) - 8 + Int32(amax > thr) + sb = cutlass.min(cutlass.max(sb, Int32(0)), Int32(254)) + + inv = ((Int32(254) - sb) << 23).bitcast(Float32) + for j in cutlass.range(32, vectorize=True): + af[j] *= inv + + for j in cutlass.range_constexpr(8): + q_i32[j] = fp32x4_to_fp8x4(*[af[4 * j + t] for t in range(4)]) + + kb = bid_n * (BN // 32) + c + sf_off = ( + (bid_m * NUM_KT + kb // 4) * 512 + + (m % 32) * 16 + + ((m // 32) % 4) * 4 + + kb % 4 + ) + if m < M: + cute.copy(f8x32_atom, q_rmem, Q_[(None, c), (m, bid_n)]) + SFQ_tensor[sf_off] = Uint8(sb) + + tmem_stage ^= 1 + if tmem_stage == 0: + tmem_parity ^= 1 + + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + if warp_id == 0: + _tcgen05.dealloc(cta_group=2) + + @cache + @staticmethod + def compile(GROUP_M: int, use_pdl: bool = True): + M = cute.sym_int() + K = cute.sym_int() + N2 = cute.sym_int() + + A = make_fake_tensor( + Float8E4M3FN, + (M, K), + (cute.sym_int64(divisibility=16), 1), + assumed_align=16, + ) + B = make_fake_tensor( + Float8E4M3FN, + (N2, K), + (cute.sym_int64(divisibility=16), 1), + assumed_align=16, + ) + SFA = make_fake_tensor( + Float8E8M0FNU, (cute.sym_int(divisibility=512),), (1,), assumed_align=16 + ) + SFB = make_fake_tensor( + Float8E8M0FNU, (cute.sym_int(divisibility=512),), (1,), assumed_align=16 + ) + Q = make_fake_tensor( + Float8E4M3FN, + (M, cute.sym_int(divisibility=256)), + (cute.sym_int64(divisibility=256), 1), + assumed_align=32, + ) + SFQ = make_fake_tensor( + Uint8, (cute.sym_int(divisibility=512),), (1,), assumed_align=16 + ) + + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = Sm100DenseGemm1Mxfp8(GROUP_M, use_pdl) + return cute.compile( + kernel, + A, + B, + SFA, + SFB, + Q, + SFQ, + Float32(7.0), + Float32(1.702), + Float32(1.0), + 148, + stream, + options="--enable-tvm-ffi", + ) + + +# Below this the unfused GEMM -> act -> quant chain is as fast (TP4 shared +# expert, I=768: break-even at ~1024 tokens). +MIN_TOKENS = 1024 + + +def _linear_kernel(linear): + method = getattr(linear, "scheme", getattr(linear, "quant_method", None)) + return getattr(method, "kernel", None) + + +def can_fuse_gemm1(gate_up_proj, down_proj, act_fn) -> bool: + """Whether fused_gemm1 can replace gate_up_proj + act_fn for down_proj. + + Called at init, before weight loading: kernels are selected at + create_weights time, so the decision does not depend on the weights. + """ + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, + ) + from vllm.model_executor.layers.activation import SiluAndMulWithClamp + + if not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ): + return False + if not isinstance(act_fn, SiluAndMulWithClamp): + return False + # The kernel reads gate_up's fp8 weight and F8_128x4-swizzled weight + # scale directly, so it depends on these kernels' weight layout. + if not isinstance( + _linear_kernel(gate_up_proj), + (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel), + ): + return False + down_kernel = _linear_kernel(down_proj) + if down_kernel is None or down_kernel.input_quant_key() != kMxfp8Dynamic: + return False + # I % 256 == 0 (output row stride), K % 128 == 0 (BK). + n = gate_up_proj.output_size_per_partition // 2 + k = gate_up_proj.input_size_per_partition + return n % 256 == 0 and k % 128 == 0 + + +def fused_gemm1(x: torch.Tensor, gate_up_proj, act_fn) -> QuantizedActivation: + """act_fn(gate_up_proj(x)), MXFP8-quantized for an MXFP8 down_proj. + + Requires can_fuse_gemm1(gate_up_proj, down_proj, act_fn). + """ + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, + ) + + assert x.ndim == 2 and x.dtype == torch.bfloat16 + weight = gate_up_proj.weight + # FlashInferCutedsl stores the weight as a column-major [K, 2I] view. + if isinstance(_linear_kernel(gate_up_proj), FlashInferCutedslMxfp8LinearKernel): + weight = weight.t() + + M = x.shape[0] + N = weight.shape[0] // 2 + q = torch.empty((M, N), dtype=torch.float8_e4m3fn, device=x.device) + sf = torch.empty( + ((M + 127) // 128) * 128 * (N // 32), dtype=torch.uint8, device=x.device + ) + + xq, xsf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True) + run( + xq, + weight, + xsf.view(-1).view(torch.float8_e8m0fnu), + gate_up_proj.weight_scale.view(torch.float8_e8m0fnu), + q, + sf, + act_fn.swiglu_limit, + act_fn.alpha, + act_fn.beta, + ) + return QuantizedActivation( + data=q, + scale=sf, + orig_dtype=x.dtype, + orig_shape=x.shape[:-1] + (N,), + quant_key=kMxfp8Dynamic, + ) + + +def run( + xq: torch.Tensor, + weight: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + q: torch.Tensor, + sf: torch.Tensor, + limit: float, + alpha: float, + beta: float, + grid: int | None = None, +) -> None: + """Functional interface to the fused kernel (hides compile + grid sizing). + + xq: [M, K] fp8e4m3; weight: [2N, K] fp8e4m3 row-major; x_scale/weight_scale: + F8_128x4-swizzled E8M0 (flat, any uint8-viewable dtype); q: [M, N] fp8e4m3; + sf: flat uint8 scale buffer (round_up(M,128) * N/32 bytes; rows >= M are + not written). + """ + M, K = xq.shape + _, N = q.shape + if grid is None: + grid_m2 = ((M + 127) // 128 + 1) // 2 * 2 + num_sms = torch.cuda.get_device_properties(xq.device).multi_processor_count + grid = min(num_sms, grid_m2 * (N // 128)) + + # Grouped (m-fastest) rasterization only pays when B (n2*k fp8 bytes) + # exceeds half of L2; below that the flat order is faster (no div/mod + # overhead in compute_bid) + GROUP_M = 32 if N * 2 * K > (63 << 20) else 0 + kernel = Sm100DenseGemm1Mxfp8.compile(GROUP_M) + kernel(xq, weight, x_scale, weight_scale, q, sf, limit, alpha, beta, grid) diff --git a/vllm/models/minimax_m3/nvidia/ops/finalize_allreduce_norm.py b/vllm/models/minimax_m3/nvidia/ops/finalize_allreduce_norm.py new file mode 100644 index 000000000000..9c75e71e3249 --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/ops/finalize_allreduce_norm.py @@ -0,0 +1,781 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# The consumer is adapted from FlashInfer's low-latency MNNVL CuTe DSL +# all-reduce (flashinfer/comm/mnnvl_cutedsl/kernel_ll, +# _LamportResidualRMSNormDeviceKernel), as DSV4.1's AllReduceMHC is. + +"""MoE top-k finalize + TP all-reduce + Gemma RMSNorm + MXFP8 quantization. + +FlashInfer's Lamport all-reduce: the finalize publish kernel sums each token's +routed expert rows and its shared-expert row, then multicasts that TP-partial +row into a three-generation mailbox; a PDL-chained consumer spins until every +rank's fragment lands, reduces them, adds the residual and applies the Gemma +RMSNorm. Instead of a BF16 row, the consumer writes the normalized row as MXFP8 +(E4M3 values and UE8M0 scales of 32-element blocks in the 128x4 swizzled +layout), so the MXFP8 linear that follows skips its own input quantization. +""" + +import math + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import make_fake_compact_tensor +from cutlass.cutlass_dsl import T, dsl_user_op + +from vllm.config import VllmConfig +from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group +from vllm.model_executor.layers.fused_moe.moe_output import ( + MoEOutput, + UnfinalizedMoEOutput, +) +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic +from vllm.models.common.ops.moe_tail_fusion.primitives import ( + ACTIVE_STAGE, + LAMPORT_GENERATIONS, + NEXT_STAGE, + VEC_BF16, + WARP_SIZE, + bf16x8_to_packed_u32x4, + current_cu_stream, + fragment_has_negative_zero, + group_leader_block_sum, + load_global_u32x4, + load_volatile_u32, + make_fake_dynamic_compact_tensor, + map_shared_to_peer, + packed_u32x4_to_bf16x8, + store_global_u32, + store_global_u32x4, + store_lamport_sentinel_u32x4, + store_shared_cluster_f32, + to_cute, + to_cute_dynamic, +) +from vllm.models.deepseek_v41.nvidia.ops.cute_dsl.all_reduce_mhc import ( + _QuadFinalizePublishDeviceKernel, +) +from vllm.platforms import current_platform + +MXFP8_BLOCK = 32 +# Fragments (of VEC_BF16 values) per MXFP8 block, held by adjacent lanes. +_FRAGMENTS_PER_BLOCK = MXFP8_BLOCK // VEC_BF16 +# The swizzled scale layout tiles 128 rows by 4 block columns. +_SCALE_ROW_TILE = 128 +_SCALE_COL_TILE = 4 + + +@dsl_user_op +def _ue8m0_ceil(value: Float32, *, loc=None, ip=None) -> Uint32: + """FlashInfer's float_to_ue8m0_fast: the biased exponent rounded up, 0 for + non-positive input, as its MXFP8 quantization kernels compute it.""" + return Uint32( + llvm.inline_asm( + T.i32(), + [Float32(value).ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n\t" + ".reg .u32 bits, exp_biased, mantissa, bump, result;\n\t" + "setp.le.f32 p_zero, $1, 0f00000000;\n\t" + "mov.b32 bits, $1;\n\t" + "shr.b32 exp_biased, bits, 23;\n\t" + "and.b32 exp_biased, exp_biased, 255;\n\t" + "and.b32 mantissa, bits, 0x7FFFFF;\n\t" + "setp.ne.u32 p_has_mant, mantissa, 0;\n\t" + "selp.u32 bump, 1, 0, p_has_mant;\n\t" + "setp.eq.u32 p_exp_zero, exp_biased, 0;\n\t" + "setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n\t" + "and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n\t" + "@p_tiny_sub mov.u32 bump, 0;\n\t" + "add.u32 result, exp_biased, bump;\n\t" + "setp.gt.u32 p_ovf, result, 254;\n\t" + "selp.u32 result, 254, result, p_ovf;\n\t" + "selp.u32 $0, 0, result, p_zero;\n\t" + "}", + "=r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ue8m0_inverse(exponent: Uint32, *, loc=None, ip=None) -> Float32: + """2^(127 - exponent), or 0 for exponent 0.""" + return Float32( + llvm.inline_asm( + T.f32(), + [Uint32(exponent).ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .s32 new_exp;\n\t" + ".reg .pred p_zero;\n\t" + "setp.eq.u32 p_zero, $1, 0;\n\t" + "sub.s32 new_exp, 254, $1;\n\t" + "max.s32 new_exp, new_exp, 0;\n\t" + "shl.b32 $0, new_exp, 23;\n\t" + "@p_zero mov.b32 $0, 0;\n\t" + "}", + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _f32x4_to_e4m3x4( + a: Float32, b: Float32, c: Float32, d: Float32, *, loc=None, ip=None +) -> Uint32: + """Four values to E4M3, ``a`` in the low byte, clamped to +-448 first so + that inf saturates like FlashInfer's MXFP8 kernels and torch.""" + return Uint32( + llvm.inline_asm( + T.i32(), + [Float32(v).ir_value(loc=loc, ip=ip) for v in (a, b, c, d)], + "{\n\t" + ".reg .f32 f0, f1, f2, f3;\n\t" + ".reg .b16 lo, hi;\n\t" + "min.f32 f0, $1, 0f43E00000;\n\t" + "max.f32 f0, f0, 0fC3E00000;\n\t" + "min.f32 f1, $2, 0f43E00000;\n\t" + "max.f32 f1, f1, 0fC3E00000;\n\t" + "min.f32 f2, $3, 0f43E00000;\n\t" + "max.f32 f2, f2, 0fC3E00000;\n\t" + "min.f32 f3, $4, 0f43E00000;\n\t" + "max.f32 f3, f3, 0fC3E00000;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 lo, f1, f0;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, f3, f2;\n\t" + "mov.b32 $0, {lo, hi};\n\t" + "}", + "=r,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _store_global_u32x2( + address: Int64, low: Uint32, high: Uint32, *, loc=None, ip=None +) -> None: + llvm.inline_asm( + None, + [ + address.ir_value(loc=loc, ip=ip), + low.ir_value(loc=loc, ip=ip), + high.ir_value(loc=loc, ip=ip), + ], + "st.global.v2.u32 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _swizzled_scale_offset(row: Int32, col: Int32, padded_cols: Int32) -> Int32: + """Byte offset of block ``col`` of ``row`` in the 128x4 swizzled layout.""" + return ( + col % Int32(4) + + (col // Int32(4)) * Int32(512) + + (row % Int32(32)) * Int32(16) + + ((row % Int32(128)) // Int32(32)) * Int32(4) + + (row // Int32(128)) * Int32(128) * padded_cols + ) + + +class _LamportGemmaRMSNormMxfp8Kernel: + """FlashInfer's ``_LamportResidualRMSNormDeviceKernel`` (residual add and + copy, one rank lane) with a Gemma RMSNorm whose BF16 result is quantized to + MXFP8 instead of stored. + + A 32-element scale block is four consecutive fragments, which adjacent lanes + hold, so two shuffles reduce its amax. The scales of the padding rows below + the next multiple of 128 are zeroed, as FlashInfer's quantization does. + """ + + def __init__( + self, + *, + hidden: int, + tp: int, + capacity_m: int, + cluster_size: int, + threads: int, + enable_pdl: bool, + ) -> None: + if hidden <= 0 or hidden % (MXFP8_BLOCK * _SCALE_COL_TILE): + raise ValueError("hidden must be a positive multiple of 128") + if threads <= 0 or threads % WARP_SIZE: + raise ValueError("threads must be a positive multiple of the warp size") + self.hidden = hidden + self.tp = tp + self.capacity_m = capacity_m + self.cluster_size = cluster_size + self.threads = threads + self.enable_pdl = enable_pdl + self.fragments = hidden // VEC_BF16 + self.fragment_stride = cluster_size * threads + self.trips = math.ceil(self.fragments / self.fragment_stride) + self.warps = threads // WARP_SIZE + self.scale_cols = hidden // MXFP8_BLOCK + + def smem_size_in_bytes(self) -> int: + return (self.warps + self.cluster_size) * 4 + + @cute.jit + def __call__( + self, + contribution_mailbox: cute.Tensor, + residual_source: cute.Tensor, + gamma: cute.Tensor, + residual_output: cute.Tensor, + quant_output: cute.Tensor, + scale_output: cute.Tensor, + stage_state: cute.Tensor, + eps: Float32, + m: Int32, + stream: cuda.CUstream, + ) -> None: + self.kernel( + contribution_mailbox, + residual_source, + gamma, + residual_output, + quant_output, + scale_output, + stage_state, + eps, + m, + ).launch( + grid=(m, self.cluster_size, 1), + block=(self.threads, 1, 1), + cluster=(1, self.cluster_size, 1), + smem=self.smem_size_in_bytes(), + stream=stream, + use_pdl=self.enable_pdl, + ) + + @cute.kernel + def kernel( + self, + contribution_mailbox: cute.Tensor, + residual_source: cute.Tensor, + gamma: cute.Tensor, + residual_output: cute.Tensor, + quant_output: cute.Tensor, + scale_output: cute.Tensor, + stage_state: cute.Tensor, + eps: Float32, + m: Int32, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + token, _, _ = cute.arch.block_idx() + cluster_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + base_fragment = cluster_rank * self.threads + tidx + + prenorm_fragments = cute.make_rmem_tensor( + cute.make_layout((self.trips, VEC_BF16), stride=(VEC_BF16, 1)), + BFloat16, + ) + prenorm_fragments.fill(BFloat16(0.0)) + gamma_fragments = cute.make_rmem_tensor( + cute.make_layout((self.trips, VEC_BF16), stride=(VEC_BF16, 1)), + BFloat16, + ) + gamma_fragments.fill(BFloat16(0.0)) + residual_fragments = cute.make_rmem_tensor( + cute.make_layout((self.trips, VEC_BF16), stride=(VEC_BF16, 1)), + BFloat16, + ) + residual_fragments.fill(BFloat16(0.0)) + + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + if fragment < self.fragments: + gamma_element = Int64(fragment) * VEC_BF16 + gamma_pointer = cute.make_ptr( + BFloat16, + (gamma.iterator + gamma_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + gamma_fragments[trip, None].store( + packed_u32x4_to_bf16x8(load_global_u32x4(gamma_pointer)) + ) + residual_element = Int64(token) * self.hidden + gamma_element + residual_pointer = cute.make_ptr( + BFloat16, + (residual_source.iterator + residual_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + residual_fragments[trip, None].store( + packed_u32x4_to_bf16x8(load_global_u32x4(residual_pointer)) + ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() + + active_stage = load_volatile_u32(stage_state.iterator + ACTIVE_STAGE) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + lane_packed = cute.make_rmem_tensor(cute.make_layout((self.tp, 4)), Uint32) + lane_packed.fill(Uint32(0)) + dirty = fragment < self.fragments + while dirty: + dirty = False + for source_rank in cutlass.range_constexpr(self.tp): + if fragment < self.fragments: + source_element = ( + (Int64(active_stage) * self.tp + Int64(source_rank)) + * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + source_pointer = cute.make_ptr( + BFloat16, + (contribution_mailbox.iterator + source_element).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + packed = load_global_u32x4(source_pointer, volatile=True) + dirty = dirty | fragment_has_negative_zero(packed) + for word in cutlass.range_constexpr(4): + lane_packed[source_rank, word] = packed[word] + + if fragment < self.fragments: + prenorm = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + prenorm.fill(Float32(0.0)) + for source_rank in cutlass.range_constexpr(self.tp): + packed = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) + for word in cutlass.range_constexpr(4): + packed[word] = lane_packed[source_rank, word] + prenorm.store( + prenorm.load() + + packed_u32x4_to_bf16x8(packed.load()).to(Float32) + ) + prenorm_bf16 = ( + prenorm.load() + residual_fragments[trip, None].load().to(Float32) + ).to(BFloat16) + prenorm_fragments[trip, None].store(prenorm_bf16) + output_element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16 + store_global_u32x4( + Int64((residual_output.iterator + output_element).toint()), + bf16x8_to_packed_u32x4(prenorm_bf16), + ) + + if token == 0 and cluster_rank == 0 and tidx == 0: + store_global_u32( + (stage_state.iterator + NEXT_STAGE).toint(), + (active_stage + Uint32(1)) % Uint32(LAMPORT_GENERATIONS), + ) + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() + + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + for source_rank in cutlass.range_constexpr(self.tp): + if fragment < self.fragments: + source_element = ( + (Int64(active_stage) * self.tp + Int64(source_rank)) + * self.capacity_m + + Int64(token) + ) * self.hidden + Int64(fragment) * VEC_BF16 + store_lamport_sentinel_u32x4( + Int64((contribution_mailbox.iterator + source_element).toint()) + ) + + thread_sum = Float32(0.0) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + if fragment < self.fragments: + values = prenorm_fragments[trip, None].load().to(Float32) + thread_sum = thread_sum + (values * values).reduce( + cute.ReductionOp.ADD, + init_val=Float32(0.0), + reduction_profile=0, + ) + + smem = cutlass.utils.SmemAllocator() + warp_sums = smem.allocate_array(Float32, self.warps) + cluster_sums = smem.allocate_array(Float32, self.cluster_size) + cta_sum = group_leader_block_sum(thread_sum, warp_sums, self.warps, 1) + if tidx < self.cluster_size: + local_slot = cluster_sums + cluster_rank + remote_slot = map_shared_to_peer(local_slot, Int32(tidx)) + store_shared_cluster_f32(remote_slot, cta_sum) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + full_sum = Float32(0.0) + for peer in cutlass.range_constexpr(self.cluster_size): + full_sum = full_sum + cute.arch.load( + (cluster_sums + peer).llvm_ptr, + Float32, + ) + inv_rms = cute.math.rsqrt( + full_sum / Float32(self.hidden) + eps, + fastmath=True, + ) + padded_cols = Int32(self.scale_cols) + for trip in cutlass.range_constexpr(self.trips): + fragment = base_fragment + trip * self.fragment_stride + # Gemma RMSNorm scales by (1 + weight); the unfused norm returns BF16. + normed = ( + ( + prenorm_fragments[trip, None].load().to(Float32) + * inv_rms + * (gamma_fragments[trip, None].load().to(Float32) + Float32(1.0)) + ) + .to(BFloat16) + .to(Float32) + ) + values = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32) + values.store(normed) + amax = Float32(0.0) + for element in cutlass.range_constexpr(VEC_BF16): + amax = cute.arch.fmax( + amax, cute.arch.fmax(values[element], -values[element]) + ) + # Out-of-range fragments are zero; every lane joins the shuffles. + for offset in cutlass.range_constexpr(1, _FRAGMENTS_PER_BLOCK): + if cutlass.const_expr((offset & (offset - 1)) == 0): + amax = cute.arch.fmax( + amax, + cute.arch.shuffle_sync_bfly( + amax, offset=offset, mask=-1, mask_and_clamp=31 + ), + ) + exponent = _ue8m0_ceil(amax * Float32(1.0 / 448.0)) + inverse = _ue8m0_inverse(exponent) + if fragment < self.fragments: + low = _f32x4_to_e4m3x4( + values[0] * inverse, + values[1] * inverse, + values[2] * inverse, + values[3] * inverse, + ) + high = _f32x4_to_e4m3x4( + values[4] * inverse, + values[5] * inverse, + values[6] * inverse, + values[7] * inverse, + ) + quant_element = Int64(token) * self.hidden + Int64(fragment) * VEC_BF16 + _store_global_u32x2( + Int64((quant_output.iterator + quant_element).toint()), low, high + ) + if fragment % _FRAGMENTS_PER_BLOCK == 0: + col = fragment // _FRAGMENTS_PER_BLOCK + scale_output[_swizzled_scale_offset(token, col, padded_cols)] = ( + exponent.to(Uint8) + ) + + # Zero the scales of the last row tile's padding rows. A 4-byte word of + # a 512-byte block holds the 4 blocks of one row, so every thread of + # the grid clears whole words. + tile_words = Int32(self.scale_cols * _SCALE_ROW_TILE // 4) + last_tile = (m - Int32(1)) // Int32(_SCALE_ROW_TILE) + word = (token * Int32(self.cluster_size) + cluster_rank) * Int32( + self.threads + ) + tidx + while word < tile_words: + block_word = word % Int32(128) + row = ( + last_tile * Int32(_SCALE_ROW_TILE) + + block_word // Int32(4) + + block_word % Int32(4) * Int32(32) + ) + if row >= m: + byte = (Int64(last_tile) * Int64(tile_words) + Int64(word)) * 4 + store_global_u32( + cute.make_ptr( + Uint32, + (scale_output.iterator + byte).llvm_ptr, + cute.AddressSpace.gmem, + assumed_align=4, + ).toint(), + Uint32(0), + ) + word = word + m * Int32(self.cluster_size * self.threads) + + +def mxfp8_scale_numel(num_tokens: int, hidden_size: int) -> int: + """Bytes of the 128x4 swizzled UE8M0 scales of ``num_tokens`` rows.""" + rows = math.ceil(num_tokens / _SCALE_ROW_TILE) * _SCALE_ROW_TILE + cols = math.ceil(hidden_size // MXFP8_BLOCK / _SCALE_COL_TILE) * _SCALE_COL_TILE + return rows * cols + + +class FinalizeAllReduceGemmaRMSNormMxfp8: + """MoE top-k finalize, TP all-reduce, residual add, Gemma RMSNorm and MXFP8 + quantization in two PDL-chained launches. + + ``__call__`` takes an unfinalized MoE output plus its TP-partial shared-expert + output and returns the MXFP8 normalized rows with the new residual. Each + instance owns a Lamport mailbox for ``max_num_tokens`` rows that all layers + share; construction is collective over the TP group. + """ + + def __init__( + self, + *, + hidden_size: int, + max_num_tokens: int, + top_k: int, + fp32_weights: bool, + device: torch.device, + ) -> None: + group = get_tp_group().device_group + self.hidden_size = hidden = hidden_size + self.capacity = capacity = max_num_tokens + self.top_k = top_k + self.fp32_weights = fp32_weights + tp = dist.get_world_size(group) + rank = dist.get_rank(group) + device = torch.device(device) + + # A cluster of 128-thread CTAs covers a token's fragments in one trip + # (six for hidden 6144); the finalize stages all top_k routes at once. + cluster_size = math.ceil(hidden // VEC_BF16 / 128) + with torch.accelerator.device_index(device.index): + self._publish = cute.compile( + _QuadFinalizePublishDeviceKernel( + hidden=hidden, + top_k=top_k, + tp=tp, + rank=rank, + capacity_m=capacity, + threads=128, + # The routing weights already carry routed_scaling_factor. + routed_scaling_factor=1.0, + include_shared_expert=True, + load_shared_expert_before_pdl=False, + enable_pdl=True, + prefetch_group=top_k, + fp32_weights=fp32_weights, + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_dynamic_compact_tensor( + Float32 if fp32_weights else BFloat16, + alignment=4 if fp32_weights else 2, + divisibility=top_k, + ), + make_fake_dynamic_compact_tensor( + Int32, alignment=4, divisibility=top_k + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Int64(0), + Int32(capacity), + current_cu_stream(), + ) + self._collective = cute.compile( + _LamportGemmaRMSNormMxfp8Kernel( + hidden=hidden, + tp=tp, + capacity_m=capacity, + cluster_size=cluster_size, + threads=128, + enable_pdl=True, + ), + make_fake_compact_tensor( + BFloat16, + (LAMPORT_GENERATIONS * tp * capacity * hidden,), + assumed_align=16, + ), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_compact_tensor(BFloat16, (hidden,), assumed_align=16), + make_fake_dynamic_compact_tensor( + BFloat16, alignment=16, divisibility=hidden + ), + make_fake_dynamic_compact_tensor( + Uint8, alignment=16, divisibility=hidden + ), + make_fake_dynamic_compact_tensor( + Uint8, alignment=16, divisibility=_SCALE_ROW_TILE + ), + make_fake_compact_tensor(Int32, (2,), assumed_align=4), + Float32(0.0), + Int32(capacity), + current_cu_stream(), + ) + + self._mailbox = symm_mem.empty( + (LAMPORT_GENERATIONS, tp, capacity, hidden), + dtype=torch.bfloat16, + device=device, + ) + self._mailbox_handle = symm_mem.rendezvous(self._mailbox, group) + multicast = int(self._mailbox_handle.multicast_ptr or 0) + if not multicast or multicast % 16: + raise RuntimeError("NVLink multicast mapping is unavailable") + self._multicast = multicast + # Every slot starts as the Lamport sentinel, BF16 negative zero. + self._mailbox.view(torch.int16).fill_(-32768) + self._stage_state = torch.zeros(2, dtype=torch.int32, device=device) + torch.accelerator.synchronize() + dist.barrier(group=group) + + def __call__( + self, + moe_output: MoEOutput, + residual: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + ) -> tuple[QuantizedActivation, torch.Tensor]: + """Return (MXFP8 normalized rows, new residual). + + Args: + moe_output: Unfinalized TP-partial routed output and the TP-partial + shared-expert output, with routed_scaling_factor in the weights. + residual: [M, hidden] BF16 residual stream. + norm_weight: [hidden] BF16 Gemma RMSNorm weight (scale is 1 + w). + norm_eps: RMSNorm epsilon. + + """ + routed = moe_output.routed + shared_output = moe_output.shared_output + assert isinstance(routed, UnfinalizedMoEOutput) + assert shared_output is not None + assert moe_output.routed_scaling_factor == 1.0 + hidden, top_k = self.hidden_size, self.top_k + m = residual.shape[0] + if not 1 <= m <= self.capacity: + raise ValueError(f"M={m} is outside [1, {self.capacity}]") + if residual.shape != (m, hidden) or shared_output.shape != (m, hidden): + raise ValueError("residual and shared_output must be [M, hidden]") + weights_dtype = torch.float32 if self.fp32_weights else torch.bfloat16 + if ( + routed.expert_weights.dtype != weights_dtype + or routed.expert_weights.numel() != m * top_k + ): + raise ValueError(f"expert_weights must be {weights_dtype} [M, top_k]") + self._publish( + to_cute_dynamic(routed.gemm2_permuted.flatten(), 16, divisibility=hidden), + to_cute_dynamic( + routed.expert_weights.flatten(), + 4 if self.fp32_weights else 2, + divisibility=top_k, + ), + to_cute_dynamic( + routed.expanded_idx_to_permuted_idx.flatten(), 4, divisibility=top_k + ), + to_cute_dynamic(shared_output.flatten(), 16, divisibility=hidden), + to_cute(self._stage_state, 4), + Int64(self._multicast), + Int32(m), + current_cu_stream(), + ) + residual_output = torch.empty_like(residual) + quantized = torch.empty( + (m, hidden), dtype=torch.float8_e4m3fn, device=residual.device + ) + scales = torch.empty( + mxfp8_scale_numel(m, hidden), dtype=torch.uint8, device=residual.device + ) + self._collective( + to_cute(self._mailbox.flatten(), 16), + to_cute_dynamic(residual.flatten(), 16, divisibility=hidden), + to_cute(norm_weight, 16), + to_cute_dynamic(residual_output.flatten(), 16, divisibility=hidden), + to_cute_dynamic( + quantized.view(torch.uint8).flatten(), 16, divisibility=hidden + ), + to_cute_dynamic(scales, 16, divisibility=_SCALE_ROW_TILE), + to_cute(self._stage_state, 4), + Float32(norm_eps), + Int32(m), + current_cu_stream(), + ) + return ( + QuantizedActivation( + quantized, scales, residual.dtype, quantized.shape, kMxfp8Dynamic + ), + residual_output, + ) + + +_fused_moe_finalize: FinalizeAllReduceGemmaRMSNormMxfp8 | None = None + + +def supports_fused_moe_finalize(vllm_config: VllmConfig) -> bool: + """Whether to build the fused MoE finalize kernel: tensor parallel without + expert parallel, a BF16 model on SM100-family GPUs, and NVLink multicast on + every TP rank. Every TP rank calls it; the multicast check is collective. + """ + from vllm.distributed.device_communicators.custom_all_reduce import ( + _all_ranks_true, + _has_local_multicast_support, + ) + + if ( + get_tensor_model_parallel_world_size() <= 1 + or vllm_config.parallel_config.enable_expert_parallel + or not current_platform.is_cuda() + or not current_platform.is_device_capability_family(100) + or vllm_config.model_config.dtype != torch.bfloat16 + ): + return False + device = torch.device( + current_platform.device_type, torch.accelerator.current_device_index() + ) + return _all_ranks_true( + get_tp_group().cpu_group, _has_local_multicast_support(device) + ) + + +def init_fused_moe_finalize( + *, hidden_size: int, max_num_tokens: int, top_k: int, fp32_weights: bool +) -> None: + """Build the kernel every deferred MoE output is finalized with. + + Collective over the TP group, so every rank calls it, and only when + ``supports_fused_moe_finalize`` holds. + """ + global _fused_moe_finalize + _fused_moe_finalize = FinalizeAllReduceGemmaRMSNormMxfp8( + hidden_size=hidden_size, + max_num_tokens=max_num_tokens, + top_k=top_k, + fp32_weights=fp32_weights, + device=torch.device( + current_platform.device_type, torch.accelerator.current_device_index() + ), + ) + + +def finalize_allreduce_gemma_rms_norm_mxfp8( + moe_output: MoEOutput, + residual: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, +) -> tuple[QuantizedActivation, torch.Tensor]: + """Finalize a deferred MoE output into the next layer's MXFP8 input and + new residual.""" + assert _fused_moe_finalize is not None, "init_fused_moe_finalize was not called" + return _fused_moe_finalize(moe_output, residual, norm_weight, norm_eps) diff --git a/vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py b/vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py new file mode 100644 index 000000000000..0c9f477180ff --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/ops/gather_nvfp4_index_k.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gather NVFP4 index-K pages into compact FP8 pages for the Q8KV8 prefill indexer. + +A packed NVFP4 page ``[128, 72]`` holds the E2M1 codes of all 128 tokens (64 B +per token, low nibble first) followed by their E4M3 16-element block scales +(``token * 8 + group``), with no global scale. Each element dequantizes to +``e4m3_rn_satfinite(e2m1 * scale)``, the value the Q8KV4 decode indexer scores, +so prefill and decode rank pages from identical keys. + +Where Triton can emit it, each four elements are one QMUL4 +(``mul.rn.satfinite.e4m3x4.e2m1x4.e4m3x4``), the instruction the decode kernel +dequantizes with; otherwise the product is formed exactly in fp32 and rounded +once, which gives the same bytes. +""" + +import functools + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +PAGE_SIZE = 128 +HEAD_DIM = 128 +E4M3_MAX = 448.0 +# QMUL4 needs PTX ISA 9.4 (ptxas 13.4) and an sm_100a/sm_103a target. +QMUL4_MIN_PTX_VERSION = 94 + + +@functools.cache +def _qmul4_supported(device_index: int) -> bool: + """Whether Triton emits PTX that can use QMUL4 on this device. + + Triton takes the PTX ISA version from its Blackwell ptxas: the bundled one + unless ``TRITON_PTXAS_BLACKWELL_PATH`` points at a CUDA 13.4+ ptxas. + """ + capability = current_platform.get_device_capability(device_index) + if capability is None or capability.to_int() not in (100, 103): + return False + try: + from triton.backends.nvidia.compiler import get_ptxas, ptx_get_version + + ptxas_version = get_ptxas(capability.to_int()).version + return ptx_get_version(ptxas_version) >= QMUL4_MIN_PTX_VERSION + except Exception: + # Triton-internal helpers; without them keep the portable path. + return False + + +@triton.jit +def _e2m1_to_f32(code): + # |code| 0..7 -> 0, 0.5, 1, 1.5, 2, 3, 4, 6; bit 3 is the sign. + exponent = (code >> 1) & 3 + mantissa = (code & 1).to(tl.float32) + normal = (2.0 + mantissa) * (1 << exponent).to(tl.float32) * 0.25 + value = tl.where(exponent == 0, 0.5 * mantissa, normal) + # Multiply (not negate) so that code 8 stays -0, as in the decode kernel. + return value * tl.where((code & 8) != 0, -1.0, 1.0) + + +@triton.jit +def _gather_nvfp4_index_k_fp8_kernel( + kv_cache, + block_table, + num_pages, + page_offsets, + out, + kv_page_stride, + block_table_stride, + PAGE_SIZE: tl.constexpr, + HEAD_DIM: tl.constexpr, + E4M3_MAX: tl.constexpr, + USE_QMUL4: tl.constexpr, +): + req = tl.program_id(0) + col = tl.program_id(1) + if col >= tl.load(num_pages + req): + return + page = tl.load(block_table + req * block_table_stride + col).to(tl.int64) + dst_page = (tl.load(page_offsets + req) + col).to(tl.int64) + + src = kv_cache + page * kv_page_stride + dst = out + dst_page * (PAGE_SIZE * HEAD_DIM) + tokens = tl.arange(0, PAGE_SIZE)[:, None] + scales = src + PAGE_SIZE * (HEAD_DIM // 2) + tokens * (HEAD_DIM // 16) + if USE_QMUL4: + # A 16-bit word holds four elements of one group; multiply it by the + # group's scale byte broadcast to E4M3x4. + words = tl.arange(0, HEAD_DIM // 4)[None, :] + packed4 = tl.load( + src.to(tl.pointer_type(tl.uint16)) + tokens * (HEAD_DIM // 4) + words + ) + scale4 = tl.load(scales + words // 4).to(tl.uint32) * 0x01010101 + values4 = tl.inline_asm_elementwise( + "{ .reg .b16 v; cvt.u16.u32 v, $1; " + "mul.rn.satfinite.e4m3x4.e2m1x4.e4m3x4 $0, v, $2; }", + "=r,r,r", + [packed4.to(tl.uint32), scale4], + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + tl.store( + dst.to(tl.pointer_type(tl.uint32)) + tokens * (HEAD_DIM // 4) + words, + values4, + ) + return + + packed_cols = tl.arange(0, HEAD_DIM // 2)[None, :] + packed = tl.load(src + tokens * (HEAD_DIM // 2) + packed_cols) + # Packed byte b holds elements 2b and 2b+1, both in group b // 8. + scale = tl.load(scales + packed_cols // 8) + scale = scale.to(tl.float8e4nv, bitcast=True).to(tl.float32) + # e2m1 * e4m3 is exact in fp32; clamp, then round once to e4m3. + low = _e2m1_to_f32(packed & 0xF) * scale + high = _e2m1_to_f32(packed >> 4) * scale + values = tl.interleave(low, high) + values = tl.minimum(tl.maximum(values, -E4M3_MAX), E4M3_MAX) + tl.store( + dst + tokens * HEAD_DIM + tl.arange(0, HEAD_DIM)[None, :], + values.to(tl.float8e4nv), + ) + + +def minimax_m3_gather_nvfp4_index_k_fp8( + kv_cache: torch.Tensor, + block_table: torch.Tensor, + num_pages: torch.Tensor, + page_offsets: torch.Tensor, + max_num_pages: int, + out: torch.Tensor, +) -> torch.Tensor: + """Dequantize each request's leading NVFP4 pages into contiguous FP8 pages. + + Args: + kv_cache: ``[num_blocks, 128, 72]`` uint8 packed NVFP4 pages. + block_table: ``[num_reqs, max_blocks]`` int32 physical page ids. + num_pages: ``[num_reqs]`` int32 pages to gather per request (its first + ``num_pages[r]`` logical pages). + page_offsets: ``[num_reqs]`` int32 first output page of each request. + max_num_pages: Host upper bound of ``num_pages``. + out: ``[>= sum(num_pages), 128, 128]`` contiguous float8_e4m3fn. + + Returns: + ``out``. + + """ + assert kv_cache.dtype == torch.uint8 and kv_cache.shape[1:] == ( + PAGE_SIZE, + HEAD_DIM // 2 + HEAD_DIM // 16, + ) + assert kv_cache.stride(1) == kv_cache.shape[2] and kv_cache.stride(2) == 1 + assert out.dtype == torch.float8_e4m3fn and out.is_contiguous() + assert out.shape[1:] == (PAGE_SIZE, HEAD_DIM) + assert block_table.stride(1) == 1 + num_reqs = block_table.shape[0] + if num_reqs == 0 or max_num_pages == 0: + return out + _gather_nvfp4_index_k_fp8_kernel[(num_reqs, max_num_pages)]( + kv_cache, + block_table, + num_pages, + page_offsets, + out, + kv_cache.stride(0), + block_table.stride(0), + PAGE_SIZE=PAGE_SIZE, + HEAD_DIM=HEAD_DIM, + E4M3_MAX=E4M3_MAX, + USE_QMUL4=_qmul4_supported(kv_cache.device.index or 0), + num_warps=4, + ) + return out diff --git a/vllm/models/minimax_m3/nvidia/ops/qkv_index_fused_epilogue.py b/vllm/models/minimax_m3/nvidia/ops/qkv_index_fused_epilogue.py new file mode 100644 index 000000000000..e531f4f9e792 --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/ops/qkv_index_fused_epilogue.py @@ -0,0 +1,1101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MXFP8 QKV+indexer GEMM with a fused QK-norm/RoPE + KV-insert epilogue. + +Replaces mm_mxfp8 + fused_minimax_m3_qknorm_rope_kv_insert for the MiniMax-M3 +qkv_proj with one persistent SM100 kernel. Per 128-col head tile, the +epilogue applies Gemma RMSNorm + partial NeoX RoPE to the f32 accumulator +(no bf16 round trip) and writes q (bf16, optionally fp8) and index_q (fp8), +or NVFP4-quantizes k / v / index_k into the paged caches. The qkv projection +is never materialized. +""" + +from dataclasses import dataclass +from functools import cache + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import ( + BFloat16, + Float8E4M3FN, + Float8E8M0FNU, + Float32, + Int32, + Int64, + Uint8, + Uint16, + Uint32, + Uint64, + cute, +) +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import make_fake_tensor + +from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem +from vllm.cute_utils.cvt import ( + bf16x2_to_fp32x2, + fp32x8_to_fp4x8, +) +from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import VllmCuTeDSLJitKernel + +HEAD_DIM = 128 # one 128-col N-tile per head +# epilogue outputs are consumed by other kernels: stores bypass L1 +_NO_L1_ALLOC = dict(level1_eviction_priority="evict_no_allocate") + + +@cute.jit +def _rms_norm(acc, sW, eps): + """Gemma RMS-norm the 128-col head row in acc in place; sW is the head + type's bf16 weight as (8, 16) tiles.""" + acc1 = cute.make_tensor(acc.iterator, cute.make_layout(HEAD_DIM)) + # 4 independent fma chains; starting from 0 makes every step + # fma(x, x, ss) (x0 * x0 + x1 * x1 would let ptxas fuse either product) + ss = cute.full(4, 0.0, Float32) + for j in cutlass.range_constexpr(HEAD_DIM // 4): + x = cute.local_tile(acc1, (4,), (j,)).load() + ss = ss + x * x + t = (ss[0] + ss[1] + ss[2] + ss[3]) * Float32(1.0 / HEAD_DIM) + eps + # t >= eps, so rsqrt.approx needs no subnormal handling + rms = cute.math.rsqrt(t, approx=True) + + # load 8 norm weights at a time -> 16-byte loads + for j in cutlass.range_constexpr(HEAD_DIM // 8): + w_f32 = bf16x2_to_fp32x2(sW[None, j].load()) + x = cute.local_tile(acc1, (8,), (j,)) + n = x.load() * rms + x.store(n + n * w_f32) + + +def _rope64(acc, CS, pos): + """NeoX-rope cols [0, 64) of acc in place, rotating chunk 0 (x) with + chunk 1 (y); CS is the bf16 cos | sin cache as [32, (pos, 2)].""" + cos = bf16x2_to_fp32x2(CS[None, (pos, 0)].load()) + sin = bf16x2_to_fp32x2(CS[None, (pos, 1)].load()) + x = acc[None, 0].load() + y = acc[None, 1].load() + acc[None, 0].store(x * cos - y * sin) + acc[None, 1].store(y * cos + x * sin) + + +def _nvfp4_cache_views(C, bs, nslots): + """Views of the [num_pages, page_bytes] NVFP4 cache C. A page holds nslots + slots of bs tokens' E2M1 data (64B each), then their E4M3 scales (8B + each). Returns the data as u32 words [8, half, token, slot, page], the + linear scales as [8, token, slot, page], and the V-swizzled scales as + [(2, 4), token, slot, page]: byte h of token t at + (t // 4 * 4 + h // 2) * 8 + h % 2 * 4 + t % 4.""" + slot_bytes = bs * (HEAD_DIM // 2 + HEAD_DIM // 16) + sf_off = bs * HEAD_DIM // 2 + nb, page = C.shape[0], C.stride[0] + data = cute.make_tensor( + C.iterator, + cute.make_layout((32, 2, bs, nslots, nb), stride=(1, 32, 64, slot_bytes, page)), + ) + sf = cute.make_tensor( + C.iterator + sf_off, + cute.make_layout((8, bs, nslots, nb), stride=(1, 8, slot_bytes, page)), + ) + vsf = cute.make_tensor( + C.iterator + sf_off, + cute.make_layout( + ((2, 4), (4, bs // 4), nslots, nb), + stride=((4, 8), (1, 32), slot_bytes, page), + ), + ) + return cute.recast_tensor(data, Uint32), sf, vsf + + +def _nvfp4_quant32(vals, words, sf, sf_scale): + """NVFP4-quantize 32 f32 (two 16-value groups) into 4 data words and the + 2 scale bytes `sf`, following + fused_minimax_m3_qknorm_rope_kv_insert's quantizer. sf_scale is + 1 / (global dequant scale), or None for unit scale.""" + # E4M3(amax / 6) floored at 2^-9 (satfinite caps it at 448); the unfused + # op's 1e-12 amax floor lies below that and is omitted. Both groups' + # scales share one e4m3x2 cvt. + amax_6 = cute.make_rmem_tensor(2, Float32) + for g in range(2): + gmax = Float32(0.0) + for i in range(16 * g, 16 * g + 16): + gmax = cute.arch.fmax(gmax, cute.math.abs(vals[i])) + a = gmax * Float32(1.0 / 6.0) + if sf_scale is not None: + a = a * sf_scale + amax_6[g] = cute.arch.fmax(a, Float32(2.0**-9)) + + sf8 = amax_6.load().to(Float8E4M3FN) + sf.store(sf8.bitcast(Uint8)) + sf_f32 = sf8.to(Float32) + + q = cute.make_rmem_tensor(32, Float32) + for g in range(2): + # rcp.approx.ftz is exact on every positive E4M3 value + r = cute.arch.rcp_approx(sf_f32[g]) + if sf_scale is not None: + r = r * sf_scale + # rz: r >= 1 / sf, so a rounded product could step past an exact E2M1 + # midpoint and round away from even; truncation lands back on it (raw + # v accumulators do hit exact midpoints, rarely) + for i in range(16 * g, 16 * g + 16, 2): + q[i], q[i + 1] = cute.arch.mul_packed_f32x2( + (vals[i], vals[i + 1]), (r, r), rnd="rz" + ) + + # clear the sign of zero-magnitude nibbles (+7 carries into bit 3 exactly + # when the 3-bit magnitude is nonzero) + m = Uint32(0x77777777) + for j in range(4): + w = fp32x8_to_fp4x8(q, 8 * j) + words[j] = w & (((w & m) + m) | m) + + +def _nvfp4_quant64(vals, sf_scale=None): + """NVFP4-quantize 64 f32 (a 64-col half, one 32-col chunk at a time); + returns the 8 data words and the 4 scale bytes.""" + words = cute.make_rmem_tensor(8, Uint32) + sf = cute.make_rmem_tensor(4, Uint8) + for c in range(2): + _nvfp4_quant32( + cute.local_tile(vals, (32,), (c,)), + cute.local_tile(words, (4,), (c,)), + cute.local_tile(sf, (2,), (c,)), + sf_scale, + ) + return words.load(), sf.load() + + +class QkvIndexFusedEpilogueKernel: + """Persistent MXFP8 GEMM with a BM x BN = 128 x (128 * cta_group) output + tile per CTA (one head per epilogue warpgroup). 2-SM (cta_group=2): a CTA + pair runs one m256n256 MMA per k-step over two adjacent M-tiles, each CTA + loading half of the BN rows of B. 1-SM (cta_group=1): one m128n128 MMA per + CTA (finer tiles for small M). Warps: 4 * cta_group epilogue, then MMA and + TMA; cta_group=2 adds 2 padding warps so setmaxnreg can move registers to + the epilogue (at cta_group=1, 6 warps already fit 255 regs each). + """ + + def __init__( + self, + nkv: int, + block_size: int, + cta_group: int = 2, + group_m: int = 1, + use_pdl: bool = True, + ) -> None: + assert nkv in (1, 2, 4) + assert block_size % 4 == 0 # V scales swizzle over 4-token quads + assert cta_group in (1, 2) + self.cta_group = cta_group + self.group_m = group_m + self.num_warps = 12 if cta_group == 2 else 6 + self.tma_group = ( + tcgen05.CtaGroup.TWO if cta_group == 2 else tcgen05.CtaGroup.ONE + ) + self.use_pdl = use_pdl + self.nkv = nkv + self.block_size = block_size + + # The N dim is a list of 128-col heads, [q x 16nkv | k x nkv | + # v x nkv | index_q x nkv | index_k x 1]; first head of each role + # (q starts at 0). An odd head count leaves a phantom head in the + # last tile at cta_group=2 (TMA zero-fill, never stored). + self.k_start = 16 * nkv + self.v_start = self.k_start + nkv + self.iq_start = self.v_start + nkv + self.ik_start = self.iq_start + nkv + + BM, BN, BK = 128, cta_group * HEAD_DIM, 128 + self.cta_tile = (BM, BN, BK) + # per CTA and stage: A + this CTA's share of B (cta_group=2 splits B + # across the pair) + SFA + full SFB (each CTA multicasts its half) + self.stage_size = (BM + BN // cta_group) * BK + (BM + BN) * (BK // 32) + self.num_stages = ( + cutlass.memory.SmemAllocator.capacity_in_bytes() // self.stage_size + ) + + @cute.jit + def prepare_AB(self, A: cute.Tensor, BM: int, BK: int) -> cpasync.TmaInfo: + tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=self.tma_group) + swizzle_128B = cute.make_swizzle(3, 4, 3) + s_layout = cute.make_layout((BM, BK, self.num_stages), stride=(BK, 1, BM * BK)) + s_layout = cute.make_composed_layout(swizzle_128B, 0, s_layout) + return cpasync.make_tiled_tma_atom(tma_op, A, s_layout, (BM, BK)) + + @cute.jit + def prepare_SF( + self, + SF: cute.Tensor, + R: Int32, + Ksf: Int32, + tma_op: cpasync.CopyG2STileBaseOp, + ) -> cpasync.TmaInfo: + # SF grid [R, Ksf] padded to R % 128 == 0 and permuted as + # [R/128, Ksf/4, 32, 4, 4]; the [32, 4, 4] atom is tcgen05's opaque + # 512B block, so treat it as bytes. One box = one atom (128 rows x + # BK / 32 = 4 K-groups) = 64 Int64 x 1 row-block. + g_layout = cute.make_layout((Ksf * 128, cute.ceil_div(R, 128))) + SF_i64 = cute.recast_tensor(cute.make_tensor(SF.iterator, g_layout), Int64) + s_layout = cute.make_layout((64, 1, self.num_stages)) + return cpasync.make_tiled_tma_atom(tma_op, SF_i64, s_layout, (64, 1)) + + @cute.jit + def __call__( + self, + A: cute.Tensor, + B: cute.Tensor, + SFA: cute.Tensor, + SFB: cute.Tensor, + Q: cute.Tensor | None, + WQ: cute.Tensor, + WK: cute.Tensor, + WIQ: cute.Tensor, + WIK: cute.Tensor, + COS_SIN: cute.Tensor, + POS: cute.Tensor, + KVC: cute.Tensor, + IDXC: cute.Tensor, + SLOT: cute.Tensor, + IDX_SLOT: cute.Tensor, + IQFP8: cute.Tensor, + QFP8: cute.Tensor | None, + eps: Float32, + q_fp8_scale: Float32, + k_scale: Float32, + v_scale: Float32, + grid_size: Int32, + stream: CUstream, + ) -> None: + M, K = A.shape + N, _ = B.shape + BM, BN, BK = self.cta_tile + + A_tma = self.prepare_AB(A, BM, BK) + B_tma = self.prepare_AB(B, BN // self.cta_group, BK) + SFA_tma = self.prepare_SF( + SFA, M, K // 32, cpasync.CopyBulkTensorTileG2SOp(self.tma_group) + ) + sfb_op = ( + cpasync.CopyBulkTensorTileG2SMulticastOp(self.tma_group) + if self.cta_group == 2 + else cpasync.CopyBulkTensorTileG2SOp(self.tma_group) + ) + SFB_tma = self.prepare_SF(SFB, N, K // 32, sfb_op) + + self.kernel( + A_tma, + B_tma, + SFA_tma, + SFB_tma, + Q, + WQ, + WK, + WIQ, + WIK, + COS_SIN, + POS, + KVC, + IDXC, + SLOT, + IDX_SLOT, + IQFP8, + QFP8, + eps, + q_fp8_scale, + k_scale, + v_scale, + ).launch( + grid=(grid_size, 1, 1), + block=(self.num_warps * 32, 1, 1), + cluster=(self.cta_group, 1, 1), + min_blocks_per_mp=1, + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + A_tma: cpasync.TmaInfo, + B_tma: cpasync.TmaInfo, + SFA_tma: cpasync.TmaInfo, + SFB_tma: cpasync.TmaInfo, + Q_tensor: cute.Tensor | None, + WQ_tensor: cute.Tensor, + WK_tensor: cute.Tensor, + WIQ_tensor: cute.Tensor, + WIK_tensor: cute.Tensor, + COS_SIN_tensor: cute.Tensor, + POS_tensor: cute.Tensor, + KVC_tensor: cute.Tensor, + IDXC_tensor: cute.Tensor, + SLOT_tensor: cute.Tensor, + IDX_SLOT_tensor: cute.Tensor, + IQFP8_tensor: cute.Tensor, + QFP8_tensor: cute.Tensor | None, + eps: Float32, + q_fp8_scale: Float32, + k_scale: Float32, + v_scale: Float32, + ) -> None: + tid, _, _ = cute.arch.thread_idx() + raw_bid, _, _ = cute.arch.block_idx() + num_bids, _, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(tid // 32) + cg = self.cta_group + cta_rank = raw_bid % cg + + BM, BN, BK = self.cta_tile + num_stages = self.num_stages + + # tmem (512 cols): double-buffered accumulators, then SFA/SFB0/SFB1. + # At cta_group=2, 2 x 256 + SF cols don't fit, so they overlap by 32 + # cols (even tiles at [0, 256), odd at [224, 480)). + acc_odd = BN - 32 if cg == 2 else BN + sfa_tmem = 496 + sfb0_tmem = sfa_tmem + 4 + sfb1_tmem = sfb0_tmem + 4 + # q / k / index_q / index_k norm weight index in sW + W_Q, W_K, W_IQ, W_IK = range(4) + bs = self.block_size + + BAR_EPI = 1 + BAR_TMEM_ALLOC = 2 + cta_pair = Uint16(3) if cg == 2 else None # multicast / commit mask + + smem = cutlass.memory.SmemAllocator() + sA = smem.allocate_tensor( + Float8E4M3FN, + A_tma.smem_layout.outer, + byte_alignment=128, + swizzle=A_tma.smem_layout.inner, + ) + sB = smem.allocate_tensor( + Float8E4M3FN, + B_tma.smem_layout.outer, + byte_alignment=128, + swizzle=B_tma.smem_layout.inner, + ) + sSFA = smem.allocate_tensor(Int64, SFA_tma.smem_layout, byte_alignment=128) + # SFB slots: CTA r's multicast copy (head tile 2 * bid_n + r) fills slot + # r of both CTAs; each slot matches the atom's (64, 1):(1, 64) tile + sSFB = smem.allocate_tensor( + Int64, + cute.make_layout((64, 1, cg, num_stages), stride=(1, 64, 64, 64 * cg)), + byte_alignment=128, + ) + # norm weights as (8, 16) tiles per head type + sW = smem.allocate_tensor( + BFloat16, cute.make_layout((8, 16, 4)), byte_alignment=16 + ) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + tmem_full_mbar = smem.allocate_array(Int64, 2) + tmem_partial_mbar = smem.allocate_array(Int64, 1) + tmem_empty_mbar = smem.allocate_array(Int64, 2) + taddr = smem.allocate(Int32, 4) + + M, K = A_tma.tma_tensor.shape + N, _ = B_tma.tma_tensor.shape + + def num_tiles(M, N, cg): + # per-CTA M-tiles (a multiple of cg) x BN-wide N-tiles. Each role + # recomputes these after setmaxnreg: a value kept live across it + # is spilled. + grid_n = cute.ceil_div(N, 128 * cg) + grid_mp = cute.ceil_div(M, 128 * cg) + return grid_mp, grid_n, grid_mp * cg * grid_n + + def tile_coords(bid, grid_mp, grid_n, cg, G): + # CTA group p = bid // cg covers M-tiles cg * m + rank (one per + # CTA) and N-tile n. Rasterized in groups of G M-tile groups, m + # fastest within each N-tile of a group (G = 1: n fastest). + p = bid // cg + if cutlass.const_expr(G == 1): + return (p // grid_n) * cg + bid % cg, p % grid_n + group = p // (G * grid_n) + first = group * G + size = cutlass.min(grid_mp - first, G) + local = p - group * (G * grid_n) + return (first + local % size) * cg + bid % cg, local // size + + if warp_id == 0: + # counts: tma_full cg CTAs; tmem_empty cg CTAs x 128 * cg epilogue + # threads; tmem_partial 2 CTAs x 128 (the overlap warpgroup) + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, cg) + cute.arch.mbarrier_init(tma_empty_mbar + i, 1) + for i in cutlass.range_constexpr(2): + cute.arch.mbarrier_init(tmem_full_mbar + i, 1) + cute.arch.mbarrier_init(tmem_empty_mbar + i, 128 * cg * cg) + cute.arch.mbarrier_init(tmem_partial_mbar, 256) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(A_tma.atom) + cpasync.prefetch_descriptor(B_tma.atom) + cpasync.prefetch_descriptor(SFA_tma.atom) + cpasync.prefetch_descriptor(SFB_tma.atom) + elif warp_id < 4: + # norm weights (static, so before the PDL wait): 4 x 256B, one 16B + # cp.async per thread of warps 2-3 + t = tid - 64 + cp_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(), BFloat16, num_bits_per_copy=128 + ) + for i, Wt in enumerate((WQ_tensor, WK_tensor, WIQ_tensor, WIK_tensor)): + if t // 16 == i: + cute.copy( + cp_atom, + cute.zipped_divide(Wt, (8,))[None, t % 16], + sW[None, t % 16, i], + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() # sW visible (cluster_arrive is relaxed) + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + # warps: 4 * cg epilogue (one warpgroup per head), then MMA, TMA (and 2 + # padding warps at cg=2, completing the producer warpgroup) + MMA_WARP = 4 * cg + TMA_WARP = MMA_WARP + 1 + # setmaxnreg (cg=2): the producer warpgroup drops to 40 regs, the + # epilogue rises to 232 (2*128*232 + 128*40 <= 64K) + REG_DEC = 40 + REG_INC = 232 + + if warp_id > TMA_WARP: + # padding warps + if cutlass.const_expr(cg == 2): + cute.arch.setmaxregister_decrease(REG_DEC) + + elif warp_id == TMA_WARP: + # TMA warp; both CTAs load, copies arrive on CTA0's mbarrier. + # SFB is multicast: the pair's MMA needs all 256 rows in each CTA. + if cutlass.const_expr(cg == 2): + cute.arch.setmaxregister_decrease(REG_DEC) + grid_mp, grid_n, total_tiles = num_tiles(M, N, cg) + tma_stage = 0 + tma_empty_parity = 1 + tma_full_mbar_ = to_cta0_smem(tma_full_mbar) + + gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK)) + gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN // cg, BK)) + # SF boxes: one [32, 4, 4] atom = (64, 1) Int64 (see prepare_SF) + gSFA_tiles = cute.zipped_divide(SFA_tma.tma_tensor, (64, 1)) + gSFB_tiles = cute.zipped_divide(SFB_tma.tma_tensor, (64, 1)) + + for bid in range(raw_bid, total_tiles, num_bids): + bid_m, bid_n = tile_coords(bid, grid_mp, grid_n, cg, self.group_m) + + for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1): + mbar = tma_full_mbar_ + tma_stage + cute.arch.mbarrier_wait( + tma_empty_mbar + tma_stage, tma_empty_parity + ) + with cute.arch.elect_one(): + mbarrier.arrive_expect_tx(mbar, self.stage_size, "cluster") + simple_tma_copy( + A_tma.atom, + gA_tiles[None, (bid_m, iter_k)], + sA[None, None, tma_stage], + mbar, + ) + simple_tma_copy( + B_tma.atom, + gB_tiles[None, (cg * bid_n + cta_rank, iter_k)], + sB[None, None, tma_stage], + mbar, + ) + simple_tma_copy( + SFA_tma.atom, + gSFA_tiles[None, (iter_k, bid_m)], + sSFA[None, None, tma_stage], + mbar, + ) + simple_tma_copy( + SFB_tma.atom, + gSFB_tiles[None, (iter_k, cg * bid_n + cta_rank)], + sSFB[None, None, cta_rank, tma_stage], + mbar, + mcast_mask=cta_pair, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_empty_parity ^= 1 + + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + elif warp_id == MMA_WARP: + # MMA warp (2-SM collective; only cta_rank 0 issues). tmem is + # allocated by epilogue warp 0. + if cutlass.const_expr(cg == 2): + cute.arch.setmaxregister_decrease(REG_DEC) + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=2 * 32) + + grid_mp, grid_n, total_tiles = num_tiles(M, N, cg) + tma_stage = 0 + tma_full_parity = 0 + + sdesc = _tcgen05.make_sdesc_128B_swizzle(0) + # SF smem descriptor: no swizzle, SBO = 8 * 16 bytes + sdesc_sf = Uint64((8 << 32) | (1 << 46)) + idesc = _tcgen05.make_mxfp8_idesc(cg * BM, BN) + + if cta_rank == 0: + tmem_stage = 0 + tmem_parity = 1 + for _ in range(raw_bid, total_tiles, num_bids): + acc_base = acc_odd * tmem_stage + cute.arch.mbarrier_wait(tmem_empty_mbar + tmem_stage, tmem_parity) + # the previous tile's overlap-region read (cols [224, 256)); + # the single partial barrier phases once per tile + if cutlass.const_expr(cg == 2): + cute.arch.mbarrier_wait(tmem_partial_mbar, tmem_stage ^ 1) + _tcgen05.fence_after_thread_sync() + + for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1): + a_desc = sdesc | ( + sA[None, None, tma_stage].iterator.toint() >> 4 + ) + b_desc = sdesc | ( + sB[None, None, tma_stage].iterator.toint() >> 4 + ) + sfa_desc = sdesc_sf | ( + sSFA[None, None, tma_stage].iterator.toint() >> 4 + ) + sfb0_desc = sdesc_sf | ( + sSFB[None, None, 0, tma_stage].iterator.toint() >> 4 + ) + + cute.arch.mbarrier_wait( + tma_full_mbar + tma_stage, tma_full_parity + ) + _tcgen05.fence_after_thread_sync() + + _tcgen05.cp(sfa_tmem, sfa_desc, "32x128b", "warpx4", cg) + _tcgen05.cp(sfb0_tmem, sfb0_desc, "32x128b", "warpx4", cg) + if cutlass.const_expr(cg == 2): + sfb1_desc = sdesc_sf | ( + sSFB[None, None, 1, tma_stage].iterator.toint() >> 4 + ) + _tcgen05.cp(sfb1_tmem, sfb1_desc, "32x128b", "warpx4", 2) + + for k in cutlass.range_constexpr(BK // 32): + enable_input_d = iter_k > 0 or k > 0 + # select this k-block's SF byte inside the staged + # 4-col tmem atoms (idesc SF-id fields) + idesc_k = idesc + ((k << 4) | (k << 29)) + _tcgen05.mma_mxfp8( + acc_base, + a_desc, + b_desc, + idesc_k, + sfa_tmem, + sfb0_tmem, + enable_input_d, + cg, + ) + a_desc += 32 >> 4 + b_desc += 32 >> 4 + _tcgen05.commit(tma_empty_mbar + tma_stage, cta_pair, cg) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_full_parity ^= 1 + + _tcgen05.commit(tmem_full_mbar + tmem_stage, cta_pair, cg) + tmem_stage ^= 1 + if tmem_stage == 0: + tmem_parity ^= 1 + + else: + # epilogue warps + if cutlass.const_expr(cg == 2): + cute.arch.setmaxregister_increase(REG_INC) + if warp_id == 0: + _tcgen05.alloc(taddr, cta_group=cg) + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=2 * 32) + + # warp w reads tmem rows (w % 4) * 32 of head cg * bid_n + w // 4; + # thread t owns row t % 128. Padding rows (m >= M) compute garbage + # from clamped metadata and skip only the stores. + warp_e = warp_id % 4 + wg_id = warp_id // 4 + tid_ = tid % 128 + + # Q / QFP8 are each optional (None -> compiled out) + has_q = cutlass.const_expr(Q_tensor is not None) + has_qfp8 = cutlass.const_expr(QFP8_tensor is not None) + # [32, (pos, cos | sin)] + CS_CHUNKS = cute.zipped_divide(COS_SIN_tensor, (1, 32))[(0, None), None] + # KV head h's K / V go to slots 2h / 2h + 1; index_k has one slot + KV_DATA, KV_SF, KV_VSF = _nvfp4_cache_views(KVC_tensor, bs, 2 * self.nkv) + IDX_DATA, IDX_SF, _ = _nvfp4_cache_views(IDXC_tensor, bs, 1) + + partial_ = to_cta0_smem(tmem_partial_mbar) + empty_ = to_cta0_smem(tmem_empty_mbar) + + grid_mp, grid_n, total_tiles = num_tiles(M, N, cg) + tmem_stage = 0 + tmem_parity = 0 + for bid in range(raw_bid, total_tiles, num_bids): + bid_m, bid_n = tile_coords(bid, grid_mp, grid_n, cg, self.group_m) + acc_base = acc_odd * tmem_stage + + if warp_id == 0: + cute.arch.mbarrier_wait(tmem_full_mbar + tmem_stage, tmem_parity) + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=cg * 128) + _tcgen05.fence_after_thread_sync() + + head = cg * bid_n + wg_id # this warpgroup's head + m = bid_m * BM + tid_ + + # Drain this warp's 32 rows of its head into registers + row = warp_e * 32 + base = acc_base + wg_id * HEAD_DIM + acc = cute.make_rmem_tensor((32, 4), Float32) # 4 x 32 cols + owns_overlap = False + if cutlass.const_expr(cg == 2): + owns_overlap = wg_id == 1 - tmem_stage + if owns_overlap: + # This warpgroup holds cols [224, 256), which the next + # tile's MMA reuses: load and release them first. They are + # chunk 0 of the head on odd tiles, chunk 3 on even. + if tmem_stage == 1: + acc[None, 0].store(_tcgen05.ld(row, base, "32x32b", 32)) + _tcgen05.wait_ld() + mbarrier.arrive(partial_, "cluster") + for c in cutlass.range_constexpr(1, 4): + col = base + 32 * c + acc[None, c].store(_tcgen05.ld(row, col, "32x32b", 32)) + else: + acc[None, 3].store(_tcgen05.ld(row, base + 96, "32x32b", 32)) + _tcgen05.wait_ld() + mbarrier.arrive(partial_, "cluster") + for c in cutlass.range_constexpr(3): + col = base + 32 * c + acc[None, c].store(_tcgen05.ld(row, col, "32x32b", 32)) + else: + acc.store(_tcgen05.ld(row, base, "32x32b", HEAD_DIM)) + _tcgen05.wait_ld() + _tcgen05.fence_before_thread_sync() + mbarrier.arrive(empty_ + tmem_stage, "cluster") + + # rope dim is 64: store the unrotated half (cols 64-127) first + # to free its registers before rope. Only the stores are + # predicated on the row / slot. + halves = cute.make_tensor(acc.iterator, cute.make_layout((64, 2))) + m_safe = cutlass.min(m, M - 1) + + if head < self.k_start: + # q: bf16 and / or scaled e4m3 + _rms_norm(acc, sW[None, None, W_Q], eps) + if cutlass.const_expr(has_qfp8): + qfp8_inv = cute.math.rcp(q_fp8_scale) + + for h in (1, 0): # NoPE, then RoPE + if cutlass.const_expr(h == 0): + _rope64(acc, CS_CHUNKS, POS_tensor[m_safe].to(Int32)) + + x = halves[None, h].load() + col = HEAD_DIM * head + 64 * h + if cutlass.const_expr(has_q): + assert Q_tensor is not None + q = x.to(BFloat16) + if m < M: + dst = Q_tensor[m, None].iterator + col + cute.arch.store(dst, q, **_NO_L1_ALLOC) + # fp8 is stored as u32 words: fp8 vectors don't lower + # in cute.arch.store (NVIDIA/cutlass#3712) + if cutlass.const_expr(has_qfp8): + assert QFP8_tensor is not None + q = (x * qfp8_inv).to(Float8E4M3FN).bitcast(Uint32) + if m < M: + dst = QFP8_tensor[m, None].iterator + col + dst = cute.recast_ptr(dst, dtype=Uint32) + cute.arch.store(dst, q, **_NO_L1_ALLOC) + + elif head < self.v_start: + # k: NVFP4 into its KV head's K slot, linear scale bytes + _rms_norm(acc, sW[None, None, W_K], eps) + kv_slot = 2 * (head - self.k_start) + slot = SLOT_tensor[m_safe] + page, token = slot // bs, slot % bs + k_sf = cute.math.rcp(k_scale) + sf = cute.make_rmem_tensor(8, Uint8) + + for h in (1, 0): # NoPE, then RoPE + if cutlass.const_expr(h == 0): + _rope64(acc, CS_CHUNKS, POS_tensor[m_safe].to(Int32)) + + words, sf_h = _nvfp4_quant64(halves[None, h], k_sf) + cute.local_tile(sf, (4,), (h,)).store(sf_h) + if m < M and slot >= 0: + dst = KV_DATA[None, h, token, kv_slot, page].iterator + cute.arch.store(dst, words, **_NO_L1_ALLOC) + + if m < M and slot >= 0: + dst = KV_SF[None, token, kv_slot, page].iterator + cute.arch.store(dst, sf.load(), **_NO_L1_ALLOC) + + elif head < self.iq_start: + # v: raw accumulator, NVFP4 into its KV head's V slot, + # swizzled scale bytes + kv_slot = 2 * (head - self.v_start) + 1 + slot = SLOT_tensor[m_safe] + page, token = slot // bs, slot % bs + v_sf = cute.math.rcp(v_scale) + + for h in (0, 1): + words, sf = _nvfp4_quant64(halves[None, h], v_sf) + if m < M and slot >= 0: + dst = KV_DATA[None, h, token, kv_slot, page].iterator + cute.arch.store(dst, words, **_NO_L1_ALLOC) + vsf = KV_VSF[None, token, kv_slot, page] + cute.local_tile(vsf, (4,), (h,)).store(sf) + + elif head < self.ik_start: + # index_q: e4m3 + _rms_norm(acc, sW[None, None, W_IQ], eps) + iq_col = HEAD_DIM * (head - self.iq_start) + + for h in (1, 0): # NoPE, then RoPE + if cutlass.const_expr(h == 0): + _rope64(acc, CS_CHUNKS, POS_tensor[m_safe].to(Int32)) + + # u32 words: see the q branch + q = halves[None, h].load().to(Float8E4M3FN).bitcast(Uint32) + if m < M: + dst = IQFP8_tensor[m, None].iterator + iq_col + 64 * h + dst = cute.recast_ptr(dst, dtype=Uint32) + cute.arch.store(dst, q, **_NO_L1_ALLOC) + + elif head == self.ik_start: + # index_k: NVFP4 into the index cache (one slot per page, + # linear scale bytes) + _rms_norm(acc, sW[None, None, W_IK], eps) + slot = IDX_SLOT_tensor[m_safe] + page, token = slot // bs, slot % bs + sf = cute.make_rmem_tensor(8, Uint8) + + for h in (1, 0): # NoPE, then RoPE + if cutlass.const_expr(h == 0): + _rope64(acc, CS_CHUNKS, POS_tensor[m_safe].to(Int32)) + + words, sf_h = _nvfp4_quant64(halves[None, h]) + cute.local_tile(sf, (4,), (h,)).store(sf_h) + if m < M and slot >= 0: + dst = IDX_DATA[None, h, token, 0, page].iterator + cute.arch.store(dst, words, **_NO_L1_ALLOC) + + if m < M and slot >= 0: + dst = IDX_SF[None, token, 0, page].iterator + cute.arch.store(dst, sf.load(), **_NO_L1_ALLOC) + + # else: phantom head (odd head count), nothing to store + + tmem_stage ^= 1 + if tmem_stage == 0: + tmem_parity ^= 1 + + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + if warp_id == 0: + _tcgen05.dealloc(cta_group=cg) + + +class _QkvIndexFusedEpilogueJit( + VllmCuTeDSLJitKernel["_QkvIndexFusedEpilogueJit.CompileKey"] +): + """Compile keys, warmup (fake) inputs and the compiled executors of + QkvIndexFusedEpilogueKernel.""" + + @dataclass(frozen=True) + class CompileKey: + has_q: bool + has_qfp8: bool + nkv: int + block_size: int + group_m: int + cta_group: int + use_pdl: bool = True + + @staticmethod + def kernel(key: CompileKey) -> QkvIndexFusedEpilogueKernel: + return QkvIndexFusedEpilogueKernel( + key.nkv, key.block_size, key.cta_group, key.group_m, key.use_pdl + ) + + def get_warmup_keys(self, **fields) -> list[CompileKey]: + # a layer's keys differ only in the tile config, picked per batch size + return [self.CompileKey(cta_group=cg, **fields) for cg in (1, 2)] + + def warmup_inputs(self, key: CompileKey) -> tuple: + M = cute.sym_int() + N = cute.sym_int() + K = cute.sym_int() + P = cute.sym_int() + slot_bytes = key.block_size * (HEAD_DIM // 2 + HEAD_DIM // 16) + + A = make_fake_tensor( + Float8E4M3FN, (M, K), (cute.sym_int64(divisibility=16), 1), assumed_align=16 + ) + B = make_fake_tensor( + Float8E4M3FN, (N, K), (cute.sym_int64(divisibility=16), 1), assumed_align=16 + ) + SFA = make_fake_tensor( + Float8E8M0FNU, (cute.sym_int(divisibility=512),), (1,), assumed_align=16 + ) + SFB = make_fake_tensor( + Float8E8M0FNU, (cute.sym_int(divisibility=512),), (1,), assumed_align=16 + ) + q_cols = 16 * key.nkv * HEAD_DIM + Q = ( + make_fake_tensor(BFloat16, (M, q_cols), (q_cols, 1), assumed_align=32) + if key.has_q + else None + ) + WQ, WK, WIQ, WIK = ( + make_fake_tensor(BFloat16, (HEAD_DIM,), (1,), assumed_align=16) + for _ in range(4) + ) + COS_SIN = make_fake_tensor(BFloat16, (P, 64), (64, 1), assumed_align=16) + POS = make_fake_tensor(Int64, (M,), (1,), assumed_align=16) + # caches as [nb, page_bytes] with a dynamic page stride: no dynamic dim + # may exceed int32, and cross-layer layouts space pages wider than a page + KVC = make_fake_tensor( + Uint8, + (cute.sym_int(), 2 * key.nkv * slot_bytes), + (cute.sym_int64(divisibility=32), 1), + assumed_align=32, + ) + IDXC = make_fake_tensor( + Uint8, + (cute.sym_int(), slot_bytes), + (cute.sym_int64(divisibility=32), 1), + assumed_align=32, + ) + SLOT = make_fake_tensor(Int64, (M,), (1,), assumed_align=16) + IDX_SLOT = make_fake_tensor(Int64, (M,), (1,), assumed_align=16) + IQFP8 = make_fake_tensor( + Float8E4M3FN, + (M, key.nkv * HEAD_DIM), + (key.nkv * HEAD_DIM, 1), + assumed_align=32, + ) + QFP8 = ( + make_fake_tensor(Float8E4M3FN, (M, q_cols), (q_cols, 1), assumed_align=32) + if key.has_qfp8 + else None + ) + + return ( + A, + B, + SFA, + SFB, + Q, + WQ, + WK, + WIQ, + WIK, + COS_SIN, + POS, + KVC, + IDXC, + SLOT, + IDX_SLOT, + IQFP8, + QFP8, + # runtime scalars (eps, q/k/v scales, grid): placeholder values + Float32(1e-6), + Float32(1.0), + Float32(1.0), + Float32(1.0), + 148, + ) + + +_JIT = _QkvIndexFusedEpilogueJit() + + +@cache +def _device_info(device: int) -> tuple[int, int]: + """(SM count, L2 bytes) of a CUDA device.""" + props = torch.cuda.get_device_properties(device) + return props.multi_processor_count, props.L2_cache_size + + +def _group_m(weight_bytes: int, l2_bytes: int) -> int: + # Group M-tiles 16 at a time when the weight is a large fraction of L2, so + # a wave reads a few N-panels instead of all of B (GB300: helps TP1's 60MB + # weight by 1.5-5%; slows TP2 / TP4, whose weights stay L2-resident anyway) + return 16 if weight_bytes >= l2_bytes // 4 else 1 + + +def register_qkv_index_fused_epilogue_warmup( + *, has_q: bool, has_qfp8: bool, nkv: int, block_size: int, weight_bytes: int +) -> None: + """Compile this layer's kernels (both tile configs) during JIT warmup.""" + _, l2_bytes = _device_info(torch.accelerator.current_device_index()) + _JIT.register_warmup( + has_q=has_q, + has_qfp8=has_qfp8, + nkv=nkv, + block_size=block_size, + group_m=_group_m(weight_bytes, l2_bytes), + ) + + +def minimax_m3_qkv_index_fused_epilogue( + xq: torch.Tensor, + xsf: torch.Tensor, + wq: torch.Tensor, + wsf: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + index_q_norm_weight: torch.Tensor, + index_k_norm_weight: torch.Tensor, + eps: float, + *, + slot_mapping: torch.Tensor, + index_slot_mapping: torch.Tensor, + kv_cache: torch.Tensor, + index_cache: torch.Tensor, + index_q_out: torch.Tensor, + q_out: torch.Tensor | None = None, + q_fp8_out: torch.Tensor | None = None, + q_fp8_scale: float = 1.0, + k_scale: float = 1.0, + v_scale: float = 1.0, + cta_group: int | None = None, + use_pdl: bool = True, + grid: int | None = None, +) -> None: + """QKV+indexer projection with fused epilogue: writes q (bf16 and / or + fp8; at least one) and index_q, and NVFP4-inserts k / v / index_k into the + paged caches. + nkv in {1, 2, 4} KV heads per rank is inferred from N = (19 * nkv + 1) * 128. + + Args: + xq: MXFP8 activation [M, K]. + xsf: xq's F8_128x4-swizzled E8M0 scales. + wq: K-major MXFP8 weight [N, K]: [q | k | v | index_q | index_k]. + wsf: wq's F8_128x4-swizzled E8M0 scales. + positions: [M] int64 token positions. + cos_sin_cache: [max_pos, 64] bf16 RoPE cache (cos | sin). + q_norm_weight: [128] bf16 Gemma RMSNorm weight; likewise the next three. + k_norm_weight: see q_norm_weight. + index_q_norm_weight: see q_norm_weight. + index_k_norm_weight: see q_norm_weight. + eps: RMSNorm epsilon. + slot_mapping: [M] int64 KV-cache slots (-1 = skip). + index_slot_mapping: [M] int64 index-cache slots (-1 = skip). + kv_cache: NVFP4 [nb, 2 * nkv, block_size, 72] uint8; packed pages at + any page stride. block_size % 4 == 0. + index_cache: NVFP4 [nb, block_size, 72] uint8; packed pages at any + page stride. + index_q_out: [M, nkv * 128] e4m3 normed + roped index_q. + q_out: optional [M, 16 * nkv * 128] bf16 normed + roped q. + q_fp8_out: optional [M, 16 * nkv * 128] e4m3 normed + roped q. + q_fp8_scale: dequantization scale of q_fp8_out. + k_scale: dequantization (global) scale of the NVFP4 k entries. + v_scale: dequantization (global) scale of the NVFP4 v entries. + cta_group: 2 (2-SM: CTA pairs, 128 x 256 tiles), 1 (1-SM: 128 x 128 + tiles) or None (pick by wave count). + use_pdl: launch with programmatic dependent launch. + grid: persistent grid size override (testing). + + """ + M, K = xq.shape + N = wq.shape[0] + nkv = (N // HEAD_DIM - 1) // 19 + assert nkv in (1, 2, 4) and (19 * nkv + 1) * HEAD_DIM == N, N + assert K % 128 == 0 and wq.shape[1] == K and wq.stride() == (K, 1) + assert index_q_out.shape == (M, nkv * HEAD_DIM) + assert index_q_out.dtype == torch.float8_e4m3fn and index_q_out.is_contiguous() + assert q_out is not None or q_fp8_out is not None, "no q output" + for t in (positions, slot_mapping, index_slot_mapping): + assert t.shape == (M,) and t.dtype == torch.int64 + assert cos_sin_cache.dtype == torch.bfloat16 and cos_sin_cache.is_contiguous() + assert cos_sin_cache.dim() == 2 and cos_sin_cache.shape[1] == 64 + for w in (q_norm_weight, k_norm_weight, index_q_norm_weight, index_k_norm_weight): + assert w.shape == (HEAD_DIM,) and w.dtype == torch.bfloat16 + for q, dtype in ((q_out, torch.bfloat16), (q_fp8_out, torch.float8_e4m3fn)): + if q is not None: + assert q.shape == (M, 16 * nkv * HEAD_DIM) + assert q.dtype == dtype and q.is_contiguous() + assert kv_cache.dtype == torch.uint8 and index_cache.dtype == torch.uint8 + block_size, token_bytes = kv_cache.shape[2:] + assert token_bytes == HEAD_DIM // 2 + HEAD_DIM // 16 + assert kv_cache.shape[1] == 2 * nkv + assert kv_cache.stride()[1:] == (block_size * token_bytes, token_bytes, 1) + assert index_cache.shape[1:] == (block_size, token_bytes) + assert index_cache.stride()[1:] == (token_bytes, 1) + for c in (kv_cache, index_cache): # 32B NVFP4 data stores + assert c.data_ptr() % 32 == 0 and c.stride(0) % 32 == 0 + num_sms, l2_bytes = _device_info(xq.device.index) + + def num_tiles(cg): + tile = 128 * cg # M rows per CTA group, N cols per tile + return -(-M // tile) * cg * -(-N // tile) + + if cta_group is None: + # a 128 x 128 tile costs ~2/3 of a 128 x 256 one (GB300): take 1-SM + # when that wins after wave quantization (odd M-tile counts, few tiles) + waves1, waves2 = (-(-num_tiles(cg) // num_sms) for cg in (1, 2)) + cta_group = 1 if 2 * waves1 < 3 * waves2 else 2 + if grid is None: + grid = min(num_sms, num_tiles(cta_group)) + # caches as [nb, page_bytes], keeping the page stride (a flat 1-D view + # exceeds the DSL's int32 shapes at production sizes) + slot_bytes = block_size * token_bytes + kvc = kv_cache.as_strided( + (kv_cache.size(0), 2 * nkv * slot_bytes), (kv_cache.stride(0), 1) + ) + idxc = index_cache.as_strided( + (index_cache.size(0), slot_bytes), (index_cache.stride(0), 1) + ) + key = _JIT.CompileKey( + has_q=q_out is not None, + has_qfp8=q_fp8_out is not None, + nkv=nkv, + block_size=block_size, + group_m=_group_m(wq.numel(), l2_bytes), + cta_group=cta_group, + use_pdl=use_pdl, + ) + _JIT._get_or_compile(key)( + xq, + wq, + xsf.view(-1).view(torch.float8_e8m0fnu), + wsf.view(-1).view(torch.float8_e8m0fnu), + q_out, + q_norm_weight, + k_norm_weight, + index_q_norm_weight, + index_k_norm_weight, + cos_sin_cache, + positions, + kvc, + idxc, + slot_mapping, + index_slot_mapping, + index_q_out, + q_fp8_out, + eps, + q_fp8_scale, + k_scale, + v_scale, + grid, + ) diff --git a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py index a5598c719bfa..7443b449f776 100644 --- a/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py +++ b/vllm/models/minimax_m3/nvidia/sparse_attention_msa.py @@ -1,16 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MSA (SM100/Blackwell) block-sparse attention for MiniMax M3. +"""MSA (SM100 family) block-sparse attention for MiniMax M3. -Prefill attends with ``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``). -Decode uses Triton split-K by default, with an opt-in CUTLASS ``fmha_sm100`` -path for regular decode and speculative verification. NVFP4 KV caches read the -vLLM packed pages directly in both prefill and CUTLASS decode. +``MiniMaxM3SparseMSAImpl`` serves BF16 and FP8 KV caches: prefill attends with +``fmha_sm100`` (``build_k2q_csr`` + ``sparse_atten_func``), and decode uses the +CUTLASS ``fmha_sm100`` path by default for regular decode and speculative +verification, with Triton split-K as the fallback. + +``MiniMaxM3SparseMSANvfp4Impl`` serves NVFP4 KV caches, reading vLLM's packed +pages directly with only the E4M3 query: decode and prefill both run the CUTLASS +Q8KV4 kernels. """ +import functools import os -from dataclasses import dataclass, replace -from typing import ClassVar +from dataclasses import dataclass, fields, replace +from typing import TYPE_CHECKING, ClassVar import torch @@ -29,8 +34,10 @@ MiniMaxM3SparseImpl, MiniMaxM3SparseMetadata, MiniMaxM3SparseMetadataBuilder, + MiniMaxM3SparsePrefillMetadata, ) from vllm.models.minimax_m3.nvidia.msa_cutlass_sparse_decode import ( + NVFP4_BLOCK_SCALE_SHIFT, MSACutlassDecodeMetadata, MSACutlassDecodePlanCache, is_nvfp4_kv_cache, @@ -39,7 +46,9 @@ prepare_decode_metadata, should_prepare_decode_metadata, supports_cutlass_sparse_decode, + uses_q8kv4_decode, ) +from vllm.platforms import current_platform from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim from vllm.v1.attention.backend import ( AttentionLayer, @@ -47,19 +56,40 @@ ) from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout +if TYPE_CHECKING: + from vllm.third_party.fmha_sm100.prefill_q8kv4 import PrefillWorkspace + logger = init_logger(__name__) +@functools.cache +def _q8kv4_prefill_blocker( + num_q_heads: int, num_kv_heads: int, page_size: int, topk_blocks: int +) -> str | None: + """Why the Q8KV4 kernel cannot prefill this layout on this device, or None.""" + from vllm.third_party.fmha_sm100.q8kv4_prefill_adapter import _plan_blocker + + return _plan_blocker( + num_qo_heads=num_q_heads, + num_kv_heads=num_kv_heads, + page_size=page_size, + kv_block_num=topk_blocks, + causal=True, + output_maxscore=False, + device=None, + ) + + def _get_msa_cutlass_min_batch_size() -> int | None: value = os.getenv("VLLM_MINIMAX_M3_MSA_CUTLASS_MIN_BATCH") return int(value) if value is not None else None def _dequantize_query( - query_fp8: torch.Tensor, q_scale: float, out: torch.Tensor + query_fp8: torch.Tensor, q_scale: float, dtype: torch.dtype ) -> torch.Tensor: - """Write ``query_fp8 * q_scale`` into ``out`` (E4M3 is exact in BF16).""" - out.copy_(query_fp8) + """``query_fp8 * q_scale`` in ``dtype`` (E4M3 is exact in BF16).""" + out = query_fp8.to(dtype) if q_scale != 1.0: out.mul_(q_scale) return out @@ -118,6 +148,12 @@ class MiniMaxM3SparseMSADecodeMetadata(MiniMaxM3SparseDecodeMetadata): msa_cutlass: MSACutlassDecodeMetadata | None = None +@dataclass +class MiniMaxM3SparseMSAPrefillMetadata(MiniMaxM3SparsePrefillMetadata): + # The builder's Q8KV4 prefill split workspace; None off Rubin. + workspace: "PrefillWorkspace | None" = None + + class MiniMaxM3SparseMSAMetadataBuilder(MiniMaxM3SparseMetadataBuilder): """Prepare MSA plans only for decode shapes supported by ``fmha_sm100``.""" @@ -138,8 +174,29 @@ def __init__( # format to distinguish E4M3 (supported) from E5M2 before planning. self.kv_cache_dtype = vllm_config.cache_config.cache_dtype self.decode_backend = vllm_config.attention_config.minimax_m3_msa_decode_backend + self.q8kv4_decode = uses_q8kv4_decode( + kv_cache_dtype=self.kv_cache_dtype, + num_q_heads=self.num_q_heads, + num_kv_heads=self.num_kv_heads, + ) self.msa_cutlass_plan_cache = MSACutlassDecodePlanCache() self.msa_cutlass_min_batch_size = _get_msa_cutlass_min_batch_size() + # Rubin only: Q8KV4 prefill split partials striped over the memory + # locality domains, reserved for the largest batch. + self.prefill_workspace: "PrefillWorkspace | None" = None # noqa: UP037 + if is_nvfp4_kv_cache( + self.kv_cache_dtype + ) and current_platform.is_device_capability(107): + from vllm.third_party.fmha_sm100.prefill_q8kv4 import PrefillWorkspace + + self.prefill_workspace = PrefillWorkspace( + device, block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT + ) + self.prefill_workspace.split_buffers( + self.topk_blocks, + vllm_config.scheduler_config.max_num_batched_tokens, + self.num_q_heads, + ) def build( self, @@ -152,6 +209,12 @@ def build( common_attn_metadata, fast_build, ) + prefill = metadata.prefill + if prefill is not None and self.prefill_workspace is not None: + metadata.prefill = MiniMaxM3SparseMSAPrefillMetadata( + **{f.name: getattr(prefill, f.name) for f in fields(prefill)}, + workspace=self.prefill_workspace, + ) decode = metadata.decode if decode is None: return metadata @@ -179,6 +242,7 @@ def build( num_kv_heads=self.num_kv_heads, page_size=SPARSE_BLOCK_SIZE, topk_blocks=self.topk_blocks, + q8kv4=self.q8kv4_decode, plan_cache=self.msa_cutlass_plan_cache, ) metadata.decode = MiniMaxM3SparseMSADecodeMetadata( @@ -190,8 +254,58 @@ def build( return metadata -class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl): - """MSA block-sparse attention with guarded CUTLASS sparse decode.""" +class MiniMaxM3SparseMSABaseImpl(MiniMaxM3SparseImpl): + """What the MSA impls share: CUTLASS decode plans and the prefill CSR.""" + + use_cutlass_decode: bool + # The prefill kernels ``fmha_sm100.warmup`` prebuilds for this impl. + prefill_backend: str + + def should_use_msa_decode(self, layer_name: str) -> bool: + if not self.use_cutlass_decode: + return False + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return False + main_md = attn_metadata[layer_name] + if not isinstance(main_md, MiniMaxM3SparseMetadata): + return False + decode = main_md.decode + return ( + isinstance(decode, MiniMaxM3SparseMSADecodeMetadata) + and decode.msa_cutlass is not None + ) + + def _build_prefill_csr( + self, + topk: torch.Tensor, + nd: int, + num_tokens: int, + p: MiniMaxM3SparsePrefillMetadata, + ): + from vllm.third_party.fmha_sm100.sparse import build_k2q_csr + + # [H, prefill, MK] transposed view; build_k2q_csr consumes the strided + # view directly (topK stays innermost-contiguous). + return build_k2q_csr( + topk[nd:num_tokens].transpose(0, 1), + p.cu_seqlens_q, + p.cu_seqlens_k, + SPARSE_BLOCK_SIZE, + total_k=0, + max_seqlen_k=p.max_seq_len, + max_seqlen_q=p.max_query_len, + total_rows=p.total_kv_blocks, + qhead_per_kv=self.num_heads // self.num_kv_heads, + return_schedule=True, + ) + + +class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseMSABaseImpl): + """MSA block-sparse attention over a BF16 or FP8 KV cache, with guarded + CUTLASS sparse decode.""" + + prefill_backend = "cute_dsl" def __init__( self, @@ -203,7 +317,7 @@ def __init__( *, topk_blocks: int, sparse_block_size: int, - msa_decode_backend: MiniMaxM3MSADecodeBackend = "triton", + msa_decode_backend: MiniMaxM3MSADecodeBackend = "cutlass", ) -> None: super().__init__( num_heads, @@ -214,7 +328,6 @@ def __init__( topk_blocks=topk_blocks, sparse_block_size=sparse_block_size, ) - self.use_nvfp4_kv = is_nvfp4_kv_cache(kv_cache_dtype) self.use_cutlass_decode = supports_cutlass_sparse_decode( decode_backend=msa_decode_backend, num_q_heads=self.num_heads, @@ -223,31 +336,18 @@ def __init__( page_size=self.block_size, topk_blocks=self.topk_blocks, ) - if self.use_nvfp4_kv and not self.use_cutlass_decode: - raise ValueError( - "MiniMax M3 NVFP4 KV cache requires CUTLASS MSA sparse decode " - f"(num_heads={self.num_heads}, num_kv_heads={self.num_kv_heads}, " - f"block_size={self.block_size}, topk_blocks={self.topk_blocks})" - ) logger.info_once( "MiniMax M3 MSA sparse decode selected %s", "CUTLASS" if self.use_cutlass_decode else "Triton", ) - def should_use_msa_decode(self, layer_name: str) -> bool: + def _bf16_query( + self, q: torch.Tensor, q_scale: float, dtype: torch.dtype + ) -> torch.Tensor: + # With CUTLASS decode the query is E4M3; the other kernels read BF16. if not self.use_cutlass_decode: - return False - attn_metadata = get_forward_context().attn_metadata - if not isinstance(attn_metadata, dict): - return False - main_md = attn_metadata[layer_name] - if not isinstance(main_md, MiniMaxM3SparseMetadata): - return False - decode = main_md.decode - return ( - isinstance(decode, MiniMaxM3SparseMSADecodeMetadata) - and decode.msa_cutlass is not None - ) + return q + return _dequantize_query(q, q_scale, dtype) def forward( self, @@ -255,8 +355,6 @@ def forward( query: torch.Tensor, kv_cache: torch.Tensor, output: torch.Tensor, - *, - query_fp8: torch.Tensor | None = None, ) -> torch.Tensor: attn_metadata = get_forward_context().attn_metadata if not isinstance(attn_metadata, dict): @@ -273,16 +371,8 @@ def forward( hd = self.head_size q = query[:num_tokens].view(-1, self.num_heads, hd) out = output[:num_tokens].view(-1, self.num_heads, hd) - # Given query_fp8, it is the only q the fused insert wrote: CUTLASS - # decode and prefill read it directly, and Triton decode reads it - # dequantized into ``query``. - q_fp8 = ( - None - if query_fp8 is None - else query_fp8[:num_tokens].view(-1, self.num_heads, hd) - ) q_scale_float = getattr(layer, "_q_scale_float", 1.0) - if self.use_fp8_kv and not self.use_nvfp4_kv: + if self.use_fp8_kv: kv_cache = kv_cache.view(self.kv_cache_fp8_dtype) k_scale = getattr(layer, "_k_scale", None) if self.use_fp8_kv else None v_scale = getattr(layer, "_v_scale", None) if self.use_fp8_kv else None @@ -290,16 +380,11 @@ def forward( # Decode [:nd]: CUTLASS for planned shapes, otherwise Triton. if main_md.num_decodes > 0: d = main_md.decode - assert d is not None - msa_metadata = ( - d.msa_cutlass - if isinstance(d, MiniMaxM3SparseMSADecodeMetadata) - else None - ) + assert isinstance(d, MiniMaxM3SparseMSADecodeMetadata) + msa_metadata = d.msa_cutlass if self.use_cutlass_decode and msa_metadata is not None: - assert q_fp8 is not None msa_cutlass_sparse_decode( - q_fp8[:nd], + q[:nd], kv_cache, topk[:nd], out[:nd], @@ -308,20 +393,10 @@ def forward( q_scale_float=q_scale_float, k_scale_float=getattr(layer, "_k_scale_float", 1.0), v_scale_float=getattr(layer, "_v_scale_float", 1.0), - k_scale=k_scale if self.use_nvfp4_kv else None, - v_scale=v_scale if self.use_nvfp4_kv else None, - ) - elif self.use_nvfp4_kv: - raise RuntimeError( - "MiniMax M3 NVFP4 KV cache has no Triton decode fallback; " - f"CUTLASS decode was not planned for {main_md.num_decodes} " - f"requests with query length {d.decode_query_len}" ) else: minimax_m3_sparse_attn_decode( - q[:nd] - if q_fp8 is None - else _dequantize_query(q_fp8[:nd], q_scale_float, q[:nd]), + self._bf16_query(q[:nd], q_scale_float, output.dtype), kv_cache, topk[:nd].transpose(0, 1), d.block_table, @@ -336,61 +411,18 @@ def forward( # Prefill [nd:]: MSA sparse FMHA over the selected blocks. if main_md.num_prefills > 0: - from vllm.third_party.fmha_sm100.sparse import ( - build_k2q_csr, - sparse_atten_func, - sparse_atten_nvfp4_kv_func, - ) + from vllm.third_party.fmha_sm100.sparse import sparse_atten_func p = main_md.prefill assert p is not None - # [H, prefill, MK] transposed view; build_k2q_csr consumes the - # strided view directly (topK stays innermost-contiguous). - prefill_topk = topk[nd:num_tokens].transpose(0, 1) + k2q_row_ptr, k2q_q_indices, schedule = self._build_prefill_csr( + topk, nd, num_tokens, p + ) # An E4M3 query attends as is, its scale folded into softmax_scale. - if q_fp8 is None: - qp, softmax_scale = q[nd:], self.scale - else: - qp, softmax_scale = q_fp8[nd:], self.scale * q_scale_float - k2q_row_ptr, k2q_q_indices, schedule = build_k2q_csr( - prefill_topk, - p.cu_seqlens_q, - p.cu_seqlens_k, - SPARSE_BLOCK_SIZE, - total_k=0, - max_seqlen_k=p.max_seq_len, - max_seqlen_q=p.max_query_len, - total_rows=p.total_kv_blocks, - qhead_per_kv=qp.shape[1] // self.num_kv_heads, - return_schedule=True, + qp = q[nd:] + softmax_scale = ( + self.scale * q_scale_float if self.use_cutlass_decode else self.scale ) - if self.use_nvfp4_kv: - k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache) - sparse_atten_nvfp4_kv_func( - qp, - k_data, - v_data, - k_sf, - v_sf, - k_scale, - v_scale, - k2q_row_ptr, - k2q_q_indices, - self.topk_blocks, - cu_seqlens_q=p.cu_seqlens_q, - cu_seqlens_k=p.cu_seqlens_k, - max_seqlen_q=p.max_query_len, - max_seqlen_k=p.max_seq_len, - blk_kv=SPARSE_BLOCK_SIZE, - causal=True, - softmax_scale=softmax_scale, - page_table=p.block_table, - seqused_k=p.seq_lens, - schedule=schedule, - out=out[nd:], - kv_layout="vllm", - ) - return output k_cache, v_cache = kv_cache.split(self.head_size, dim=-1) sparse_atten_func( qp, @@ -412,3 +444,158 @@ def forward( out=out[nd:], ) return output + + +class MiniMaxM3SparseMSANvfp4Impl(MiniMaxM3SparseMSABaseImpl): + """MSA block-sparse attention over an NVFP4 KV cache, SM100 family only. + + Decode and prefill run the CUTLASS Q8KV4 kernels, which read the E4M3 + query; there is no BF16-query kernel and no Triton fallback. + """ + + use_cutlass_decode = True + prefill_backend = "q8kv4" + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int | None = None, + kv_cache_dtype: str = "nvfp4", + *, + topk_blocks: int, + sparse_block_size: int, + msa_decode_backend: MiniMaxM3MSADecodeBackend = "cutlass", + ) -> None: + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + kv_cache_dtype, + topk_blocks=topk_blocks, + sparse_block_size=sparse_block_size, + ) + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ): + raise ValueError("MiniMax M3 NVFP4 KV cache requires an SM100-family GPU") + # NVFP4 decodes only on CUTLASS, whatever msa_decode_backend says. + if not supports_cutlass_sparse_decode( + decode_backend="cutlass", + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + kv_cache_dtype=self.kv_cache_dtype, + page_size=self.block_size, + topk_blocks=self.topk_blocks, + ): + raise ValueError( + "MiniMax M3 NVFP4 KV cache requires CUTLASS MSA sparse decode " + f"(num_heads={self.num_heads}, num_kv_heads={self.num_kv_heads}, " + f"block_size={self.block_size}, topk_blocks={self.topk_blocks})" + ) + blocker = _q8kv4_prefill_blocker( + self.num_heads, self.num_kv_heads, self.block_size, self.topk_blocks + ) + if blocker is not None: + raise ValueError( + "MiniMax M3 NVFP4 KV cache requires the CUTLASS Q8KV4 prefill, " + f"which cannot serve this layout or device: {blocker}" + ) + logger.info_once("MiniMax M3 NVFP4 MSA: CUTLASS Q8KV4 decode and prefill") + # NVFP4_BLOCK_SCALE_SHIFT assumes unit global scales; checked on the + # first forward, once the checkpoint's KV scales are loaded. + self._check_global_scales = True + + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + kv_cache: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + assert query.dtype == torch.float8_e4m3fn + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return output # profiling run; caches unbound + main_md = attn_metadata[layer.layer_name] # type: ignore[attr-defined] + assert isinstance(main_md, MiniMaxM3SparseMetadata) + + nd = main_md.num_decode_tokens + num_tokens = main_md.num_actual_tokens + topk = layer.topk_indices_buffer # type: ignore[attr-defined] + assert topk is not None + hd = self.head_size + q_fp8 = query[:num_tokens].view(-1, self.num_heads, hd) + out = output[:num_tokens].view(-1, self.num_heads, hd) + q_scale_float = getattr(layer, "_q_scale_float", 1.0) + k_scale = getattr(layer, "_k_scale", None) + v_scale = getattr(layer, "_v_scale", None) + if self._check_global_scales: + k_global = getattr(layer, "_k_scale_float", 1.0) + v_global = getattr(layer, "_v_scale_float", 1.0) + if k_global != 1.0 or v_global != 1.0: + raise ValueError( + "MiniMax M3 reads its NVFP4 KV cache on the Q8KV4 kernels with " + f"block_scale_shift={NVFP4_BLOCK_SCALE_SHIFT}, which needs unit " + f"global scales; got k_scale={k_global}, v_scale={v_global}" + ) + self._check_global_scales = False + + if main_md.num_decodes > 0: + d = main_md.decode + assert isinstance(d, MiniMaxM3SparseMSADecodeMetadata) + msa_metadata = d.msa_cutlass + if msa_metadata is None: + raise RuntimeError( + "MiniMax M3 NVFP4 KV cache has no Triton decode fallback; " + f"CUTLASS decode was not planned for {main_md.num_decodes} " + f"requests with query length {d.decode_query_len}" + ) + msa_cutlass_sparse_decode( + q_fp8[:nd], + kv_cache, + topk[:nd], + out[:nd], + msa_metadata, + scale=self.scale, + q_scale_float=q_scale_float, + k_scale_float=getattr(layer, "_k_scale_float", 1.0), + v_scale_float=getattr(layer, "_v_scale_float", 1.0), + k_scale=k_scale, + v_scale=v_scale, + ) + + if main_md.num_prefills > 0: + from vllm.third_party.fmha_sm100.prefill_q8kv4 import run_prefill + + p = main_md.prefill + assert p is not None + k2q_row_ptr, _, schedule = self._build_prefill_csr(topk, nd, num_tokens, p) + k_data, k_sf, v_data, v_sf = nvfp4_kv_cache_views(kv_cache) + run_prefill( + q_fp8[nd:], + (k_data, v_data), + (k_sf, v_sf), + kv_indices=p.block_table, + kv_indptr=None, + cu_seqlens_q=p.cu_seqlens_q, + cu_seqlens_k=p.cu_seqlens_k, + k2q_row_ptr=k2q_row_ptr, + schedule=schedule, + sm_scale=self.scale * q_scale_float, + topk=self.topk_blocks, + seqused_k=p.seq_lens, + k_global_scale=k_scale, + v_global_scale=v_scale, + block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT, + out=out[nd:], + workspace=( + p.workspace + if isinstance(p, MiniMaxM3SparseMSAPrefillMetadata) + else None + ), + ) + return output diff --git a/vllm/reasoning/minimax_m3_reasoning_parser.py b/vllm/reasoning/minimax_m3_reasoning_parser.py index d49f9311f460..b6a8e48e49ef 100644 --- a/vllm/reasoning/minimax_m3_reasoning_parser.py +++ b/vllm/reasoning/minimax_m3_reasoning_parser.py @@ -39,6 +39,9 @@ def __init__(self, tokenizer, *args, **kwargs): super().__init__(tokenizer, *args, **kwargs) self._start_token_ids = self._encode_marker(self.start_token) self._end_token_ids = self._encode_marker(self.end_token) + # The chat template opens every message, including the generation + # prompt, with this token. + self._message_begin_token_id = self.vocab.get("]~b]") chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} self._initial_in_reasoning = chat_kwargs.get("thinking_mode") == "enabled" self._reasoning_ended_streaming = False @@ -117,6 +120,16 @@ def _ends_with_token_sequence_prefix( return True return False + def _last_message_ids(self, input_ids: Sequence[int]) -> Sequence[int]: + # The system message's thinking instructions and earlier turns contain + # reasoning markers; only the last message (the generation prompt) + # tells whether generation starts inside reasoning. Output-only ids + # contain no message-begin token and are returned whole. + for i in range(len(input_ids) - 1, -1, -1): + if input_ids[i] == self._message_begin_token_id: + return input_ids[i + 1 :] + return input_ids + @staticmethod def _strip_partial_marker_suffix(text: str, marker: str) -> str: max_len = min(len(text), len(marker) - 1) @@ -204,6 +217,7 @@ def is_reasoning_end_streaming( return False delta_ids = tuple(delta_ids) + input_ids = self._last_message_ids(input_ids) if self._contains_token_sequence(delta_ids, self._end_token_ids): return True if self._contains_token_sequence(input_ids, self._end_token_ids): @@ -311,6 +325,7 @@ def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: return count def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: + input_ids = self._last_message_ids(input_ids) start_index = self._rfind_token_sequence(input_ids, self._start_token_ids) end_index = self._rfind_token_sequence(input_ids, self._end_token_ids) if end_index < 0: diff --git a/vllm/transformers_utils/processors/minimax_m3.py b/vllm/transformers_utils/processors/minimax_m3.py index 9d922f8e4c64..8d4212cee395 100644 --- a/vllm/transformers_utils/processors/minimax_m3.py +++ b/vllm/transformers_utils/processors/minimax_m3.py @@ -1,16 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MiniMax M3 VL HuggingFace-compatible Processor / ImageProcessor / +"""MiniMax M3.1 VL HuggingFace-compatible Processor / ImageProcessor / VideoProcessor, vendored into vLLM so the model loads without ``--trust-remote-code`` (the released checkpoint only ships these classes as remote code via ``auto_map``). -Adapted verbatim from the ``MiniMaxAI/Minimax-M3-preview`` repository files -``image_processor.py``, ``video_processor.py`` and ``processing_minimax.py`` -(revision ``db01c0fe``). Both image and video processors use Qwen-style -``smart_resize`` (bound by total pixels). The original async frame-sampling -helpers are intentionally omitted: vLLM performs its own frame loading and -feeds decoded frames to the processor. +Both image and video processors use the MiniMax M3.1 aligned-size resize +(``compute_aligned_size``): tiny inputs are upscaled to a minimum area / +short side, then both dims are ceil-aligned to the patch factor and capped +at a per-modality long side (selectable via the OpenAI-style ``detail`` +kwarg or an explicit ``max_long_side_pixel``). The original async +frame-sampling helpers are intentionally omitted: vLLM performs its own +frame loading and feeds decoded frames to the processor. """ import math @@ -35,125 +36,158 @@ from transformers.video_processing_utils import BaseVideoProcessor from transformers.video_utils import group_videos_by_shape, reorder_videos -# Maximum allowed aspect ratio before smart_resize rejects the input. +# Maximum allowed aspect ratio before compute_aligned_size rejects the input. MAX_RATIO = 200 -# Fixed (non-configurable) bounds for the long-side resize logic, per the -# MiniMax-M3 size spec. ``min_short_side_pixel`` is the floor the short edge is -# enlarged to; ``*_MAX_TOTAL_PIXELS`` is the hard area cap that, once exceeded, -# aborts processing instead of downscaling. -MIN_SHORT_SIDE_PIXEL = 112 -IMAGE_MAX_TOTAL_PIXELS = 12_845_056 # 3584 ** 2 (width * height) -VIDEO_MAX_TOTAL_PIXELS = 301_056_000 # width * height * frames +# Fixed (non-configurable) bounds of the aligned-size resize, per the +# MiniMax M3.1 size spec. ``MIN_SHORT_SIDE_PIXEL`` is the floor the short +# edge is enlarged to; ``VIDEO_MAX_TOTAL_PIXELS`` is the hard volumetric cap +# (width * height * frames) that, once exceeded, aborts video processing. +MIN_SHORT_SIDE_PIXEL = 168 +MIN_PIXELS = MIN_SHORT_SIDE_PIXEL**2 +VIDEO_MAX_TOTAL_PIXELS = 301_056_000 +# OpenAI-style ``detail`` -> long-side pixel cap, per modality. +IMAGE_DETAIL_MAX_LONG_SIDE_PIXEL = {"low": 672, "auto": 2016, "high": 3584} +VIDEO_DETAIL_MAX_LONG_SIDE_PIXEL = {"low": 504, "auto": 672, "high": 2188} -def round_by_factor(number: int | float, factor: int) -> int: - return round(number / factor) * factor - -def ceil_by_factor(number: int | float, factor: int) -> int: - return math.ceil(number / factor) * factor - - -def floor_by_factor(number: int | float, factor: int) -> int: - return math.floor(number / factor) * factor - - -def _smart_resize_by_long_side( +def compute_aligned_size( height: int, width: int, factor: int, + *, + min_pixels: int = MIN_PIXELS, max_long_side_pixel: int, - min_short_side_pixel: int, - max_total_pixels: int | None, + min_short_side: int = MIN_SHORT_SIDE_PIXEL, ) -> tuple[int, int]: - """Long-side based resize (MiniMax-M3 size spec). - - (a) if the long side exceeds ``max_long_side_pixel`` → shrink so the long - side equals ``max_long_side_pixel``; - (b) else if the short side is below ``min_short_side_pixel`` → enlarge so the - short side equals ``min_short_side_pixel``; - (c) if the resulting area still exceeds ``max_total_pixels`` → raise. - - (a) and (b) are mutually exclusive (they branch on the *original* long side). - Both sides are then rounded to a multiple of ``factor``. For videos the - ``max_total_pixels`` cap is volumetric (width * height * frames) and is - enforced by the caller, so pass ``max_total_pixels=None`` here. - """ - long_side = max(height, width) - short_side = min(height, width) - - scaled_height: float = height - scaled_width: float = width - if long_side > max_long_side_pixel: - beta = max_long_side_pixel / long_side - scaled_height = height * beta - scaled_width = width * beta - elif short_side < min_short_side_pixel: - beta = min_short_side_pixel / short_side - scaled_height = height * beta - scaled_width = width * beta - - h_bar = max(factor, round_by_factor(scaled_height, factor)) - w_bar = max(factor, round_by_factor(scaled_width, factor)) - - if max_total_pixels is not None and h_bar * w_bar > max_total_pixels: - raise ValueError( - f"image area {h_bar * w_bar} exceeds max_total_pixels " - f"{max_total_pixels} after resizing" - ) - return h_bar, w_bar + """Compute the resize target ``(resized_height, resized_width)``. + Three layers (mirrors sglang's M3.1 ``compute_aligned_size``): -def smart_resize( - height: int, - width: int, - factor: int = 28, - min_pixels: int = 4 * 28 * 28, - max_pixels: int = 451584, - max_long_side_pixel: int | None = None, - min_short_side_pixel: int = MIN_SHORT_SIDE_PIXEL, - max_total_pixels: int | None = None, -) -> tuple[int, int]: - """Rescale (height, width) so each side is a multiple of ``factor``. + 1. Aspect-ratio guard: reject inputs whose ratio exceeds ``MAX_RATIO``. + 2. Floor layers: upscale proportionally when the area is below + ``min_pixels``, or the short side is below ``min_short_side``. + 3. Alignment layer: ``_align_to_factor`` ceil-aligns both dims to + ``factor`` and caps them at ``max_long_side_pixel``. + + Args: + height: Original height in pixels, must be > 0. + width: Original width in pixels, must be > 0. + factor: Alignment factor, usually ``patch_size * merge_size``. + min_pixels: Area floor; smaller inputs are upscaled to at least + this area before alignment. + max_long_side_pixel: Per-dim cap applied by the alignment layer + (used for both dims). + min_short_side: Short-side pixel floor. + + Returns: + ``(resized_height, resized_width)``, both multiples of ``factor`` + and at most ``max_long_side_pixel``. + + Raises: + ValueError: If the aspect ratio exceeds ``MAX_RATIO``. - When ``max_long_side_pixel`` is set, use the MiniMax-M3 long-side resize - spec (see :func:`_smart_resize_by_long_side`). Otherwise fall back to the - Qwen-VL area bound, keeping the total area within ``[min_pixels, max_pixels]``. """ if max(height, width) / min(height, width) > MAX_RATIO: raise ValueError( f"absolute aspect ratio must be smaller than {MAX_RATIO}, " f"got {max(height, width) / min(height, width)}" ) - if max_long_side_pixel is not None: - return _smart_resize_by_long_side( - height, - width, - factor=factor, - max_long_side_pixel=max_long_side_pixel, - min_short_side_pixel=min_short_side_pixel, - max_total_pixels=max_total_pixels, - ) - h_bar = max(factor, round_by_factor(height, factor)) - w_bar = max(factor, round_by_factor(width, factor)) - if h_bar * w_bar > max_pixels: - beta = math.sqrt((height * width) / max_pixels) - h_bar = floor_by_factor(height / beta, factor) - w_bar = floor_by_factor(width / beta, factor) - elif h_bar * w_bar < min_pixels: + if height * width < min_pixels: beta = math.sqrt(min_pixels / (height * width)) - h_bar = ceil_by_factor(height * beta, factor) - w_bar = ceil_by_factor(width * beta, factor) - return h_bar, w_bar + height = math.ceil(height * beta) + width = math.ceil(width * beta) + if min(height, width) < min_short_side: + beta = min_short_side / min(height, width) + height = math.ceil(height * beta) + width = math.ceil(width * beta) + return _align_to_factor( + height, width, factor, (max_long_side_pixel, max_long_side_pixel) + ) + + +def _align_to_factor( + height: int, + width: int, + factor: int, + max_size: tuple[int, int], +) -> tuple[int, int]: + """Ceil-align both dims to ``factor``, then cap at per-dim ``max_size``. + + When either aligned dim exceeds its cap, scale both dims back down + proportionally (integer arithmetic, preserving the aspect ratio) and + ceil-align again. + + Returns: + ``(height, width)`` with both dims multiples of ``factor`` and at + most the corresponding ``max_size`` dim. + + """ + max_w, max_h = max_size + # Round the caps down to factor multiples so alignment never exceeds them. + max_w = max_w - (max_w % factor) + max_h = max_h - (max_h % factor) + new_w = width if width % factor == 0 else width + (factor - width % factor) + new_h = height if height % factor == 0 else height + (factor - height % factor) + if new_w > max_w or new_h > max_h: + new_w_ = min((new_w * max_w) // new_w, (new_w * max_h) // new_h) + new_h_ = min((new_h * max_w) // new_w, (new_h * max_h) // new_h) + new_w = new_w_ + new_h = new_h_ + new_w = new_w if new_w % factor == 0 else new_w + (factor - new_w % factor) + new_h = new_h if new_h % factor == 0 else new_h + (factor - new_h % factor) + assert new_w <= max_w and new_h <= max_h, ( + f"_align_to_factor: result (h={new_h}, w={new_w}) exceeds max_size " + f"{max_size}; input (h={height}, w={width}, factor={factor})" + ) + return new_h, new_w + + +def resolve_max_long_side_pixel( + detail: "str | None", + max_long_side_pixel: "int | None", + *, + modality: str, +) -> int: + """Resolve the long-side pixel cap from ``detail`` / explicit override. + + ``detail`` (default ``"auto"``) maps to a per-modality cap. An explicit + ``max_long_side_pixel`` takes priority over a non-auto ``detail``; the + ``"auto"`` tier value itself is treated as "unset" so the processor class + defaults (2016 for images, 672 for videos) do not shadow ``detail``. + + Raises: + ValueError: If ``detail`` or ``modality`` is not a known value. + + """ + if modality == "image": + detail_map = IMAGE_DETAIL_MAX_LONG_SIDE_PIXEL + elif modality == "video": + detail_map = VIDEO_DETAIL_MAX_LONG_SIDE_PIXEL + else: + raise ValueError(f"unknown modality {modality!r}") + detail = detail or "auto" + if detail not in detail_map: + raise ValueError( + f"invalid detail {detail!r} for {modality}; " + f"expected one of {sorted(detail_map)}" + ) + if max_long_side_pixel is not None and max_long_side_pixel != detail_map["auto"]: + return max_long_side_pixel + return detail_map[detail] class MiniMaxM3VLImageProcessorKwargs(ImagesKwargs, total=False): # type: ignore[call-arg] patch_size: int temporal_patch_size: int merge_size: int + detail: "str | None" + max_long_side_pixel: "int | None" + # Legacy keys from released preprocessor_config.json files; accepted for + # backward compatibility but ignored by the aligned-size resize. + min_pixels: int max_pixels: int - max_long_side_pixel: int class MiniMaxM3VLImageProcessor(TorchvisionBackend): @@ -171,12 +205,11 @@ class MiniMaxM3VLImageProcessor(TorchvisionBackend): patch_size = 14 temporal_patch_size = 2 merge_size = 2 - max_pixels = 451584 # 672 * 672 - # Long-side resize spec (opt-in via ``max_long_side_pixel``). The latter two - # are fixed per the spec and are not exposed as configurable kwargs. - max_long_side_pixel = None + # Aligned-size resize bounds (M3.1). ``detail`` selects the long-side cap + # ("low"/"auto"/"high"); an explicit ``max_long_side_pixel`` wins. + detail = "auto" + max_long_side_pixel = 2016 min_short_side_pixel = MIN_SHORT_SIDE_PIXEL - max_total_pixels = IMAGE_MAX_TOTAL_PIXELS valid_kwargs = MiniMaxM3VLImageProcessorKwargs model_input_names = ["pixel_values", "image_grid_thw"] @@ -202,7 +235,7 @@ def _preprocess( patch_size: int, temporal_patch_size: int, merge_size: int, - max_pixels: int, + detail: str, max_long_side_pixel: "int | None", disable_grouping: "bool | None", return_tensors: "str | TensorType | None", @@ -213,17 +246,18 @@ def _preprocess( ) resized_images_grouped = {} factor = patch_size * merge_size + max_long_side = resolve_max_long_side_pixel( + detail, max_long_side_pixel, modality="image" + ) for shape, stacked_images in grouped_images.items(): height, width = stacked_images.shape[-2:] if do_resize: - resized_height, resized_width = smart_resize( + resized_height, resized_width = compute_aligned_size( height, width, factor=factor, - max_pixels=max_pixels, - max_long_side_pixel=max_long_side_pixel, - min_short_side_pixel=self.min_short_side_pixel, - max_total_pixels=self.max_total_pixels, + max_long_side_pixel=max_long_side, + min_short_side=self.min_short_side_pixel, ) stacked_images = self.resize( stacked_images, @@ -308,19 +342,19 @@ def get_number_of_image_patches(self, height: int, width: int, images_kwargs=Non images_kwargs = images_kwargs or {} patch_size = images_kwargs.get("patch_size", self.patch_size) merge_size = images_kwargs.get("merge_size", self.merge_size) - max_pixels = images_kwargs.get("max_pixels", self.max_pixels) + detail = images_kwargs.get("detail", self.detail) max_long_side_pixel = images_kwargs.get( "max_long_side_pixel", self.max_long_side_pixel ) - - resized_height, resized_width = smart_resize( + max_long_side = resolve_max_long_side_pixel( + detail, max_long_side_pixel, modality="image" + ) + resized_height, resized_width = compute_aligned_size( height, width, factor=patch_size * merge_size, - max_pixels=max_pixels, - max_long_side_pixel=max_long_side_pixel, - min_short_side_pixel=self.min_short_side_pixel, - max_total_pixels=self.max_total_pixels, + max_long_side_pixel=max_long_side, + min_short_side=self.min_short_side_pixel, ) grid_h, grid_w = resized_height // patch_size, resized_width // patch_size return grid_h * grid_w @@ -330,13 +364,16 @@ class MiniMaxM3VLVideoProcessorKwargs(VideosKwargs, total=False): # type: ignor patch_size: int temporal_patch_size: int merge_size: int - min_pixels: int - max_pixels: int - max_long_side_pixel: int - total_pixels: int + detail: "str | None" + max_long_side_pixel: "int | None" min_frames: int max_frames: int fps: "float | int" + # Legacy keys from released preprocessor_config.json files; accepted for + # backward compatibility but ignored by the aligned-size resize. + min_pixels: int + max_pixels: int + total_pixels: int class MiniMaxM3VLVideoProcessor(BaseVideoProcessor): @@ -354,13 +391,12 @@ class MiniMaxM3VLVideoProcessor(BaseVideoProcessor): patch_size = 14 temporal_patch_size = 2 merge_size = 2 - min_pixels = 4 * 28 * 28 - max_pixels = 768 * 28 * 28 # 602,112 - total_pixels = int(64000 * 28 * 28 * 0.9) # ~45M, ~64k tokens budget - # Long-side resize spec (opt-in via ``max_long_side_pixel``). The video + # Aligned-size resize bounds (M3.1). ``detail`` selects the long-side cap + # ("low"/"auto"/"high"); an explicit ``max_long_side_pixel`` wins. The # ``max_total_pixels`` cap is volumetric (width * height * frames) and is # enforced in ``_preprocess`` once the frame count is known. - max_long_side_pixel = None + detail = "auto" + max_long_side_pixel = 672 min_short_side_pixel = MIN_SHORT_SIDE_PIXEL max_total_pixels = VIDEO_MAX_TOTAL_PIXELS fps = 1.0 @@ -386,8 +422,7 @@ def _preprocess( patch_size: int, temporal_patch_size: int, merge_size: int, - min_pixels: int, - max_pixels: int, + detail: str = "auto", max_long_side_pixel: "int | None" = None, return_tensors: "str | TensorType | None" = None, **kwargs, @@ -395,27 +430,21 @@ def _preprocess( grouped_videos, grouped_videos_index = group_videos_by_shape(videos) resized_videos_grouped = {} factor = patch_size * merge_size + max_long_side = resolve_max_long_side_pixel( + detail, max_long_side_pixel, modality="video" + ) for shape, stacked_videos in grouped_videos.items(): batch_size, num_frames, channels, height, width = stacked_videos.shape resized_height, resized_width = height, width if do_resize: - resized_height, resized_width = smart_resize( + resized_height, resized_width = compute_aligned_size( height, width, factor=factor, - min_pixels=min_pixels, - max_pixels=max_pixels, - max_long_side_pixel=max_long_side_pixel, - min_short_side_pixel=self.min_short_side_pixel, - # Per-frame raise disabled; the video cap is volumetric and - # is enforced below once num_frames is known. - max_total_pixels=None, + max_long_side_pixel=max_long_side, + min_short_side=self.min_short_side_pixel, ) - if ( - max_long_side_pixel is not None - and resized_height * resized_width * num_frames - > self.max_total_pixels - ): + if resized_height * resized_width * num_frames > self.max_total_pixels: raise ValueError( f"video area {resized_height * resized_width * num_frames} " f"(width * height * frames) exceeds max_total_pixels " @@ -424,8 +453,10 @@ def _preprocess( stacked_videos = stacked_videos.view( batch_size * num_frames, channels, height, width ) + # Resize in float32, matching the training/SFT pipeline (and + # sglang): uint8 resize rounds intermediates and shifts pixels. stacked_videos = self.resize( - stacked_videos, + stacked_videos.float(), size=SizeDict(height=resized_height, width=resized_width), resample=resample, ) diff --git a/vllm/transformers_utils/repo_utils.py b/vllm/transformers_utils/repo_utils.py index 43ef07676e3a..8f4573780395 100644 --- a/vllm/transformers_utils/repo_utils.py +++ b/vllm/transformers_utils/repo_utils.py @@ -304,6 +304,30 @@ def get_model_path(model: str | Path, revision: str | None = None): ) +def get_subfolder_path( + model: str, + subfolder: str, + revision: str | None = None, + cache_dir: str | None = None, +) -> str: + """Return the local path of `subfolder` in a local or Hub checkpoint, + downloading only that subfolder from the Hub.""" + if os.path.isdir(model): + return os.path.join(model, subfolder) + # Avoid circular import + from vllm.model_executor.model_loader.weight_utils import get_lock + + with get_lock(model, cache_dir): + snapshot = hf_api().snapshot_download( + repo_id=model, + revision=revision, + cache_dir=cache_dir, + allow_patterns=f"{subfolder}/*", + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ) + return os.path.join(snapshot, subfolder) + + _WEIGHT_FILE_PATTERNS = [ "*.safetensors", "*.bin", diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 725e2f8039a6..88f7e99891b7 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -29,6 +29,13 @@ logger = init_logger(__name__) +def is_flashinfer_autotuning() -> bool: + """Whether a FlashInfer autotune(True) context is active.""" + from flashinfer.autotuner import AutoTuner + + return AutoTuner.get().is_tuning_mode + + def flashinfer_bf16_mm( a: torch.Tensor, b: torch.Tensor, @@ -1310,6 +1317,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: __all__ = [ "has_flashinfer", "flashinfer_bf16_mm", + "is_flashinfer_autotuning", "has_flashinfer_bf16_gemm", "has_flashinfer_concat_mla_kv_quant_fp8", "is_flashinfer_bf16_gemm_supported", diff --git a/vllm/v1/attention/backends/flashinfer_ts.py b/vllm/v1/attention/backends/flashinfer_ts.py new file mode 100644 index 000000000000..248eb7f65013 --- /dev/null +++ b/vllm/v1/attention/backends/flashinfer_ts.py @@ -0,0 +1,512 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer task-scheduled (``prims_ts``) paged decode. + +Wraps ``flashinfer.attention.prims_ts.BatchDecodePagedTSWrapper`` for uniform +query-block decode: every request contributes exactly ``max_query_len`` query +tokens over a paged KV cache. That is the DFlash/DSpark drafter shape, and the +reason this backend exists is that ``mask_type="dense"`` expresses a +*non-causal* draft block directly -- the trtllm-gen paged decode kernel is +bottom-right causal and silently ignores a supplied mask. + +Layout. No KV cache change: K and V are the two halves of vLLM's ordinary +``[pages, Hkv, page_size, 2 * D]`` page, so their token stride is ``2 * D``. +The kernel builds its TMA descriptors from the tensors' own strides +(``_validate_native_hnd_tma_strides``), so the only requirement is a +contiguous head_dim. Verified numerically equal to the same data repacked +into compact planes. + +Sliding window. ``prims_ts`` rejects ``window_left >= 0`` unless the mask is +causal, so a non-causal windowed layer cannot be expressed directly. The query +block always sits at the end of the sequence, so the symmetric window reduces +to a suffix, and we hand the kernel a block table restricted to the window's +pages (see :func:`_restrict_to_window`). Because the restriction is page +aligned, each query token sees up to ``page_size - 1`` extra tokens at the old +edge, and all query tokens share one window start rather than each having its +own (``max_query_len - 1`` tokens of drift). +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionImpl, + AttentionLayer, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, + CommonAttentionMetadata, + MultipleOf, + max_decode_query_len, +) +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec + +if TYPE_CHECKING: + from vllm.config.cache import CacheDType + from vllm.platforms.interface import DeviceCapability + +logger = init_logger(__name__) + +# flashinfer/attention/prims_ts/decode.py: _SUPPORTED_PAGE_SIZES, +# _SUPPORTED_HEAD_DIMS, _SUPPORTED_COMPUTE_CAPABILITIES. (10, 7) is SM107a / +# Rubin, which the kernel additionally gates on a CuTe DSL that can emit for +# the sm_100f family (>= 4.8, else CUTE_DSL_ARCH=sm_100f in the environment). +_TS_PAGE_SIZES = (16, 32, 64, 128) +_TS_HEAD_DIMS = (64, 128, 256) +_TS_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) +_TS_MAX_HEAD_RATIO = 128 + + +def _require_prims_ts(): + """Import the prims_ts wrapper, or explain which flashinfer is needed. + + 0.7.0.post1 ships prims_ts but calls ``Task._run_task_body_impl`` with a + ``context`` kwarg that cutlass-dsl 4.8.0 removed, so it raises TypeError + deep inside kernel codegen. + """ + try: + from flashinfer.attention.prims_ts import BatchDecodePagedTSWrapper + except ImportError as exc: + raise RuntimeError( + "FLASHINFER_TS needs flashinfer-python>=0.7.1 " + "(flashinfer.attention.prims_ts is missing)." + ) from exc + import flashinfer + + if tuple(int(p) for p in flashinfer.__version__.split(".")[:3]) < (0, 7, 1): + raise RuntimeError( + f"FLASHINFER_TS needs flashinfer-python>=0.7.1, found " + f"{flashinfer.__version__}: earlier builds of prims_ts are " + "incompatible with cutlass-dsl>=4.8." + ) + return BatchDecodePagedTSWrapper + + +def _ts_output_dtype(q_dtype: torch.dtype, model_dtype: torch.dtype) -> torch.dtype: + """The output dtype the kernel can write for this query dtype. + + An FP8 query only yields FP16 or FP8 out, so a bf16 model stages through + FP16 and converts; every other combination writes the model dtype directly. + """ + if q_dtype is torch.float8_e4m3fn and model_dtype is not torch.float16: + return torch.float16 + return model_dtype + + +def _kv_kernel_dtype(spec: AttentionSpec) -> torch.dtype: + """An FP8 cache is stored as uint8 and viewed as e4m3 at use.""" + if spec.dtype is torch.uint8: + return current_platform.fp8_dtype() + return spec.dtype + + +@dataclass +class FlashInferTSMetadata(AttentionMetadata): + # Per-request K/V lengths the kernel reads, already restricted to the + # window when one applies. + seq_lens: torch.Tensor + # [B, C] int32 page ids, likewise restricted. + block_table: torch.Tensor + num_reqs: int + num_actual_tokens: int + query_len: int + # Plan bound at build time; the same object must be replayed by a captured + # graph, which is why plans are keyed only on (batch size, query length). + wrapper: Any + + +class FlashInferTSMetadataBuilder(AttentionMetadataBuilder[FlashInferTSMetadata]): + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: KVCacheSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ): + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + assert isinstance(kv_cache_spec, AttentionSpec) + self.page_size = kv_cache_spec.block_size + self.num_kv_heads = kv_cache_spec.num_kv_heads + self.head_size = kv_cache_spec.head_size + # Q and K/V dtypes are kept matched: FP8 query against the FP8 cache, + # or model-dtype against an unquantized one. The FP8 query is produced + # upstream (MiniMax-M3's fused QK-norm/RoPE kernel writes it through + # q_fp8_out), so the attention layer does not quantize it again. + self.kv_dtype = _kv_kernel_dtype(kv_cache_spec) + self.q_dtype = self.kv_dtype + self.out_dtype = vllm_config.model_config.dtype + self.num_qo_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + # FullAttentionSpec records the window when the hybrid KV cache manager + # is disabled; SlidingWindowSpec carries it directly. + self.sliding_window: int | None = getattr(kv_cache_spec, "sliding_window", None) + self.query_len = max_decode_query_len(vllm_config) + self._plans: dict[tuple[int, int, bool], Any] = {} + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) + + # CUDA-graph-safe metadata. This builder advertises UNIFORM_BATCH, so + # the drafter's attention is captured and replayed; the kernel records + # the ADDRESS of seq_lens and the page table at capture time. Building + # fresh tensors per step would leave a replayed graph reading freed + # memory, so every build writes into these buffers in place. + max_reqs = max( + vllm_config.scheduler_config.max_num_seqs, + vllm_config.scheduler_config.max_num_batched_tokens + // max(self.query_len, 1), + ) + self._buf_seq_lens = torch.zeros(max_reqs, dtype=torch.int32, device=device) + self._buf_block_table = torch.zeros( + (max_reqs, self._window_pages(self.query_len)), + dtype=torch.int32, + device=device, + ) + + def _window_pages(self, query_len: int) -> int: + """Page-table width a restricted (windowed) batch needs.""" + if self.sliding_window is None: + return cdiv(self.vllm_config.model_config.max_model_len, self.page_size) + return cdiv(self.sliding_window + query_len - 1, self.page_size) + 1 + + def _max_kv_len(self) -> int: + """Constant K/V capacity every plan compiles for. + + Keeping it constant is what makes one plan per (batch size, query + length) safe to replay from a CUDA graph: ``build()`` can never pick a + differently-compiled plan for a batch the graph already captured. + """ + if self.sliding_window is not None: + # _restrict_to_window leaves at most one partial page in front of + # the window plus the query block itself. + return self.sliding_window + self.query_len + self.page_size + return self.vllm_config.model_config.max_model_len + + def _get_plan(self, num_reqs: int, query_len: int, causal: bool) -> Any: + key = (num_reqs, query_len, causal) + plan = self._plans.get(key) + if plan is not None: + return plan + + BatchDecodePagedTSWrapper = _require_prims_ts() + + window = self.sliding_window + wrapper = BatchDecodePagedTSWrapper(kv_layout="HND") + wrapper.plan( + device=self.device, + batch_size=num_reqs, + num_qo_heads=self.num_qo_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_size, + page_size=self.page_size, + max_kv_len=self._max_kv_len(), + max_seq_len_q=query_len, + packed_query=False, + q_data_type=self.q_dtype, + kv_data_type=self.kv_dtype, + o_data_type=_ts_output_dtype(self.q_dtype, self.out_dtype), + # A dense plan is what makes a non-causal draft block expressible; + # it cannot carry a window, which _restrict_to_window applies to + # the page table instead. + mask_type="causal" if causal else "dense", + window_left=window - 1 if (causal and window is not None) else -1, + ) + self._plans[key] = wrapper + return wrapper + + def _restrict_to_window( + self, seq_lens: torch.Tensor, block_table: torch.Tensor, query_len: int + ) -> tuple[torch.Tensor, torch.Tensor]: + """Cut the page table down to the pages the window can reach. + + The query block ends the sequence, so a symmetric window around it is + the suffix ``[L - w - query_len + 1, L)``. Round that start down to a + page so the slice is expressible as a page table, then renumber the + lengths relative to the new first page. + """ + assert self.sliding_window is not None + span = self.sliding_window + query_len - 1 + num_pages = self._buf_block_table.shape[1] + num_reqs = seq_lens.shape[0] + + start_page = torch.clamp_min( + (seq_lens - span).div(self.page_size, rounding_mode="floor"), 0 + ) + cols = torch.arange(num_pages, device=block_table.device, dtype=torch.int32) + idx = start_page.unsqueeze(1).to(torch.int64) + cols.unsqueeze(0) + idx = idx.clamp_max(block_table.shape[1] - 1) + + assert num_reqs <= self._buf_seq_lens.shape[0], ( + f"FLASHINFER_TS metadata buffer holds {self._buf_seq_lens.shape[0]} " + f"requests, got {num_reqs}" + ) + out_lens = self._buf_seq_lens[:num_reqs] + out_table = self._buf_block_table[:num_reqs] + out_lens.copy_(seq_lens - start_page.to(seq_lens.dtype) * self.page_size) + out_table.copy_(block_table.gather(1, idx)) + return out_lens, out_table + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> FlashInferTSMetadata: + num_reqs = common_attn_metadata.num_reqs + num_tokens = common_attn_metadata.num_actual_tokens + causal = common_attn_metadata.causal + if isinstance(causal, torch.Tensor): + raise NotImplementedError( + "FLASHINFER_TS has no per-sequence causality; it plans one " + "mask for the whole batch." + ) + # The drafter contributes the same number of query tokens per request, + # so the block length is the quotient; max_query_len is only an upper + # bound. Startup profiling also builds ragged batches that no real + # draft step produces (e.g. 5 tokens over 2 requests). Rounding down + # covers a whole number of blocks and leaves the remainder untouched, + # which keeps those builds in bounds instead of failing the engine. + query_len = num_tokens // num_reqs if num_reqs else 0 + num_tokens = num_reqs * query_len + # The model runner builds metadata for every KV cache group on every + # batch, so this is also called with the target's prefill shapes for + # the drafter's own group -- metadata the speculator never uses, since + # it builds its own. A drafter block is never longer than + # max_decode_query_len, so anything wider is not a batch this backend + # runs; planning it would compile a prefill-shaped plan and fail + # validation (seq_len_q > max_kv_len). + if query_len > self.query_len: + query_len = 0 + + seq_lens = common_attn_metadata.seq_lens[:num_reqs] + block_table = common_attn_metadata.block_table_tensor[:num_reqs] + if self.sliding_window is not None and not causal: + seq_lens, block_table = self._restrict_to_window( + seq_lens, block_table, query_len + ) + else: + num_pages = min(block_table.shape[1], self._buf_block_table.shape[1]) + out_lens = self._buf_seq_lens[:num_reqs] + out_table = self._buf_block_table[:num_reqs, :num_pages] + out_lens.copy_(seq_lens) + out_table.copy_(block_table[:, :num_pages]) + seq_lens, block_table = out_lens, out_table + + if query_len == 0: + return FlashInferTSMetadata( + seq_lens=common_attn_metadata.seq_lens[:num_reqs], + block_table=common_attn_metadata.block_table_tensor[:num_reqs], + num_reqs=num_reqs, + num_actual_tokens=0, + query_len=0, + wrapper=None, + ) + + return FlashInferTSMetadata( + seq_lens=seq_lens, + block_table=block_table, + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + query_len=query_len, + wrapper=self._get_plan(num_reqs, query_len, causal), + ) + + +class FlashInferTSImpl(AttentionImpl[FlashInferTSMetadata]): + forward_includes_kv_cache_update: bool = False + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int | None = None, + alibi_slopes: list[float] | None = None, + sliding_window: int | None = None, + kv_cache_dtype: str = "auto", + logits_soft_cap: float | None = None, + attn_type: str = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + sinks: torch.Tensor | None = None, + ) -> None: + if sinks is not None: + raise NotImplementedError("FLASHINFER_TS does not support attention sinks.") + if alibi_slopes is not None: + raise NotImplementedError("FLASHINFER_TS does not support alibi.") + if logits_soft_cap: + raise NotImplementedError("FLASHINFER_TS does not support softcap.") + if kv_sharing_target_layer_name is not None: + raise NotImplementedError("FLASHINFER_TS does not support KV sharing.") + if attn_type != AttentionType.DECODER: + raise NotImplementedError( + f"FLASHINFER_TS does not support attention type {attn_type}." + ) + self.num_heads = num_heads + self.head_size = head_size + self.scale = float(scale) + self.num_kv_heads = num_kv_heads or num_heads + self.sliding_window = sliding_window + self.kv_cache_dtype = kv_cache_dtype + # The kernel wants the query in the cache's dtype. For FP8 the caller + # supplies it already quantized (query_pre_quantized on the layer), so + # advertise support rather than making the layer quantize a second time. + self.supports_quant_query_input = kv_cache_dtype.startswith("fp8") + + def _kv_views(self, kv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """HND K/V views of vLLM's ``[pages, Hkv, page, 2 * D]`` page. + + Their token stride is ``2 * D`` rather than ``D``; the kernel builds its + TMA descriptors from the tensor's own strides, so no repacking and no + layout change are needed. + """ + if self.kv_cache_dtype.startswith("fp8"): + kv_cache = kv_cache.view(current_platform.fp8_dtype()) + return kv_cache[..., : self.head_size], kv_cache[..., self.head_size :] + + def do_kv_cache_update( + self, + layer: AttentionLayer, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + from vllm._custom_ops import reshape_and_cache_flash + + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + self.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) + + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: FlashInferTSMetadata, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + if attn_metadata is None: + # Profiling run: the KV cache is not bound yet. + return output + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError( + "FLASHINFER_TS does not support fused output quantization." + ) + + n = attn_metadata.num_actual_tokens + if n == 0: + return output + b, sq = attn_metadata.num_reqs, attn_metadata.query_len + k_cache, v_cache = self._kv_views(kv_cache) + q = query[:n].view(b, sq, self.num_heads, self.head_size) + + out = output[:n].view(b, sq, self.num_heads, self.head_size) + kernel_dtype = _ts_output_dtype(q.dtype, out.dtype) + needs_cast = kernel_dtype is not out.dtype + kernel_out = torch.empty_like(out, dtype=kernel_dtype) if needs_cast else out + + attn_metadata.wrapper.run( + q=q, + paged_kv_cache=(k_cache, v_cache), + seq_lens=attn_metadata.seq_lens, + block_tables=attn_metadata.block_table, + bmm1_scale=self.scale * layer._q_scale_float * layer._k_scale_float, + bmm2_scale=layer._v_scale_float, + out=kernel_out, + validate=False, + ) + if needs_cast: + out.copy_(kernel_out) + return output + + +class FlashInferTSBackend(AttentionBackend): + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = [ + "auto", + "float16", + "bfloat16", + "fp8", + "fp8_e4m3", + ] + forward_includes_kv_cache_update: bool = False + + @staticmethod + def get_name() -> str: + return "FLASHINFER_TS" + + @staticmethod + def get_impl_cls() -> type[FlashInferTSImpl]: + return FlashInferTSImpl + + @staticmethod + def get_builder_cls() -> type[FlashInferTSMetadataBuilder]: + return FlashInferTSMetadataBuilder + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return list(_TS_HEAD_DIMS) + + @staticmethod + def get_supported_kernel_block_sizes( + kv_cache_spec=None, + ) -> list[int | MultipleOf]: + return list(_TS_PAGE_SIZES) + + @classmethod + def supports_sliding_window(cls) -> bool: + # Causal layers use the kernel's own window_left; non-causal ones get a + # page-restricted block table (see the module docstring). + return True + + @classmethod + def supports_non_causal(cls) -> bool: + return True + + @classmethod + def supports_attn_type(cls, attn_type: str) -> bool: + return attn_type == AttentionType.DECODER + + @classmethod + def supports_compute_capability(cls, capability: "DeviceCapability") -> bool: + return (capability.major, capability.minor) in _TS_CAPABILITIES + + @classmethod + def supports_combination( + 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, + device_capability: "DeviceCapability", + ) -> str | None: + if use_mla or use_sparse or use_mm_prefix: + return "FLASHINFER_TS supports dense non-MLA attention only" + if has_sink: + return "FLASHINFER_TS does not support attention sinks" + return None diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 1e40298c7d69..eca068d40b34 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -69,6 +69,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): XPU_MLA_SPARSE = "vllm.v1.attention.backends.mla.xpu_mla_sparse.XPUMLASparseBackend" TORCH_SDPA = "" # this tag is only used for ViT FLASHINFER = "vllm.v1.attention.backends.flashinfer.FlashInferBackend" + FLASHINFER_TS = "vllm.v1.attention.backends.flashinfer_ts.FlashInferTSBackend" TRITON_FLASHINFER = ( "vllm.v1.attention.backends.triton_flashinfer.TritonFlashInferBackend" ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index fa5b3365561e..d12a8aa1e9d1 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -133,9 +133,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): @property def attn_vllm_config(self) -> VllmConfig: # The draft's attention differs from the target's in causality. - config = copy.copy(super().attn_vllm_config) + config = super().attn_vllm_config config.attention_config = replace( - self.vllm_config.attention_config, + config.attention_config, use_non_causal=self.requires_non_causal, ) return config diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 5d29b207b7fc..bcd6bfd608e4 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import replace @@ -11,6 +12,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.config.compilation import CUDAGraphMode +from vllm.config.speculative import _DRAFT_VLLM_CONFIG_OVERRIDES from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -317,7 +319,16 @@ def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None: def attn_vllm_config(self) -> VllmConfig: """Config for the draft's attention metadata builders. Overridden by speculators whose attention mode differs from the target's.""" - return self.vllm_config + # Shallow copies instead of apply_draft_overrides: its replace() + # rebuilds each sub-config and drops engine-resolved init=False state + # such as cache_config.kv_cache_layout. + config = copy.copy(self.vllm_config) + for src, config_name, dst in _DRAFT_VLLM_CONFIG_OVERRIDES: + if (value := getattr(self.speculative_config, src)) is not None: + sub_config = copy.copy(getattr(config, config_name)) + setattr(sub_config, dst, value) + setattr(config, config_name, sub_config) + return config def set_attn( self, From abaec97b4222957a1e9a8b6c1f672c84df82b6cf Mon Sep 17 00:00:00 2001 From: Alexander Matveev <59768536+alexm-redhat@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:12:58 -0400 Subject: [PATCH 116/406] [ROCm] Fuse tiny-M MXFP8 quant into the block32 GEMM; guard lm_head skinny-GEMM by weight size (#63179) Signed-off-by: Alexander Matveev --- .../quantization/test_rocm_mxfp8_linear.py | 220 ++++++++++++ .../layers/test_lmhead_skinny_gemm_guard.py | 306 ++++++++++++++++ vllm/envs.py | 18 + .../kernels/linear/mxfp8/rocm_block32_gemm.py | 338 +++++++++++++++++- .../kernels/linear/mxfp8/rocm_native.py | 27 +- vllm/model_executor/layers/utils.py | 99 +++++ 6 files changed, 986 insertions(+), 22 deletions(-) create mode 100644 tests/model_executor/layers/test_lmhead_skinny_gemm_guard.py diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py index 03cbc96428a1..b784551928a6 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -8,6 +8,13 @@ counts). These tests pin the two properties those workarounds stood in for -- finite output and MXFP8-level accuracy -- across the V4.1 attention shapes at TP=4 and the small token counts decode actually runs. + +They also cover the tiny-M fused quant+GEMM fastpath of +``rocm_mxfp8_block32_gemm_quant`` (at M <= 8 the activation quant runs inside +the packed GEMM's K loop): bit-exactness with the stock two-launch +``mxfp8_e4m3_quantize`` + ``rocm_mxfp8_block32_gemm`` path at the serving +shapes and under CUDA-graph capture/replay, agreement with a high-precision +QDQ fp32 reference, and bit-identical fallbacks for ineligible inputs. """ import pytest @@ -26,6 +33,17 @@ reason="native MX linear requires CDNA4 (gfx95x)", ) +from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( # noqa: E402 + Mxfp8LinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm import ( # noqa: E402 + rocm_mxfp8_block32_gemm, + rocm_mxfp8_block32_gemm_quant, +) +from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( # noqa: E402 + RocmDotScaledMxfp8LinearKernel, +) + # (name, N, K) for the DeepSeek V4.1 attention linears at tensor_parallel=4: # fused wq_a|wkv is replicated, wq_b/wo_a are column-parallel, wo_b is # row-parallel (so K, not N, is sharded). @@ -386,3 +404,205 @@ def test_rocm_mxfp8_quantizer_matches_torch_on_degenerate_blocks(num_rows): torch.testing.assert_close( q_triton.view(torch.uint8), q_torch.view(torch.uint8), rtol=0, atol=0 ) + + +# --- tiny-M fused quant+GEMM (rocm_mxfp8_block32_gemm_quant) ----------------- +# +# Serving block32 linears at TP=4 that hit the fused path at decode, +# covering the K-structure variety: the multi-band K=5120 shape with the +# in-launch split-K tier (shared_gate_up), a second K=5120 tier, the +# masked-K-tail K=1280 (EVEN_K=False), a small-K tile, and a draft-model +# shape absent from the tuned table (_default_config tile). +FUSED_QUANT_SERVING_SHAPES = [ + ("shared_gate_up_tp4", 1152, 5120), + ("fused_wqa_wkv", 1792, 5120), + ("wq_b_tp4", 8192, 1280), + ("shared_down_tp4", 5120, 576), + ("default_config", 1024, 2048), +] + +# Decode single, the DSpark k=5 hot M, and the M <= 8 tier bound. +FUSED_QUANT_MS = [1, 5, 8] + + +def _make_block32_weight(n: int, k: int, device: str, seed: int = 0): + """e4m3 weight [n, k] with compact [n//32, k//32] E8M0 block scales.""" + torch.manual_seed(seed) + blocks = (torch.randn(n, k, device=device) * 0.05).view(n // 32, 32, k // 32, 32) + amax = blocks.abs().amax(dim=(1, 3)).clamp(min=torch.finfo(torch.float32).tiny) + bits = (torch.ceil(torch.log2(amax / 448.0)) + 127.0).clamp(0, 254) + weight = (blocks * torch.exp2(127.0 - bits)[:, None, :, None]).view(n, k) + return ( + weight.to(torch.float8_e4m3fn).contiguous(), + bits.to(torch.uint8).contiguous(), + ) + + +def _stock(x, weight, weight_scale, out_dtype=torch.bfloat16): + """The pre-fusion two-launch path, verbatim.""" + x_q, x_scale = mxfp8_e4m3_quantize(x) + return rocm_mxfp8_block32_gemm(x_q, x_scale, weight, weight_scale, out_dtype) + + +def _fused_quant_hit(m, n, k) -> bool: + from vllm.model_executor.kernels.linear.mxfp8 import rocm_block32_gemm as g + + return (m, n, k) in g._FUSED_QUANT_HIT_SHAPES + + +@pytest.mark.parametrize("name,n,k", FUSED_QUANT_SERVING_SHAPES, ids=lambda v: str(v)) +@pytest.mark.parametrize("m", FUSED_QUANT_MS) +@torch.inference_mode() +def test_rocm_mxfp8_fused_quant_bit_exact_serving_shapes(name, n, k, m): + """The fused output equals the stock two-launch path byte-for-byte at the + serving shapes, and a high-precision QDQ fp32 reference within tolerance. + + The reference quantizes the activation with the same quantizer and then + dequantizes both operands and contracts in fp32, so it is independent of + the block32 GEMM kernel under test. + """ + torch.manual_seed(m) + weight, weight_scale = _make_block32_weight(n, k, "cuda") + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5 + + expected = _stock(x, weight, weight_scale) + out = rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, torch.bfloat16) + out2 = rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, torch.bfloat16) + + assert _fused_quant_hit(m, n, k), f"{name}: fused fastpath did not engage at M={m}" + assert torch.equal(out, expected), f"{name}: fused output differs from stock" + assert torch.equal(out, out2), f"{name}: fused output not deterministic" + + x_q, x_scale = mxfp8_e4m3_quantize(x) + row_scale = weight_scale.repeat_interleave(32, dim=0) + reference = ( + dequant_mxfp8_to_bf16(x_q, x_scale).float() + @ dequant_mxfp8_to_bf16(weight, row_scale).float().T + ) + rel = (out.float() - reference).norm() / reference.norm() + assert rel < 5e-3, f"{name}: fused output vs QDQ reference, rel err {rel:.4f}" + + +@torch.inference_mode() +def test_rocm_mxfp8_fused_quant_bit_exact_degenerate_blocks(): + """All-zero / denormal / huge / negative blocks quantize identically. + + An all-zero 32-block clamps the E8M0 scale to byte 0 (2**-127, + subnormal in fp32): the reciprocal-form rescale must keep it NaN-free, + exactly as the separate quantizer does. + """ + n, k, m = 512, 1024, 6 + weight, weight_scale = _make_block32_weight(n, k, "cuda", seed=7) + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5 + x[0] = 0.0 # all-zero row + x[1, :32] = 0.0 # single zero block + x[1, 32:64] = 1e-38 # denormal block + x[2] = 3.0e30 # huge: scale byte clamps at 254 + x[3] = -x[3] # negative mirror + x[4, ::32] = 0.0 # zero at every block start + + expected = _stock(x, weight, weight_scale) + out = rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, torch.bfloat16) + assert _fused_quant_hit(m, n, k) + assert torch.isfinite(out.float()).all(), "non-finite fused output" + assert torch.equal(out, expected), "degenerate blocks differ from stock" + + +@torch.inference_mode() +def test_rocm_mxfp8_fused_quant_fallbacks_match_stock(): + """Ineligible inputs take the internal stock path, bit-identical.""" + n, k = 1024, 512 + weight, weight_scale = _make_block32_weight(n, k, "cuda", seed=11) + + # Non-contiguous bf16 x (stride-2 view). + base = torch.randn(6, 2 * k, device="cuda", dtype=torch.bfloat16) * 0.5 + x_nc = base[:, ::2] + assert not x_nc.is_contiguous() + out = rocm_mxfp8_block32_gemm_quant(x_nc, weight, weight_scale, torch.bfloat16) + assert not _fused_quant_hit(6, n, k), ( + "non-contiguous x must not take the fused path" + ) + assert torch.equal(out, _stock(x_nc.contiguous(), weight, weight_scale)) + + # fp32 x falls back (quantizer loads any float dtype, fused does not). + x_f32 = torch.randn(6, k, device="cuda", dtype=torch.float32) * 0.5 + out = rocm_mxfp8_block32_gemm_quant(x_f32, weight, weight_scale, torch.float32) + assert not _fused_quant_hit(6, n, k), "fp32 x must not take the fused path" + assert torch.equal(out, _stock(x_f32, weight, weight_scale, torch.float32)) + + # M > 8 (prefill-class) falls back. + x_big = torch.randn(16, k, device="cuda", dtype=torch.bfloat16) * 0.5 + out = rocm_mxfp8_block32_gemm_quant(x_big, weight, weight_scale, torch.bfloat16) + assert not _fused_quant_hit(16, n, k), "M=16 must not take the fused path" + assert torch.equal(out, _stock(x_big, weight, weight_scale)) + + # M = 0 returns an empty output. + x0 = torch.randn(0, k, device="cuda", dtype=torch.bfloat16) + out = rocm_mxfp8_block32_gemm_quant(x0, weight, weight_scale, torch.bfloat16) + assert out.shape == (0, n) + + +@torch.inference_mode() +def test_rocm_mxfp8_fused_quant_per_row_scale_unaffected(): + """Regression guard: the per-row-scale branch of + RocDotScaledMxfp8LinearKernel.apply_weights (weight_scale [N, K//32], + NOT compacted to block32) must keep working when the block32 branch + dispatches to the fused op. The numerics of this branch are owned by + test_rocm_mxfp8_linear_matches_dequant_reference above; this test pins + the dispatch interaction only. + """ + n, k, m = 1152, 5120, 7 + torch.manual_seed(23) + weight_bf16 = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.05 + weight, weight_scale = _mxfp8_e4m3_quantize_torch(weight_bf16) + assert weight_scale.shape == (n, k // 32), "expected per-row scales" + + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(weight, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k))) + kernel.process_weights_after_loading(layer) + # Non-uniform per-row scales are not compacted to [N//32, K//32]. + assert layer.weight_scale.shape[0] == n + + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5 + out = kernel.apply_weights(layer, x) + assert out.shape == (m, n) + assert torch.isfinite(out.float()).all() + # This branch never reaches the fused op. + assert not _fused_quant_hit(m, n, k), ( + "per-row-scale layer must not hit the fused path" + ) + + +@pytest.mark.parametrize("warm_counters", [True, False]) +@torch.inference_mode() +def test_rocm_mxfp8_fused_quant_cuda_graph_capture(warm_counters): + """Capture+replay is bit-exact with and without warm split-K counters. + + shared_gate_up TP4 at M=6 uses the in-launch split-K config + (fused_splits=3): with warmup the counters exist and the splits meet in + one launch; captured first-use on a fresh stream has no counters and + must fall back to the two-launch split-K (partials + reduce). + """ + n, k, m = 1152, 5120, 6 + weight, weight_scale = _make_block32_weight(n, k, "cuda", seed=17) + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5 + expected = _stock(x, weight, weight_scale) + + stream = torch.cuda.Stream() + if warm_counters: + with torch.cuda.stream(stream): + warm = rocm_mxfp8_block32_gemm_quant( + x, weight, weight_scale, torch.bfloat16 + ) + stream.synchronize() + assert torch.equal(warm, expected) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = rocm_mxfp8_block32_gemm_quant(x, weight, weight_scale, torch.bfloat16) + for _ in range(3): + graph.replay() + torch.accelerator.synchronize() + assert torch.equal(out, expected), "replay differs from stock" diff --git a/tests/model_executor/layers/test_lmhead_skinny_gemm_guard.py b/tests/model_executor/layers/test_lmhead_skinny_gemm_guard.py new file mode 100644 index 000000000000..32f3ec0b679e --- /dev/null +++ b/tests/model_executor/layers/test_lmhead_skinny_gemm_guard.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""lm_head skinny-GEMM weight-size guard: core contracts. + +wvSplitK_hf_sml's cost at n activation rows scales as ~n passes over the +weight, so past the last-level (memory-side) cache its effective +bandwidth collapses. At the DSV4.1-Flash TP=4 serving point the draft +vocab-parallel lm_head shard ([32320, 5120] bf16, ~315 MiB, one call per +decode step) hit this collapse. The guard keeps every weight above +VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES on the generic GEMM fall-through at +0 < n <= 5; the default cap is the memory-side cache size, read once +from the hardware (256 MiB on MI355X, 128 MiB on MI300X-class parts), +with a 256 MiB fallback where that cannot be read. The swap is NOT +bit-exact (different GEMM kernel, different accumulation order) -- these +tests check the decision table, dispatch selection, cap tuning and +hardware-derivation, and numerical agreement within bf16 GEMM tolerance. +""" + +import os + +import pytest +import torch + +import vllm.envs as envs +from vllm.model_executor.layers.utils import ( + _amd_l3_cache_bytes, + _skinny_gemm_max_weight_bytes, + _skinny_guard_warned, + lmhead_skinny_guard_permits, + rocm_unquantized_gemm_impl, +) + +try: + from vllm import _custom_ops as ops + + WVSPLITK_AVAILABLE = hasattr(ops, "wvSplitK") +except ImportError: + WVSPLITK_AVAILABLE = False + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="GEMM kernels run on GPU"), + pytest.mark.skipif( + not WVSPLITK_AVAILABLE, + reason="requires the ROCm skinny-GEMM custom ops", + ), +] + +_MAX_BYTES_VAR = "VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES" + +# The serving-site shape: the DSpark draft vocab-parallel lm_head shard at +# TP=4 (draft hidden size 5120), above the 256 MiB default cap. +BIG_SHAPE = (32320, 5120) # 315.5 MiB bf16 +# A dense-projection-class weight: far below the cap. +SMALL_SHAPE = (2048, 7168) # 29 MiB bf16 + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + # These tests assert exact dispatch selection, so pin every env var + # that feeds it: state leaked by unrelated tests earlier in a long + # suite run must not change the outcome here. + monkeypatch.delenv(_MAX_BYTES_VAR, raising=False) + monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "1") + monkeypatch.delenv("VLLM_BATCH_INVARIANT", raising=False) + _skinny_guard_warned.clear() + envs.disable_envs_cache() + _skinny_gemm_max_weight_bytes.cache_clear() + yield + _skinny_guard_warned.clear() + envs.disable_envs_cache() + _skinny_gemm_max_weight_bytes.cache_clear() + + +def _weight(shape, seed=0): + torch.manual_seed(seed) + return torch.randn(*shape, dtype=torch.bfloat16, device="cuda") + + +def _x(n, k, seed=1): + torch.manual_seed(seed) + return torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + +def _set_cap(max_bytes): + # None -> the hardware-derived default (memory-side cache size, with + # the 256 MiB fallback); an int -> that exact cap. The accessor caches + # its one evaluation, so clear it after every change. + if max_bytes is None: + os.environ.pop(_MAX_BYTES_VAR, None) + else: + os.environ[_MAX_BYTES_VAR] = str(max_bytes) + envs.disable_envs_cache() + _skinny_gemm_max_weight_bytes.cache_clear() + + +# Event names that are not compute-kernel records: transfers, runtime API +# wrappers, framework overhead. Every fn profiled here is a GEMM, so a +# capture holding only these (in the observed failure: nothing at all) +# recorded no kernels. +_NOT_KERNEL = ("Memcpy", "Memset", "cuda", "hip", "aten::", "at::", "ProfilerStep") + + +def _launched_kernels(fn, iters=3, attempts=3): + """Kernel names recorded for fn(), retrying failed captures. + + Mid-suite, the ROCm activity tracer can drop every kernel record of a + capture while the neighboring captures in the same test stay healthy + (seen once in the full model_executor suite, next to a "ROCTracer + produced duplicate flow start" warning). A GEMM always launches + device kernels, so a capture without any kernel record is a failed + capture -- retry it rather than read it as dispatch evidence. + """ + from torch.profiler import ProfilerActivity, profile + + fn() + torch.accelerator.synchronize() + kernels = [] + for _ in range(attempts): + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + fn() + torch.accelerator.synchronize() + kernels = [e.key for e in prof.key_averages()] + if any(not k.startswith(_NOT_KERNEL) for k in kernels): + break + return kernels + + +def _has_wvsplitk_hf(names): + """Any wvSplitK_hf skinny-kernel variant ran. + + The csrc dispatch (WVSPLITK_CFG in csrc/rocm/skinny_gemms.cu) picks + between the wvSplitK_hf_sml_, wvSplitK_hf_ and wvSplitK_hf_big_ + variants by LDS fit (Kbp*N vs the part's shared-memory size) and + M-tile divisibility, so the same shape lands on different variants + on different parts (MI300/gfx942 runs wvSplitK_hf_ for (2048, 7168) + at n=5 where MI355X/gfx950 runs wvSplitK_hf_sml_). The guard's + contract is about the skinny path as a whole, so match the family. + """ + return any("wvSplitK_hf" in k for k in names) + + +def _has_cijk(names): + return any("Cijk" in k for k in names) + + +def test_guard_decision(): + """Decision table: bytes-gated at 0 < n <= 5; n=1 (no pathology) + always permits; the cap is tunable. The default (auto) cap gates the + 315.5 MiB site shape on every part it can run on: the memory-side + cache is 256 MiB (MI355X) / 128 MiB (MI300X-class), and the fallback + is 256 MiB.""" + small = torch.zeros(SMALL_SHAPE, dtype=torch.bfloat16) + big = torch.zeros(BIG_SHAPE, dtype=torch.bfloat16) + assert lmhead_skinny_guard_permits(small, 5) + for n in (2, 3, 4, 5): + assert not lmhead_skinny_guard_permits(big, n) + assert lmhead_skinny_guard_permits(big, 1), "n=1 carve-out missing" + + _set_cap(1024 * 1024 * 1024) + assert lmhead_skinny_guard_permits(big, 5), "raised cap must permit" + + _set_cap(128 * 1024 * 1024) + assert not lmhead_skinny_guard_permits(big, 5), "lowered cap must gate" + + _set_cap(None) + assert not lmhead_skinny_guard_permits(big, 5), "auto cap must gate" + + +# Canned amdsmi cache lists (amdsmi_get_gpu_cache_info(...)["cache"]), +# trimmed to the fields the parse keys on; sizes are KB as amdsmi +# reports them. +_AMD_SMI_MI355X = [ + # L1 data / L1 inst / L2: the levels below the boundary. + { + "cache_properties": ["DATA_CACHE", "SIMD_CACHE"], + "cache_level": 1, + "cache_size": 32, + }, + { + "cache_properties": ["INST_CACHE", "SIMD_CACHE"], + "cache_level": 1, + "cache_size": 64, + }, + { + "cache_properties": ["DATA_CACHE", "SIMD_CACHE"], + "cache_level": 2, + "cache_size": 4096, + }, + # The memory-side cache (MALL): 256 MiB, the level the guard keys on. + { + "cache_properties": ["DATA_CACHE", "SIMD_CACHE"], + "cache_level": 3, + "cache_size": 262144, + }, +] +_AMD_SMI_MI300X_CLASS = [ + {"cache_properties": ["DATA_CACHE"], "cache_level": 2, "cache_size": 8192}, + {"cache_properties": ["DATA_CACHE"], "cache_level": 3, "cache_size": 131072}, +] + + +def test_l3_parse_selects_level3_data_cache(): + """The parse takes the level-3 DATA_CACHE (the memory-side cache), + not the L2 below it, in amdsmi's KB units.""" + assert _amd_l3_cache_bytes(_AMD_SMI_MI355X) == 256 * 1024 * 1024 + assert _amd_l3_cache_bytes(_AMD_SMI_MI300X_CLASS) == 128 * 1024 * 1024 + # Empty list, or instruction caches only: no L3 to key on. + assert _amd_l3_cache_bytes([]) is None + assert ( + _amd_l3_cache_bytes( + [{"cache_properties": ["INST_CACHE"], "cache_level": 3, "cache_size": 64}] + ) + is None + ) + + +def test_cap_env_hardware_fallback_precedence(monkeypatch): + """Unset env: the cap is the L3 size the library reports; an + unreadable L3 falls back to the documented 256 MiB; an explicit env + value overrides the hardware.""" + _set_cap(None) + monkeypatch.setattr( + "vllm.model_executor.layers.utils._amd_l3_cache_bytes", + lambda cache_info=None: 128 * 1024 * 1024, + ) + _skinny_gemm_max_weight_bytes.cache_clear() + assert _skinny_gemm_max_weight_bytes() == 128 * 1024 * 1024 + + monkeypatch.setattr( + "vllm.model_executor.layers.utils._amd_l3_cache_bytes", + lambda cache_info=None: None, + ) + _skinny_gemm_max_weight_bytes.cache_clear() + assert _skinny_gemm_max_weight_bytes() == 256 * 1024 * 1024 + + _set_cap(1024 * 1024 * 1024) + assert _skinny_gemm_max_weight_bytes() == 1024 * 1024 * 1024 + + +def test_census_warning_fires_once_per_site(caplog): + w = torch.zeros(BIG_SHAPE, dtype=torch.bfloat16) + with caplog.at_level("WARNING", logger="vllm.model_executor.layers.utils"): + lmhead_skinny_guard_permits(w, 5) + lmhead_skinny_guard_permits(w, 5) # same (shape, n): once + lmhead_skinny_guard_permits(w, 4) # new n: second line + warns = [r for r in caplog.records if "skinny-GEMM guard engaged" in r.message] + assert len(warns) == 2 + assert "32320" in warns[0].message and "n=5" in warns[0].message + + +@pytest.mark.parametrize("cap_gib,expect_skinny", [(None, False), (1, True)]) +def test_n5_big_weight_dispatch(cap_gib, expect_skinny): + """The draft-lm_head shape (n=5, >256 MiB): under the default cap the + guard redirects to the generic hipBLASLt GEMM (the same kernel class + its n=6 verify twin dispatches to); a raised cap keeps the stock + wvSplitK path.""" + _set_cap(cap_gib * 1024 * 1024 * 1024 if cap_gib else None) + w = _weight(BIG_SHAPE) + x = _x(5, BIG_SHAPE[1]) + ref = torch.nn.functional.linear(x, w) + kernels = _launched_kernels(lambda: rocm_unquantized_gemm_impl(x, w, None)) + assert _has_wvsplitk_hf(kernels) == expect_skinny, f"kernels={kernels}" + if not expect_skinny: + assert _has_cijk(kernels), ( + f"fall-through did not reach the hipBLASLt GEMM; kernels={kernels}" + ) + out = rocm_unquantized_gemm_impl(x, w, None) + torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-1) + + +@pytest.mark.parametrize( + "shape,n,expect_skinny", + [ + (BIG_SHAPE, 6, False), # n=6 never matches the skinny branch + (SMALL_SHAPE, 5, True), # small weight: bytes-gated only + ], +) +def test_guard_no_op_where_it_must_be(shape, n, expect_skinny): + """The guard must not touch dispatch outside its window, under the + default cap AND a raised one: n=6 (above the skinny branch) and small + weights (below the byte cap).""" + w = _weight(shape) + x = _x(n, shape[1]) + ref = torch.nn.functional.linear(x, w) + for cap_gib in (None, 1): + _set_cap(cap_gib * 1024 * 1024 * 1024 if cap_gib else None) + _skinny_guard_warned.clear() + kernels = _launched_kernels(lambda: rocm_unquantized_gemm_impl(x, w, None)) + assert _has_wvsplitk_hf(kernels) == expect_skinny, ( + f"cap={cap_gib} GiB: skinny dispatch changed for " + f"shape={shape} n={n}; kernels={kernels}" + ) + out = rocm_unquantized_gemm_impl(x, w, None) + torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-1) + + +def test_n1_big_weight_dispatch_unchanged(): + """n=1 must keep the stock wvSplitK_hf skinny kernel under the + guard (default ON) -- the carve-out keeps n=1 bit-identical to stock.""" + w = _weight(BIG_SHAPE) + x = _x(1, BIG_SHAPE[1]) + kernels = _launched_kernels(lambda: rocm_unquantized_gemm_impl(x, w, None)) + assert _has_wvsplitk_hf(kernels), ( + f"n=1 was redirected by the guard; the carve-out is missing; kernels={kernels}" + ) diff --git a/vllm/envs.py b/vllm/envs.py index 145cc8b0e8e5..46e11e549234 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -309,6 +309,7 @@ VLLM_DEBUG_WORKSPACE: bool = False VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION: bool = False + VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES: int = 0 VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256 VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024 VLLM_COHERE2_ATTN_MOE_DUAL_STREAM_TOKEN_THRESHOLD: int = 256 @@ -2136,6 +2137,23 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION": lambda: bool( int(os.getenv("VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION", "0")) ), + # Size cap for the lm_head skinny-GEMM (wvSplitK) weight guard: weights + # larger than this skip the skinny path at 1 < n <= 5 (see + # vllm.model_executor.layers.utils.lmhead_skinny_guard_permits). + # wvSplitK's cost at n activation rows scales as ~n passes over the + # weight, so once the weight exceeds the last-level (memory-side) + # cache those extra passes are HBM round trips. 0 (the default) + # derives the cap from that cache size once at first use (256 MiB on + # MI355X, 128 MiB on MI300X-class parts), with a 256 MiB fallback + # where the size cannot be read — see the guard for how it's read. + # The portable device-properties API cannot be used: it exposes only + # the far smaller L2 (4 MiB on gfx950), which would disable the + # fastpath on every real lm_head shard. Set a value explicitly to + # override the hardware-derived cap; set it very large (e.g. 2**60) + # to disable the guard. + "VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES": lambda: int( + os.getenv("VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES", "0") + ), # Limits when we run shared_experts in a separate stream. # We found out that for large batch sizes, the separate stream # execution is not beneficial (most likely because of the input clone) diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py index 3d5e9c8357ca..9effa2b50bd2 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py @@ -50,9 +50,15 @@ import torch +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, +) from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import direct_register_custom_op +logger = init_logger(__name__) + BLOCK_ROWS = 32 # gfx950 hands out workgroups to its XCDs round-robin by program id, and 8 is @@ -272,6 +278,124 @@ def _block32_packed_kernel( ) +# _block32_packed_kernel with the MXFP8 activation quantizer folded into +# the K loop: x_ptr is the RAW bf16/fp16 [M, K] activation (row-major +# contiguous), quantized per 32-element block in registers with the exact +# _mxfp8_quant_triton_kernel arithmetic (amax -> biased E8M0 scale -> +# reciprocal rescale -> RN cast to e4m3), so the separate quant launch and +# the x_q/x_scale global round-trip disappear from the decode critical +# path. Tiling, K-panel packing, split-K reduction and stores are +# byte-for-byte the packed kernel, so the output is bit-identical to +# quantize-then-GEMM by construction (same quant arithmetic on the same +# values, same dot order). +@triton.jit(do_not_specialize=["M"]) +def _block32_packed_quant_kernel( + x_ptr, + w_ptr, + ws_ptr, + out_ptr, + slot_ptr, + count_ptr, + M, + N: tl.constexpr, + K: tl.constexpr, + stride_om, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + K_PACK: tl.constexpr, + K_PER_SPLIT: tl.constexpr, + EVEN_K: tl.constexpr, + FUSED_SPLITS: tl.constexpr, + FP8_MAX: tl.constexpr, + TINY: tl.constexpr, +): + if FUSED_SPLITS > 1: + pid_m, pid_n, tile, pid_k = _split_tile(N, BLOCK_N, FUSED_SPLITS) + if pid_m * BLOCK_M >= M: + return + k_start = pid_k * K_PER_SPLIT + k_stop = tl.minimum(k_start + K_PER_SPLIT, K) + else: + pid_m, pid_n, pid_k = tl.program_id(0), tl.program_id(1), tl.program_id(2) + if K_PER_SPLIT >= K: + k_start = 0 + k_stop = K + else: + k_start = pid_k * K_PER_SPLIT + k_stop = tl.minimum(k_start + K_PER_SPLIT, K) + # Tile row r is token r // K_PACK reading K panel r % K_PACK (and likewise + # for weight rows), so one dot covers K_PACK consecutive BLOCK_K slices. + rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M * K_PACK) // K_PACK + x_panel = tl.arange(0, BLOCK_M * K_PACK) % K_PACK + cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N * K_PACK) // K_PACK + w_panel = tl.arange(0, BLOCK_N * K_PACK) % K_PACK + ks = tl.arange(0, BLOCK_K) + gs = tl.arange(0, BLOCK_K // 32) + row_ok = rows < M + col_ok = cols < N + acc = tl.zeros((BLOCK_M * K_PACK, BLOCK_N * K_PACK), dtype=tl.float32) + for base in range(k_start, k_stop, BLOCK_K * K_PACK): + xk = base + x_panel[:, None] * BLOCK_K + ks[None, :] + wk = base + w_panel[:, None] * BLOCK_K + ks[None, :] + wg = base // 32 + w_panel[:, None] * (BLOCK_K // 32) + gs[None, :] + x_mask = row_ok[:, None] + w_mask = col_ok[:, None] + ws_mask = col_ok[:, None] + if not EVEN_K: + x_mask = x_mask & (xk < K) + w_mask = w_mask & (wk < K) + ws_mask = ws_mask & (wg < K // 32) + # Fused MXFP8 quant, verbatim _mxfp8_quant_triton_kernel arithmetic. + # K is a multiple of 32 and every 32-block of the tile is 32-aligned, + # so each block is either fully inside K or fully masked (loaded 0.0, + # scale byte 0) -- the same values the separate quantizer produces + # for these rows. Masked rows likewise load 0.0 and quantize to + # scale byte 0 / value 0, matching the kernel's other=0 loads. + x = tl.load(x_ptr + rows[:, None] * K + xk, mask=x_mask, other=0.0).to( + tl.float32 + ) + xb = tl.reshape(x, (BLOCK_M * K_PACK, BLOCK_K // 32, 32)) + amax = tl.maximum(tl.max(tl.abs(xb), axis=2), TINY) + sb = tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0 + sb = tl.minimum(tl.maximum(sb, 0.0), 254.0) + rescale = tl.exp2(127.0 - sb) + xq = tl.reshape(xb * rescale[:, :, None], (BLOCK_M * K_PACK, BLOCK_K)).to( + w_ptr.dtype.element_ty + ) + xs = sb.to(tl.uint8) + w = tl.load(w_ptr + cols[:, None] * K + wk, mask=w_mask, other=0.0) + ws = tl.load( + ws_ptr + (cols[:, None] // 32) * (K // 32) + wg, mask=ws_mask, other=0 + ) + acc = tl.dot_scaled(xq, xs, "e4m3", w.T, ws, "e4m3", acc=acc) + # Keep only products of matching K panels: the block diagonal. + blocks = acc.reshape(BLOCK_M, K_PACK, BLOCK_N, K_PACK).trans(0, 2, 1, 3) + pair = tl.arange(0, K_PACK) + diagonal = tl.where( + pair[None, None, :, None] == pair[None, None, None, :], blocks, 0.0 + ) + y = tl.sum(tl.sum(diagonal, 3), 2) + out_rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + out_cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + out_ptrs = out_ptr + out_rows[:, None] * stride_om + out_cols[None, :] + out_mask = (out_rows[:, None] < M) & (out_cols[None, :] < N) + if FUSED_SPLITS > 1: + _sum_splits( + y, + out_ptrs, + out_mask, + slot_ptr + tile * (FUSED_SPLITS * BLOCK_M * BLOCK_N), + count_ptr + tile, + pid_k, + FUSED_SPLITS, + ) + else: + tl.store( + out_ptrs + pid_k * M * N, y.to(out_ptr.dtype.element_ty), mask=out_mask + ) + + @triton.jit(do_not_specialize=["M"]) def _block32_splitk_reduce_kernel( partial_ptr, @@ -331,9 +455,9 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf ) -# (N, K) -> [(max M, config), ...] for the DeepSeek-V4.1-Flash MXFP8 linears at -# TP2/TP4, tuned on MI355X under HIP graphs with weights cold in the LLC. A -# tier serves every M up to its bound; larger M reuse the last tier. +# (N, K) -> [(max M, config), ...] for the DeepSeek-V4.1-Flash MXFP8 +# linears at TP2/TP4, tuned on MI355X under HIP graphs. A tier serves +# every M up to its bound; larger M reuse the last tier. _TUNED: dict[tuple[int, int], list[tuple[int, _Config]]] = { # shared expert gate_up, TP4 (1152, 5120): [ @@ -563,6 +687,49 @@ def _config(M: int, N: int, K: int) -> _Config: return tiers[-1][1] +# Retuned tiling/scheduling for the fused-quant kernel, whose per-CTA +# quant prologue shifts the best (BLOCK_M, BLOCK_N, num_warps, num_stages, +# waves_per_eu). Entries override ONLY those five axes of the stock tier +# active at M; the K-structure axes (block_k, k_split/k_pack, +# fused_splits, matrix_instr_nonkdim) stay at that tier's values so no +# output element's K-accumulation order changes and the fused output stays +# bit-identical to the two-launch path at every M. Consulted only on the +# quant_in_kernel path, never by the stock two-launch path. +_TUNED_FUSED_QUANT: dict[ + tuple[int, int], list[tuple[int, tuple[int, int, int, int, int]]] +] = { + # (N, K): [(max_m, (block_m, block_n, warps, stages, waves)), ...] + # DeepSeek-V4.1-Flash TP4 decode shapes: + (1152, 5120): [(8, (8, 32, 4, 2, 0))], # shared_gate_up + (1792, 5120): [(8, (8, 16, 4, 3, 1))], # fused_wqa_wkv + (4096, 1280): [(8, (8, 16, 4, 2, 0))], # indexer_wq_b + (5120, 576): [(8, (8, 16, 4, 2, 0))], # shared_down + (5120, 2048): [(8, (8, 32, 8, 3, 0))], # wo_b + (8192, 1280): [(8, (8, 32, 2, 3, 1))], # wq_b + # DSpark draft shapes (stock falls to _default_config): + (1024, 2048): [(8, (8, 16, 4, 2, 1))], # draft + (768, 1536): [(8, (8, 16, 4, 3, 0))], # draft +} + + +def _fused_quant_config(M: int, N: int, K: int) -> _Config: + """The stock tier active at M with its tiling/scheduling retuned for the + in-kernel quant prologue (see _TUNED_FUSED_QUANT); the stock config + where the shape has no fused-tier entry. K-structure axes are always + the stock tier's own, so the output stays bit-exact at every M.""" + base = _config(M, N, K) + for max_m, (bm, bn, warps, stages, waves) in _TUNED_FUSED_QUANT.get((N, K), ()): + if max_m >= M: + return base._replace( + block_m=bm, + block_n=bn, + num_warps=warps, + num_stages=stages, + waves_per_eu=waves, + ) + return base + + _split_counters_by_stream: dict[tuple[int, int], torch.Tensor] = {} @@ -599,7 +766,11 @@ def _launch( weight_scale: torch.Tensor, out: torch.Tensor, cfg: _Config, + quant_in_kernel: bool = False, ) -> None: + """Launch cfg's kernel. quant_in_kernel=True (packed cfgs only) makes + x the raw bf16/fp16 activation and x_scale None: the kernel quantizes + x per 32-block in registers (see _block32_packed_quant_kernel).""" M, K = x.shape N = weight.shape[0] opts = dict( @@ -638,21 +809,46 @@ def _launch( grid = (grid_m, grid_n, splits) else: grid = (grid_n, grid_m, splits) - args = (x, x_scale, weight, weight_scale, target, slots, counters, M, N, K) + args: tuple + if quant_in_kernel: + if not cfg.packed: + raise ValueError("quant_in_kernel needs a packed tile config") + args = (x, weight, weight_scale, target, slots, counters, M, N, K) + else: + args = (x, x_scale, weight, weight_scale, target, slots, counters, M, N, K) stride = out.stride(0) if splits == 1 or fused > 1 else N if cfg.packed: - _block32_packed_kernel[grid]( - *args, - stride, - BLOCK_M=cfg.block_m, - BLOCK_N=cfg.block_n, - BLOCK_K=cfg.block_k, - K_PACK=cfg.k_split, - K_PER_SPLIT=k_per_split, - EVEN_K=K % step == 0, - FUSED_SPLITS=fused, - **opts, - ) + if quant_in_kernel: + # Same output contract as the packed kernel below (partials at + # pid_k * M * N when FUSED_SPLITS == 1 with multiple pid_k), so + # the split-K reduce after this branch applies unchanged. + _block32_packed_quant_kernel[grid]( + *args, + stride, + BLOCK_M=cfg.block_m, + BLOCK_N=cfg.block_n, + BLOCK_K=cfg.block_k, + K_PACK=cfg.k_split, + K_PER_SPLIT=k_per_split, + EVEN_K=K % step == 0, + FUSED_SPLITS=fused, + FP8_MAX=_FP8_MAX, + TINY=_FP32_TINY, + **opts, + ) + else: + _block32_packed_kernel[grid]( + *args, + stride, + BLOCK_M=cfg.block_m, + BLOCK_N=cfg.block_n, + BLOCK_K=cfg.block_k, + K_PACK=cfg.k_split, + K_PER_SPLIT=k_per_split, + EVEN_K=K % step == 0, + FUSED_SPLITS=fused, + **opts, + ) else: _block32_tiled_kernel[grid]( *args, @@ -735,3 +931,113 @@ def rocm_mxfp8_block32_gemm( return torch.ops.vllm.rocm_mxfp8_block32_gemm( x, x_scale, weight, weight_scale, out_dtype ) + + +# --- tiny-M fused quant+GEMM ------------------------------------------------ +# At M <= 8 (decode-sized batches) the fused kernel below computes the +# activation quant inside the GEMM's K loop, saving one launch and its +# global round-trip per call. + +_TINY_M_FUSED_QUANT_MAX_M = 8 +# Shared with _mxfp8_quant_triton_kernel (the bit-exactness contract). +_FP8_MAX = float(torch.finfo(torch.float8_e4m3fn).max) +_FP32_TINY = float(torch.finfo(torch.float32).tiny) + +_FUSED_QUANT_HIT_SHAPES: set = set() + + +def _rocm_mxfp8_block32_gemm_quant_impl( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + """``x @ weight.T`` for raw bf16/fp16 ``x`` and a 32x32 block-scaled + MXFP8 weight, quantizing ``x`` inside the GEMM at tiny M. + + M <= 8 with a packed tile config (every tuned M <= 8 tier is packed) + runs one launch whose K loop computes the per-32-block E8M0 scale and + the e4m3 cast exactly as ``_mxfp8_quant_triton_kernel`` and feeds + ``tl.dot_scaled`` from registers -- bit-identical output to + ``mxfp8_e4m3_quantize`` + ``rocm_mxfp8_block32_gemm`` with one launch + fewer on the decode critical path. Everything else (larger M, + non-packed tiles, or non-contiguous or non-bf16/fp16 x) runs the stock + two-launch path with identical results. The branch is inside this + opaque op so a compiled caller cannot freeze the token-count decision. + """ + M, K = x.shape + out = torch.empty((M, weight.shape[0]), dtype=out_dtype, device=x.device) + if M <= 0: + return out + cfg = _config(M, weight.shape[0], K) + eligible = ( + cfg.packed + and M <= _TINY_M_FUSED_QUANT_MAX_M + and x.dtype in (torch.bfloat16, torch.float16) + and x.is_contiguous() + ) + if eligible: + # Retuned tiling for the in-kernel quant prologue (same K-structure + # as the stock tier at M, so bit-exactness holds); the stock path + # below keeps the stock config untouched. + fused_cfg = _fused_quant_config(M, weight.shape[0], K) + shape = (M, weight.shape[0], K) + if shape not in _FUSED_QUANT_HIT_SHAPES: + _FUSED_QUANT_HIT_SHAPES.add(shape) + # warning_once keeps the engagement log one-shot per shape; + # WARNING, not INFO, so it reaches the server log. + logger.warning_once( + "tiny-M fused MXFP8 quant+GEMM HIT (M=%d, N=%d, K=%d, packed " + "tile %dx%dx%d k_pack=%d%s)", + M, + weight.shape[0], + K, + fused_cfg.block_m, + fused_cfg.block_n, + fused_cfg.block_k, + fused_cfg.k_split, + ( + f", fused_splits={fused_cfg.fused_splits}" + if fused_cfg.fused_splits > 1 + else "" + ), + ) + _launch(x, None, weight, weight_scale, out, fused_cfg, quant_in_kernel=True) + else: + x_q, x_scale = mxfp8_e4m3_quantize(x) + _launch(x_q, x_scale, weight, weight_scale, out, cfg) + return out + + +def _rocm_mxfp8_block32_gemm_quant_fake( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + return x.new_empty((x.shape[0], weight.shape[0]), dtype=out_dtype) + + +# An opaque op for the same reason as rocm_mxfp8_block32_gemm: the fused vs +# two-launch choice branches on the token count. +direct_register_custom_op( + op_name="rocm_mxfp8_block32_gemm_quant", + op_func=_rocm_mxfp8_block32_gemm_quant_impl, + fake_impl=_rocm_mxfp8_block32_gemm_quant_fake, +) + + +def rocm_mxfp8_block32_gemm_quant( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + """``x @ weight.T`` for a raw bf16/fp16 ``x`` [M, K] and a 32x32 + block-scaled MXFP8 ``weight`` [N, K] with ``weight_scale`` + [N / 32, K / 32]. Tiny-M batches quantize inside the packed GEMM (see + ``_rocm_mxfp8_block32_gemm_quant_impl``). + """ + return torch.ops.vllm.rocm_mxfp8_block32_gemm_quant( + x, weight, weight_scale, out_dtype + ) diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 800e0d1aa6cd..878a9ee1cd0e 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -34,7 +34,11 @@ from vllm.triton_utils import tl, triton from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig -from .rocm_block32_gemm import BLOCK_ROWS, rocm_mxfp8_block32_gemm +from .rocm_block32_gemm import ( + BLOCK_ROWS, + rocm_mxfp8_block32_gemm, + rocm_mxfp8_block32_gemm_quant, +) _FP8_MAX = float(torch.finfo(MXFP8_VALUE_DTYPE).max) _TINY = float(torch.finfo(torch.float32).tiny) @@ -437,11 +441,22 @@ def apply_weights( out_shape = (*in_shape[:-1], layer.weight.shape[0]) if layer.weight_scale.shape[0] != layer.weight.shape[0]: # One scale row per 32 weight rows (see process_weights_after_loading). - if qa is None: - x_q, x_scale = mxfp8_e4m3_quantize(x2d) - out = rocm_mxfp8_block32_gemm( - x_q, x_scale, layer.weight, layer.weight_scale, out_dtype - ) + if qa is not None: + out = rocm_mxfp8_block32_gemm( + x_q, x_scale, layer.weight, layer.weight_scale, out_dtype + ) + else: + # The op quantizes the activation itself: at tiny M (the + # decode / DSpark shapes) the quant is fused into the packed + # GEMM kernel -- one launch, bit-exact with + # mxfp8_e4m3_quantize, kill switch + # VLLM_DISABLE_FUSED_MXFP8_QUANT_GEMM=1 -- while larger M + # keeps the separate launches. The branch lives inside the + # opaque op so a compiled caller cannot freeze the + # token-count decision. + out = rocm_mxfp8_block32_gemm_quant( + x2d, layer.weight, layer.weight_scale, out_dtype + ) elif qa is not None: out = _mxfp8_dot_scaled_gemm( x_q, x_scale, layer.weight, layer.weight_scale, out_dtype diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index ef4594e11374..fa1b598967d9 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -328,6 +328,104 @@ def maybe_pad_weight_avoid_cache_aliasing(weight: torch.Tensor) -> torch.Tensor: return torch.nn.functional.pad(weight, (0, pad))[..., :-pad] +_skinny_guard_warned: set[tuple[tuple[int, ...], int]] = set() + +# Cap used when the L3 size cannot be read: the MI355X value, the part +# this guard was measured on (see envs.py for the cap). +_SKINNY_GEMM_FALLBACK_WEIGHT_BYTES = 256 * 1024 * 1024 + + +def _amd_l3_cache_bytes(cache_info: list | None = None) -> int | None: + """The GPU's L3 (memory-side) cache size in bytes, or None. + + The skinny-GEMM weight cap tracks this cache: past it a weight cannot + stay resident, and the skinny kernel's cost at n activation rows (~n + weight passes) pays HBM round trips instead. The device-properties + API exposes only the far smaller L2 (4 MiB on gfx950), so query the + L3 through AMD's amdsmi library: 256 MiB on MI355X, 128 MiB on + MI300X-class parts. `cache_info` (an amdsmi cache list) lets tests + drive the parse; the default queries the library. + """ + if cache_info is None: + try: + import amdsmi + + amdsmi.amdsmi_init() + try: + handles = amdsmi.amdsmi_get_processor_handles() + cache_info = amdsmi.amdsmi_get_gpu_cache_info(handles[0])["cache"] + finally: + amdsmi.amdsmi_shut_down() + except Exception: # noqa: BLE001 -- optional lib, any failure -> fallback + return None + for cache in cache_info: + if cache.get("cache_level") == 3 and "DATA_CACHE" in cache.get( + "cache_properties", [] + ): + return cache["cache_size"] * 1024 # amdsmi reports KB + return None + + +@functools.lru_cache(maxsize=1) +def _skinny_gemm_max_weight_bytes() -> int: + """Effective skinny-GEMM weight cap in bytes, evaluated once. + + VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES > 0 overrides; 0 (the default) + derives the cap from the hardware L3 (memory-side) cache size, + falling back to _SKINNY_GEMM_FALLBACK_WEIGHT_BYTES where that cannot + be read. lmhead_skinny_guard_permits runs per decode step, so this is + computed on the first call only (env access in hot paths: #17067). + Tests that change the env var or mock the hardware read must call + cache_clear(). + """ + env_bytes = envs.VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES + if env_bytes > 0: + return env_bytes + return _amd_l3_cache_bytes() or _SKINNY_GEMM_FALLBACK_WEIGHT_BYTES + + +def lmhead_skinny_guard_permits(weight: torch.Tensor, n: int) -> bool: + """Should this weight still take the wvSplitK skinny path at n? + + wvSplitK_hf_sml's cost at n activation rows scales as ~n passes over + the weight (MI355X measurements in the PR description), so past the + last-level (memory-side) cache its effective bandwidth collapses: + every pass beyond the first is an HBM round trip. Weights above + VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES (default: the memory-side cache + size, read from the hardware once at first use, 256 MiB on MI355X; + see envs.py for what the cap controls and how to tune it) therefore + skip the skinny path at 1 < n <= 5 and fall through to the generic + GEMM. + + n == 1 always keeps wvSplitK: the n=1 kernel variant shows no + pathology on large weights, and its fall-through would be a numerics + change with no speed win, so single-row calls stay on the stock path. + + Returns True (keep wvSplitK) when n <= 1 or the weight is under the + cap; otherwise returns False after emitting a census WARNING (once per + (shape, n) per process) naming the redirected site. + """ + if n <= 1: + return True + max_bytes = _skinny_gemm_max_weight_bytes() + if weight.nbytes <= max_bytes: + return True + key = (tuple(weight.shape), n) + if key not in _skinny_guard_warned: + _skinny_guard_warned.add(key) + logger.warning( + "lm_head skinny-GEMM guard engaged: weight %s (%d bytes) exceeds " + "the weight cap %d (VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES, or the " + "hardware memory-side cache size when unset) at n=%d; skipping " + "wvSplitK for the generic GEMM fall-through", + tuple(weight.shape), + weight.nbytes, + max_bytes, + n, + ) + return False + + def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: @@ -396,6 +494,7 @@ def rocm_unquantized_gemm_impl( and (bias is None or bias.is_contiguous()) and (m == 1 or m > 8) and 0 < n <= 5 + and lmhead_skinny_guard_permits(weight, n) ) use_llmm1 = ( From fa3b9d1377219247060d3c8b49eae199c9e7e154 Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Thu, 29 Oct 2026 00:30:51 -0600 Subject: [PATCH 117/406] [Bugfix][AMD] A GPU-CPU KV transfer fault in OffloadingConnector takes the engine down instead of degrading the cache. (#52838) Signed-off-by: okorzh-amd Signed-off-by: Oxana Korzh Co-authored-by: okorzh-amd Co-authored-by: Claude Opus 5 (1M context) --- .../offloading_connector/test_scheduler.py | 161 +++++++++++ .../unit/offloading_connector/test_worker.py | 65 +++++ .../test_worker_metadata.py | 7 + .../unit/offloading_connector/utils.py | 14 +- tests/v1/kv_connector/unit/utils.py | 3 + tests/v1/kv_offload/cpu/test_gpu_worker.py | 270 ++++++++++++++++++ .../kv_connector/v1/offloading/common.py | 6 +- .../kv_connector/v1/offloading/scheduler.py | 29 +- .../kv_connector/v1/offloading/worker.py | 20 +- .../kv_connector/v1/offloading_connector.py | 13 +- vllm/v1/kv_offload/cpu/gpu_worker.py | 112 ++++++-- 11 files changed, 667 insertions(+), 33 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 1573101b244a..1df5786c46a7 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -3512,6 +3512,167 @@ def test_complete_store_waits_for_all_worker_acks( assert job_id not in runner.connector_scheduler._jobs +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_failed_store_is_not_published_as_cache(request_runner, async_scheduling: bool): + """A store that failed on any worker must not become loadable. + + Its destination blocks were never written, so publishing them would serve + whatever the recycled blocks last held -- potentially another request's KV. + The failure is reported in an earlier step than the one that completes the + job, which is the case a per-step failure lookup would miss. + """ + tokens_per_block = 4 + blocks_per_chunk = 3 + runner = request_runner( + blocks_per_chunk=blocks_per_chunk, + block_size=tokens_per_block, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + worker_count=3, + ) + runner.new_request(token_ids=[0] * (tokens_per_block * blocks_per_chunk)) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run(decoded_tokens=[0, 0], complete_transfers=False) + job_id = next(iter(runner.connector_scheduler._jobs)) + runner.manager.complete_store.reset_mock() + + runner.connector_scheduler.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 1}, failed_jobs={job_id} + ) + ) + ) + assert runner.manager.complete_store.call_count == 0 + + runner.connector_scheduler.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 2} + ) + ) + ) + assert runner.manager.complete_store.call_count == 1 + assert runner.manager.complete_store.call_args.kwargs["success"] is False + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_successful_store_is_published(request_runner, async_scheduling: bool): + """The success path still publishes, so the guard cannot pass vacuously.""" + tokens_per_block = 4 + blocks_per_chunk = 3 + runner = request_runner( + blocks_per_chunk=blocks_per_chunk, + block_size=tokens_per_block, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + worker_count=3, + ) + runner.new_request(token_ids=[0] * (tokens_per_block * blocks_per_chunk)) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run(decoded_tokens=[0, 0], complete_transfers=False) + job_id = next(iter(runner.connector_scheduler._jobs)) + runner.manager.complete_store.reset_mock() + + runner.connector_scheduler.update_connector_output( + KVConnectorOutput( + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 3} + ) + ) + ) + assert runner.manager.complete_store.call_count == 1 + assert runner.manager.complete_store.call_args.kwargs["success"] is True + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +@pytest.mark.parametrize("recompute", [False, True], ids=["fail", "recompute"]) +@pytest.mark.parametrize("hybrid", [False, True], ids=["one_group", "full_and_swa"]) +def test_failed_load_ends_or_reruns_its_request( + request_runner, recompute: bool, hybrid: bool, async_scheduling: bool +): + """A failed load ends or reruns its request, never the engine. + + Its destination blocks hold undefined KV, so the request is reported through + failed_recving, which the scheduler handles for any number of KV cache + groups. A rerun must not read the offloaded cache again, or a persistent + fault would reload the same chunks forever. + """ + block_size = 4 + kv_cache_groups = None + expected_loaded: tuple[int | tuple[int, int], ...] = (0, 1, 2) + if hybrid: + # The sliding-window group loads only the blocks inside its window. + expected_loaded = ((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)) + kv_cache_groups = [ + KVCacheGroupSpec( + ["layer0"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["layer1"], + SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=2 * block_size, + ), + ), + ] + runner = request_runner( + block_size=block_size, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + kv_cache_groups=kv_cache_groups, + ) + runner.scheduler.recompute_kv_load_failures = recompute + runner.new_request(token_ids=[0] * block_size * 3) + 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)) + + handler = runner.offloading_spec.handler + + def num_loads() -> int: + return sum( + isinstance(dst_spec, GPULoadStoreSpec) + for _, dst_spec in handler.transfer_specs.values() + ) + + def at_most_one_load() -> None: + assert num_loads() <= 1, "the failed load was retried" + + runner.scheduler.reset_prefix_cache() + runner.new_request(token_ids=[0] * (block_size * 3 + 1)) + request = runner.scheduler.requests[str(runner.req_id)] + runner.manager.lookup.return_value = LookupResult.HIT + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output([]) + ) + handler.fail_loads = True + runner.run(decoded_tokens=[], expected_loaded=expected_loaded) + + if not recompute: + assert request.status == RequestStatus.FINISHED_ERROR + return + assert request.num_computed_tokens == 0 + runner.run(decoded_tokens=[EOS_TOKEN_ID], post_step_fn=at_most_one_load) + assert num_loads() == 1 + assert request.status == RequestStatus.FINISHED_STOPPED + assert request.num_output_tokens == 1 + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_max_offload_tokens_validation(request_runner, async_scheduling: bool): """Validates max_offload_tokens: type coercion, boundary values, and capping. diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 7e050e01f737..1e800abfd735 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -309,6 +309,71 @@ def test_load_completion(sync: bool): assert meta.transfer_stats.store.is_empty() +def _kv_cache_config(num_groups: int) -> KVCacheConfig: + return KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[MagicMock() for _ in range(num_groups)], + ) + + +def _failed(job_id: int) -> list[TransferResult]: + return [TransferResult(job_id=job_id, success=False)] + + +def test_failed_store_is_dropped_and_still_completes(): + """A failed store is a future cache miss, not an error. + + It must still be marked completed, or the scheduler's per-job pending + count never reaches zero, and marked failed so it is not published. + """ + worker, _ = _make_worker(_kv_cache_config(1)) + worker.prepare_store_kv(_store_metadata(11)) + worker.start_kv_transfers(_empty_metadata()) + worker.worker.get_finished.return_value = _failed(11) + + finished_sending, finished_recving = worker.get_finished(set()) + + assert finished_sending == set() + assert finished_recving == set() + assert worker.take_failed_recving() == set() + meta = worker.build_connector_worker_meta() + assert meta is not None + assert meta.completed_jobs == {11: 1} + assert meta.failed_jobs == {11} + + +def test_failed_load_fails_its_request(): + """A failed load leaves undefined KV in its destination blocks. + + The request is reported as finished, per the base contract, and as failed, + which the scheduler handles for any number of KV cache groups. + """ + worker, _ = _make_worker(_kv_cache_config(1)) + worker.start_kv_transfers(_load_metadata(12)) + worker.worker.get_finished.return_value = _failed(12) + + _, finished_recving = worker.get_finished(set()) + + assert finished_recving == {"req"} + assert worker.take_failed_recving() == {"req"} + assert worker.take_failed_recving() == set() + meta = worker.build_connector_worker_meta() + assert meta is not None + assert meta.completed_jobs == {12: 1} + assert meta.failed_jobs == {12} + + +def test_failed_sync_load_raises(): + """A sync load's request already ran on the KV it was meant to load.""" + worker, _ = _make_worker(_kv_cache_config(1)) + worker.start_kv_transfers(_load_metadata(13, sync=True)) + worker.worker.get_finished.return_value = _failed(13) + + with pytest.raises(RuntimeError, match="sync load job 13 failed"): + worker.get_finished(set()) + + def test_offloading_connector_worker_accepts_plugin_spec_default_layout(): from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import ( OffloadingConnectorWorker, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py index ad2505abcefc..7afa93313613 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker_metadata.py @@ -37,6 +37,13 @@ def test_aggregate_multiple_workers(): assert result.completed_jobs == {42: 3, 43: 2, 7: 2, 8: 2} +def test_aggregate_keeps_a_failure_reported_by_any_worker(): + ok = OffloadingWorkerMetadata(completed_jobs={1: 1}) + failed = OffloadingWorkerMetadata(completed_jobs={1: 1}, failed_jobs={1}) + assert ok.aggregate(failed).failed_jobs == {1} + assert failed.aggregate(ok).failed_jobs == {1} + + def test_aggregate_transfer_stats(): meta1 = OffloadingWorkerMetadata( transfer_stats=TransferStats( diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index f8740e0bab87..30437ad3a111 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -84,6 +84,9 @@ def __init__(self): self.transfer_specs: dict[int, tuple[LoadStoreSpec, LoadStoreSpec]] = {} self.completed_transfers: list[TransferResult] = [] self.waiting_jobs: set[int] = set() + # When set, loads submitted from now on complete as failed. + self.fail_loads = False + self._failing_jobs: set[int] = set() self.completed_jobs: list[int] = [] self.flushed_jobs: set[int] = set() @@ -108,6 +111,8 @@ def submit_load( ) -> bool: self.transfer_specs[job_id] = (src_spec, dst_spec) self.waiting_jobs.add(job_id) + if self.fail_loads: + self._failing_jobs.add(job_id) return True def complete_jobs(self, job_ids: set[int]) -> None: @@ -117,7 +122,7 @@ def complete_jobs(self, job_ids: set[int]) -> None: self.completed_jobs.append(job_id) result = TransferResult( job_id=job_id, - success=True, + success=job_id not in self._failing_jobs, transfer_size=None, transfer_time=None, ) @@ -535,7 +540,7 @@ def _run( self.offloading_spec.complete_transfers() self.worker_connector.wait_for_save() - finished_sending, finished_recving = self.worker_connector.get_finished( + transfer_results = self.worker_connector.get_transfer_results( scheduler_output.finished_req_ids ) worker_meta = ( @@ -547,10 +552,11 @@ def _run( model_runner_output = create_model_runner_output( reqs=self.scheduler.running, - finished_sending=finished_sending, - finished_recving=finished_recving, + finished_sending=transfer_results.finished_sending, + finished_recving=transfer_results.finished_recving, token_id=token_id or 0, kv_connector_worker_meta=worker_meta, + failed_recving=transfer_results.failed_recving, ) prev_token_id = token_id diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 903ef4e0e327..e165d95932cf 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -289,6 +289,7 @@ def create_model_runner_output( use_eos: bool = False, token_id: int = 0, kv_connector_worker_meta: KVConnectorWorkerMetadata | None = None, + failed_recving: set[str] | None = None, ) -> ModelRunnerOutput: """Make dummy model runner output for testing.""" # Make request data. @@ -306,12 +307,14 @@ def create_model_runner_output( and finished_recving is None and invalid_block_ids is None and kv_connector_worker_meta is None + and failed_recving is None ) else KVConnectorOutput( finished_sending=finished_sending, finished_recving=finished_recving, invalid_block_ids=invalid_block_ids or set(), kv_connector_worker_meta=kv_connector_worker_meta, + failed_recving=failed_recving or set(), ) ) diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index aad2c2f4fbfb..b2ad143c6af4 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib import logging import random import time @@ -204,6 +205,202 @@ def test_handler_shutdown_skips_transfers_after_event_sync_failure() -> None: assert not handler.dst_tensors +def _bare_handler(gpu_to_cpu: bool = True): + """Handler with only the state the failure paths touch.""" + handler = gpu_worker.SingleDirectionOffloadingHandler.__new__( + gpu_worker.SingleDirectionOffloadingHandler + ) + handler.gpu_to_cpu = gpu_to_cpu + handler._transfers = gpu_worker.deque() + handler._transfer_events = {} + handler._stream_pool = [] + handler._event_pool = [] + handler._buffer_pool = [] + return handler + + +def _submittable_handler(monkeypatch: pytest.MonkeyPatch, gpu_to_cpu: bool = True): + """Bare handler wired up far enough to run transfer_async off-device.""" + handler = _bare_handler(gpu_to_cpu) + handler.layer_refs_per_group = [MagicMock()] + handler.src_blocks_per_chunk = 1 + handler.dst_blocks_per_chunk = 1 + handler._estimate_max_copy_ops = lambda group_sizes: 4 + handler._fill_group_ops = lambda g_idx, **kwargs: (1, 128) + handler._swap_blocks_batch = MagicMock() + handler._buffer_pool = [tuple(torch.zeros(4, dtype=torch.int64) for _ in range(3))] + + platform = MagicMock() + platform.stream.return_value = contextlib.nullcontext() + platform.is_cuda_alike.return_value = False + monkeypatch.setattr(gpu_worker, "current_platform", platform) + return handler + + +def _submit(handler, job_id: int) -> bool: + spec = GPULoadStoreSpec(block_ids=[0], group_sizes=[1], block_indices=[0]) + return handler.transfer_async(job_id, spec, spec) + + +def test_submit_failure_is_not_reported_until_stream_drains( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A submit fault leaves copies in flight, so the job completes only later. + + swap_blocks_batch submits the descriptor list in chunks and the driver's + batch memcpy stops at the first failing copy, so earlier copies of the same + job may still be touching its blocks. Reporting the job before the stream + drains would let the scheduler reuse blocks that are still being written. + """ + handler = _submittable_handler(monkeypatch) + stream, start_event, end_event = MagicMock(), MagicMock(), MagicMock() + handler._stream_pool = [stream] + handler._event_pool = [end_event, start_event] + handler._swap_blocks_batch.side_effect = RuntimeError("failed at index 1") + clear = MagicMock() + monkeypatch.setattr(gpu_worker, "_clear_pending_device_error", clear) + + assert _submit(handler, 7) is True + + clear.assert_called_once_with() + end_event.record.assert_called_once_with(stream) + (transfer,) = handler._transfers + assert (transfer.job_id, transfer.failed) == (7, True) + assert handler._transfer_events == {7: end_event} + assert not handler._stream_pool + assert not handler._event_pool + assert not handler._buffer_pool + + end_event.query.return_value = False + assert handler.get_finished() == [] + + end_event.query.return_value = True + results = handler.get_finished() + + assert [(r.job_id, r.success) for r in results] == [(7, False)] + assert not handler._transfers + assert not handler._transfer_events + assert not handler._stream_pool + assert not handler._event_pool + assert not handler._buffer_pool + start_event.elapsed_time.assert_not_called() + + +def test_submit_failure_still_fences_wait_and_next_submission( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The faulted job keeps ordering the work that follows it. + + wait() is the store-flush fence and every submission chains onto the + previous transfer's end event; a faulted job that left neither would let + both run ahead of its in-flight copies. + """ + handler = _submittable_handler(monkeypatch) + failed_stream, failed_start, failed_end = MagicMock(), MagicMock(), MagicMock() + handler._stream_pool = [failed_stream] + handler._event_pool = [failed_end, failed_start] + handler._swap_blocks_batch.side_effect = RuntimeError("failed at index 1") + _submit(handler, 7) + + handler.wait({7}) + failed_end.synchronize.assert_called_once_with() + + next_stream = MagicMock() + handler._stream_pool = [next_stream] + handler._event_pool = [MagicMock(), MagicMock()] + handler._buffer_pool = [tuple(torch.zeros(4, dtype=torch.int64) for _ in range(3))] + handler._swap_blocks_batch.side_effect = None + _submit(handler, 8) + + next_stream.wait_event.assert_called_once_with(failed_end) + + +def test_submit_failure_propagates_when_the_end_event_cannot_be_recorded( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Without an end event the in-flight copies cannot be bounded. + + The device is unrecoverable at that point, so failing loudly beats + reporting a completion the blocks have not actually reached. + """ + handler = _submittable_handler(monkeypatch) + end_event = MagicMock() + end_event.record.side_effect = RuntimeError("device lost") + handler._stream_pool = [MagicMock()] + handler._event_pool = [end_event, MagicMock()] + handler._swap_blocks_batch.side_effect = RuntimeError("failed at index 1") + + with pytest.raises(RuntimeError, match="device lost"): + _submit(handler, 7) + + +def test_get_finished_reports_failure_when_event_query_raises() -> None: + """A device error latched after submit surfaces at poll time, not as a raise.""" + handler = _bare_handler() + failing = MagicMock() + failing.end_event.query.side_effect = RuntimeError("device lost") + failing.job_id = 3 + ready = MagicMock() + ready.job_id = 4 + ready.failed = False + ready.end_event.query.return_value = True + ready.start_event.elapsed_time.return_value = 2.0 + ready.num_bytes = 512 + handler._transfers.extend([failing, ready]) + handler._transfer_events = {3: failing.end_event, 4: ready.end_event} + + results = handler.get_finished() + + assert [(r.job_id, r.success) for r in results] == [(3, False), (4, True)] + assert not handler._transfers + assert not handler._transfer_events + assert handler._stream_pool == [ready.stream] + assert handler._buffer_pool == [ + (ready.batch_src, ready.batch_dst, ready.batch_sizes) + ] + + +def test_get_finished_propagates_non_runtime_errors() -> None: + """Only RuntimeError, which device faults raise, becomes a failed transfer. + + Anything else, such as a missing op, must stay loud rather than silently + degrade the cache. + """ + handler = _bare_handler() + broken = MagicMock() + broken.job_id = 5 + broken.end_event.query.side_effect = AttributeError("no such op") + handler._transfers.append(broken) + + with pytest.raises(AttributeError, match="no such op"): + handler.get_finished() + + +def test_submit_propagates_non_runtime_errors(monkeypatch: pytest.MonkeyPatch) -> None: + handler = _submittable_handler(monkeypatch) + handler._stream_pool = [MagicMock()] + handler._event_pool = [MagicMock(), MagicMock()] + handler._swap_blocks_batch.side_effect = AttributeError("no such op") + clear = MagicMock() + monkeypatch.setattr(gpu_worker, "_clear_pending_device_error", clear) + + with pytest.raises(AttributeError, match="no such op"): + _submit(handler, 7) + clear.assert_not_called() + + +def test_wait_swallows_synchronize_failure() -> None: + """wait() must not raise; the failure is reported by the next poll.""" + handler = _bare_handler() + event = MagicMock() + event.synchronize.side_effect = RuntimeError("device lost") + handler._transfer_events[9] = event + + handler.wait({9}) + + event.synchronize.assert_called_once_with() + + @pytest.mark.parametrize("device_sync_fails", [False, True]) def test_worker_syncs_before_cleanup_after_handler_failure( caplog_vllm, monkeypatch: pytest.MonkeyPatch, device_sync_fails: bool @@ -710,3 +907,76 @@ def test_load_waits_for_pending_compute_stream_writes(default_vllm_config) -> No torch.testing.assert_close(gpu_tensor[block_id].cpu(), expected) finally: worker.shutdown() + + +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="faults a real hipMemcpyBatchAsync call", +) +@pytest.mark.parametrize("gpu_to_cpu", [True, False]) +@pytest.mark.parametrize( + "num_blocks", [16, 9000], ids=["one_batch", "fault_in_second_batch"] +) +def test_device_fault_does_not_fail_the_next_kernel_launch( + default_vllm_config, gpu_to_cpu: bool, num_blocks: int +) -> None: + """A caught copy fault must not leave its error pending on the thread. + + The runtime keeps a failed call's error until it is read, and the next + kernel launch check would raise it, killing the model step the catch was + meant to protect. + """ + page_size_bytes = 4096 + gpu_tensor = torch.zeros( + (num_blocks, page_size_bytes), dtype=torch.int8, device=DEVICES[0] + ) + worker = CPUOffloadingWorker( + kv_caches=CanonicalKVCaches( + tensors=[ + CanonicalKVCacheTensor( + tensor=gpu_tensor, page_size_bytes=page_size_bytes + ) + ], + group_data_refs=[ + [CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=page_size_bytes)] + ], + ), + blocks_per_chunk=1, + num_cpu_chunks=num_blocks, + ) + handler = worker._store_handler if gpu_to_cpu else worker._load_handler + swap_blocks_batch = handler._swap_blocks_batch + faulted = [] + + def zero_one_copy(src, dst, sizes, **kwargs): + sizes[len(sizes) - 500 if num_blocks > 8192 else len(sizes) // 2] = 0 + try: + swap_blocks_batch(src, dst, sizes, **kwargs) + except RuntimeError: + faulted.append(True) + raise + + handler._swap_blocks_batch = zero_one_copy + gpu_spec = GPULoadStoreSpec( + list(range(num_blocks)), group_sizes=(num_blocks,), block_indices=(0,) + ) + cpu_spec = CPULoadStoreSpec(list(range(num_blocks))) + try: + if gpu_to_cpu: + assert worker.submit_store(1, gpu_spec, cpu_spec) + else: + assert worker.submit_load(1, cpu_spec, gpu_spec) + if not faulted: + pytest.skip("this runtime accepts zero-size copies") + + torch.arange(4096, device=DEVICES[0]).cumsum(0) + torch.accelerator.synchronize() + + deadline = time.time() + 10 + finished: list[TransferResult] = [] + while time.time() < deadline and not finished: + finished = worker.get_finished() + assert [(r.job_id, r.success) for r in finished] == [(1, False)] + finally: + gpu_worker._clear_pending_device_error() + worker.shutdown() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py index 61cbcf99c1a5..afcd787dbaa7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py @@ -85,11 +85,14 @@ class OffloadingWorkerMetadata(KVConnectorWorkerMetadata): """ completed_jobs: dict[int, int] = field(default_factory=dict) + failed_jobs: set[int] = field(default_factory=set) transfer_stats: TransferStats = field(default_factory=TransferStats) - def mark_completed(self, job_id: int) -> None: + def mark_completed(self, job_id: int, success: bool = True) -> None: """Record a transfer job completion from this worker.""" self.completed_jobs[job_id] = 1 + if not success: + self.failed_jobs.add(job_id) def aggregate( self, other: "KVConnectorWorkerMetadata" @@ -102,5 +105,6 @@ def aggregate( return OffloadingWorkerMetadata( completed_jobs=merged, + failed_jobs=self.failed_jobs | other.failed_jobs, transfer_stats=self.transfer_stats.aggregate(other.transfer_stats), ) 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 5bf8a77d351d..63e4d8cfaac6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -87,6 +87,7 @@ class TransferJobStatus: is_store: bool # Unlike an async load, store jobs may be created next to it. is_sync_load: bool = False + failed: bool = False # Store source blocks fenced after the request finishes. deferred_fence_block_ids: list[int] | None = None # Store source blocks fenced when the transfer is created. @@ -374,6 +375,9 @@ class RequestOffloadState: load_key_ranges: list[tuple[int, int, OffloadKey]] = field(default_factory=list) # True once on_request_finished has been signaled to the manager. finished_signaled: bool = False + # A load for this request failed; all its later lookups, including after + # preemption, skip the offloaded cache so a faulting load is not retried. + load_failed: bool = False def __post_init__(self) -> None: self.group_states = tuple( @@ -1299,7 +1303,7 @@ def get_num_new_matched_tokens( req_status.num_locally_computed_tokens = num_computed_tokens num_hit_tokens: int | None - if request.skip_reading_external_kv: + if request.skip_reading_external_kv or req_status.load_failed: num_hit_tokens = 0 else: lookup_start = time.monotonic() @@ -2325,6 +2329,8 @@ def _update_connector_output(self, connector_output: KVConnectorOutput): ) continue job_status = self._jobs[job_id] + if job_id in meta.failed_jobs: + job_status.failed = True job_status.pending_count -= count if job_status.pending_count > 0: continue @@ -2332,8 +2338,27 @@ def _update_connector_output(self, connector_output: KVConnectorOutput): req_status = self._req_status[job_status.req_id] if job_status.is_store: - self.manager.complete_store(job_status.keys, req_status.req_context) + if job_status.failed: + logger.warning( + "KV offload store job %d failed; dropping %d chunk(s) " + "from the CPU cache rather than publishing them.", + job_id, + len(job_status.keys), + ) + self.manager.complete_store( + job_status.keys, + req_status.req_context, + success=not job_status.failed, + ) else: + if job_status.failed: + logger.warning( + "KV offload load job %d failed; request %s will not " + "load from the offloaded cache again.", + job_id, + job_status.req_id, + ) + req_status.load_failed = True self.manager.complete_load(job_status.keys, req_status.req_context) if self._chunks_being_loaded: self._chunks_being_loaded.difference_update(job_status.keys) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 616b20e4d085..8698deee8733 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -60,6 +60,7 @@ def __init__( # job_id -> req_id for in-flight loads (None for sync loads). self._load_jobs: dict[int, ReqId | None] = {} + self._failed_recving: set[ReqId] = set() self._unsubmitted_store_jobs: list[ tuple[int, GPULoadStoreSpec, LoadStoreSpec] ] = [] @@ -296,10 +297,19 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: assert self.worker is not None finished_recving: set[str] = set() for transfer_result in self.worker.get_finished(): - # we currently do not support job failures job_id = transfer_result.job_id - assert transfer_result.success is_load = job_id in self._load_jobs + if not transfer_result.success: + self._connector_worker_meta.mark_completed(job_id, success=False) + req_id = self._load_jobs.pop(job_id, None) + if is_load and req_id is None: + # Its request already ran this step on the unloaded KV, + # which failed_recving cannot undo. + raise RuntimeError(f"KV offload sync load job {job_id} failed") + if req_id is not None: + self._failed_recving.add(req_id) + finished_recving.add(req_id) + continue if ( transfer_result.transfer_time is not None and transfer_result.transfer_size is not None @@ -320,6 +330,11 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: return set(), finished_recving + def take_failed_recving(self) -> set[ReqId]: + """Drain the requests whose load failed since the last call.""" + failed_recving, self._failed_recving = self._failed_recving, set() + return failed_recving + def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None: """Return completed transfer job IDs since the last call.""" if not self._connector_worker_meta.completed_jobs: @@ -331,6 +346,7 @@ def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None: def shutdown(self) -> None: self._unsubmitted_store_jobs.clear() self._load_jobs.clear() + self._failed_recving.clear() self._connector_worker_meta = OffloadingWorkerMetadata() if self.worker is not None: self.worker.shutdown() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index af3e113b5f09..087d1e88d644 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -13,7 +13,10 @@ KVConnectorRole, SupportsHMA, ) -from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorMetadata, + KVConnectorTransferResults, +) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( KVConnectorPromMetrics, KVConnectorStats, @@ -142,6 +145,14 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: assert isinstance(self._connector_metadata, OffloadingConnectorMetadata) return self.connector_worker.get_finished(finished_req_ids) + def get_transfer_results( + self, finished_req_ids: set[str] + ) -> KVConnectorTransferResults: + results = super().get_transfer_results(finished_req_ids) + assert self.connector_worker is not None + results.failed_recving = self.connector_worker.take_failed_recving() + return results + def build_connector_worker_meta(self) -> OffloadingWorkerMetadata | None: if self.connector_worker is not None: return self.connector_worker.build_connector_worker_meta() diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 9d20b7825af8..199632d053a3 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -11,6 +11,7 @@ import torch from vllm import _custom_ops as ops +from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, triton @@ -36,6 +37,16 @@ logger = init_logger(__name__) +def _clear_pending_device_error() -> None: + """Clear the error a failed copy left pending for the next kernel launch.""" + if not current_platform.is_cuda_alike(): + return + try: + CudaRTLibrary().cudaGetLastError() + except (AssertionError, AttributeError, OSError): + logger.warning("Could not clear the pending device error", exc_info=True) + + def _select_swap_blocks_fn( layer_refs_per_group: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, @@ -75,6 +86,7 @@ class Transfer: batch_src: torch.Tensor batch_dst: torch.Tensor batch_sizes: torch.Tensor + failed: bool = False def compute_sub_block_ptrs( @@ -704,16 +716,40 @@ def transfer_async( # writing; we must keep STREAM ordering so source reads are gated # by the transfer stream's wait_stream(compute) barrier. is_src_access_order_any = not self.gpu_to_cpu - with current_platform.stream(stream): - start_event.record(stream) - if op_idx > 0: - self._swap_blocks_batch( - src, - dst, - sizes, - is_src_access_order_any=is_src_access_order_any, - ) + try: + with current_platform.stream(stream): + start_event.record(stream) + if op_idx > 0: + self._swap_blocks_batch( + src, + dst, + sizes, + is_src_access_order_any=is_src_access_order_any, + ) + end_event.record(stream) + except RuntimeError: + _clear_pending_device_error() + logger.exception( + "KV offload %s transfer failed for job %d", + "GPU->CPU" if self.gpu_to_cpu else "CPU->GPU", + job_id, + ) end_event.record(stream) + self._transfer_events[job_id] = end_event + self._transfers.append( + Transfer( + job_id=job_id, + stream=stream, + start_event=start_event, + end_event=end_event, + num_bytes=0, + batch_src=batch_src, + batch_dst=batch_dst, + batch_sizes=batch_sizes, + failed=True, + ) + ) + return True self._transfer_events[job_id] = end_event self._transfers.append( @@ -734,33 +770,63 @@ def transfer_async( def get_finished(self) -> list[TransferResult]: results: list[TransferResult] = [] - while self._transfers and self._transfers[0].end_event.query(): - transfer = self._transfers.popleft() - transfer_time = ( - transfer.start_event.elapsed_time(transfer.end_event) * 1e-3 - ) # elapsed_time is in milliseconds - result = TransferResult( - job_id=transfer.job_id, - success=True, - transfer_size=transfer.num_bytes, - transfer_time=transfer_time, - ) + while self._transfers: + transfer = self._transfers[0] + try: + if not transfer.end_event.query(): + break + transfer_time = ( + None + if transfer.failed + # elapsed_time is in milliseconds + else transfer.start_event.elapsed_time(transfer.end_event) * 1e-3 + ) + except RuntimeError: + logger.exception( + "KV offload %s transfer failed for job %d while polling " + "for completion", + "GPU->CPU" if self.gpu_to_cpu else "CPU->GPU", + transfer.job_id, + ) + self._transfers.popleft() + self._transfer_events.pop(transfer.job_id, None) + results.append(TransferResult(job_id=transfer.job_id, success=False)) + continue + + self._transfers.popleft() + self._transfer_events.pop(transfer.job_id, None) + if transfer.failed: + results.append(TransferResult(job_id=transfer.job_id, success=False)) + continue - results.append(result) + results.append( + TransferResult( + job_id=transfer.job_id, + success=True, + transfer_size=transfer.num_bytes, + transfer_time=transfer_time, + ) + ) self._stream_pool.append(transfer.stream) self._event_pool.append(transfer.end_event) self._event_pool.append(transfer.start_event) self._buffer_pool.append( (transfer.batch_src, transfer.batch_dst, transfer.batch_sizes) ) - del self._transfer_events[transfer.job_id] return results def wait(self, job_ids: set[int]): for job_id in job_ids: event = self._transfer_events.get(job_id) if event is not None: - event.synchronize() + try: + event.synchronize() + except RuntimeError: + logger.exception( + "KV offload transfer for job %d failed while waiting " + "for it to complete", + job_id, + ) def shutdown(self) -> None: """Drain this direction and release its transfer-side resources.""" From 34091a07c7d78230e54d53d92b66b125b0f21a26 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Thu, 29 Oct 2026 16:55:05 +0800 Subject: [PATCH 118/406] [XPU] use multimodal_rotary_embedding xpu kernel for mrope (#59938) Signed-off-by: Yan Ma Co-authored-by: Kunshang Ji Co-authored-by: vllm-agent --- .../layers/rotary_embedding/mrope.py | 80 +++++++++++++------ 1 file changed, 57 insertions(+), 23 deletions(-) diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 9bcc5a95e16d..e1dbd1169f09 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -436,38 +436,18 @@ def forward_native( key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) return query, key - def forward_cuda( + def _forward_text_rope( self, positions: torch.Tensor, query: torch.Tensor, - key: torch.Tensor | None = None, - offsets: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - assert positions.ndim == 1 or positions.ndim == 2 - assert key is not None - + key: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: cos_sin_cache = self._match_cos_sin_cache_dtype(query) num_tokens = positions.shape[-1] cos_sin = cos_sin_cache[positions] cos, sin = cos_sin.chunk(2, dim=-1) query_shape = query.shape key_shape = key.shape - if positions.ndim == 2: - assert self.mrope_section - - q, k = triton_mrope( - query, - key, - cos, - sin, - self.mrope_section, - self.head_size, - self.rotary_dim, - self.mrope_interleaved, - self.is_neox_style, - ) - - return q.reshape(query_shape), k.reshape(key_shape) query = query.view(num_tokens, -1, self.head_size) query_rot = query[..., : self.rotary_dim] @@ -490,6 +470,40 @@ def forward_cuda( key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) return query, key + def forward_cuda( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None = None, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert positions.ndim == 1 or positions.ndim == 2 + assert key is not None + + if positions.ndim == 2: + assert self.mrope_section + cos_sin_cache = self._match_cos_sin_cache_dtype(query) + cos_sin = cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + query_shape = query.shape + key_shape = key.shape + + q, k = triton_mrope( + query, + key, + cos, + sin, + self.mrope_section, + self.head_size, + self.rotary_dim, + self.mrope_interleaved, + self.is_neox_style, + ) + + return q.reshape(query_shape), k.reshape(key_shape) + + return self._forward_text_rope(positions, query, key) + def forward_hip( self, positions: torch.Tensor, @@ -536,6 +550,26 @@ def forward_xpu( offsets: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: return self.forward_cuda(positions, query, key, offsets) + if positions.ndim == 2: + assert key is not None + assert self.mrope_section is not None + import vllm._xpu_ops # noqa: F401 + + cos_sin_cache = self._match_cos_sin_cache_dtype(query) + torch.ops._xpu_C.multimodal_rotary_embedding( + positions, + query, + key, + self.head_size, + cos_sin_cache, + self.is_neox_style, + self.mrope_section, + self.mrope_interleaved, + ) + return query, key + assert positions.ndim == 1 + assert key is not None + return self._forward_text_rope(positions, query, key) @staticmethod def get_next_input_positions( From a0e31b7353cc46232e8a9152b883b14726a024ab Mon Sep 17 00:00:00 2001 From: qli88 Date: Thu, 29 Oct 2026 04:27:22 -0500 Subject: [PATCH 119/406] [CI][ROCM][Bug] Elastic DP doesn't work with AITER-on (#63053) Signed-off-by: Qiang Li --- .../device_communicators/cuda_communicator.py | 8 ++++- .../distributed/elastic_ep/elastic_execute.py | 32 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index bba4792fea87..85f7541062b4 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -165,9 +165,15 @@ def __init__( register_graph_buffers=register, ) - # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine + # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine. + # AITER requires a torch.distributed-registered group and exchanges IPC + # handles via the default store; stateless groups satisfy neither. + # Only StatelessGroupCoordinator passes tcp_store_group, so this checks + # the property AITER depends on rather than the feature that currently + # creates stateless groups (elastic EP). if ( "dp" in unique_name + and tcp_store_group is None and self.world_size in (2, 4, 8) and current_platform.is_rocm() and rocm_aiter_ops.is_custom_all_reduce_enabled() diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 182a4b9bb4f3..f064e2e86af0 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -12,6 +12,7 @@ import torch.nn as nn from torch.distributed import P2POp +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.compilation.cuda_graph import CUDAGraphWrapper from vllm.compilation.wrapper import reset_compile_wrapper from vllm.config import ParallelConfig, set_current_vllm_config @@ -460,6 +461,37 @@ def _release_cuda_graphs(self) -> None: gc.collect() torch.accelerator.synchronize() torch.accelerator.empty_cache() + self._switch_to_fresh_graph_pool() + + def _switch_to_fresh_graph_pool(self) -> None: + # Libraries may keep capture-time tensors alive after their graphs are + # released (e.g. FlyDSL retains the last launch's DLPack arguments), so + # the old pool cannot be erased. Re-capturing into that pool trips the + # allocator's "use_count > 0" assert; capture into a fresh one instead. + platform_cls = type(current_platform) + if platform_cls._global_graph_pool is None: + return + new_pool = current_platform.graph_pool_handle() + platform_cls._global_graph_pool = new_pool + + model_runner = self.worker.model_runner + manager = getattr(model_runner, "cudagraph_manager", None) + speculator = getattr(model_runner, "speculator", None) + managers = [manager] + [ + attr + for attr in (vars(speculator).values() if speculator else ()) + if isinstance(attr, CudaGraphManager) + ] + for m in managers: + if m is not None and m.pool is not None: + m.pool = new_pool + + wrappers: list[CUDAGraphWrapper | BreakableCUDAGraphWrapper] = [ + *CUDAGraphWrapper._all_instances, + *BreakableCUDAGraphWrapper._all_instances, + ] + for wrapper in wrappers: + wrapper.graph_pool = new_pool def switch_and_remove(self) -> None: # Removing ranks skipped preparation, so wait for prior cleanup here. From 0d6f92ae7c83f1d338d1b6a742ebcd79ab0fc2bb Mon Sep 17 00:00:00 2001 From: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:46:03 -0400 Subject: [PATCH 120/406] [Bugfix][Spec Decode] Don't specialize MRV2 logits-stats kernels on num_logits (#64137) Signed-off-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> --- vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index af086b0508f0..bf86dcf0ea5e 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -213,7 +213,8 @@ def _is_bonus_logit(expanded_local_pos_ptr, logit_idx, num_logits): return next_local_pos == 0 -@triton.jit +# num_logits only bounds the bonus-logit lookahead; one kernel serves every value. +@triton.jit(do_not_specialize=["num_logits"]) def _compute_local_logits_stats_kernel( # [num_logits, num_blocks] target_local_argmax_ptr, @@ -413,7 +414,7 @@ def _compute_cumulative_log_p_kernel( tl.store(cumulative_log_p_ptr + logit_idx, log_p) -@triton.jit +@triton.jit(do_not_specialize=["num_logits"]) def _compute_local_residual_mass_kernel( # [num_logits, num_blocks] local_residual_mass_ptr, From fa1cf1d436a0364a230c56f8e2c53a3d5043d231 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 29 Oct 2026 16:59:54 +0700 Subject: [PATCH 121/406] [Perf][GDN] Allocate GDN core output uninitialized, write FlashInfer prefill in place (#64419) Signed-off-by: Thien Tran Co-authored-by: Claude Opus 5.5 Co-authored-by: Kimi Code Co-authored-by: vllm-agent --- .../mamba/test_gdn_forward_core_split.py | 1 + tests/kernels/mamba/test_gdn_fused_mtp.py | 1 + .../mamba/test_gdn_prefill_flashinfer.py | 7 +- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 166 +++++++++++++----- .../warmup/qwen_triton_warmup.py | 5 +- 5 files changed, 129 insertions(+), 51 deletions(-) diff --git a/tests/kernels/mamba/test_gdn_forward_core_split.py b/tests/kernels/mamba/test_gdn_forward_core_split.py index 8a99e223be3b..7e8b9ef238cb 100644 --- a/tests/kernels/mamba/test_gdn_forward_core_split.py +++ b/tests/kernels/mamba/test_gdn_forward_core_split.py @@ -123,6 +123,7 @@ def _build_layer( layer.kv_cache = (conv_state, ssm_state) with set_current_vllm_config(vllm_config): layer.chunk_gated_delta_rule = ChunkGatedDeltaRule() + layer.gdn_prefill_backend = layer.chunk_gated_delta_rule.gdn_prefill_backend for name in ( "rearrange_mixed_qkv", "_forward_core", diff --git a/tests/kernels/mamba/test_gdn_fused_mtp.py b/tests/kernels/mamba/test_gdn_fused_mtp.py index 24ea20325ef8..690d3a0983f6 100644 --- a/tests/kernels/mamba/test_gdn_fused_mtp.py +++ b/tests/kernels/mamba/test_gdn_fused_mtp.py @@ -128,6 +128,7 @@ def _build_layer( ) with set_current_vllm_config(vllm_config): layer.chunk_gated_delta_rule = ChunkGatedDeltaRule() + layer.gdn_prefill_backend = layer.chunk_gated_delta_rule.gdn_prefill_backend for name in ( "rearrange_mixed_qkv", "_forward_core", diff --git a/tests/kernels/mamba/test_gdn_prefill_flashinfer.py b/tests/kernels/mamba/test_gdn_prefill_flashinfer.py index be6ab89b3c43..9c694573a027 100644 --- a/tests/kernels/mamba/test_gdn_prefill_flashinfer.py +++ b/tests/kernels/mamba/test_gdn_prefill_flashinfer.py @@ -28,7 +28,8 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID # noqa: E402 -def test_flashinfer_gdn_prefill_uses_int64_cu_seqlens(monkeypatch): +def test_flashinfer_gdn_prefill_passes_cu_seqlens_through(monkeypatch): + """FlashInfer accepts int32 cu_seqlens; the wrapper must not cast them.""" captured_cu_seqlens = None def fake_chunk_gated_delta_rule(**kwargs): @@ -57,7 +58,7 @@ def fake_chunk_gated_delta_rule(**kwargs): ) assert captured_cu_seqlens is not None - assert captured_cu_seqlens.dtype == torch.int64 + assert captured_cu_seqlens is cu_seqlens assert output.shape == q.shape assert final_state is None @@ -99,7 +100,7 @@ def test_flashinfer_checkpoints_export_the_state_at_each_offset(use_cp, ssm_dtyp torch.randn(1, num_tokens, num_qk_heads, head_dim, device=device), dim=-1 ).bfloat16() v = torch.randn(1, num_tokens, num_v_heads, head_dim, device=device).bfloat16() - g = F.logsigmoid(torch.randn(1, num_tokens, num_v_heads, device=device)) + g = F.logsigmoid(torch.randn(1, num_tokens, num_v_heads, device=device)).exp() beta = torch.sigmoid(torch.randn(1, num_tokens, num_v_heads, device=device)) initial_state = 0.1 * torch.randn( len(query_lens), num_v_heads, head_dim, head_dim, device=device diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 3c6c15b961a6..110929f1f8fb 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -324,6 +324,7 @@ def fi_chunk_gated_delta_rule( checkpoint_cu_starts: torch.Tensor | None = None, checkpoint_every_n_tokens: int = 0, use_cp: Literal["auto"] | bool = "auto", + output: torch.Tensor | None = None, ): from flashinfer.gdn_prefill import ( chunk_gated_delta_rule as chunk_gated_delta_rule_fi, @@ -341,19 +342,21 @@ def fi_chunk_gated_delta_rule( g = g.squeeze(0).contiguous() beta = beta.squeeze(0).contiguous() fi_state = initial_state.to(torch.float32) + # Like FlashInfer's own API, g is the already-exponentiated gate; + # fused_post_conv_prep emits it directly (output_g_exp=True). fi_g = g.to(torch.float32) fi_beta = beta.to(torch.float32) - if cu_seqlens is not None: - cu_seqlens = cu_seqlens.to(torch.int64) + # cu_seqlens goes through as is: FlashInfer accepts int32 and int64. result = chunk_gated_delta_rule_fi( q=q, k=k, v=v, - g=torch.exp(fi_g), + g=fi_g, beta=fi_beta, initial_state=fi_state, output_final_state=output_final_state, cu_seqlens=cu_seqlens, + output=output, state_checkpoints=state_checkpoints, checkpoint_cu_starts=checkpoint_cu_starts, checkpoint_every_n_tokens=checkpoint_every_n_tokens, @@ -414,6 +417,9 @@ def forward_cuda( core_attn_out: torch.Tensor | None = None, aiter_prefill_metadata: object | None = None, ): + # g is in log space in this dispatch layer (shared with the Triton/FLA + # backends); FlashInfer takes exp(g), so exp here. + g = torch.exp(g.to(torch.float32)) o, final_state = fi_chunk_gated_delta_rule( q=q, k=k, @@ -1192,9 +1198,10 @@ def forward_cuda( # ============================================================ # Part 2: Core Attention (Custom Op) # ============================================================ - # Note: we should not use torch.empty here like other attention backends, - # see discussions in https://github.com/vllm-project/vllm/pull/28182 - core_attn_out = torch.zeros( + # Uninitialized like the other attention backends: _forward_core + # writes every row up to num_actual_tokens. Padding rows behind it + # are only read by row-wise ops and never reach sampled outputs. + core_attn_out = torch.empty( (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), dtype=hidden_states.dtype, device=hidden_states.device, @@ -1491,6 +1498,9 @@ def _warmup_flashinfer_prefill_variants( output. A single sequence under ``use_cp="auto"`` takes the CP route wherever the heuristic can pick it; ``use_cp=False`` forces the other. """ + # g arrives in log space (from _warmup_prefill_kernels); FlashInfer + # takes exp(g). + g = torch.exp(g.to(torch.float32)) checkpoint_buffers: list[torch.Tensor | None] = [None] if self._exports_checkpoints: (state_checkpoints,) = current_workspace_manager().get_simultaneous( @@ -1766,7 +1776,7 @@ def _forward_core( head_k_dim=self.head_k_dim, head_v_dim=self.head_v_dim, apply_l2norm=True, - output_g_exp=False, + output_g_exp=self.gdn_prefill_backend == "flashinfer", ) query_non_spec = query_non_spec.unsqueeze(0) key_non_spec = key_non_spec.unsqueeze(0) @@ -1832,6 +1842,7 @@ def _forward_core( core_attn_out_decode = None # 2.3: Process the remaining part (prefill chunk, or non-spec decode-only) + prefill_in_place = False if attn_metadata.num_prefills > 0: # State indices, initial-state mask and cu_seqlens for the chunk # kernel are precomputed by the metadata builder (the prefill tail @@ -1841,10 +1852,31 @@ def _forward_core( prefill_has_initial_state = attn_metadata.prefill_has_initial_state assert prefill_state_indices is not None assert prefill_has_initial_state is not None - initial_state = ssm_state[prefill_state_indices] - initial_state[~prefill_has_initial_state, ...] = 0 - plan = attn_metadata.flashinfer_checkpoint_plan - if plan is None: + if self.gdn_prefill_backend == "flashinfer": + # Without spec decodes the prefill rows follow the peeled + # decodes contiguously, so FlashInfer writes them straight + # into core_attn_out. Spec batches scatter by token index. + prefill_out = None + if spec_sequence_masks is None: + prefill_out = core_attn_out[num_decode_tokens:num_actual_tokens] + core_attn_out_non_spec, last_recurrent_state = ( + self._forward_flashinfer_prefill( + query_non_spec, + key_non_spec, + value_non_spec, + g_non_spec, + beta_non_spec, + attn_metadata, + ssm_state=ssm_state, + conv_state=conv_state, + conv_input=conv_input_non_spec, + output=prefill_out, + ) + ) + prefill_in_place = prefill_out is not None + else: + initial_state = ssm_state[prefill_state_indices] + initial_state[~prefill_has_initial_state, ...] = 0 ( core_attn_out_non_spec, last_recurrent_state, @@ -1862,44 +1894,12 @@ def _forward_core( use_qk_l2norm_in_kernel=False, aiter_prefill_metadata=attn_metadata.aiter_prefill_metadata, ) - else: - # Only FlashInfer's prefill emits checkpoints. - (state_checkpoints,) = current_workspace_manager().get_simultaneous( - self._checkpoint_buffer_spec - ) - state_checkpoints = state_checkpoints[: plan.num_checkpoints] - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = fi_chunk_gated_delta_rule( - q=query_non_spec, - k=key_non_spec, - v=value_non_spec, - g=g_non_spec, - beta=beta_non_spec, - initial_state=initial_state, - output_final_state=True, - cu_seqlens=attn_metadata.prefill_query_start_loc, - use_qk_l2norm_in_kernel=False, - state_checkpoints=state_checkpoints, - checkpoint_cu_starts=plan.cu_starts, - checkpoint_every_n_tokens=plan.every_n_tokens, - ) - assert attn_metadata.checkpoint is not None - assert non_spec_query_start_loc is not None - self._checkpoint_exporter.export( - attn_metadata.checkpoint, - raw_qkv=conv_input_non_spec, - conv_state=conv_state, - recurrent_checkpoint=state_checkpoints, - recurrent_state=ssm_state, - cu_seqlens=non_spec_query_start_loc, - recurrent_rows=plan.rows, - ) # Init cache ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype) - if split_non_spec: + if split_non_spec and prefill_in_place: + core_attn_out[:num_decode_tokens] = core_attn_out_decode.squeeze(0) + elif split_non_spec: # Stitch the peeled decode outputs in front of the prefill # outputs (decode-first order). core_attn_out_non_spec = torch.cat( @@ -1936,9 +1936,81 @@ def _forward_core( ) elif spec_sequence_masks is not None: core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0) - else: + elif not prefill_in_place: core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0) + def _forward_flashinfer_prefill( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + attn_metadata: GDNAttentionMetadata, + *, + ssm_state: torch.Tensor, + conv_state: torch.Tensor, + conv_input: torch.Tensor | None, + output: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """FlashInfer chunked prefill of the batch's prefill rows. + + Writes the output into ``output`` when given and exports the + batch's internal prefill checkpoints. Returns the output and the + rows' final states. + """ + has_initial_state = attn_metadata.prefill_has_initial_state + assert has_initial_state is not None + initial_state = ssm_state[attn_metadata.prefill_state_indices] + initial_state[~has_initial_state, ...] = 0 + plan = attn_metadata.flashinfer_checkpoint_plan + if plan is None: + return fi_chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=attn_metadata.prefill_query_start_loc, + use_qk_l2norm_in_kernel=False, + output=output, + ) + + (state_checkpoints,) = current_workspace_manager().get_simultaneous( + self._checkpoint_buffer_spec + ) + state_checkpoints = state_checkpoints[: plan.num_checkpoints] + core_out, final_state = fi_chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=attn_metadata.prefill_query_start_loc, + use_qk_l2norm_in_kernel=False, + output=output, + state_checkpoints=state_checkpoints, + checkpoint_cu_starts=plan.cu_starts, + checkpoint_every_n_tokens=plan.every_n_tokens, + ) + assert attn_metadata.checkpoint is not None + assert attn_metadata.non_spec_query_start_loc is not None + assert conv_input is not None + self._checkpoint_exporter.export( + attn_metadata.checkpoint, + raw_qkv=conv_input, + conv_state=conv_state, + recurrent_checkpoint=state_checkpoints, + recurrent_state=ssm_state, + cu_seqlens=attn_metadata.non_spec_query_start_loc, + recurrent_rows=plan.rows, + ) + return core_out, final_state + def _forward_core_decode_aiter( self, qkvz: torch.Tensor, diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index dc512eee02d9..00a76899d7f1 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -49,6 +49,8 @@ class _QwenGDNWarmupConfig: state_stride_token: int state_dtype: torch.dtype fused_spec_decode: bool = False + # The FlashInfer prefill takes exp(g) straight from fused_post_conv_prep. + post_conv_g_exp: bool = False @property def conv_dim(self) -> int: @@ -141,6 +143,7 @@ def _qwen_gdn_warmup_config( fused_spec_decode=bool( getattr(layer, "enable_fused_gdn_spec_decode", False) ), + post_conv_g_exp=getattr(layer, "gdn_prefill_backend", None) == "flashinfer", ) if found_layer: @@ -235,7 +238,7 @@ def _warm_fused_post_conv_kernel( config.k, config.v, apply_l2norm=True, - output_g_exp=False, + output_g_exp=config.post_conv_g_exp, ) From c1b7dffa1457f4ffe7eafb08092eb2505a9651e4 Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Thu, 29 Oct 2026 19:01:26 +0800 Subject: [PATCH 122/406] [Bugfix][KV Offload] Refresh CPU LRU for GPU-hit prefixes in eager SimpleCPUOffload (#62444) Signed-off-by: whx-sjtu Co-authored-by: Cursor --- tests/v1/simple_kv_offload/test_scheduler.py | 102 +++++++++++++++++++ vllm/v1/simple_kv_offload/manager.py | 21 ++++ 2 files changed, 123 insertions(+) diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index dd2028c6eaca..69912a818814 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -31,6 +31,7 @@ from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager from vllm.v1.core.kv_cache_utils import ( + KVCacheBlock, get_request_block_hasher, init_none_hash, make_block_hash_with_group_id, @@ -1329,6 +1330,68 @@ def test_touched_blocks_survive_eviction() -> None: assert cached is not None, f"req_a block {i} should survive (was touched/MRU)" +def test_gpu_prefix_hit_refreshes_cpu_lru() -> None: + """Eager mode: a prefix that hits on GPU keeps its CPU copy at the MRU end.""" + fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False) + sched = fix.scheduler + + # Fill all 4 usable CPU blocks: req_a is LRU, req_b is MRU. + req_a = make_request(num_blocks=2) + req_b = make_request(num_blocks=2) + kv_a = _alloc_and_register(fix, req_a, 2) + kv_b = _alloc_and_register(fix, req_b, 2) + sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0) + sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0) + meta = sched.build_connector_meta( + make_scheduler_output( + {req_a.request_id: 2 * BLOCK_SIZE, req_b.request_id: 2 * BLOCK_SIZE}, + new_reqs={ + req_a.request_id: kv_a.get_block_ids(), + req_b.request_id: kv_b.get_block_ids(), + }, + ) + ) + simulate_store_completion(sched, meta.store_event) + + # req_a's prefix is reused and hits entirely on GPU. + req_a2 = Request( + request_id="req-a2", + prompt_token_ids=req_a.prompt_token_ids, + sampling_params=req_a.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256), + ) + assert sched.get_num_new_matched_tokens(req_a2, 2 * BLOCK_SIZE) == (0, False) + sched.update_state_after_alloc( + req_a2, KVCacheBlocks(blocks=([],)), num_external_tokens=0 + ) + + # Storing 2 new blocks evicts the 2 LRU blocks, which are now req_b's. + req_c = make_request(num_blocks=2) + kv_c = _alloc_and_register(fix, req_c, 2) + sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0) + meta2 = sched.build_connector_meta( + make_scheduler_output( + {req_c.request_id: 2 * BLOCK_SIZE}, + new_reqs={req_c.request_id: kv_c.get_block_ids()}, + ) + ) + simulate_store_completion(sched, meta2.store_event) + + cache_map = sched.cpu_block_pool.cached_block_hash_to_block + + def cpu_cached(req: Request) -> list[bool]: + return [ + cache_map.get_one_block(make_block_hash_with_group_id(h, 0)) is not None + for h in req.block_hashes[:2] + ] + + assert cpu_cached(req_a) == [True, True] + assert cpu_cached(req_b) == [False, False] + assert cpu_cached(req_c) == [True, True] + + # --------------------------------------------------------------------------- # Test 5: Preemption no CPU block leak # --------------------------------------------------------------------------- @@ -2966,6 +3029,45 @@ def test_finished_eager_store_does_not_duplicate_completed_partial_tail() -> Non assert len(finish_meta.store_cpu_blocks) == len(stored) +def test_gpu_prefix_hit_refreshes_cpu_mamba_boundary() -> None: + """A hybrid GPU prefix hit refreshes the CPU attention blocks and the Mamba + state at the hit boundary, not older Mamba states, once admitted.""" + block_size = 4 * BLOCK_SIZE + fix = _make_hybrid_attention_mamba_scheduler(block_size=block_size) + sched = fix.scheduler + cpu_pool = sched.cpu_block_pool + req = _make_cp_request(num_blocks=2, virtual_block_size=block_size) + other = _make_cp_request(num_blocks=2, virtual_block_size=block_size) + + def cache_on_cpu(request: Request, group_id: int) -> list[KVCacheBlock]: + blocks = cpu_pool.get_new_blocks(2) + cpu_pool.cache_full_blocks( + request=request, + blocks=blocks, + num_cached_blocks=0, + num_full_blocks=2, + block_size=sched.group_block_sizes[group_id], + kv_cache_group_id=group_id, + ) + return blocks + + attention, mamba = cache_on_cpu(req, 0), cache_on_cpu(req, 1) + newer = cache_on_cpu(other, 0) + cpu_pool.free_blocks([*attention, *mamba]) + cpu_pool.free_blocks(newer) + + lru_before = cpu_pool.free_block_queue.get_all_free_blocks() + assert sched.get_num_new_matched_tokens(req, 2 * block_size) == (0, False) + assert cpu_pool.free_block_queue.get_all_free_blocks() == lru_before + sched.update_state_after_alloc( + req, KVCacheBlocks(blocks=([], [])), num_external_tokens=0 + ) + + lru_order = cpu_pool.free_block_queue.get_all_free_blocks() + assert lru_order[-3:] == [mamba[1], attention[1], attention[0]] + assert lru_order.index(mamba[0]) < lru_order.index(newer[0]) + + def test_external_lookup_rejects_misaligned_local_prefix(caplog_vllm) -> None: scheduler_block_size = 4 * BLOCK_SIZE fix = _make_hybrid_attention_mamba_scheduler(block_size=scheduler_block_size) diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 6794412402dc..a877f0a53002 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -245,6 +245,9 @@ def __init__( self._pending_cpu_hits: dict[ str, tuple[tuple[list[KVCacheBlock], ...], int, int] ] = {} + # Eager mode: local prefix length from get_num_new_matched_tokens(), + # whose CPU copy update_state_after_alloc() refreshes. + self._pending_local_hits: dict[str, int] = {} # Store metadata self._lazy_mode = lazy_offload @@ -373,10 +376,16 @@ def get_num_new_matched_tokens( # is dropped first. if stale := self._pending_cpu_hits.pop(request.request_id, None): self._free_pending_cpu_hit(stale) + self._pending_local_hits.pop(request.request_id, None) if request.skip_reading_external_kv: return 0, False + if num_computed_tokens and not self._lazy_mode: + # Refreshed at admission: a request waiting for blocks repeats + # this lookup every step. + self._pending_local_hits[request.request_id] = num_computed_tokens + if num_computed_tokens % self.block_size != 0: # Transfers are whole-block copies, so an external suffix cannot # start in the middle of a destination block. When the GPU-side @@ -449,6 +458,16 @@ def update_state_after_alloc( num_stored_blocks=[0] * num_groups, ) + if num_local_tokens := self._pending_local_hits.pop(req_id, None): + # Eager CPU is inclusive: a GPU prefix hit also uses the CPU copy, + # but eager stores skip it, so move it to the MRU end here. + prefix_hit, _, _ = self.cpu_coordinator.find_longest_cache_hit( + request.block_hashes, num_local_tokens + ) + prefix_blocks = [b for grp in prefix_hit for b in grp if not b.is_null] + self.cpu_block_pool.touch(prefix_blocks) + self.cpu_block_pool.free_blocks(reversed(prefix_blocks)) + # Pop the CPU hit cached by get_num_new_matched_tokens(). The # found blocks were pinned there to survive LRU eviction in the window # between get_num_new_matched_tokens() and this matching call. @@ -1232,6 +1251,7 @@ def request_finished( pending = self._pending_cpu_hits.pop(req_id, None) if pending is not None: self._free_pending_cpu_hit(pending) + self._pending_local_hits.pop(req_id, None) # Handle load: defer cleanup if load is in-flight load_state = self._reqs_to_load.get(req_id) @@ -1435,6 +1455,7 @@ def reset(self) -> bool: for pending in self._pending_cpu_hits.values(): self._free_pending_cpu_hit(pending) self._pending_cpu_hits.clear() + self._pending_local_hits.clear() self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks) for transfer in self._pending_finished_stores: From 3c39dda2a5da9e4b71b5eb61db686bf669374569 Mon Sep 17 00:00:00 2001 From: Sherif Waly Date: Thu, 29 Oct 2026 11:06:03 +0000 Subject: [PATCH 123/406] [Performance][KV Offload] Avoid zero-filling the SimpleCPUOffload host KV pool (#64345) Signed-off-by: Sherif Waly Co-authored-by: Mistral Vibe --- vllm/v1/simple_kv_offload/worker.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index 6a43c67ed133..af3f3312f5eb 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -244,7 +244,8 @@ def _init_cpu_mode( # Allocate non-pinned first, then pin via cudaHostRegister to # bypass PyTorch's CUDACachingHostAllocator which rounds up to # the next power of 2 (e.g. 100 GB -> 128 GB). - tensor = torch.zeros(cpu_shape, dtype=gpu_tensor.dtype, device="cpu") + # CPU blocks are loaded only after a completed GPU-to-CPU store. + tensor = torch.empty(cpu_shape, dtype=gpu_tensor.dtype, device="cpu") if pin_memory: pin_tensor(tensor) self.cpu_kv_caches[name] = tensor From 3765000a372042125f9a17c203558368a21f0803 Mon Sep 17 00:00:00 2001 From: namgyu-youn Date: Thu, 29 Oct 2026 20:32:08 +0900 Subject: [PATCH 124/406] [Perf][Qwen3.8] Add SM100 skinny GEMM plans for TP=1 (#64482) Signed-off-by: namgyu-youn --- .../qwen4_exp/nvidia/low_latency_gemm.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py index ea56750702fc..679f145cc0fd 100644 --- a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -118,7 +118,7 @@ } -# B200 plans; the winning configs differ from the B300 table. +# B200 plans for TP=4 and TP=1; the winning configs differ from the B300 table. QWEN4_EXP_SM100_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = { # GDN fused B/A projection, TP=4. (24, 2560): { @@ -165,6 +165,28 @@ (62080, 2560): { 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4), }, + # GDN fused B/A projection, TP=1. + (96, 2560): { + 1: SkinnyGemmConfig(1, 128, 4, k_unroll=2, vector_width=2, static_k=2560), + }, + # Router, TP=1. + (512, 2560): { + 1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, vector_width=2, static_k=2560), + }, + # Shared-expert fused gate/up projection, TP=1. + (1280, 2560): { + 1: SkinnyGemmConfig(1, 32, 1, k_unroll=2, static_k=2560), + }, + # Shared-expert down projection, TP=1. + (2560, 640): { + 1: SkinnyGemmConfig(1, 160, 4, k_unroll=2, vector_width=4), + 2: SkinnyGemmConfig(2, 160, 4, vector_width=4), + }, + # GDN and QSA output projections, TP=1. + (2560, 6144): { + 1: SkinnyGemmConfig(1, 256, 2, k_unroll=2, static_k=6144), + 2: SkinnyGemmConfig(2, 128, 2, static_k=6144), + }, } From ee1056b3179adfd4523dfdceb87b0a079801b6f4 Mon Sep 17 00:00:00 2001 From: freyfwt <110148731+freyfwt@users.noreply.github.com> Date: Thu, 29 Oct 2026 19:38:42 +0800 Subject: [PATCH 125/406] [Doc] Fix MiniMax-M3 docstring cross-references (#64474) Signed-off-by: freyfwt Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: vllm-agent --- vllm/models/minimax_m3/amd/mono/kernels/index_score.py | 4 ++-- vllm/models/minimax_m3/amd/mono/runner.py | 2 +- .../minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/models/minimax_m3/amd/mono/kernels/index_score.py b/vllm/models/minimax_m3/amd/mono/kernels/index_score.py index aaad81f44796..ac95896137e4 100644 --- a/vllm/models/minimax_m3/amd/mono/kernels/index_score.py +++ b/vllm/models/minimax_m3/amd/mono/kernels/index_score.py @@ -373,8 +373,8 @@ def score_task(tk, b, page): def q_frag_from(iq, heads=ONE_INDEX_HEAD): - """(``q_frag``, ``q_frag_head``) reading index q from memory: [token][head] - [HEAD_DIM] bf16, ``heads.count`` heads a token.""" + """(`q_frag`, `q_frag_head`) reading index q from memory: + `[token][head][HEAD_DIM]` bf16, `heads.count` heads a token.""" def q_frag_head(tk, s, head): lane = fx.thread_idx.x % 64 diff --git a/vllm/models/minimax_m3/amd/mono/runner.py b/vllm/models/minimax_m3/amd/mono/runner.py index 048a4aaba6a8..89b9b84afd8b 100644 --- a/vllm/models/minimax_m3/amd/mono/runner.py +++ b/vllm/models/minimax_m3/amd/mono/runner.py @@ -235,7 +235,7 @@ def run_sparse_layer( positions: torch.Tensor, res: torch.Tensor, ) -> torch.Tensor: - """Sparse layer ``i``: (ars[i % 2][:n], res) -> (ars[(i + 1) % 2][:n], the + """Sparse layer `i`: (`ars[i % 2][:n]`, res) -> (`ars[(i + 1) % 2][:n]`, the returned residual).""" n = res.shape[0] block_table, seq_lens, slots, q_len = tables diff --git a/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py index 390dbc075c0f..da68e37f49ca 100644 --- a/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py +++ b/vllm/models/minimax_m3/amd/ops/nvfp4_index_prefill/kernel_shared.py @@ -79,7 +79,7 @@ def _fp8_pair(w, sv): def _mfma_fp8_32x32x64(a8, b8, acc16): """v_mfma_f32_32x32x64_f8f6f4 on e4m3; lane (r, h) gets - D[8 (i // 4) + 4 h + i % 4][r].""" + `D[8 (i // 4) + 4 h + i % 4][r]`.""" return Vec( fx.rocdl.mfma_scale_f32_32x32x64_f8f6f4( Vec.make_type(16, fx.Float32), [a8, b8, acc16, 0, 0, 0, 0, 0, 0] From 5082c45853ca93ae809c58255e6126cfe20c415f Mon Sep 17 00:00:00 2001 From: Rex Mingjiang Lee Date: Thu, 29 Oct 2026 19:44:30 +0800 Subject: [PATCH 126/406] [Docs] Add vllm-iluvatar to custom_op.md (#62500) Signed-off-by: mingjiang.li --- docs/design/custom_op.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design/custom_op.md b/docs/design/custom_op.md index f99c8043fa04..bedefbcdf9b2 100644 --- a/docs/design/custom_op.md +++ b/docs/design/custom_op.md @@ -264,7 +264,7 @@ Currently, thanks to [vLLM's hardware-plugin mechanism](./plugin_system.md), the - **Official device plugins:** [vllm-ascend](https://github.com/vllm-project/vllm-ascend) (for Huawei Ascend NPU), [vllm-spyre](https://github.com/vllm-project/vllm-spyre) (for Spyre), [vllm-gaudi](https://github.com/vllm-project/vllm-gaudi) (for Intel Gaudi), [vllm-neuron](https://github.com/vllm-project/vllm-neuron) (for AWS Neuron), [vllm-meta](https://github.com/vllm-project/vllm-metal) (for Apple Silicon), etc. -- **Non-official device plugins:** [vllm-metax](https://github.com/MetaX-MACA/vLLM-metax) (for MetaX GPU), [vllm-kunlun](https://github.com/baidu/vLLM-Kunlun) (for Baidu Kunlun XPU), [vllm-musa](https://github.com/MooreThreads/vllm-musa) (for Moore Threads GPU), etc. +- **Non-official device plugins:** [vllm-metax](https://github.com/MetaX-MACA/vLLM-metax) (for MetaX GPU), [vllm-kunlun](https://github.com/baidu/vLLM-Kunlun) (for Baidu Kunlun XPU), [vllm-musa](https://github.com/MooreThreads/vllm-musa) (for Moore Threads GPU), [vllm-iluvatar](https://github.com/Deep-Spark/vllm-iluvatar) (for Iluvatar CoreX GPU), etc. In this case, `CustomOp` can enable these hardware manufacturers to seamlessly replace vLLM's operations with their deep-optimized kernels for specific devices at runtime, by just registering an OOT `CustomOp` and implementing the `forward_oot()` method. From bfa4d6d2ad7f39afd7c9baf8aa6435303f252f08 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 29 Oct 2026 20:06:55 +0800 Subject: [PATCH 127/406] [KV Connector] Support sleep mode with MooncakeStoreConnector (#59934) Signed-off-by: aoshen02 Signed-off-by: Dakai An Co-authored-by: Claude Opus 5.5 Co-authored-by: Dakai An Co-authored-by: Dakai An <77474977+andakai@users.noreply.github.com> --- tests/v1/kv_connector/unit/test_config.py | 18 +++ .../unit/test_mooncake_store_worker.py | 113 ++++++++++++++++++ .../v1/mooncake/store/connector.py | 25 +++- .../kv_connector/v1/mooncake/store/worker.py | 22 ++++ 4 files changed, 177 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index 2c413708b6aa..f9164dfa9e62 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/tests/v1/kv_connector/unit/test_config.py @@ -3,6 +3,8 @@ """Tests for KV cache offloading configuration.""" +import sys + import pytest from vllm.config import ( @@ -185,6 +187,22 @@ def test_sleep_mode_refuses_unsupported_kv_connector( ) +def test_mooncake_store_supports_sleep_with_explicit_rdma(monkeypatch, tmp_path): + """Disabling auto-discovery keeps explicit RDMA despite NVLink overrides.""" + monkeypatch.setitem(sys.modules, "mooncake", None) + config_path = tmp_path / "mooncake.json" + config_path.write_text('{"protocol": "rdma", "device_name": "mlx5_0"}') + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(config_path)) + monkeypatch.setenv("MC_MS_AUTO_DISC", "0") + monkeypatch.delenv("MC_USE_TENT", raising=False) + monkeypatch.delenv("MC_USE_TEV1", raising=False) + monkeypatch.setenv("MC_FORCE_MNNVL", "1") + monkeypatch.setenv("MC_INTRANODE_NVLINK", "1") + config = KVTransferConfig(kv_connector="MooncakeStoreConnector", kv_role="kv_both") + connector_cls = KVConnectorFactory.get_connector_class(config) + assert connector_cls.supports_sleep_mode(config) + + def test_kv_connector_allows_other_alloc_conf(monkeypatch): """Other PYTORCH_CUDA_ALLOC_CONF values must not be rejected.""" monkeypatch.setenv( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index ee3961171b88..20ca558f12a4 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -3720,6 +3720,8 @@ def _make_bare_worker( worker.kv_recv_threads = [] worker.num_recv_threads = 1 worker.recv_request_queue = queue.Queue() + worker._device_regions = {} + worker._released_regions = set() worker.finished_store_req = set() worker.tp_size = 1 worker.store_tp_size = None @@ -4774,6 +4776,117 @@ def test_register_kv_caches_separate_head_groups(): worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes) +# --------------------------------------------------------------------------- +# Sleep mode: release / restore of the KV cache registration +# --------------------------------------------------------------------------- + + +class _FakeRegistrationStore: + """Keeps the registered regions like the store; `fail` fails one region.""" + + def __init__(self): + self.registered: dict[int, int] = {} + self.fail: tuple[str, int] | None = None + + def register_buffer(self, addr: int, length: int) -> int: + if self.fail == ("register", addr) or addr in self.registered: + return -1 + self.registered[addr] = length + return 0 + + def unregister_buffer(self, addr: int) -> int: + if self.fail == ("unregister", addr) or addr not in self.registered: + return -1 + del self.registered[addr] + return 0 + + +def _registered_worker(kind: str = "dense"): + """A worker that registered its KV caches with a fake store, plus the device + and host regions it registered.""" + worker = _make_bare_worker(num_gpu_blocks=2) + if kind == "dense": # two device regions + worker.token_dbs[0].store_layout = MagicMock() + caches = {n: torch.zeros(2, 4) for n in ("layer0", "__cross_layer__")} + else: # one host region and one device region + from vllm.v1.kv_cache_interface import KVCacheTensor + + spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float16 + ) + config = KVCacheConfig( + num_blocks=2, + kv_cache_tensors=[ + KVCacheTensor( + size=4 * 32, + layers=["source"], + layer_stride=4 * 32, + block_stride=32, + host_resident=True, + ), + KVCacheTensor( + size=2 * 32, layers=["layer0"], layer_stride=2 * 32, block_stride=32 + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["source"], spec, host_resident=True), + KVCacheGroupSpec(["layer0"], spec), + ], + hisparse_host_num_blocks=4, + ) + worker._kv_cache_config = config + worker._kv_cache_groups = list(config.transfer_groups) + worker.token_dbs = [ + ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16) + for _ in worker._kv_cache_groups + ] + caches = { + "source": torch.zeros(4, 32, dtype=torch.uint8), + "layer0": torch.zeros(2, 32, dtype=torch.uint8), + } + worker.store = _FakeRegistrationStore() + _register_with_mocked_threads(worker, caches) + worker.kv_send_thread = SimpleNamespace(request_queue=queue.Queue()) + host = {c.data_ptr(): c.nbytes for n, c in caches.items() if n == "source"} + device = {c.data_ptr(): c.nbytes for n, c in caches.items() if n != "source"} + return worker, device, host + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("kind", ["dense", "hisparse"]) +def test_sleep_cycle_releases_and_restores_device_regions(kind: str): + """Asleep, only host regions stay registered; waking registers the same + device addresses and lengths again. Both calls are idempotent.""" + worker, device, host = _registered_worker(kind) + assert device and worker.store.registered == device | host + worker.restore_kv_caches() # a wake-up without a sleep + for _ in range(2): + worker.unregister_kv_caches() + worker.unregister_kv_caches() # sleeping after a discard + assert worker.store.registered == host + worker.restore_kv_caches() + worker.restore_kv_caches() # a retried wake-up + assert worker.store.registered == device | host + + +@pytest.mark.cpu_test +@pytest.mark.parametrize("op", ["unregister", "register"]) +def test_failed_store_call_raises_and_retry_converges(op: str): + """A failed region keeps its state and is the only one a retry touches.""" + worker, device, _ = _registered_worker() + if op == "register": + worker.unregister_kv_caches() + step = worker.restore_kv_caches if op == "register" else worker.unregister_kv_caches + _, failing = device + worker.store.fail = (op, failing) + with pytest.raises(RuntimeError, match=f"{op}_buffer failed"): + step() + assert (failing in worker.store.registered) == (op == "unregister") + worker.store.fail = None + step() + assert worker.store.registered == (device if op == "register" else {}) + + # --------------------------------------------------------------------------- # Dual-mode (embedded / standalone-store) config validation tests # --------------------------------------------------------------------------- diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index a1e28b97d06a..a91167234da4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -11,6 +11,7 @@ enabling prefix caching via hash-based deduplication. """ +import os from collections import Counter from collections.abc import Iterable, Sequence from contextlib import AbstractContextManager @@ -38,6 +39,7 @@ PromMetric, PromMetricT, ) +from vllm.distributed.mooncake_store import MooncakeStoreConfig from vllm.forward_context import ForwardContext from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata @@ -129,7 +131,20 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): @classmethod def supports_sleep_mode(cls, kv_transfer_config: KVTransferConfig) -> bool: - return False + config = MooncakeStoreConfig.load_from_config() + if ( + config.protocol != "rdma" + or not config.device_name + or os.getenv("MC_MS_AUTO_DISC", "0") != "0" + or any(v in os.environ for v in ("MC_USE_TENT", "MC_USE_TEV1")) + ): + raise ValueError( + "MooncakeStoreConnector sleep requires protocol='rdma', " + "device_name set, MC_MS_AUTO_DISC=0 or unset, and " + "MC_USE_TENT/MC_USE_TEV1 unset. " + "Sleep with NVLink Fabric is not supported yet." + ) + return True @staticmethod def _validate_kv_cache_config( @@ -342,6 +357,14 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): assert self.connector_worker is not None self.connector_worker.register_kv_caches(kv_caches) + def unregister_kv_caches(self) -> None: + assert self.connector_worker is not None + self.connector_worker.unregister_kv_caches() + + def restore_kv_caches(self) -> None: + assert self.connector_worker is not None + self.connector_worker.restore_kv_caches() + def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None: assert self.connector_worker is not None metadata = self._get_connector_metadata() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 27bbd80a4f65..42a97001ff5e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1492,6 +1492,10 @@ def __init__( self.kv_recv_threads: list[KVCacheStoreRecvingThread] = [] self.num_recv_threads = max(1, envs.VLLM_MOONCAKE_LOAD_RECV_THREADS) self.recv_request_queue: queue.Queue[ReqMeta] = queue.Queue() + # Device KV regions (addr -> len) registered with the store. + self._device_regions: dict[int, int] = {} + # Device regions unregistered for sleep and not yet registered again. + self._released_regions: set[int] = set() self.finished_store_req: set[str] = set() self._kv_connector_stats_lock = threading.Lock() self.kv_connector_stats = MooncakeStoreConnectorStats() @@ -1774,6 +1778,8 @@ def _layer_views(v: torch.Tensor | list[torch.Tensor]) -> list[torch.Tensor]: region_len, ret, ) + else: + self._device_regions[base_addr] = region_len for group, db in zip(self._kv_cache_groups, self.token_dbs, strict=True): # In a hybrid model each KV cache group registers only its own # layers. A group with no layers on this PP rank registers all. @@ -1842,6 +1848,8 @@ def _layer_views(v: torch.Tensor | list[torch.Tensor]) -> list[torch.Tensor]: f"{base_addr:#x} len {region_len} ({mem_kind}): {ret}" ) registered_buffers[base_addr] = region_len + if not is_host_resident: + self._device_regions[base_addr] = region_len elif registered_len != region_len: raise ValueError( f"KV cache views at {base_addr:#x} expose inconsistent " @@ -1942,6 +1950,20 @@ def _layer_views(v: torch.Tensor | list[torch.Tensor]) -> list[torch.Tensor]: "Started %d Mooncake KV-load receive thread(s)", self.num_recv_threads ) + def unregister_kv_caches(self) -> None: + """Unregister device KV buffers after the engine has drained transfers.""" + for addr in self._device_regions.keys() - self._released_regions: + if self.store.unregister_buffer(addr) != 0: + raise RuntimeError(f"Mooncake unregister_buffer failed for {addr:#x}") + self._released_regions.add(addr) + + def restore_kv_caches(self) -> None: + """Register the released regions again: same addresses, new pages.""" + for addr in list(self._released_regions): + if self.store.register_buffer(addr, self._device_regions[addr]) != 0: + raise RuntimeError(f"Mooncake register_buffer failed for {addr:#x}") + self._released_regions.discard(addr) + def start_load_kv(self, metadata: MooncakeStoreConnectorMetadata): """Issue async loads. From d8e6fbe461f54ec17274fc49000e94d332347e6a Mon Sep 17 00:00:00 2001 From: Andrew Sansom Date: Thu, 29 Oct 2026 07:30:28 -0500 Subject: [PATCH 128/406] [Bugfix] Consume prompt embedding sentinel IDs before model execution (#64367) Signed-off-by: Andrew Sansom Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/features/prompt_embeds.md | 2 + ...test_chat_completion_with_prompt_embeds.py | 56 +++++++++++++++++++ tests/renderers/test_hf.py | 45 ++++++++++++++- vllm/renderers/hf.py | 2 + 4 files changed, 104 insertions(+), 1 deletion(-) diff --git a/docs/features/prompt_embeds.md b/docs/features/prompt_embeds.md index 14b4e56afa5c..4721294a16c9 100644 --- a/docs/features/prompt_embeds.md +++ b/docs/features/prompt_embeds.md @@ -59,6 +59,8 @@ Prompt embeddings can be included as content parts in chat messages, interleaved Each `prompt_embeds` content part contains a `data` field with a base64-encoded `torch.Tensor` of shape `(num_tokens, hidden_size)`. Multiple `prompt_embeds` parts can appear in any message, in any position relative to text parts. The server expands each part into the correct number of placeholder tokens during chat template rendering, then splices the pre-computed embeddings into the model's input at the corresponding positions. +For multimodal models, `prompt_embeds` parts can also be mixed with image content parts. Their placeholder token IDs are replaced with safe dummy IDs before model execution; the supplied embeddings replace the dummy embeddings at those positions. + Unlike the Completions API, a `prompt_embeds` content part should encode **only** the content, not a templated conversation. The server wraps the chat template around the embedded content at request time, the same way it would for a plain text `content` string. Embedding a full templated conversation here would double-apply the template and produce incorrect inputs to the model. Assistant messages can also pass a single `prompt_embeds` part as `reasoning` instead of a string: diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py index 80b512699a33..0696b3ac70d5 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py @@ -12,8 +12,11 @@ import pytest_asyncio import torch from openai import BadRequestError +from PIL import Image +from transformers import AutoConfig from tests.utils import VLLM_PATH, RemoteOpenAIServer +from vllm.multimodal.utils import encode_image_url from vllm.platforms import current_platform MODEL_NAME = "facebook/opt-125m" @@ -71,6 +74,59 @@ def _encode_embeds(embeds: torch.Tensor) -> str: return base64.b64encode(buf.getvalue()).decode("utf-8") +@pytest.mark.asyncio +async def test_prompt_embeds_with_image() -> None: + """An unpadded vocabulary must accept mixed embeddings and image inputs.""" + model = "HuggingFaceTB/SmolVLM-256M-Instruct" + config = AutoConfig.from_pretrained(model) + embeds = _encode_embeds( + torch.zeros(1, config.text_config.hidden_size, dtype=SERVER_DTYPE) + ) + with RemoteOpenAIServer( + model, + [ + "--enable-prompt-embeds", + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--gpu-memory-utilization", + "0.2", + "--enforce-eager", + ], + ) as remote_server: + async with remote_server.get_async_client() as client: + response = await client.chat.completions.create( + model=model, + temperature=0, + max_tokens=1, + messages=[ + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": embeds}, + { + "type": "image_url", + "image_url": { + "url": encode_image_url(Image.new("RGB", (32, 32))) + }, + }, + {"type": "text", "text": "Describe the image."}, + ], + } + ], + ) + assert response.choices[0].finish_reason in ("length", "stop") + assert response.usage is not None + assert response.usage.completion_tokens == 1 + with pytest.raises(BadRequestError, match="out of vocabulary"): + await client.completions.create( + model=model, + prompt=[config.text_config.vocab_size + 1], + max_tokens=1, + ) + + @pytest.fixture(scope="module") def prompt_embeds_b64(hf_runner) -> list[str]: """Pre-compute embeddings for two short prompts and return as base64.""" diff --git a/tests/renderers/test_hf.py b/tests/renderers/test_hf.py index 491dad801de3..77e114188b0a 100644 --- a/tests/renderers/test_hf.py +++ b/tests/renderers/test_hf.py @@ -3,16 +3,22 @@ import jinja2 import pytest +import torch -from vllm.config import ModelConfig +from vllm.config import ModelConfig, VllmConfig from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.exceptions import VLLMValidationError +from vllm.inputs.engine import MultiModalInput +from vllm.multimodal.inputs import MultiModalKwargsItems, PlaceholderRange from vllm.renderers.chat_utils import ConversationMessage, load_chat_template from vllm.renderers.hf import ( + HfRenderer, + _build_prompt_embeds_updates, _consolidate_system_messages, _convert_developer_to_system, _detect_content_format, _detect_developer_role_support, + _ensure_prompt_embeds_placeholder_token, _get_hf_base_chat_template_params, _template_error_reason, _try_extract_ast, @@ -83,6 +89,43 @@ ASSISTANT_MESSAGE_TO_CONTINUE = {"role": "assistant", "content": "The capital of"} +def test_apply_prompt_embeds_keeps_engine_token_ids_in_vocabulary() -> None: + """Consume OOV sentinels without changing text IDs or embedding offsets.""" + model = "HuggingFaceTB/SmolVLM-256M-Instruct" + model_config = ModelConfig(model, enable_prompt_embeds=True) + renderer = HfRenderer(VllmConfig(model_config=model_config), get_tokenizer(model)) + try: + assert renderer.tokenizer is not None + sentinel = _ensure_prompt_embeds_placeholder_token(renderer.tokenizer) + vocab_size = model_config.get_vocab_size() + assert sentinel >= vocab_size + tensors = [ + torch.zeros(length, model_config.get_hidden_size()) for length in (2, 1) + ] + engine_input: MultiModalInput = { + "type": "multimodal", + "prompt_token_ids": [1, sentinel, sentinel, 2, sentinel, 3], + "mm_kwargs": MultiModalKwargsItems(), + "mm_hashes": {}, + "mm_placeholders": {}, + } + + renderer._apply_prompt_embeds_to_engine_input( + engine_input, tensors, _build_prompt_embeds_updates(tensors, sentinel) + ) + + assert all( + 0 <= token_id < vocab_size for token_id in engine_input["prompt_token_ids"] + ) + assert engine_input["prompt_token_ids"] == [1, 0, 0, 2, 0, 3] + assert engine_input["mm_placeholders"]["prompt_embeds"] == [ + PlaceholderRange(offset=1, length=2), + PlaceholderRange(offset=4, length=1), + ] + finally: + renderer.shutdown() + + def test_load_chat_template(): # Testing chatml template template_content = load_chat_template(chat_template=chatml_jinja_path) diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index fa8a1db71689..b1e096cd2230 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -1479,6 +1479,7 @@ def _apply_prompt_embeds_to_engine_input( for tensor, (start, length) in zip( prompt_embeds_tensors, positions, strict=True ): + token_ids[start : start + length] = [0] * length pe_kwargs_items.append( MultiModalKwargsItem( { @@ -1500,6 +1501,7 @@ def _apply_prompt_embeds_to_engine_input( PlaceholderRange(offset=start, length=length, is_embed=None) ) + engine_input["prompt_token_ids"] = token_ids cast( MultiModalKwargsItems[MultiModalKwargsItem | None], engine_input["mm_kwargs"], From fb7def9ef9fa27fb794f655fb993a41e0b9bfe0c Mon Sep 17 00:00:00 2001 From: Felix Li Date: Thu, 29 Oct 2026 20:39:55 +0800 Subject: [PATCH 129/406] [ROCm][Perf] Reuse compiled FlyDSL host launchers (#64008) Signed-off-by: Felix Li Co-authored-by: OpenAI Codex Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> --- vllm/models/deepseek_v4/amd/mono/runner.py | 12 ++++++++---- vllm/models/deepseek_v41/amd/mono/runner.py | 17 +++++++++++------ .../specialized_decode/mono/attention/kda.py | 6 ++++-- .../amd/specialized_decode/mono/layer.py | 6 ++++-- .../attention/ops/flydsl_kernels/__init__.py | 18 ++++++++++++++++++ .../attention/ops/flydsl_turboquant_decode.py | 4 +++- .../attention/ops/flydsl_ultraquant_decode.py | 4 +++- 7 files changed, 51 insertions(+), 16 deletions(-) diff --git a/vllm/models/deepseek_v4/amd/mono/runner.py b/vllm/models/deepseek_v4/amd/mono/runner.py index 47bea49be826..0468f0acb84c 100644 --- a/vllm/models/deepseek_v4/amd/mono/runner.py +++ b/vllm/models/deepseek_v4/amd/mono/runner.py @@ -21,6 +21,8 @@ import torch +from vllm.v1.attention.ops.flydsl_kernels import _run_compiled + from .common.plan import BLOCKS from .dims import ( EXPERTS, @@ -359,7 +361,8 @@ def forward( comp_block_table.stride(0), topk_indices, ) - k1( + _run_compiled( + k1, residual.data_ptr(), x.data_ptr(), post_mix.data_ptr(), @@ -400,10 +403,11 @@ def forward( scratch.data_ptr(), self.epoch.data_ptr(), self.tl[0].data_ptr() if self.tl is not None else 0, - stream=st, + st, ) select(self.normed[:M], self.qr[:M]) - k2( + _run_compiled( + k2, self.q.data_ptr(), swa_cache.data_ptr(), swa_cache.stride(0), @@ -458,6 +462,6 @@ def forward( self.rank, self.epoch.data_ptr(), self.tl[1].data_ptr() if self.tl is not None else 0, - stream=st, + st, ) return outs diff --git a/vllm/models/deepseek_v41/amd/mono/runner.py b/vllm/models/deepseek_v41/amd/mono/runner.py index 6cf144db52f5..46743f63cf34 100644 --- a/vllm/models/deepseek_v41/amd/mono/runner.py +++ b/vllm/models/deepseek_v41/amd/mono/runner.py @@ -22,6 +22,8 @@ import torch +from vllm.v1.attention.ops.flydsl_kernels import _run_compiled + from .attention.front import FRONT_POINTS from .attention.plan import HEAD_DIM, HIDDEN, KEYS, Q_RANK, WINDOW, Dims from .common.arch import GFX942 @@ -386,7 +388,8 @@ def forward( else: comp, comp_stride, comp_block = self._dummy, 0, 1 bt, bt_stride, topk = self._dummy, 0, self._dummy - k1( + _run_compiled( + k1, residual.data_ptr(), x.data_ptr(), post_mix.data_ptr(), @@ -427,7 +430,7 @@ def forward( self.scratch(M).data_ptr(), self.epoch.data_ptr(), self.tl1.data_ptr() if self.timeline else 0, - stream=st, + st, ) if select is not None: select(self.normed[:M], self.qr[:M]) @@ -450,7 +453,8 @@ def forward( comp_block_table=bt, comp_block=comp_block, ) - k2( + _run_compiled( + k2, self.q.data_ptr(), swa_cache.data_ptr(), swa_cache.stride(0), @@ -505,7 +509,7 @@ def forward( self.rank, self.epoch.data_ptr(), self.tl2.data_ptr() if self.timeline else 0, - stream=st, + st, ) return outs @@ -533,7 +537,8 @@ def ffn( if outs is None: outs = self._outs(M, residual) out, res_out, post_out, comb_out, pre_out = outs - self.ffn_kernel(M, w.gate_w.shape[0], topk)( + _run_compiled( + self.ffn_kernel(M, w.gate_w.shape[0], topk), part.data_ptr(), residual.data_ptr(), post_mix.data_ptr(), @@ -564,6 +569,6 @@ def ffn( self.rank, self.epoch.data_ptr(), self.tl3.data_ptr() if self.timeline else 0, - stream=torch.cuda.current_stream(), + torch.cuda.current_stream(), ) return outs diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py index 43a624f7c723..f9b038496086 100644 --- a/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py +++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/attention/kda.py @@ -73,6 +73,7 @@ publish, sreg, ) +from vllm.v1.attention.ops.flydsl_kernels import _run_compiled HIDDEN = 7168 NA = 56 # attn CTAs @@ -987,7 +988,8 @@ def kda_pre( layer (tags are a layer's); ``epoch`` a device int32 bumped once a step and ``layer`` this launch's slot below ``TAG_SLOTS``.""" f = build(key) - f( + _run_compiled( + f, prefix.data_ptr(), ptr(delta), blocks.data_ptr(), @@ -1017,5 +1019,5 @@ def kda_pre( epoch.data_ptr(), health.data_ptr(), alarm, - stream=torch.cuda.current_stream(), + torch.cuda.current_stream(), ) diff --git a/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py b/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py index c89c2b309bcd..9b56e4ddf5fc 100644 --- a/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py +++ b/vllm/models/kimi_k3/amd/specialized_decode/mono/layer.py @@ -64,6 +64,7 @@ ) from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import gemv as gemv from vllm.models.kimi_k3.amd.specialized_decode.mono.stages import moe as k2b +from vllm.v1.attention.ops.flydsl_kernels import _run_compiled O0, O1 = NA, BLOCKS # o_proj CTAs: every CTA past AttnRes's O_N = O1 - O0 # 200 @@ -540,7 +541,8 @@ def k2_launch( """Host launcher: ``core`` [S, 1536] (o_proj's input), ``prefix`` updated in place, ``out`` the MoE output [S, 7168]; MoE weights as ``moe``.""" f = build(key) - f( + _run_compiled( + f, core.data_ptr(), w_o.data_ptr(), prefix.data_ptr(), @@ -569,5 +571,5 @@ def k2_launch( layer, health.data_ptr(), alarm, - stream=torch.cuda.current_stream(), + torch.cuda.current_stream(), ) diff --git a/vllm/v1/attention/ops/flydsl_kernels/__init__.py b/vllm/v1/attention/ops/flydsl_kernels/__init__.py index 208f01a7cb5e..74addfc33e6d 100644 --- a/vllm/v1/attention/ops/flydsl_kernels/__init__.py +++ b/vllm/v1/attention/ops/flydsl_kernels/__init__.py @@ -1,2 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +def _run_compiled(launcher, *args) -> None: + """Reuse the compiled callable of a factory-style host launcher.""" + import flydsl.compiler as flyc + from flydsl.utils import env + + if env.compile.compile_only: + flyc.compile(launcher, *args) + return + + compiled = getattr(launcher, "_cf", None) + if compiled is not None: + compiled(*args) + return + + # flyc.compile executes once; do not launch its result again on this call. + launcher._cf = flyc.compile(launcher, *args) diff --git a/vllm/v1/attention/ops/flydsl_turboquant_decode.py b/vllm/v1/attention/ops/flydsl_turboquant_decode.py index e7e9ab3b9944..25523091bd55 100644 --- a/vllm/v1/attention/ops/flydsl_turboquant_decode.py +++ b/vllm/v1/attention/ops/flydsl_turboquant_decode.py @@ -29,6 +29,7 @@ from vllm.logger import init_logger from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.flydsl_kernels import _run_compiled logger = init_logger(__name__) @@ -698,7 +699,8 @@ def flydsl_turboquant_decode_attention( num_partitions * tile_groups_per_partition * kv_compute_block, worst_case_max_seq_len, ) - launch( + _run_compiled( + launch, segm_out, segm_sum, segm_max, diff --git a/vllm/v1/attention/ops/flydsl_ultraquant_decode.py b/vllm/v1/attention/ops/flydsl_ultraquant_decode.py index d2da66ffc8c4..022fff8a1243 100644 --- a/vllm/v1/attention/ops/flydsl_ultraquant_decode.py +++ b/vllm/v1/attention/ops/flydsl_ultraquant_decode.py @@ -15,6 +15,7 @@ from vllm.logger import init_logger from vllm.triton_utils import triton +from vllm.v1.attention.ops.flydsl_kernels import _run_compiled logger = init_logger(__name__) @@ -773,7 +774,8 @@ def flydsl_ultraquant_decode_attention( num_partitions * tile_groups_per_partition * kv_compute_block, worst_case_max_seq_len, ) - launch( + _run_compiled( + launch, segm_out, segm_sum, segm_max, From a1e212757c409ce56a15109a7b1342eef4d814b1 Mon Sep 17 00:00:00 2001 From: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:46:35 -0400 Subject: [PATCH 130/406] [Bugfix][Rust Frontend] Rust vllm-bench accepts --max-concurrency 0 and hangs forever (Python rejects it since #54887) (#64390) Signed-off-by: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com> Co-authored-by: vllm-agent --- rust/src/bench/src/config.rs | 65 ++++++++++++++++++++++++++++++ rust/src/bench/src/mm_processor.rs | 3 ++ rust/src/bench/src/sweep.rs | 31 +++++++++++++- 3 files changed, 97 insertions(+), 2 deletions(-) diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 3e180cfb3085..ed8f8f7d36cd 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -246,6 +246,17 @@ impl BenchConfig { )); } + if args.max_concurrency == Some(0) { + return Err(BenchError::Config( + "--max-concurrency must be greater than 0".into(), + )); + } + if args.multi_turn_concurrency == Some(0) { + return Err(BenchError::Config( + "--multi-turn-concurrency must be greater than 0".into(), + )); + } + let base_url = args.resolve_base_url(); let api_url = args.resolve_api_url(); @@ -1398,6 +1409,60 @@ mod tests { assert!(BenchConfig::from_args(&args).is_err()); } + + #[test] + fn test_zero_max_concurrency_fails() { + let args = parse_args(vec![ + "vllm-bench", + "--model", + "test-model", + "--max-concurrency", + "0", + ]); + let err = BenchConfig::from_args(&args).expect_err("0 concurrency must be rejected"); + assert!( + err.to_string().contains("--max-concurrency must be greater than 0"), + "{}", + err + ); + } + + #[test] + fn test_positive_max_concurrency_passes() { + for value in ["1", "16"] { + let args = parse_args(vec![ + "vllm-bench", + "--model", + "test-model", + "--max-concurrency", + value, + ]); + assert!( + BenchConfig::from_args(&args).is_ok(), + "value {value} should pass" + ); + } + } + + #[test] + fn test_zero_multi_turn_concurrency_fails() { + let args = parse_args(vec![ + "vllm-bench", + "--model", + "test-model", + "--multi-turn", + "--multi-turn-concurrency", + "0", + ]); + let err = + BenchConfig::from_args(&args).expect_err("0 multi-turn concurrency must be rejected"); + assert!( + err.to_string().contains("--multi-turn-concurrency must be greater than 0"), + "{}", + err + ); + } + #[test] fn test_range_ratio_parse_float() { let rr = RangeRatio::parse("0.2").unwrap(); diff --git a/rust/src/bench/src/mm_processor.rs b/rust/src/bench/src/mm_processor.rs index 3454dc7993c7..814b693af90b 100644 --- a/rust/src/bench/src/mm_processor.rs +++ b/rust/src/bench/src/mm_processor.rs @@ -135,6 +135,9 @@ pub struct MmProcessorArgs { /// Entry point for `vllm-bench mm-processor`. pub async fn run_mm_processor(args: MmProcessorArgs, timing: RequestTimingStats) -> Result<()> { + if args.max_concurrency == 0 { + anyhow::bail!("--max-concurrency must be greater than 0"); + } if args.engine.data_parallel_size_local == Some(0) { anyhow::bail!( "`--data-parallel-size-local 0` is not supported; the benchmark requires \ diff --git a/rust/src/bench/src/sweep.rs b/rust/src/bench/src/sweep.rs index 367e7adfc34e..7fb6a6129e00 100644 --- a/rust/src/bench/src/sweep.rs +++ b/rust/src/bench/src/sweep.rs @@ -376,9 +376,15 @@ impl SummaryColumn { pub fn parse_concurrency_values(s: &str) -> Result> { s.split(',') .map(|v| { - v.trim().parse::().map_err(|_| { + let value = v.trim().parse::().map_err(|_| { BenchError::Config(format!("Invalid concurrency value: '{}'", v.trim())) - }) + })?; + if value == 0 { + return Err(BenchError::Config( + "--sweep-max-concurrency values must be greater than 0, got 0".into(), + )); + } + Ok(value) }) .collect() } @@ -561,4 +567,25 @@ mod tests { "expected at least 2 dashes, got {dash_count}: {row}" ); } + + #[test] + fn parse_concurrency_values_rejects_zero() { + for input in ["0", "0,1", "1,0", "1,0,16"] { + let err = parse_concurrency_values(input).expect_err("0 must be rejected"); + let msg = err.to_string(); + assert!( + msg.contains("--sweep-max-concurrency") && msg.contains("greater than 0"), + "{input}: {msg}" + ); + } + } + + #[test] + fn parse_concurrency_values_accepts_positive() { + assert_eq!( + parse_concurrency_values("1,10,50").unwrap(), + vec![1, 10, 50] + ); + assert_eq!(parse_concurrency_values(" 8 , 16 ").unwrap(), vec![8, 16]); + } } From 5c19391400c58ae10feca2483f55c547df211715 Mon Sep 17 00:00:00 2001 From: Vegetog <2542856720@qq.com> Date: Thu, 29 Oct 2026 20:48:29 +0800 Subject: [PATCH 131/406] [Frontend] Support thinking_token_budget in the Responses API (#63896) Signed-off-by: Vegetog <110553275+Vegetog@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Chauncey --- docs/features/reasoning_outputs.md | 2 ++ .../test_thinking_token_budget.py | 34 +++++++++++++++++++ .../openai/responses/test_sampling_params.py | 3 ++ vllm/entrypoints/openai/responses/protocol.py | 3 ++ 4 files changed, 42 insertions(+) diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index 5e59644def9b..729325e29c36 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -308,6 +308,8 @@ curl http://localhost:8000/v1/chat/completions \ }' ``` +The Responses API (`/v1/responses`) accepts the same `thinking_token_budget` field. + ### Offline Inference ```python diff --git a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index ecff3f8ee6be..e0aafae13edf 100644 --- a/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -14,6 +14,7 @@ import pytest import pytest_asyncio +import vllm.envs as envs from tests.utils import RemoteOpenAIServer, multi_gpu_only, requires_fp8 from vllm.platforms import current_platform from vllm.tokenizers import get_tokenizer @@ -212,6 +213,39 @@ async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI ) +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="the Rust frontend does not serve /v1/responses", +) +@pytest.mark.parametrize("client", ["default"], indirect=True) +async def test_thinking_token_budget_responses_api(client: openai.AsyncOpenAI): + """/v1/responses applies thinking_token_budget like /v1/chat/completions.""" + extra_body = {"thinking_token_budget": THINK_BUDGET} + chat = await client.chat.completions.create( + model=MODEL_NAME, + messages=MESSAGES, + max_tokens=100, + temperature=0.0, + extra_body=extra_body, + ) + response = await client.responses.create( + model=MODEL_NAME, + input=MESSAGES, + max_output_tokens=100, + temperature=0.0, + extra_body=extra_body, + ) + + reasoning = "".join( + part.text + for item in response.output + if item.type == "reasoning" + for part in item.content or [] + ) + assert reasoning == chat.choices[0].message.reasoning + + @pytest.mark.asyncio @multi_gpu_only(num_gpus=2) @requires_fp8 diff --git a/tests/entrypoints/openai/responses/test_sampling_params.py b/tests/entrypoints/openai/responses/test_sampling_params.py index eaa3b4cedacf..e502d3db58bf 100644 --- a/tests/entrypoints/openai/responses/test_sampling_params.py +++ b/tests/entrypoints/openai/responses/test_sampling_params.py @@ -63,6 +63,7 @@ def test_extra_sampling_params(self): stop_token_ids=[42, 43], ignore_eos=True, spaces_between_special_tokens=False, + thinking_token_budget=128, vllm_xargs={"custom": "value"}, ) @@ -74,6 +75,7 @@ def test_extra_sampling_params(self): assert sampling_params.stop_token_ids == [42, 43] assert sampling_params.ignore_eos is True assert sampling_params.spaces_between_special_tokens is False + assert sampling_params.thinking_token_budget == 128 assert sampling_params.extra_args == {"custom": "value"} def test_stop_string_conversion(self): @@ -99,6 +101,7 @@ def test_default_values(self): assert sampling_params.repetition_penalty == 1.0 # None → 1.0 assert sampling_params.stop == [] # Empty list + assert sampling_params.thinking_token_budget is None assert sampling_params.extra_args == {} # Empty dict def test_null_top_logprobs_is_same_as_omitted(self): diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 80f9bc28d3c1..fa9536fe9a98 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -86,6 +86,7 @@ RequestOutputKind, SamplingParams, StructuredOutputsParams, + ThinkingTokenBudget, ) from vllm.utils import random_uuid @@ -448,6 +449,7 @@ class ResponsesRequest(OpenAIBaseModel): stop: StopParam = [] stop_token_ids: list[int] | None = [] ignore_eos: bool = False + thinking_token_budget: ThinkingTokenBudget = None vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field( default=None, description=( @@ -645,6 +647,7 @@ def to_sampling_params( ), structured_outputs=self.extract_structured_outputs(), logit_bias=self.logit_bias, + thinking_token_budget=self.thinking_token_budget, extra_args=extra_args, skip_clone=True, # Created fresh per request, safe to skip clone skip_special_tokens=skip_special_tokens, From 61a6f8138fd81f89467af762dc1b74674f544300 Mon Sep 17 00:00:00 2001 From: Yiliu Dong Date: Thu, 29 Oct 2026 20:51:03 +0800 Subject: [PATCH 132/406] [Bugfix][Bench] Aggregate DP running requests for profiling (#61777) Signed-off-by: qianlihuang Co-authored-by: Codex Co-authored-by: vllm-agent --- rust/src/bench/src/benchmark.rs | 75 ++++++++++++++++++++++++++++----- rust/src/bench/src/cli.rs | 9 ++-- 2 files changed, 70 insertions(+), 14 deletions(-) diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs index 117e9fe72d9c..1f6b0567ac2f 100644 --- a/rust/src/bench/src/benchmark.rs +++ b/rust/src/bench/src/benchmark.rs @@ -1390,7 +1390,7 @@ pub(crate) async fn send_profile_request( Ok(resp.status().is_success()) } -/// Parse `vllm:num_requests_running` from the server's `/metrics` Prometheus endpoint. +/// Fetch total `vllm:num_requests_running` across engines from `/metrics`. pub(crate) async fn fetch_num_requests_running( client: &reqwest::Client, base_url: &str, @@ -1398,18 +1398,35 @@ pub(crate) async fn fetch_num_requests_running( let url = format!("{base_url}/metrics"); let resp = client.get(&url).send().await.ok()?; let text = resp.text().await.ok()?; + + parse_num_requests_running(&text) +} + +/// Parse total running requests across the engine series in `/metrics`. +fn parse_num_requests_running(text: &str) -> Option { + let mut total: Option = None; for line in text.lines() { - if let Some(rest) = line.strip_prefix("vllm:num_requests_running") { - // Line format: `vllm:num_requests_running{model_name="..."} 42` - // The value is after the last space. - if let Some(val_str) = rest.rsplit(' ').next() - && let Ok(v) = val_str.parse::() - { - return Some(v as usize); - } + let Some(rest) = line.trim().strip_prefix("vllm:num_requests_running") else { + continue; + }; + let sample = if rest.trim_start().starts_with('{') { + let Some((_, sample)) = rest.rsplit_once('}') else { + continue; + }; + sample + } else if rest.starts_with(char::is_whitespace) { + rest + } else { + continue; + }; + if let Some(value) = sample.split_whitespace().next().and_then(|s| s.parse::().ok()) + && value.is_finite() + && value >= 0.0 + { + total = Some(total.unwrap_or(0).saturating_add(value as usize)); } } - None + total } /// Poll `/metrics` until `num_requests_running >= threshold`, then start the @@ -1921,6 +1938,44 @@ async fn sample_verify_prompts( mod tests { use super::*; + #[test] + fn test_parse_num_requests_running_sums_dp_engines() { + let metrics: String = (0..8) + .map(|engine| { + format!("vllm:num_requests_running{{model_name=\"m\",engine=\"{engine}\"}} 16.0\n") + }) + .collect(); + + assert_eq!(parse_num_requests_running(&metrics), Some(128)); + } + + #[test] + fn test_parse_num_requests_running_matches_exact_metric_name() { + let metrics = "\ +vllm:num_requests_running_total{engine=\"0\"} 1024 +vllm:num_requests_running{engine=\"0\"} 16 +vllm:num_requests_waiting{model_name=\"vllm:num_requests_running\"} 256 +vllm:num_requests_running{engine=\"1\"} 16 +"; + + assert_eq!(parse_num_requests_running(metrics), Some(32)); + } + + #[test] + fn test_parse_num_requests_running_distinguishes_zero_and_missing() { + for (metrics, expected) in [("vllm:num_requests_running 0\n", Some(0)), ("", None)] { + assert_eq!(parse_num_requests_running(metrics), expected); + } + } + + #[test] + fn test_parse_num_requests_running_reads_sample_before_timestamp() { + let metrics = + "vllm:num_requests_running {model_name=\"model with spaces\"}\t4.2e1 123456\n"; + + assert_eq!(parse_num_requests_running(metrics), Some(42)); + } + fn make_metrics( drafts: u64, draft_tokens: u64, diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs index 850fcb35ddfb..693f7f6585ba 100644 --- a/rust/src/bench/src/cli.rs +++ b/rust/src/bench/src/cli.rs @@ -406,10 +406,11 @@ pub struct BenchServeArgs { #[arg(long, default_value_t = false)] pub profile: bool, - /// Minimum server batch size (num_requests_running) before starting the - /// profiler. When set, profiling is deferred until the /metrics endpoint - /// reports at least this many running requests, then captures for - /// --profile-duration seconds. Requires --profile. + /// Minimum total running requests across engines reported by /metrics + /// before starting the profiler. When set, profiling is deferred until + /// this threshold is reached, then captures for --profile-duration seconds. + /// Only engines exposed by the /metrics endpoint at --base-url are counted. + /// Requires --profile. #[arg(long)] pub profile_batch_threshold: Option, From d83a7df60aa34f80d88414789a452c0b85728d2e Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:55:57 -0700 Subject: [PATCH 133/406] [Model] Update RADIO vision encoder to match the current Transformers version (#58468) Signed-off-by: Shengliang Xu Co-authored-by: vllm-agent --- tests/models/multimodal/pooling/test_radio.py | 214 ++++- .../model_executor/models/nano_nemotron_vl.py | 10 +- .../models/nemotron_parse3_5.py | 36 +- vllm/model_executor/models/radio.py | 905 +++++++++--------- vllm/transformers_utils/configs/radio.py | 41 +- 5 files changed, 724 insertions(+), 482 deletions(-) diff --git a/tests/models/multimodal/pooling/test_radio.py b/tests/models/multimodal/pooling/test_radio.py index 7c039ab8b9c6..ed57e3f6cd0c 100644 --- a/tests/models/multimodal/pooling/test_radio.py +++ b/tests/models/multimodal/pooling/test_radio.py @@ -1,8 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import inspect +import tempfile +from pathlib import Path +from typing import Any + import pytest import torch import torch.nn as nn +from safetensors.torch import load_file from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory @@ -33,13 +39,8 @@ def run_radio_test( img_processor = CLIPImageProcessor.from_pretrained(model) images = [asset.pil_image for asset in image_assets] - # Input resolution must be a multiple of `self.min_resolution_step`. - # Using `self.get_nearest_supported_resolution`, for assets 432x642 the - # nearest supported resolution is 432x640. pixel_values = [ - img_processor(image, return_tensors="pt").pixel_values.to(torch_dtype)[ - :, :, :, :640 - ] + img_processor(image, return_tensors="pt").pixel_values.to(torch_dtype) for image in images ] @@ -64,6 +65,22 @@ def run_radio_test( # affects evaluation benchmarks. hf_model.make_preprocessor_external() + # Input resolution must be a multiple of the tower's ``min_resolution_step`` + # (patch_size * window_size), which differs across RADIO generations. Crop + # down to the nearest supported resolution so one test body fits every + # checkpoint; for C-RADIOv2-H this is a no-op (432x640, as before). + radio = getattr(hf_model, "radio_model", hf_model) + step = radio.min_resolution_step + pixel_values = [ + pv[ + :, + :, + : (pv.shape[-2] // step) * step, + : (min(pv.shape[-1], 640) // step) * step, + ] + for pv in pixel_values + ] + hf_outputs_per_image = [ hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values ] @@ -73,7 +90,17 @@ def run_radio_test( **hf_config.args, ) vllm_model = RadioModel(vllm_config) - vllm_model.load_weights(hf_model.state_dict()) + loaded = vllm_model.load_weights(hf_model.state_dict()) + # Guard the remote-code -> vLLM key remap: every parameter must be loaded + # from the checkpoint, except the LayerScale gains (layer_scale{1,2}.lambda1), + # which are identity-init and may be absent from the remote-code checkpoint. + expected = { + name + for name, _ in vllm_model.named_parameters() + if not name.endswith((".layer_scale1.lambda1", ".layer_scale2.lambda1")) + } + missing = expected - loaded + assert not missing, f"parameters not loaded from checkpoint: {sorted(missing)}" vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype) vllm_outputs_per_image = [ @@ -85,14 +112,22 @@ def run_radio_test( cos_similar = nn.CosineSimilarity(dim=-1) for vllm_output, hf_output in zip(vllm_outputs_per_image, hf_outputs_per_image): - assert cos_similar(vllm_output[0], hf_output[0]).mean() > 0.99 - assert cos_similar(vllm_output[1], hf_output[1]).mean() > 0.99 + summary_sim = cos_similar(vllm_output[0], hf_output[0]).mean() + feature_sim = cos_similar(vllm_output[1], hf_output[1]).mean() + print( + f"[radio] {model_id} {dtype} summary={summary_sim:.6f} " + f"features={feature_sim:.6f}" + ) + assert summary_sim > 0.99 + assert feature_sim > 0.99 @pytest.mark.parametrize( "model_id", [ "nvidia/C-RADIOv2-H", + "nvidia/C-RADIOv3-H", + "nvidia/C-RADIOv4-H", ], ) @pytest.mark.parametrize("dtype", ["half", "bfloat16"]) @@ -104,3 +139,164 @@ def test_radio( model_id, dtype=dtype, ) + + +def _radio_summary(*, teachers, cls_token_per_teacher): + """Summary tensor from a tiny RadioModel with no encoder layers, so it + builds no parallel layers and needs neither a GPU nor a checkpoint.""" + config = RadioConfig( + model_name="vit_small_patch16_224", + teachers=teachers, + cls_token_per_teacher=cls_token_per_teacher, + ) + model = RadioModel(config, num_hidden_layers_override=0) + # Synthetic encoder output: [batch, num_skip + num_patches, hidden]. + y = torch.randn(2, model.embeddings.num_skip + 4, config.hidden_size) + summary, _ = model._extract_final(y) + return summary + + +def test_summary_idxs_no_teachers_keeps_class_tokens(): + # No teachers -> summary_idxs is None -> keep the class-token summary. + summary = _radio_summary(teachers=[], cls_token_per_teacher=False) + assert summary.shape[1] > 0 + + +def test_summary_idxs_all_use_summary_false_is_empty(): + # Teachers present but none flagged use_summary -> empty selection. + summary = _radio_summary( + teachers=[{"name": "a", "use_summary": False}], + cls_token_per_teacher=True, + ) + assert summary.shape[1] == 0 + + +def _native_state_dict_from_model(model: RadioModel) -> dict[str, torch.Tensor]: + """Reverse-map a model's parameters onto native Transformers checkpoint keys + (``encoder.layer.N.*`` with split ``attention.attention.{query,key,value}`` + and ``attention.output.dense``), so loading it must round-trip back to the + original parameters.""" + state_dict: dict[str, torch.Tensor] = {} + for name, param in model.named_parameters(): + data = param.detach().clone() + # The module tree already mirrors the native ``encoder.layer.N.*`` names + # (including ``layer_scale{1,2}.lambda1``); only the fused attention + # projections differ from the native split ones. + if ".attention.qkv." in name: + base, suffix = name.split(".attention.qkv.") + query, key, value = data.chunk(3, dim=0) + for proj, shard in (("query", query), ("key", key), ("value", value)): + state_dict[f"{base}.attention.attention.{proj}.{suffix}"] = ( + shard.clone() + ) + elif ".attention.proj." in name: + base, suffix = name.split(".attention.proj.") + state_dict[f"{base}.attention.output.dense.{suffix}"] = data + else: + state_dict[name] = data + return state_dict + + +def test_native_format_weight_loading(default_vllm_config, dist_init): + # Native Transformers checkpoints nest the encoder under ``encoder.layer.N`` + # with split attention projections; loading must fuse q/k/v into the packed + # ``qkv`` in the right order and land every other tensor in place, matching + # the legacy layout the integration test covers. + config = RadioConfig(model_name="vit_small_patch16_224") + model = RadioModel(config, num_hidden_layers_override=2) + + # Every parameter (including layer_scale) must load from the native keys. + expected = { + name: param.detach().clone() for name, param in model.named_parameters() + } + native_weights = _native_state_dict_from_model(model) + # Keys both layouts intentionally drop must be skipped, not error. + native_weights["input_conditioner.norm_mean"] = torch.zeros(3) + native_weights["input_conditioner.norm_std"] = torch.ones(3) + native_weights["summary_idxs"] = torch.zeros(1, dtype=torch.long) + + # Zero every parameter so a weight that fails to load is visible as zeros. + with torch.no_grad(): + for param in model.parameters(): + param.zero_() + + loaded = model.load_weights(native_weights) + + missing = set(expected) - loaded + assert not missing, f"native weights not loaded: {sorted(missing)}" + params = dict(model.named_parameters()) + for name, original in expected.items(): + assert torch.equal(params[name], original), ( + f"parameter mismatch after native load: {name}" + ) + + +def test_native_format_from_transformers(default_vllm_config, dist_init): + # Real-checkpoint native-format round-trip: load a published (legacy) + # checkpoint through Transformers' native ``RadioModel`` (which converts it to + # the native layout), re-save it with ``save_original_format=False`` so the + # on-disk checkpoint stays native -- the default re-applies the reverse + # native->legacy mapping for backward compatibility -- then confirm vLLM's + # ``RadioModel`` loads that native checkpoint with values matching the + # reference. ``nvidia/C-RADIOv4-H`` loads on Transformers' config defaults and + # is already downloaded by ``test_radio``. + import transformers + + hf_radio_cls: Any = getattr(transformers, "RadioModel", None) + if hf_radio_cls is None: + pytest.skip("this transformers version has no native RadioModel") + + model_id = "nvidia/C-RADIOv4-H" + hf_native = hf_radio_cls.from_pretrained(model_id, dtype=torch.float32) + if ( + "save_original_format" + not in inspect.signature(hf_native.save_pretrained).parameters + ): + pytest.skip("transformers save_pretrained has no save_original_format flag") + + with tempfile.TemporaryDirectory() as tmp: + hf_native.save_pretrained(tmp, save_original_format=False) + native_weights: dict[str, torch.Tensor] = {} + for shard in sorted(Path(tmp).glob("*.safetensors")): + native_weights.update(load_file(str(shard))) + + # Guard that the saved checkpoint is genuinely native-layout (split attention + # under ``encoder.layer.N``), not the legacy ``radio_model.model.blocks.*``. + assert any(".attention.attention.query" in name for name in native_weights), ( + "expected native-layout keys in the save_original_format=False checkpoint" + ) + assert not any(name.startswith("radio_model.") for name in native_weights) + + hf_config = AutoConfig.from_pretrained(model_id, trust_remote_code=True) + vllm_config = RadioConfig(model_name=hf_config.args["model"], **hf_config.args) + vllm_model = RadioModel(vllm_config) + loaded = vllm_model.load_weights(native_weights) + + expected = { + name + for name, _ in vllm_model.named_parameters() + if not name.endswith((".layer_scale1.lambda1", ".layer_scale2.lambda1")) + } + missing = expected - loaded + assert not missing, f"native weights not loaded: {sorted(missing)}" + + # Values match the reference: the fused qkv is the concatenation of the + # native split query/key/value, and proj is the native output.dense. + for name, param in vllm_model.named_parameters(): + if ".attention.qkv." in name: + base, suffix = name.split(".attention.qkv.") + ref = torch.cat( + [ + native_weights[f"{base}.attention.attention.{proj}.{suffix}"] + for proj in ("query", "key", "value") + ], + dim=0, + ) + elif ".attention.proj." in name: + base, suffix = name.split(".attention.proj.") + ref = native_weights[f"{base}.attention.output.dense.{suffix}"] + elif name in native_weights: + ref = native_weights[name] + else: + continue + assert torch.equal(param, ref), f"value mismatch after native load: {name}" diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index b77206295d5c..f24419501ef0 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1620,7 +1620,11 @@ def is_adapter_weights(weight: tuple[str, torch.Tensor]): return weight[0].startswith("mlp1") def is_vision_weights(name: str) -> bool: - return name.startswith("vision_model.radio_model.") + # The whole RADIO tower lives under ``vision_model.``: legacy + # remote-code checkpoints nest it as ``vision_model.radio_model.*``, + # native Transformers checkpoints as ``vision_model.embeddings.*`` / + # ``vision_model.encoder.*``. RadioModel.load_weights handles both. + return name.startswith("vision_model.") def is_sound_weights(name: str) -> bool: return name.startswith("sound") @@ -1653,7 +1657,9 @@ def llm_weights_gen(): elif is_vision_weights(name): if not load_multimodal_weights: continue - # Convert: vision_model.radio_model.* → radio_model.* + # Strip the ``vision_model.`` prefix; the remainder is the + # RADIO key (legacy ``radio_model.*`` or native + # ``embeddings.*``/``encoder.*``). hf_key = name[len("vision_model.") :] vision_weights.append((hf_key, w.detach().clone())) elif is_sound_weights(name): diff --git a/vllm/model_executor/models/nemotron_parse3_5.py b/vllm/model_executor/models/nemotron_parse3_5.py index 1ff0f1d0ac02..a14ac38175a9 100644 --- a/vllm/model_executor/models/nemotron_parse3_5.py +++ b/vllm/model_executor/models/nemotron_parse3_5.py @@ -76,22 +76,22 @@ def _content_normalized_positions( patches: torch.Tensor, - patch_generator: nn.Module, + embeddings: nn.Module, image_size: tuple[int, int], content_size: tuple[int, int], ) -> torch.Tensor: """Apply Parse 3.5's content-normalized positional field to one page.""" height, width = image_size content_height, content_width = content_size - patch_size = int(patch_generator.patch_size) + patch_size = int(embeddings.patch_size) rows, cols = height // patch_size, width // patch_size if height % patch_size or width % patch_size: raise ValueError(f"image size {image_size} is not patch aligned") if not (0 < content_height <= height and 0 < content_width <= width): raise ValueError(f"content size {content_size} must be within {image_size}") - positional_grid = patch_generator.pos_embed.reshape( - 1, patch_generator.num_rows, patch_generator.num_cols, -1 + positional_grid = embeddings.position_embedding.reshape( + 1, embeddings.num_rows, embeddings.num_cols, -1 ).permute(0, 3, 1, 2) content_rows = (content_height + patch_size - 1) // patch_size content_cols = (content_width + patch_size - 1) // patch_size @@ -590,12 +590,12 @@ def forward( raise ValueError("image_sizes and content_sizes must have equal length") if pixel_values.ndim == 2: pixel_values = pixel_values.unsqueeze(0) - patch_generator = self.model_encoder.model.patch_generator - backbone = self.model_encoder.model - patch_size = int(patch_generator.patch_size) - num_skip = int(patch_generator.num_skip) - num_cls = int(patch_generator.num_cls_tokens) - patches = patch_generator.embedder(pixel_values) + embeddings = self.model_encoder.embeddings + backbone = self.model_encoder + patch_size = int(embeddings.patch_size) + num_skip = int(embeddings.num_skip) + num_cls = int(embeddings.num_cls_tokens) + patches = embeddings.patch_projection(pixel_values) cursor = 0 positioned_pages = [] for image_size, content_size in zip(image_sizes, content_sizes): @@ -607,18 +607,18 @@ def forward( cursor += patch_count positioned_pages.append( _content_normalized_positions( - page, patch_generator, image_size, content_size + page, embeddings, image_size, content_size ) ) if not positioned_pages or cursor != patches.shape[1]: raise ValueError("packed patch sequence does not match image_sizes") - hidden = patch_generator.cls_token_dynamic( + hidden = embeddings._cls_token_dynamic( torch.cat(positioned_pages, dim=1), imgs_sizes=image_sizes ) - hidden = patch_generator.patch_normalizer(hidden) + hidden = embeddings.patch_normalizer(hidden) mask_meta = None if len(image_sizes) > 1: - mask_meta = backbone.inter_image_mask_metadata( + mask_meta = backbone._inter_image_mask_metadata( image_sizes, device=hidden.device ) hidden = backbone.encoder(inputs_embeds=hidden, mask_meta=mask_meta) @@ -876,11 +876,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: synthetic_layerscales = { name: param for name, param in self.encoder.model_encoder.named_parameters() - if name.endswith((".ls1", ".ls2")) and name not in radio_loaded + if name.endswith((".layer_scale1.lambda1", ".layer_scale2.lambda1")) + and name not in radio_loaded } if synthetic_layerscales: + layerscale_markers = (".ls1", ".ls2", ".layer_scale1", ".layer_scale2") source_layerscales = { - name for name in radio_source_names if ".ls1" in name or ".ls2" in name + name + for name in radio_source_names + if any(marker in name for marker in layerscale_markers) } if radio_model_name not in identity_layerscale_models: raise RuntimeError( diff --git a/vllm/model_executor/models/radio.py b/vllm/model_executor/models/radio.py index d01e0f846418..d7f651cae740 100644 --- a/vllm/model_executor/models/radio.py +++ b/vllm/model_executor/models/radio.py @@ -1,16 +1,32 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved. -# -# NVIDIA CORPORATION and its licensors retain all intellectual property -# and proprietary rights in and to this software, related documentation -# and any modifications thereto. Any use, reproduction, disclosure or -# distribution of this software and related documentation without an express -# license agreement from NVIDIA CORPORATION is strictly prohibited. + +# adapted from https://github.com/NVlabs/RADIO and +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/radio/modeling_radio.py +# -------------------------------------------------------- +# RADIO +# Copyright (c) 2023-2026, NVIDIA CORPORATION. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# -------------------------------------------------------- +"""RADIO vision encoder, matching the ``RadioModel`` in the Transformers library. + +The module tree mirrors the Transformers ``RadioModel`` (``embeddings`` and an +``encoder.layer`` stack of ``norm1``/``attention``/``layer_scale1``/``norm2``/ +``mlp``/``layer_scale2`` blocks). The vLLM-specific deviations are the fused +``attention.qkv`` (a :class:`~vllm.model_executor.layers.linear.QKVParallelLinear` +for TP/quantization, vs the split ``query``/``key``/``value`` in HF) and running +input normalization in the processor instead of an ``input_conditioner`` module; +the attention backend and MLP reuse vLLM's :class:`MMEncoderAttention` and +:class:`~vllm.model_executor.models.intern_vit.InternMLP`. ``load_weights`` maps +both native Transformers and legacy remote-code checkpoint keys onto this tree +via :data:`hf_to_vllm_mapper`. +""" import math +from collections.abc import Iterable from dataclasses import dataclass -from itertools import accumulate +from functools import partial +from itertools import accumulate, repeat from typing import TypeAlias import torch @@ -19,185 +35,129 @@ from einops import rearrange from transformers import PreTrainedConfig -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.intern_vit import ( - InternParallelAttention, - InternVisionEncoder, - InternVisionEncoderLayer, +from vllm.distributed import ( + divide, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, ) +from vllm.logger import init_logger +from vllm.model_executor.layers.attention import MMEncoderAttention +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.models.intern_vit import InternMLP +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper +from vllm.model_executor.models.vision import is_vit_use_data_parallel -input_dim_t: TypeAlias = int | tuple[int, int] - +logger = init_logger(__name__) -def calc_seq_len(size: tuple[int, int], patch_size: int) -> int: - h, w = size - return (h // patch_size) * (w // patch_size) +input_dim_t: TypeAlias = int | tuple[int, int] +norm_t: TypeAlias = tuple[float, float, float] | torch.Tensor -def calc_seq_lens(sizes: list[tuple[int, int]], patch_size: int) -> list[int]: - return [calc_seq_len(size, patch_size) for size in sizes] +def _ntuple(n): + def parse(x): + if isinstance(x, Iterable) and not isinstance(x, str): + return tuple(x) + return tuple(repeat(x, n)) + return parse -class ClsToken(nn.Module): - def __init__( - self, - ndim: int, - num_tokens: int = 1, - enabled: bool = True, - register_multiple: int | None = None, - num_registers: int | None = None, - ): - super().__init__() - self.ndim = ndim - self.enabled = enabled - self.num_registers = 0 - self.num_tokens = num_tokens - if enabled: - if num_registers: - self.num_registers = num_registers - elif register_multiple: - self.num_registers = register_multiple - ( - num_tokens % register_multiple - ) +to_1tuple = _ntuple(1) +to_2tuple = _ntuple(2) +to_3tuple = _ntuple(3) +to_4tuple = _ntuple(4) +to_ntuple = _ntuple - scale = ndim**-0.5 - self.token = nn.Parameter( - torch.randn(num_tokens + self.num_registers, ndim) * scale - ) - else: - self.token = None - - self.num_patches = self.num_tokens + self.num_registers +def calc_seq_len(size: tuple[int, int], patch_size: int) -> int: + h, w = size + return (h // patch_size) * (w // patch_size) - def forward(self, x: torch.Tensor): - if self.token is None: - return x - token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1) - x = torch.cat( - [ - token, - x, - ], - dim=1, - ) +def calc_seq_lens(sizes: list[tuple[int, int]], patch_size: int) -> list[int]: + return [calc_seq_len(size, patch_size) for size in sizes] - return x +class RadioPatchEmbeddings(nn.Module): + """Cropped Position Embedding (CPE) patch generator, matching Transformers' + ``RadioPatchEmbeddings``, with additional support for dynamic multi-resolution + packing and video temporal compression. + """ -class ViTPatchGenerator(nn.Module): - def __init__( - self, - # config: PreTrainedConfig, - patch_size: int, - embed_dim: int, - input_dims: input_dim_t, - abs_pos: bool = True, - normalize_patches: bool = False, - cls_token: bool = False, - max_input_dims: input_dim_t | None = None, - pos_dropout: float = 0.0, - return_pos_enc: bool = False, - num_cls_tokens: int = 1, - register_multiple: int | None = None, - num_registers: int | None = None, - patch_bias: bool = False, - temporal_patch_size: int = 1, - separate_video_embedder: bool = True, - device=None, - dtype=None, - ): + def __init__(self, config: PreTrainedConfig): super().__init__() - if isinstance(input_dims, int): - input_dims = (input_dims, input_dims) - - if max_input_dims is None: - max_input_dims = input_dims - if isinstance(max_input_dims, int): - max_input_dims = (max_input_dims, max_input_dims) + patch_size = config.patch_size + embed_dim = config.hidden_size + num_channels = config.num_channels - max_input_dims = ( - int(math.ceil(max_input_dims[0] / patch_size) * patch_size), - int(math.ceil(max_input_dims[1] / patch_size) * patch_size), + max_input_dims = to_2tuple(config.max_img_size) + max_input_dims = tuple( + int(math.ceil(d / patch_size) * patch_size) for d in max_input_dims ) - - self.cpe_mode = max_input_dims != input_dims - self.pos_dropout = pos_dropout - self.return_pos_enc = return_pos_enc - - factory = dict(device=device, dtype=dtype) + input_dims = to_2tuple(config.image_size) self.patch_size = patch_size - self.abs_pos = abs_pos self.embed_dim = embed_dim - self.temporal_patch_size = temporal_patch_size + self.num_cls_tokens = config.num_cls_tokens + self.num_registers = config.num_registers + self.temporal_patch_size = config.video_temporal_patch_size + self.cpe_mode = max_input_dims != input_dims self.num_rows = max_input_dims[0] // patch_size self.num_cols = max_input_dims[1] // patch_size - self.input_dims = ( - input_dims[0] // patch_size, - input_dims[1] // patch_size, - ) + self.input_dims = tuple(d // patch_size for d in input_dims) self.num_patches = self.num_rows * self.num_cols - self.max_input_dims = max_input_dims self.im_to_patches = Im2Patches(patch_size) - self.embedder = ViTPatchLinear( - patch_size, embed_dim, bias=patch_bias, **factory + self.patch_projection = ViTPatchLinear( + patch_size, embed_dim, num_channels=num_channels, bias=False ) - if temporal_patch_size > 1: - if not separate_video_embedder: + if self.temporal_patch_size > 1: + if not config.separate_video_embedder: raise NotImplementedError( "Only separate_video_embedder=True is supported for" - " temporal compression (temporal_patch_size > 1)" + " temporal compression (video_temporal_patch_size > 1)" ) - self.video_embedder = ViTPatchLinear( + self.video_patch_projection = ViTPatchLinear( patch_size, embed_dim, - bias=patch_bias, - temporal_patch_size=temporal_patch_size, - **factory, + num_channels=num_channels, + bias=False, + temporal_patch_size=self.temporal_patch_size, ) - if abs_pos: - scale = embed_dim**-0.5 - self.pos_embed = nn.Parameter( - torch.randn(1, self.num_patches, embed_dim, **factory) * scale - ) - - self.cls_token = ClsToken( - embed_dim, - num_tokens=num_cls_tokens, - enabled=cls_token, - register_multiple=register_multiple, - num_registers=num_registers, + scale = embed_dim**-0.5 + self.position_embedding = nn.Parameter( + torch.randn(1, self.num_patches, embed_dim) * scale ) - - self.patch_normalizer = ( - nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity() + self.cls_register_token = nn.Parameter( + torch.randn(self.num_cls_tokens + self.num_registers, embed_dim) * scale ) + self.patch_normalizer = nn.Identity() + + @property + def num_skip(self) -> int: + return self.num_cls_tokens + self.num_registers def forward( self, x: torch.Tensor, imgs_sizes: list[tuple[int, int]] | None = None ) -> torch.Tensor: if imgs_sizes is not None: - patches = self.embedder(x) - patches, pos_enc = self.apply_pos_enc_dynamic( - patches, imgs_sizes=imgs_sizes - ) - patches = self.cls_token_dynamic(patches, imgs_sizes=imgs_sizes) + # Dynamic multi-resolution: ``x`` is already patchified and packed + # to [B, total_patches, C*P*P]. + patches = self.patch_projection(x) + patches = self._apply_pos_enc_dynamic(patches, imgs_sizes=imgs_sizes) + patches = self._cls_token_dynamic(patches, imgs_sizes=imgs_sizes) else: - patches = self.embed_patches(x) - patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:]) - patches = self.cls_token(patches) - patches = self.patch_normalizer(patches) - if self.return_pos_enc: - return patches, pos_enc - return patches + patches = self.patch_projection(self.im_to_patches(x)) + patches = self._apply_pos_enc(patches, input_size=x.shape[2:]) + patches = self._prepend_cls_token(patches) + return self.patch_normalizer(patches) def forward_video(self, x: torch.Tensor) -> torch.Tensor: """Process video frames with temporal compression. @@ -238,148 +198,75 @@ def forward_video(self, x: torch.Tensor) -> torch.Tensor: feat=feat_dim, ) - patches = self.video_embedder(patches) - - patches, pos_enc = self.apply_pos_enc(patches, input_size=input_size) - - patches = self.cls_token(patches) - - patches = self.patch_normalizer(patches) - if self.return_pos_enc: - return patches, pos_enc - return patches - - def apply_pos_enc_dynamic( - self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]] - ) -> tuple[torch.Tensor, torch.Tensor | None]: - if not self.abs_pos: - return patches, None - - current_length = 0 - pos_enc_list = [] - - for size in imgs_sizes: - seq_length = calc_seq_len(size, self.patch_size) - - img_patches = patches[:, current_length : current_length + seq_length, :] - pos_enc = self.get_pos_enc(patches.shape[0], input_size=size) - img_patches_with_pos = img_patches + pos_enc - - patches = torch.cat( - [ - patches[:, :current_length, :], - img_patches_with_pos, - patches[:, current_length + seq_length :, :], - ], - dim=1, - ) - pos_enc_list.append(pos_enc) - current_length += seq_length + patches = self.video_patch_projection(patches) + patches = self._apply_pos_enc(patches, input_size=input_size) + patches = self._prepend_cls_token(patches) + return self.patch_normalizer(patches) - full_pos_enc = torch.cat(pos_enc_list, dim=1) if pos_enc_list else None - return patches, full_pos_enc + def _prepend_cls_token(self, x: torch.Tensor) -> torch.Tensor: + token = self.cls_register_token.unsqueeze(0).expand(x.shape[0], -1, -1) + return torch.cat([token, x], dim=1) - def cls_token_dynamic( + def _cls_token_dynamic( self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]] ) -> torch.Tensor: - if not self.cls_token.enabled: - return patches - out = [] current_length = 0 - for seq_len in calc_seq_lens(imgs_sizes, self.patch_size): - class_token = self.cls_token.token.unsqueeze(0).expand( + class_token = self.cls_register_token.unsqueeze(0).expand( patches.shape[0], -1, -1 ) out.append(class_token) out.append(patches[:, current_length : current_length + seq_len, :]) current_length += seq_len - return torch.cat(out, dim=1) - @property - def apply_cls_token(self): - return self.cls_token.enabled - - @property - def num_cls_tokens(self): - return self.cls_token.num_tokens - - @property - def num_cls_patches(self): - return self.cls_token.num_patches - - @property - def num_registers(self): - return self.cls_token.num_registers - - @property - def num_skip(self): - return self.num_cls_tokens + self.num_registers - - def embed_patches(self, x: torch.Tensor) -> torch.Tensor: - patches = self.im_to_patches(x) - patches = self.embedder(patches) - return patches - - def apply_pos_enc( + def _apply_pos_enc( self, patches: torch.Tensor, - patch_idxs: torch.Tensor | None = None, input_size: tuple[int, int] | None = None, ) -> torch.Tensor: - if not self.abs_pos: - return patches - - pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size) + pos_enc = self._get_pos_enc(patches.shape[0], input_size=input_size) + return patches + pos_enc - if self.training and self.pos_dropout > 0: - keeps = ( - torch.rand( - patches.shape[0], 1, 1, dtype=pos_enc.dtype, device=pos_enc.device - ) - > self.pos_dropout + def _apply_pos_enc_dynamic( + self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]] + ) -> torch.Tensor: + current_length = 0 + for size in imgs_sizes: + seq_length = calc_seq_len(size, self.patch_size) + img_patches = patches[:, current_length : current_length + seq_length, :] + pos_enc = self._get_pos_enc(patches.shape[0], input_size=size) + img_patches_with_pos = img_patches + pos_enc + patches = torch.cat( + [ + patches[:, :current_length, :], + img_patches_with_pos, + patches[:, current_length + seq_length :, :], + ], + dim=1, ) - pos_enc_drop = torch.where(keeps, pos_enc, 0) - else: - pos_enc_drop = pos_enc - - return patches + pos_enc_drop, pos_enc + current_length += seq_length + return patches - def get_pos_enc( + def _get_pos_enc( self, batch_size: int, - patch_idxs: torch.Tensor | None = None, input_size: tuple[int, int] | None = None, ) -> torch.Tensor: if input_size is None: input_dims = self.input_dims else: - input_dims = ( - input_size[0] // self.patch_size, - input_size[1] // self.patch_size, - ) - - pos_embed = self._get_pos_embeddings(batch_size, input_dims) - - if patch_idxs is None: - return pos_embed - - exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1]) - - pos_embed = torch.gather( - pos_embed.expand(patch_idxs.shape[0], -1, -1), dim=1, index=exp_patch_idxs - ) - return pos_embed + input_dims = tuple(d // self.patch_size for d in input_size) + return self._get_pos_embeddings(batch_size, input_dims) def _get_pos_embeddings(self, batch_size: int, input_dims: tuple[int, int]): if (self.num_rows, self.num_cols) == input_dims: - return self.pos_embed + return self.position_embedding - pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute( - 0, 3, 1, 2 - ) + pos_embed = self.position_embedding.reshape( + 1, self.num_rows, self.num_cols, -1 + ).permute(0, 3, 1, 2) def window_select(pos_embed): if input_dims[0] < pos_embed.shape[-2]: @@ -440,12 +327,16 @@ def __init__( self, patch_size: int, embed_dim: int, + num_channels: int = 3, bias: bool = False, temporal_patch_size: int = 1, **factory, ): super().__init__( - 3 * temporal_patch_size * (patch_size**2), embed_dim, bias=bias, **factory + num_channels * temporal_patch_size * (patch_size**2), + embed_dim, + bias=bias, + **factory, ) self.patch_size = patch_size self.temporal_patch_size = temporal_patch_size @@ -457,11 +348,100 @@ class MaskMetadata: max_seqlen: torch.Tensor -class RadioParallelAttention(InternParallelAttention): +class RadioAttention(nn.Module): + """RADIO self-attention: a fused ``qkv`` projection feeding the vLLM + attention backend, with optional QK-normalization and dynamic-resolution + packed attention (via ``mask_meta``). The fused counterpart of the + Transformers ``RadioAttention`` (which keeps ``query``/``key``/``value`` and + ``output.dense`` split); ``load_weights`` bridges the two.""" + + def __init__( + self, + config: PreTrainedConfig, + quant_config: QuantizationConfig | None = None, + *, + num_dummy_heads: int = 0, + prefix: str = "", + ) -> None: + super().__init__() + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if self.head_dim * self.num_heads != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got embed_dim=" + f"{self.embed_dim} and num_heads={self.num_heads})." + ) + + use_data_parallel = is_vit_use_data_parallel() + # Disable attention TP when the head count is not divisible by tp_size. + tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size() + use_data_parallel = ( + use_data_parallel or (self.num_heads + num_dummy_heads) % tp_size != 0 + ) + self.tp_size = 1 if use_data_parallel else tp_size + self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank() + + # Dummy heads pad the head count so TP divides evenly on common GPU counts. + self.dummy_dim = (num_dummy_heads + self.num_heads) * self.head_dim + self.num_heads_per_partition = divide( + num_dummy_heads + self.num_heads, self.tp_size + ) + self.scale = self.head_dim**-0.5 + + self.qkv = QKVParallelLinear( + self.embed_dim, + self.head_dim, + num_dummy_heads + self.num_heads, + bias=config.qkv_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv", + disable_tp=use_data_parallel, + ) + self.qk_normalization = config.qk_normalization + if self.qk_normalization: + self.q_norm = RMSNorm( + self.dummy_dim, + eps=config.layer_norm_eps, + var_hidden_size=self.embed_dim, + ) + self.k_norm = RMSNorm( + self.dummy_dim, + eps=config.layer_norm_eps, + var_hidden_size=self.embed_dim, + ) + self.proj = RowParallelLinear( + self.dummy_dim, + self.embed_dim, + quant_config=quant_config, + prefix=f"{prefix}.proj", + disable_tp=use_data_parallel, + ) + self.attn = MMEncoderAttention( + self.num_heads_per_partition, + self.head_dim, + self.scale, + prefix=f"{prefix}.attn", + ) + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.tp_size > 1: + q = tensor_model_parallel_all_gather(q.contiguous()) + k = tensor_model_parallel_all_gather(k.contiguous()) + q = self.q_norm(q) + k = self.k_norm(k) + if self.tp_size > 1: + splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size) + q = splitter(q)[self.tp_rank] + k = splitter(k)[self.tp_rank] + return q, k + def forward( - self, x: torch.Tensor, mask_meta: MaskMetadata | None = None + self, hidden_states: torch.Tensor, mask_meta: MaskMetadata | None = None ) -> torch.Tensor: - qkv, _ = self.qkv(x) + qkv, _ = self.qkv(hidden_states) q, k, v = qkv.chunk(3, dim=-1) if self.qk_normalization: @@ -476,28 +456,89 @@ def forward( return out -class RadioVisionEncoderLayer(InternVisionEncoderLayer): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, attn_cls=RadioParallelAttention, **kwargs) +class RadioLayerScale(nn.Module): + """LayerScale matching the Transformers ``RadioLayerScale``: a ``lambda1`` + parameter scaling the residual branch.""" + + def __init__(self, config: PreTrainedConfig) -> None: + super().__init__() + self.lambda1 = nn.Parameter( + config.layerscale_value * torch.ones(config.hidden_size) + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states * self.lambda1 + + +class RadioLayer(nn.Module): + """RADIO transformer block, mirroring the Transformers ``RadioLayer`` + (``norm1`` / ``attention`` / ``layer_scale1`` / ``norm2`` / ``mlp`` / + ``layer_scale2``). ``drop_path`` is an inference-time ``Identity``; the MLP + reuses :class:`InternMLP` (``fc1``/``fc2``, matching ``RadioMLP``).""" + + def __init__( + self, + config: PreTrainedConfig, + quant_config: QuantizationConfig | None = None, + *, + num_dummy_heads: int = 0, + prefix: str = "", + ) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attention = RadioAttention( + config, + quant_config=quant_config, + num_dummy_heads=num_dummy_heads, + prefix=f"{prefix}.attention", + ) + self.layer_scale1 = RadioLayerScale(config) + self.drop_path = nn.Identity() + self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.mlp = InternMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp") + self.layer_scale2 = RadioLayerScale(config) def forward( self, hidden_states: torch.Tensor, mask_meta: MaskMetadata | None = None, - ): - hidden_states = ( - hidden_states - + self.attn(self.norm1(hidden_states), mask_meta=mask_meta) * self.ls1 - ) - - hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) * self.ls2 - + ) -> torch.Tensor: + attn_out = self.attention(self.norm1(hidden_states), mask_meta=mask_meta) + hidden_states = self.drop_path(self.layer_scale1(attn_out)) + hidden_states + mlp_out = self.mlp(self.norm2(hidden_states)) + hidden_states = self.drop_path(self.layer_scale2(mlp_out)) + hidden_states return hidden_states -class RadioVisionEncoder(InternVisionEncoder): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, layer_cls=RadioVisionEncoderLayer, **kwargs) +class RadioVisionEncoder(nn.Module): + """RADIO encoder stack, mirroring the Transformers ``RadioEncoder`` module + tree (a ``layer`` ModuleList of :class:`RadioLayer`).""" + + def __init__( + self, + config: PreTrainedConfig, + quant_config: QuantizationConfig | None = None, + *, + num_hidden_layers_override: int | None = None, + num_dummy_heads: int = 0, + prefix: str = "", + ) -> None: + super().__init__() + if num_hidden_layers_override is None: + num_hidden_layers = config.num_hidden_layers + else: + num_hidden_layers = num_hidden_layers_override + self.layer = nn.ModuleList( + [ + RadioLayer( + config, + quant_config=quant_config, + num_dummy_heads=num_dummy_heads, + prefix=f"{prefix}.layer.{layer_idx}", + ) + for layer_idx in range(num_hidden_layers) + ] + ) def forward( self, @@ -505,19 +546,76 @@ def forward( mask_meta: MaskMetadata | None = None, ): hidden_states = inputs_embeds - for encoder_layer in self.layers: + for encoder_layer in self.layer: hidden_states = encoder_layer(hidden_states, mask_meta=mask_meta) return hidden_states -class RadioInternVisionModel(nn.Module): +class RadioModel(nn.Module): packed_modules_mapping = { "qkv": ["qkv"], } + # Map both RADIO checkpoint layouts onto this module tree, which mirrors the + # Transformers ``RadioModel`` (``embeddings.*`` and ``encoder.layer.N.*``): + # * native Transformers checkpoints load onto it directly; only the split + # attention projections fuse into ``attention.qkv`` via + # ``orig_to_new_stacked`` and the output projection is renamed; + # * legacy remote-code checkpoints nest weights under + # ``radio_model.model.{patch_generator,blocks}.*`` with a fused + # ``attn.qkv`` (renamed to ``attention.qkv`` and loaded directly). + # ``video_patch_projection`` must precede ``embedder`` (applied in order); + # keys mapping to ``None`` are intentionally not loaded (input normalization + # runs in the processor and ``summary_idxs`` comes from the config). + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + # legacy remote-code layout onto the native module tree + "radio_model.model.patch_generator.video_embedder": ( + "embeddings.video_patch_projection" + ), + "radio_model.model.patch_generator.embedder": ( + "embeddings.patch_projection" + ), + "radio_model.model.patch_generator.pos_embed": ( + "embeddings.position_embedding" + ), + "radio_model.model.patch_generator.cls_token.token": ( + "embeddings.cls_register_token" + ), + "radio_model.model.blocks": "encoder.layer", + # legacy names onto the native tree: ``attn`` -> ``attention`` (its + # fused ``qkv``/``proj`` load directly) and bare ``ls1``/``ls2`` -> + # the ``layer_scale{1,2}`` LayerScale modules. + ".attn.": ".attention.", + ".ls1": ".layer_scale1.lambda1", + ".ls2": ".layer_scale2.lambda1", + # native attention output projection: ``attention.output.dense`` is + # current at transformers 5.x (Dinov2 naming); ``attention.o_proj`` + # covers the post-refactor spelling on transformers main. + "attention.output.dense": "attention.proj", + "attention.o_proj": "attention.proj", + # dropped in both layouts: normalization runs in the processor and + # ``summary_idxs`` comes from the config. + "input_conditioner": None, + "summary_idxs": None, + }, + orig_to_new_stacked={ + # Native split q/k/v projections fuse into the ``qkv`` + # QKVParallelLinear. ``attention.attention.{query,key,value}`` is + # current at transformers 5.x; ``attention.{q,k,v}_proj`` covers the + # post-refactor spelling on transformers main. + "attention.attention.query": ("attention.qkv", "q"), + "attention.attention.key": ("attention.qkv", "k"), + "attention.attention.value": ("attention.qkv", "v"), + "attention.q_proj": ("attention.qkv", "q"), + "attention.k_proj": ("attention.qkv", "k"), + "attention.v_proj": ("attention.qkv", "v"), + }, + ) + def __init__( self, - config: PreTrainedConfig = None, + config: PreTrainedConfig, quant_config: QuantizationConfig | None = None, *, num_hidden_layers_override: int | None = None, @@ -527,27 +625,9 @@ def __init__( super().__init__() self.config = config - self.img_size, self.grid_size, self.num_patches = self._init_img_size( - config.patch_size, config.image_size - ) - assert self.img_size is not None - max_img_size = int( - round(config.cpe_max_size / config.patch_size) * config.patch_size - ) self.temporal_patch_size = config.video_temporal_patch_size - unique_teachers = set(t["name"] for t in config.teachers) - self.patch_generator = ViTPatchGenerator( - config.patch_size, - config.hidden_size, - input_dims=self.img_size, - max_input_dims=max_img_size, - cls_token=True, - num_cls_tokens=len(unique_teachers) if config.cls_token_per_teacher else 1, - register_multiple=config.register_multiple, - temporal_patch_size=self.temporal_patch_size, - separate_video_embedder=config.separate_video_embedder, - ) + self.embeddings = RadioPatchEmbeddings(config) self.encoder = RadioVisionEncoder( config=config, quant_config=quant_config, @@ -556,200 +636,89 @@ def __init__( prefix=f"{prefix}.encoder", ) - def _init_img_size( - self, - patch_size: input_dim_t, - img_size: input_dim_t | None, - ) -> tuple[tuple[int, int] | None, tuple[int, int] | None, int | None]: - if img_size is None: - return None, None, None - if isinstance(img_size, int): - img_size = (img_size, img_size) - if isinstance(patch_size, int): - patch_size = (patch_size, patch_size) - grid_size = ( - img_size[0] // patch_size[0], - img_size[1] // patch_size[1], - ) - num_patches = grid_size[0] * grid_size[1] - return img_size, grid_size, num_patches - - def get_input_embeddings(self): - return self.embeddings - - def inter_image_mask_metadata( - self, imgs_sizes: list[tuple[int, int]], device: torch.device - ) -> MaskMetadata: - """Build mask metadata from image pixel sizes. Adds num_skip to each - sequence length (cls/register tokens) to match patch generator output.""" - patch_size = self.patch_generator.patch_size - num_skip = self.patch_generator.num_skip - - seq_lens = calc_seq_lens(imgs_sizes, patch_size) - adjusted = [s + num_skip for s in seq_lens] - return self._inter_image_mask_metadata_from_seq_lens(adjusted, device=device) - - def _inter_image_mask_metadata_from_seq_lens( - self, seq_lens: list[int], device: torch.device - ) -> MaskMetadata: - """Build mask metadata from actual sequence lengths (already including - cls/register tokens, i.e. patch_count + num_skip per item). - Use inter_image_mask_metadata() when you only have imgs_sizes.""" - assert len(seq_lens) > 0 - cu_seqlens = torch.tensor( - list(accumulate(seq_lens, initial=0)), dtype=torch.int32, device=device - ) - # Keep max_seqlen on CPU to avoid .item() sync - # See: https://github.com/vllm-project/vllm/blob/20b6b01/vllm/v1/attention/ops/vit_attn_wrappers.py#L48 - max_seqlen = torch.tensor(max(seq_lens), dtype=torch.int32) - return MaskMetadata(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) + # summary_idxs selects which class tokens to gather: None (no teachers) + # keeps all class tokens; a list (possibly empty) gathers those indices. + if config.summary_idxs is not None: + self.register_buffer( + "summary_idxs", torch.tensor(config.summary_idxs, dtype=torch.long) + ) + else: + self.summary_idxs = None def forward( self, - x: torch.Tensor, + pixel_values: torch.Tensor | None = None, + *, imgs_sizes: list[tuple[int, int]] | None = None, num_frames: int | None = None, - ) -> torch.FloatTensor: + ) -> tuple[torch.FloatTensor, torch.FloatTensor]: T = self.temporal_patch_size - # Build packed-sequence metadata for MMEncoderAttention when needed. mask_meta = None - packed_batch_size = None # Original batch size before packing + packed_batch_size = None # Original batch size before packing. if num_frames is not None and T > 1: # Conv3d video: all tubelets have the same sequence length. - # Pack [num_tubelets, seq_per_tubelet, hidden] → [1, total, hidden] - hidden_states = self.patch_generator.forward_video(x) + # Pack [num_tubelets, seq_per_tubelet, hidden] -> [1, total, hidden]. + hidden_states = self.embeddings.forward_video(pixel_values) packed_batch_size, seq_per_tubelet, hidden_dim = hidden_states.shape hidden_states = hidden_states.reshape(1, -1, hidden_dim) - mask_meta = self._inter_image_mask_metadata_from_seq_lens( + mask_meta = self._mask_metadata_from_seq_lens( [seq_per_tubelet] * packed_batch_size, device=hidden_states.device ) else: - # Images for any model, or video for non-conv3d model - hidden_states = self.patch_generator(x, imgs_sizes=imgs_sizes) + hidden_states = self.embeddings(pixel_values, imgs_sizes=imgs_sizes) if imgs_sizes is not None and len(imgs_sizes) > 1: - # Dynamic resolution w/ > 1 image, create attn mask - mask_meta = self.inter_image_mask_metadata( + # Dynamic resolution w/ > 1 image, create attn mask. + mask_meta = self._inter_image_mask_metadata( imgs_sizes, device=hidden_states.device ) - encoder_outputs = self.encoder(inputs_embeds=hidden_states, mask_meta=mask_meta) + encoder_outputs = self.encoder(hidden_states, mask_meta=mask_meta) - # Unpack back to original batch shape if we packed for video + # Unpack back to original batch shape if we packed for video. if packed_batch_size is not None: encoder_outputs = encoder_outputs.reshape( packed_batch_size, seq_per_tubelet, -1 ) - return encoder_outputs - - -class RadioModel(nn.Module): - packed_modules_mapping = { - "qkv": ["qkv"], - } - - def __init__( - self, - config: PreTrainedConfig, - quant_config: QuantizationConfig | None = None, - *, - num_hidden_layers_override: int | None = None, - num_dummy_heads: int = 0, - prefix: str = "", - ) -> None: - super().__init__() + return self._extract_final(encoder_outputs, imgs_sizes=imgs_sizes) - self.config = config - self.model = RadioInternVisionModel( - config=config, - quant_config=quant_config, - num_hidden_layers_override=num_hidden_layers_override, - num_dummy_heads=num_dummy_heads, - prefix=prefix, + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load RADIO checkpoint weights, remapping legacy remote-code and + native Transformers keys onto this module tree via + :data:`hf_to_vllm_mapper`.""" + if isinstance(weights, dict): + weights = weights.items() + # Keep only tensors belonging to the RADIO tower (legacy ``radio_model.*`` + # or native top-level keys); ignore any auxiliary keys a fuller checkpoint + # may carry (an unmapped tower key still errors). + prefixes = ( + "radio_model.", + "embeddings.", + "encoder.", + "input_conditioner.", + "summary_idxs", ) - - summary_idxs = None - if config.teachers: - summary_idxs = torch.tensor( - [i for i, t in enumerate(config.teachers) if t.get("use_summary", True)] + weights = ((name, w) for name, w in weights if name.startswith(prefixes)) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + if not loaded: + logger.warning( + "RadioModel.load_weights matched no checkpoint tensors; the " + "vision tower will keep its random initialization. Expected " + "legacy 'radio_model.*' or native 'embeddings.*'/'encoder.*' " + "keys." ) - if summary_idxs.numel() > 0: - self.register_buffer("summary_idxs", summary_idxs) - self.summary_idxs = summary_idxs - - def forward( - self, - pixel_values: torch.Tensor | None = None, - pixel_embeds: torch.Tensor | None = None, - *, - imgs_sizes: list[tuple[int, int]] | None = None, - num_frames: int | None = None, - ) -> tuple[torch.FloatTensor, torch.FloatTensor]: - y = self.model( - pixel_values, - imgs_sizes=imgs_sizes, - num_frames=num_frames, - ) - return self._extract_final(y, imgs_sizes=imgs_sizes) - - def load_weights(self, weights) -> set[str]: - loaded_params: set[str] = set() - params_dict = dict(self.named_parameters()) - - if isinstance(weights, dict): - weights_list = list(weights.items()) - else: - weights_list = list(weights) - - for name, weight in weights_list: - if not name.startswith("radio_model."): - # Skip non-radio weights - continue - - sub = name[len("radio_model.") :] # drop "radio_model." prefix - - # Skip buffers not used in vLLM - if sub in {"summary_idxs"}: - continue - if sub.startswith("input_conditioner."): - # we normalize in the input processor, - # based on norm and std values from the config - continue - - vllm_key = None - if sub.startswith("model.patch_generator."): - vllm_key = f"model.patch_generator.{sub.split('.', 2)[-1]}" - elif sub.startswith("input_conditioner."): - vllm_key = f"input_conditioner.{sub.split('.', 1)[-1]}" - elif sub.startswith("model.blocks."): - # Encoder blocks: HF 'model.blocks.{i}.' -> - # vLLM 'model.encoder.layers.{i}.' - parts = sub.split(".") - if len(parts) >= 4: - layer_idx = parts[2] - suffix = ".".join(parts[3:]) - # Skip layer-scale entries that vLLM doesn't use - if suffix in {"ls1", "ls2"} or suffix.startswith(("ls1.", "ls2.")): - continue - vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}" - - if vllm_key and vllm_key in params_dict: - param = params_dict[vllm_key] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, weight) - loaded_params.add(vllm_key) - - return loaded_params + return loaded def _extract_final( self, y: torch.Tensor, imgs_sizes: list[tuple[int, int]] | None = None ) -> tuple[torch.FloatTensor, torch.FloatTensor]: - # Remove CLS + REGISTERS tokens - num_skip = self.model.patch_generator.num_skip - patch_size = self.model.patch_generator.patch_size - num_cls_tokens = self.model.patch_generator.num_cls_tokens + # Remove CLS + REGISTER tokens. + num_skip = self.embeddings.num_skip + patch_size = self.embeddings.patch_size + num_cls_tokens = self.embeddings.num_cls_tokens if imgs_sizes is None: all_summary = y[:, :num_cls_tokens] all_feat = y[:, num_skip:] @@ -773,3 +742,31 @@ def _extract_final( else: bb_summary = all_summary return bb_summary.flatten(1), all_feat + + def get_input_embeddings(self): + return self.embeddings + + def _inter_image_mask_metadata( + self, imgs_sizes: list[tuple[int, int]], device: torch.device + ) -> MaskMetadata: + """Build mask metadata from image pixel sizes. Adds ``num_skip`` to each + sequence length (cls/register tokens) to match patch generator output.""" + patch_size = self.embeddings.patch_size + num_skip = self.embeddings.num_skip + seq_lens = calc_seq_lens(imgs_sizes, patch_size) + adjusted = [s + num_skip for s in seq_lens] + return self._mask_metadata_from_seq_lens(adjusted, device=device) + + def _mask_metadata_from_seq_lens( + self, seq_lens: list[int], device: torch.device + ) -> MaskMetadata: + """Build mask metadata from sequence lengths (already including + cls/register tokens, i.e. ``patch_count + num_skip`` per item).""" + assert len(seq_lens) > 0 + cu_seqlens = torch.tensor( + list(accumulate(seq_lens, initial=0)), dtype=torch.int32, device=device + ) + # Keep max_seqlen on CPU to avoid .item() sync. + # See: https://github.com/vllm-project/vllm/blob/20b6b01/vllm/v1/attention/ops/vit_attn_wrappers.py#L48 + max_seqlen = torch.tensor(max(seq_lens), dtype=torch.int32) + return MaskMetadata(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) diff --git a/vllm/transformers_utils/configs/radio.py b/vllm/transformers_utils/configs/radio.py index 556e1b451a24..671985fec6e8 100644 --- a/vllm/transformers_utils/configs/radio.py +++ b/vllm/transformers_utils/configs/radio.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Radio vision model configuration.""" +"""Radio vision model configuration""" from typing import Any @@ -37,6 +37,7 @@ class RadioConfig(PreTrainedConfig): norm_type: The normalization type to use. layer_norm_eps: The epsilon used by the layer normalization layers. initializer_factor: A factor for initializing all weight matrices. + layerscale_value: Initial value for the LayerScale ``lambda1`` parameters. hidden_act: The non-linear activation function in the encoder. cpe_max_size: Maximum image size for position embeddings. norm_mean: Mean values for image normalization (RGB channels). @@ -44,6 +45,8 @@ class RadioConfig(PreTrainedConfig): norm_std: Standard deviation values for image normalization (RGB channels). Defaults to (0.26862954, 0.26130258, 0.27577711)). register_multiple: Number of register tokens to use. + cpe_num_registers: Explicit CPE register-token count, used by checkpoints + that set ``register_multiple=0`` (e.g. C-RADIOv3) instead. teachers: A list of teacher model configurations. Each teacher configuration is a dict with keys like "name" and some may have "use_summary". cls_token_per_teacher: Whether to use a separate CLS token for each teacher. @@ -55,6 +58,11 @@ class RadioConfig(PreTrainedConfig): dedicated video patch embedder (3*T*P*P -> hidden) separate from the image embedder (3*P*P -> hidden). When False, a single embedder with input size 3*T*P*P is used for both (images are duplicated T times). + num_channels: Number of input image channels (RGB -> 3). + + ``num_cls_tokens``, ``num_registers``, ``summary_idxs`` and ``max_img_size`` + are derived from the fields above so the model reads a single canonical + schema (these were previously computed inside the model). """ @@ -70,15 +78,18 @@ def __init__( norm_type: str = "layer_norm", layer_norm_eps: float = 1e-6, initializer_factor: float = 1.0, + layerscale_value: float = 1.0, hidden_act: str = "gelu", cpe_max_size: int = 2048, norm_mean: tuple[float, float, float] | list = OPENAI_CLIP_MEAN, norm_std: tuple[float, float, float] | list = OPENAI_CLIP_STD, register_multiple: int | None = None, + cpe_num_registers: int | None = None, teachers: list[dict[str, Any]] | None = None, cls_token_per_teacher: bool = False, video_temporal_patch_size: int = 1, separate_video_embedder: bool = True, + num_channels: int = 3, **kwargs, ): self.model_name = model_name @@ -95,6 +106,7 @@ def __init__( self.norm_type = norm_type self.layer_norm_eps = layer_norm_eps self.initializer_factor = initializer_factor + self.layerscale_value = layerscale_value self.hidden_act = hidden_act self.cpe_max_size = cpe_max_size self.norm_mean = ( @@ -104,8 +116,35 @@ def __init__( list(norm_std) if isinstance(norm_std, (tuple, list)) else norm_std ) self.register_multiple = register_multiple + self.cpe_num_registers = cpe_num_registers self.teachers = teachers if teachers is not None else [] self.cls_token_per_teacher = cls_token_per_teacher self.video_temporal_patch_size = video_temporal_patch_size self.separate_video_embedder = separate_video_embedder + self.num_channels = num_channels + + # Fields derived from the fields above (previously computed inside the + # model), so the model consumes a single canonical schema. + unique_teachers = {t["name"] for t in self.teachers} + self.num_cls_tokens = len(unique_teachers) if self.cls_token_per_teacher else 1 + if self.cpe_num_registers: + # An explicit register count wins, matching the reference ClsToken. + # C-RADIOv3 carries its count here and sets register_multiple=0. + self.num_registers = self.cpe_num_registers + elif self.register_multiple: + self.num_registers = self.register_multiple - ( + self.num_cls_tokens % self.register_multiple + ) + else: + self.num_registers = 0 + # None (no teachers) means no explicit selection -> keep all class + # tokens; a list (possibly empty) gathers exactly those indices. + self.summary_idxs = ( + [i for i, t in enumerate(self.teachers) if t.get("use_summary", True)] + if self.teachers + else None + ) + self.max_img_size = int( + round(self.cpe_max_size / self.patch_size) * self.patch_size + ) super().__init__(**kwargs) From 5cf5c8fa6d13ea7bef6ad799e32760530beca88f Mon Sep 17 00:00:00 2001 From: "Wenfei (Charles) Qi" <74716511+Manfredss@users.noreply.github.com> Date: Thu, 29 Oct 2026 06:01:00 -0700 Subject: [PATCH 134/406] [Bugfix][Quantization] Fix Humming FP8 linear for per-tensor/per-channel FP8 weights (#60163) Signed-off-by: wenfei qi Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/quantization/test_fp8.py | 9 +++++---- .../model_executor/kernels/linear/__init__.py | 3 +++ .../model_executor/layers/quantization/fp8.py | 2 ++ .../layers/quantization/online/fp8.py | 19 ++++++++++++++----- 4 files changed, 24 insertions(+), 9 deletions(-) diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 0b59a068cafa..0638c752c512 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -659,14 +659,15 @@ def test_static_fp8_moe_input_scales_remain_scalar() -> None: ) @pytest.mark.parametrize("model_id", MODELS) @pytest.mark.parametrize( - "force_marlin", [True, False] if current_platform.is_cuda() else [False] + "linear_backend", + ["auto", "marlin", "humming"] if current_platform.is_cuda() else ["auto"], ) @pytest.mark.parametrize( "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False] ) def test_model_load_and_run( model_id: str, - force_marlin: bool, + linear_backend: str, use_rocm_aiter: bool, monkeypatch, dist_init, @@ -676,8 +677,8 @@ def test_model_load_and_run( monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") kernel_config = KernelConfig( - linear_backend="marlin" if force_marlin else "auto", - moe_backend="marlin" if force_marlin else "auto", + linear_backend=linear_backend, + moe_backend="marlin" if linear_backend == "marlin" else "auto", ) model, vllm_config = load_model_without_vllm_runner( model_id, diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index b4d7bd3ef669..8301cbe481a3 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -836,6 +836,9 @@ def fp8_scaled_mm_is_batch_invariant( kernel: FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel, ) -> bool: """Whether a row's output from the kernel is independent of the other rows.""" + # Weight-only here, and Humming honors VLLM_BATCH_INVARIANT itself. + if isinstance(kernel, HummingFP8ScaledMMLinearKernel): + return True if not isinstance( kernel, (CutlassFP8ScaledMMLinearKernel, TritonFP8ScaledMMLinearKernel) ): diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 0c4fe9889268..e63efb792fbe 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -406,6 +406,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Update layer with new values. replace_parameter(layer, "weight", weight.data) + # The helper returns (K, N); the param keeps its checkpoint dims. + layer.processed_weight_dims = (0, 1) replace_parameter(layer, "weight_scale", weight_scale.data) if input_scale is not None: diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index f4a327192108..01bf5f87033a 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -22,6 +22,9 @@ from vllm.model_executor.kernels.linear.scaled_mm import ( MarlinFP8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.humming import ( + HummingFP8ScaledMMLinearKernel, +) from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( @@ -249,6 +252,8 @@ def process_weights_after_loading(self, layer: Module) -> None: # Update layer with new values. replace_parameter(layer, "weight", qweight.t().data) + # Transposed to (K, N); the param keeps its checkpoint dims. + layer.processed_weight_dims = (0, 1) replace_parameter(layer, "weight_scale", weight_scale.data) if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"): @@ -415,13 +420,17 @@ def create_weights( out_dtype=self.out_dtype, module_name=self.__class__.__name__, ) - # PTPC requires per-token activation FP8; MarlinFP8 is W8A16 and - # would silently produce a weight-only fp8 model. - if isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel): + # PTPC requires per-token activation FP8; Marlin and Humming FP8 are + # W8A16 and would silently produce a weight-only fp8 model. + if isinstance( + self.fp8_linear, + (MarlinFP8ScaledMMLinearKernel, HummingFP8ScaledMMLinearKernel), + ): raise ValueError( "FP8 PTPC online quant requires a kernel that honors " - "per-token activation quantization; MarlinFP8 is W8A16 " - "weight-only. Requires SM89+ for Cutlass FP8 or ROCm MI3xx " + "per-token activation quantization; " + f"{type(self.fp8_linear).__name__} is W8A16 weight-only. " + "Requires SM89+ for Cutlass FP8 or ROCm MI3xx " "for rowwise scaled_mm." ) From cf3f6405b849f2b672d93af0c4cf4e5e6dc6bd1b Mon Sep 17 00:00:00 2001 From: jbottleson-nvidia Date: Thu, 29 Oct 2026 06:07:56 -0700 Subject: [PATCH 135/406] [Bugfix] Default Nemotron Omni mamba_ssm_cache_dtype to float32 (#63218) Signed-off-by: Jeremy Bottleson Signed-off-by: jbottleson-nvidia Co-authored-by: Cyrus Leung Co-authored-by: vllm-agent --- .../multimodal/test_nano_nemotron_vl.py | 23 +++++++++++++++++++ vllm/model_executor/models/config.py | 3 +++ 2 files changed, 26 insertions(+) diff --git a/tests/models/multimodal/test_nano_nemotron_vl.py b/tests/models/multimodal/test_nano_nemotron_vl.py index 41f9428e6d73..8d1b9e4b06fc 100644 --- a/tests/models/multimodal/test_nano_nemotron_vl.py +++ b/tests/models/multimodal/test_nano_nemotron_vl.py @@ -8,6 +8,7 @@ import torch from vllm import envs +from vllm.engine.arg_utils import EngineArgs from vllm.model_executor.models.nano_nemotron_vl import ( NanoNemotronVLMultiModalProcessor, NemotronH_Nano_VL_V2, @@ -256,3 +257,25 @@ def test_extract_audio_from_videos_rejects_oversized_audio(): assert audio_items == [] assert has_audio == [False] + + +@pytest.mark.core_model +@pytest.mark.parametrize( + "model", + [ + "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16", + "nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16", + ], +) +@pytest.mark.parametrize( + ("requested", "expected"), [("auto", "float32"), ("bfloat16", "bfloat16")] +) +def test_omni_mamba_ssm_cache_dtype(model: str, requested: str, expected: str): + """Omni archs get the NemotronH float32 SSM-cache default; explicit wins.""" + vllm_config = EngineArgs( + model=model, + trust_remote_code=True, + mamba_ssm_cache_dtype=requested, + max_model_len=4096, + ).create_engine_config() + assert vllm_config.cache_config.mamba_ssm_cache_dtype == expected diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index e9125baae425..946e905d0cf5 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -1326,6 +1326,9 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: "NemotronHForCausalLM": NemotronHForCausalLMConfig, "NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig, "NemotronH_Nano_VL_V2": NemotronHNanoVLV2Config, + "NemotronH_Nano_Omni_Reasoning_V3": NemotronHNanoVLV2Config, + "NemotronH_Super_Omni_Reasoning_V3": NemotronHNanoVLV2Config, + "NemotronH_Omni_Reasoning_V3": NemotronHNanoVLV2Config, "Nemotron3_5TranscribeForConditionalGeneration": ( Nemotron3_5TranscribeForConditionalGenerationConfig ), From 21db9ecf2a6271ab7e745cbd5537fcc70bd69c99 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:21:24 +0800 Subject: [PATCH 136/406] [Model] Support Ming3-ASR (inclusionAI/Ming-3.0-asr-tiny) (#63585) Signed-off-by: Yuanheng Co-authored-by: codex Co-authored-by: vllm-agent Co-authored-by: Isotr0py --- docs/contributing/model/transcription.md | 8 +- docs/models/supported_models.md | 1 + .../multimodal/audio_language_offline.py | 21 + .../models/inclusionAI/Ming-3.0-asr-tiny.yaml | 244 +++++++++ .../test_transcription_validation.py | 179 ++++++- .../test_transcription_validation_whisper.py | 1 + .../multimodal/processing/test_ming_v3_asr.py | 220 ++++++++ tests/models/registry.py | 11 + vllm/config/speech_to_text.py | 9 +- .../speech_to_text/base/serving.py | 10 +- .../speech_to_text/transcription/protocol.py | 1 + vllm/model_executor/models/bailing_moe_v3.py | 3 +- vllm/model_executor/models/ming_audio.py | 151 ++++++ vllm/model_executor/models/ming_v3_asr.py | 497 ++++++++++++++++++ .../models/ming_v3_asr_logits_processor.py | 76 +++ vllm/model_executor/models/registry.py | 4 + vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 2 + .../transformers_utils/configs/ming_v3_asr.py | 85 +++ .../transformers_utils/processors/__init__.py | 2 + .../processors/ming_audio.py | 35 ++ 21 files changed, 1547 insertions(+), 14 deletions(-) create mode 100644 recipes/models/inclusionAI/Ming-3.0-asr-tiny.yaml create mode 100644 tests/models/multimodal/processing/test_ming_v3_asr.py create mode 100644 vllm/model_executor/models/ming_audio.py create mode 100644 vllm/model_executor/models/ming_v3_asr.py create mode 100644 vllm/model_executor/models/ming_v3_asr_logits_processor.py create mode 100644 vllm/transformers_utils/configs/ming_v3_asr.py create mode 100644 vllm/transformers_utils/processors/ming_audio.py diff --git a/docs/contributing/model/transcription.md b/docs/contributing/model/transcription.md index f82bafa44d52..4fdab121cd85 100644 --- a/docs/contributing/model/transcription.md +++ b/docs/contributing/model/transcription.md @@ -66,7 +66,11 @@ This is for controlling general behavior of the API when serving your model: See [Audio preprocessing and chunking](#audio-preprocessing-and-chunking) for what each field controls. -Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns: +Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. + +For transcription requests, `stt_params.extra_args` contains the optional `vllm_xargs` dictionary, which is also passed to `SamplingParams.extra_args`. Models can read request-specific frontend or prompt options here and should validate the options they consume. An omitted dictionary is `None`. Multipart clients can send `vllm_xargs` as a JSON string. These options apply to the request; do not use them to mutate the shared model configuration. + +There are two common patterns: #### Multimodal LLM with audio embeddings (e.g., Voxtral, Gemma3n) @@ -194,7 +198,7 @@ Provide a fast duration→token estimate to improve streaming usage statistics: The API server takes care of basic audio I/O and optional chunking before building prompts: -- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate` using `AudioResampler`. +- Resampling: Input audio is resampled to `SpeechToTextConfig.sample_rate`. Models can set specific `audio_resample_method` (e.g., `pyav`, `scipy`, `torchaudio`, etc) to decode at the native sample rate, then resample with `AudioResampler` before duration calculation and chunking. This can match a model's reference audio frontend without changing other models' defaults. - Chunking: If `SpeechToTextConfig.allow_audio_chunking` is True and the duration exceeds `max_audio_clip_s`, the server splits the audio into chunks and generates a prompt per chunk. There is no overlap between chunks, overlap_chunk_second controls the size of the search window used to find the split point. - Energy-aware splitting: When `min_energy_split_window_size` is set, the server finds low-energy regions to minimize cutting within words. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 8cb133a66647..cbcc08b2a006 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -744,6 +744,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition. | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ | | `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ | +| `MingV3ASRForConditionalGeneration` | Ming-3.0-ASR | `inclusionAI/Ming-3.0-asr-tiny` | | | | `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | | `Nemotron3_5TranscribeForConditionalGeneration` | Nemotron-3.5-Transcribe | `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-FP8`, `nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B-NVFP4` | | ✅︎ | | `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ | diff --git a/examples/generate/multimodal/audio_language_offline.py b/examples/generate/multimodal/audio_language_offline.py index 8b2215d8c1eb..5866101973d0 100644 --- a/examples/generate/multimodal/audio_language_offline.py +++ b/examples/generate/multimodal/audio_language_offline.py @@ -251,6 +251,26 @@ def run_midashenglm(question: str, audio_count: int): ) +# Ming3-ASR +def run_ming_v3_asr(question: str, audio_count: int) -> ModelRequestData: + from vllm.model_executor.models.ming_v3_asr import ASR_INSTRUCTION + + assert audio_count == 1, "Ming3-ASR only supports single audio input per prompt" + engine_args = EngineArgs( + model="inclusionAI/Ming-3.0-asr-tiny", + trust_remote_code=True, + dtype="bfloat16", + max_model_len=8192, + max_num_seqs=4, + limit_mm_per_prompt={"audio": audio_count}, + enforce_eager=True, + ) + return ModelRequestData( + engine_args=engine_args, + prompt=ASR_INSTRUCTION + "", + ) + + # MiniCPM-O def run_minicpmo(question: str, audio_count: int) -> ModelRequestData: model_name = "openbmb/MiniCPM-o-2_6" @@ -503,6 +523,7 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData: "granite_speech": run_granite_speech, "kimi_audio": run_kimi_audio, "midashenglm": run_midashenglm, + "ming_v3_asr": run_ming_v3_asr, "minicpmo": run_minicpmo, "phi4_mm": run_phi4mm, "qwen2_audio": run_qwen2_audio, diff --git a/recipes/models/inclusionAI/Ming-3.0-asr-tiny.yaml b/recipes/models/inclusionAI/Ming-3.0-asr-tiny.yaml new file mode 100644 index 000000000000..0a43c7029aba --- /dev/null +++ b/recipes/models/inclusionAI/Ming-3.0-asr-tiny.yaml @@ -0,0 +1,244 @@ +meta: + title: "Ming3-ASR-tiny" + slug: "ming-3-0-asr-tiny" + provider: "inclusionAI" + description: "Chinese, English, and dialect speech recognition with topic, hotword, and dialogue context." + date_added: 2026-10-26 + date_updated: 2026-10-27 + tasks: + - multimodal + performance_headline: "Context-aware Chinese, English, and dialect ASR; single-GPU serving" + +model: + model_id: "inclusionAI/Ming-3.0-asr-tiny" + min_vllm_version: "nightly" + nightly_required: true + install: + pip: + note: "Install from main for now." + command: | + git clone --branch main https://github.com/vllm-project/vllm.git + cd vllm + uv venv --python 3.12 + source .venv/bin/activate + VLLM_USE_PRECOMPILED=1 uv pip install --editable ".[audio]" --torch-backend=auto + vllm download-kernels + docker: false + architecture: moe + parameter_count: "8.55B" + active_parameters: "1.3B" + context_length: 8192 + base_args: + - "--trust-remote-code" + base_env: {} + +dependencies: + - note: "Ming uses torchaudio's sinc resampler; install a build compatible with PyTorch." + command: "uv pip install torchaudio --torch-backend=auto" + +features: + no_repeat_ngram: + description: "Load the Model Runner V2 plugin for request-level full-history n-gram blocking." + args: + - "--logits-processors" + - "vllm.model_executor.models.ming_v3_asr_logits_processor:NoRepeatNGramLogitsProcessor" + +opt_in_features: [] + +variants: + default: + precision: bf16 + vram_minimum_gb: 22 + description: "BF16 checkpoint; the capacity estimate includes the audio tower and is not a measured minimum." + +compatible_strategies: + - single_node_tp + +hardware_overrides: {} + +strategy_overrides: + single_node_tp: + tp: 1 + +guide: | + ## Overview + + [Ming-3.0-asr-tiny](https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny) supports Chinese, English, and dialect transcription with optional topic, hotword, and dialogue context. + This guide covers single-GPU serving, transcription, streaming, Context ASR, and offline inference with vLLM. + + ## Public example audio + + Download the three clips from the official inference repository: + + ```bash + MING_AUDIO_BASE=https://raw.githubusercontent.com/inclusionAI/Ming-3.0-asr/6141e8c552815db724d8df0f72c9d865c547a1ea/examples/cases/audio + curl -fL "$MING_AUDIO_BASE/english_common_voice.wav" -o english.wav + curl -fL "$MING_AUDIO_BASE/chinese_common_voice.wav" -o chinese.wav + curl -fL "$MING_AUDIO_BASE/context_same_sentence_pcm16.wav" -o context.wav + ``` + + | Clip | Observed transcription | + | --- | --- | + | English | `same way you did` | + | Mandarin | `漓江和湘江有秦代开凿的灵渠沟通` | + | Context clip with the hints below | Contains the film title `同囚` | + + The Context clip illustrates correction of `铜球` to `同囚` with a relevant topic and hotwords. These are sample-specific observations, not a general Context ASR accuracy claim. Audio provenance and licenses are documented in the [official examples](https://github.com/inclusionAI/Ming-3.0-asr/blob/6141e8c552815db724d8df0f72c9d865c547a1ea/examples/cases/README.md). + + For accuracy comparisons, keep batch settings and attention backends consistent; greedy transcription outputs can differ across these configurations. + + ## Transcription API + + Use the audio requests below to verify speech recognition. The catalog's generic text-only cURL and benchmark shortcuts do not exercise ASR. + + ```bash + curl http://localhost:8000/v1/audio/transcriptions \ + -F model=inclusionAI/Ming-3.0-asr-tiny \ + -F file=@chinese.wav \ + -F temperature=0 \ + -F max_completion_tokens=256 \ + -F 'vllm_xargs={"no_repeat_ngram_size":10}' + ``` + + Omit `language` for automatic recognition, or use `language=zh` / `language=en` to seed the known language. These are the supported language-hint codes; dialect transcription does not introduce additional hint codes. The response contains the transcript without the generated language label. `response_format=json` and `response_format=text` are supported. + + ### OpenAI Python client + + ```python + import json + from openai import OpenAI + + client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY") + with open("english.wav", "rb") as audio: + result = client.audio.transcriptions.create( + model="inclusionAI/Ming-3.0-asr-tiny", + file=audio, + temperature=0, + extra_body={ + "max_completion_tokens": 256, + "vllm_xargs": json.dumps({"no_repeat_ngram_size": 10}), + }, + ) + print(result.text) + ``` + + `vllm_xargs` is JSON-encoded for the multipart request, as in the cURL examples. + + ### Context ASR + + ```bash + curl http://localhost:8000/v1/audio/transcriptions \ + -F model=inclusionAI/Ming-3.0-asr-tiny \ + -F file=@context.wav \ + -F language=zh \ + -F temperature=0 \ + -F max_completion_tokens=256 \ + -F 'prompt=This is an audio about 同囚.' \ + -F 'hotwords=同囚,黄国权,游学修,关楚耀,赵永洪,赵咏瑶,狱警,长镜头,吞玻璃戏,糖化玻璃,真实案例数据.' \ + -F 'vllm_xargs={"no_repeat_ngram_size":10,"context_chunks":true}' + ``` + + `prompt` accepts topic or preceding dialogue appended to the official ASR instruction, or the full official instruction followed by context. Context and hotwords must be single-line text without audio placeholders. To reproduce the vendor's full prompt exactly, put its hotword text in `prompt` and omit the separate `hotwords` field. + + ### Streaming + + ```bash + curl -N http://localhost:8000/v1/audio/transcriptions \ + -F model=inclusionAI/Ming-3.0-asr-tiny \ + -F file=@english.wav \ + -F language=en \ + -F temperature=0 \ + -F stream=true \ + -F max_completion_tokens=256 \ + -F 'vllm_xargs={"no_repeat_ngram_size":10}' + ``` + + Streaming sends transcript deltas after the complete audio upload. It is not live microphone input. The English clip produces `same way you did` without leaking the language prefix into the transcript. + + ## Chunking and repetition control + + Audio is downmixed and resampled to 16 kHz using torchaudio. The processor extracts the entire clip's log-mel features before splitting encoder chunks. Chunk settings are request-local and are included in multimodal cache keys. + + | Option in `vllm_xargs` | Behavior | + | --- | --- | + | No chunk option | Use 3,000-mel-frame encoder chunks, approximately 30 seconds each. | + | `"context_chunks": true` | Use 3,000 frames up to 30 seconds and 6,000 frames for 30–60 seconds; reject longer clips. | + | `"audio_chunk_frames": 6000` | Explicit chunk size; any integer from 1 through 6,000 is accepted. | + | `"no_repeat_ngram_size": 10` | Block repeated 10-grams over prompt plus generated tokens. | + + `context_chunks` and an explicit `audio_chunk_frames` are mutually exclusive. Ordinary chunking can handle clips beyond 60 seconds within the shared 8,192-token input/output budget and audio-upload size limit. Internal encoder chunking does not create separate transcription requests. + + The no-repeat feature loads a plugin; each request must still set a nonzero `no_repeat_ngram_size` to enable blocking. Omitted values and `0` leave ordinary sampling in effect. Blocking uses the full prompt/output history, and requires Model Runner V2. Speculative decoding and MTP are unsupported with this plugin. + + ## Offline inference + + From the same vLLM source checkout, use the shared [audio language model example](https://github.com/vllm-project/vllm/blob/main/examples/generate/multimodal/audio_language_offline.py) with its built-in audio assets: + + ```bash + python examples/generate/multimodal/audio_language_offline.py --model-type ming_v3_asr + ``` + + This uses the shared example's sampling defaults (`temperature=0.2`, `max_tokens=64`) and prints raw model output, including the language label. For the downloaded Context clip, use the native prompt builder with greedy decoding and 10-gram blocking: + + ```python + from vllm import LLM, SamplingParams + from vllm.config.speech_to_text import SpeechToTextParams + from vllm.model_executor.models.ming_v3_asr import ( + MingV3ASRForConditionalGeneration as MingASR, + ) + from vllm.multimodal.audio import AudioResampler + from vllm.multimodal.media.audio import load_audio + + def main(): + llm = LLM( + model="inclusionAI/Ming-3.0-asr-tiny", + trust_remote_code=True, + dtype="bfloat16", + max_model_len=8192, + max_num_seqs=4, + enforce_eager=True, + logits_processors=[ + "vllm.model_executor.models.ming_v3_asr_logits_processor:" + "NoRepeatNGramLogitsProcessor", + ], + ) + audio, rate = load_audio("context.wav", sr=None, mono=True) + audio = AudioResampler(16000).resample(audio, orig_sr=rate) + params = SpeechToTextParams( + audio=audio, + stt_config=MingASR.get_speech_to_text_config(llm.model_config, "transcribe"), + model_config=llm.model_config, + language="zh", + request_prompt="This is an audio about 同囚.", + hotwords=( + "同囚,黄国权,游学修,关楚耀,赵永洪,赵咏瑶,狱警,长镜头," + "吞玻璃戏,糖化玻璃,真实案例数据." + ), + extra_args={"context_chunks": True}, + ) + outputs = llm.generate( + MingASR.get_generation_prompt(params), + SamplingParams( + temperature=0, + max_tokens=256, + extra_args={"no_repeat_ngram_size": 10}, + ), + ) + print(MingASR.post_process_output(outputs[0].outputs[0].text)) + + + if __name__ == "__main__": + main() + ``` + + ## Limitations + + The [official model card](https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny#usage-notes) describes transcription text output without word-level timestamps or speaker diarization. + + The current vLLM integration does not support translation, segment timestamps, or live audio input. Translation and timestamp requests are rejected. The `context_chunks=true` policy accepts clips up to 60 seconds; ordinary chunking, including requests with context prompts, can accept longer clips subject to the model context and upload limits. Longer recordings can produce incomplete transcripts; accepting an upload does not guarantee complete transcription. + + ## References + + - [model card](https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny) + - [Official inference repository](https://github.com/inclusionAI/Ming-3.0-asr) + - [vLLM speech-to-text API](https://docs.vllm.ai/en/latest/serving/online_serving/speech_to_text.html) diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py index 2237a091ce64..6a14635ca47d 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py @@ -1,12 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# imports for structured outputs tests import json +from io import BytesIO +import httpx +import numpy as np import pytest +import soundfile as sf +from fastapi import UploadFile +from transformers import WhisperConfig from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer +from vllm.config import ModelConfig, SpeechToTextConfig +from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionRequest +from vllm.multimodal.audio import AudioResampler MISTRAL_FORMAT_ARGS = [ "--tokenizer_mode", @@ -16,6 +24,69 @@ "--load_format", "mistral", ] +MING_ASR_ARGS = [ + "--trust-remote-code", + "--max-model-len", + "8192", + *(ROCM_EXTRA_ARGS or ["--max-num-seqs", "4"]), + "--max-num-batched-tokens", + "1024", + "--logits-processors", + "vllm.model_executor.models.ming_v3_asr_logits_processor:NoRepeatNGramLogitsProcessor", +] + + +@pytest.fixture +def ming_server(model_name, request): + args = ["--enforce-eager", *MING_ASR_ARGS] + options = getattr(request, "param", {}) + if options.get("enable_prefix_caching", False): + args = [arg for arg in args if arg != "--no-enable-prefix-caching"] + args.append("--enable-prefix-caching") + with RemoteOpenAIServer(model_name, args) as remote: + yield remote + + +@pytest.fixture(scope="module") +def ming_audio(): + base_url = ( + "https://raw.githubusercontent.com/inclusionAI/Ming-3.0-asr/" + "6141e8c552815db724d8df0f72c9d865c547a1ea/examples/cases/audio/" + ) + with httpx.Client() as client: + clips = {} + for name in ( + "chinese_common_voice.wav", + "context_same_sentence_pcm16.wav", + ): + response = client.get(base_url + name) + response.raise_for_status() + clips[name] = response.content + return clips + + +@pytest.mark.parametrize("encoding", ["absent", "dict", "json"]) +def test_request_extensions_reach_prompt_and_sampling_params(tmp_path, encoding): + """Both request consumers must receive extensions, including form JSON.""" + WhisperConfig().save_pretrained(tmp_path) + config = ModelConfig(model=str(tmp_path)) + extensions = {"frontend_window": 128, "custom_tags": ["speech", 1]} + options = {} + if encoding != "absent": + options["vllm_xargs"] = ( + json.dumps(extensions) if encoding == "json" else extensions + ) + request = TranscriptionRequest(file=UploadFile(BytesIO()), **options) + params = request.build_stt_params( + audio=np.zeros(160, dtype=np.float32), + stt_config=SpeechToTextConfig(), + model_config=config, + task_type="transcribe", + ) + + expected = None if encoding == "absent" else extensions + assert params.extra_args == expected + assert (request.to_sampling_params(16).extra_args or {}) == (expected or {}) async def transcribe_and_check( @@ -23,10 +94,12 @@ async def transcribe_and_check( model_name: str, file, *, - language: str, + language: str | None, expected_text: str, expected_seconds: int | None = None, case_sensitive: bool = False, + response_format: str = "text", + extra_body: dict | None = None, ): """Run a transcription request and assert the output contains *expected_text* and optionally that usage reports *expected_seconds*. @@ -37,12 +110,16 @@ async def transcribe_and_check( model=model_name, file=file, language=language, - response_format="text", + response_format=response_format, temperature=0.0, + extra_body=extra_body, + ) + out = ( + json.loads(transcription) + if isinstance(transcription, str) + else transcription.model_dump() ) - out = json.loads(transcription) out_text = out["text"] - out_usage = out["usage"] if case_sensitive: assert expected_text in out_text, ( @@ -55,21 +132,35 @@ async def transcribe_and_check( ) if expected_seconds is not None: + out_usage = out["usage"] assert out_usage["seconds"] == expected_seconds, ( f"Expected {expected_seconds}s of audio, " f"got {out_usage['seconds']}s. Full usage: {out_usage!r}" ) + return out_text + @pytest.mark.asyncio @pytest.mark.parametrize( - "model_name", ["mistralai/Voxtral-Mini-3B-2507", "Qwen/Qwen3-ASR-0.6B"] + "model_name", + [ + "mistralai/Voxtral-Mini-3B-2507", + "Qwen/Qwen3-ASR-0.6B", + "inclusionAI/Ming-3.0-asr-tiny", + ], ) async def test_basic_audio(mary_had_lamb, model_name): server_args = ["--enforce-eager", *ROCM_EXTRA_ARGS] - + extra_body = None if model_name.startswith("mistralai"): server_args += MISTRAL_FORMAT_ARGS + elif model_name == "inclusionAI/Ming-3.0-asr-tiny": + server_args = ["--enforce-eager", *MING_ASR_ARGS] + extra_body = { + "max_completion_tokens": 256, + "vllm_xargs": json.dumps({"no_repeat_ngram_size": 10}), + } # Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb. with RemoteOpenAIServer(model_name, server_args) as remote_server: @@ -81,6 +172,7 @@ async def test_basic_audio(mary_had_lamb, model_name): language="en", expected_text="Mary had a little lamb", expected_seconds=16, + extra_body=extra_body, ) @@ -142,3 +234,76 @@ async def test_basic_audio_foscolo(foscolo, model_name): language="it", expected_text="ove il mio corpo fanciulletto", ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", ["inclusionAI/Ming-3.0-asr-tiny"]) +async def test_ming_context(model_name, ming_server, ming_audio): + """Topic and hotword hints must reach the native ASR prompt over HTTP.""" + async with ming_server.get_async_client() as client: + text = await transcribe_and_check( + client, + model_name, + ("context.wav", ming_audio["context_same_sentence_pcm16.wav"], "audio/wav"), + language=None, + expected_text="同囚", + response_format="json", + extra_body={ + "max_completion_tokens": 256, + "vllm_xargs": json.dumps({"no_repeat_ngram_size": 10}), + "prompt": "This is an audio about 同囚.", + "hotwords": ( + "同囚,黄国权,游学修,关楚耀,赵永洪,赵咏瑶,狱警,长镜头," + "吞玻璃戏,糖化玻璃,真实案例数据." + ), + }, + ) + assert "\t" not in text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", ["inclusionAI/Ming-3.0-asr-tiny"]) +@pytest.mark.parametrize( + "ming_server", [{"enable_prefix_caching": True}], indirect=True +) +async def test_ming_chunk_policy_prefix_cache(model_name, ming_server, ming_audio): + """Restoring a long request's chunk policy must reuse its cached prefix.""" + waveform, rate = sf.read( + BytesIO(ming_audio["chinese_common_voice.wav"]), dtype="float32" + ) + waveform = AudioResampler(16000).resample(waveform, orig_sr=rate) + waveform = np.pad(waveform, (0, 64 * 16000 - len(waveform))) + clip = BytesIO() + sf.write(clip, waveform, 16000, format="WAV", subtype="PCM_16") + async with httpx.AsyncClient() as metrics, ming_server.get_async_client() as client: + + async def prefix_hits(): + response = await metrics.get(ming_server.url_for("metrics")) + response.raise_for_status() + return sum( + float(line.rsplit(" ", 1)[1]) + for line in response.text.splitlines() + if line.startswith("vllm:prefix_cache_hits_total{") + ) + + for frames in (3000, 6000, 3000): + before = await prefix_hits() + text = await transcribe_and_check( + client, + model_name, + ("audio.wav", clip.getvalue(), "audio/wav"), + language="zh", + expected_text="漓江和湘江有秦代开凿的灵渠沟通", + response_format="json", + extra_body={ + "max_completion_tokens": 256, + "vllm_xargs": json.dumps( + {"audio_chunk_frames": frames, "no_repeat_ngram_size": 10} + ), + }, + ) + assert "\t" not in text + after = await prefix_hits() + assert after > before, ( + f"Restoring the chunk policy added no prefix hits: {before=} {after=}" + ) diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py index d5658e7cf0b4..17b15563fb3a 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py @@ -89,6 +89,7 @@ async def test_basic_audio(whisper_client, mary_had_lamb): language="en", response_format="text", temperature=0.0, + extra_body={"vllm_xargs": json.dumps({"frontend_window": 128})}, ) out = json.loads(transcription) out_text = out["text"] diff --git a/tests/models/multimodal/processing/test_ming_v3_asr.py b/tests/models/multimodal/processing/test_ming_v3_asr.py new file mode 100644 index 000000000000..7363e61a3c5e --- /dev/null +++ b/tests/models/multimodal/processing/test_ming_v3_asr.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The native Ming frontend must preserve the released ASR request contract.""" + +import os +from io import BytesIO + +import numpy as np +import pytest +import torch +from starlette.datastructures import UploadFile +from transformers import AutoProcessor + +from vllm.config import ModelConfig +from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionRequest +from vllm.model_executor.models.ming_audio import ( + MingAudioProjector, + get_audio_token_count, +) +from vllm.model_executor.models.ming_v3_asr import ( + ASR_INSTRUCTION, + MingASRStreamingPostProcessor, + MingV3ASRForConditionalGeneration, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.cache import MultiModalProcessorOnlyCache +from vllm.transformers_utils.configs.ming_v3_asr import MingAudioConfig + +from .test_common import _assert_inputs_equal + +MODEL = os.environ.get("VLLM_TEST_MING_ASR_MODEL", "inclusionAI/Ming-3.0-asr-tiny") + + +@pytest.fixture(scope="module") +def asr_processors(): + config = ModelConfig( + MODEL, dtype="bfloat16", mm_processor_cache_gb=1, trust_remote_code=True + ) + return ( + config, + MULTIMODAL_REGISTRY.create_processor(config), + AutoProcessor.from_pretrained(MODEL, trust_remote_code=True), + ) + + +@pytest.mark.parametrize( + "samples,language,frames", + [ + (201, None, 1), + (321, None, 3000), + (16001, "en", 3000), + (480001, "zh", 3000), + (480161, "en", 1001), + (960000, "zh", 6000), + (960161, "zh", 6000), + ], +) +def test_transcription_request_matches_native( + asr_processors, samples, language, frames +): + """Compare features and token IDs, including language seed placement and tails.""" + config, processor, native = asr_processors + audio = np.random.RandomState(42).uniform(-0.5, 0.5, samples).astype(np.float32) + prompt = ASR_INSTRUCTION + "This is an audio about 同囚." + request = TranscriptionRequest( + file=UploadFile(BytesIO()), + prompt=prompt, + language=language, + vllm_xargs={"audio_chunk_frames": frames}, + ) + stt_params = request.build_stt_params( + audio, + MingV3ASRForConditionalGeneration.get_speech_to_text_config( + config, "transcribe" + ), + config, + "transcribe", + ) + inputs = MingV3ASRForConditionalGeneration.get_generation_prompt(stt_params) + actual = processor( + inputs["prompt_token_ids"], + mm_items=processor.info.parse_mm_data(inputs["multi_modal_data"]), + hf_processor_mm_kwargs=inputs["mm_processor_kwargs"], + ) + expected = native.apply_transcription_request( + (audio, 16000), + prompt=prompt, + language_hint=language, + audio_chunk_frames=frames, + ) + assert actual["prompt_token_ids"] == expected.input_ids[0].tolist() + data = actual["mm_kwargs"].get_data() + torch.testing.assert_close( + data["audio_feats"][0], + expected.audio_feats[0], + rtol=0, + atol=0, + ) + + +@pytest.mark.parametrize( + "kernels,strides,paddings", [([3], [2], [1]), ([5, 3], [3, 2], [2, 0])] +) +def test_audio_token_count_matches_chunked_convolutions(kernels, strides, paddings): + """Count each chunk's rounding separately, including a nonempty tail.""" + config = MingAudioConfig( + n_audio_state=4, + downsampler_kernel_size_list=kernels, + downsampler_stride_list=strides, + downsampler_padding_list=paddings, + ) + encoder_conv = torch.nn.Conv1d(1, config.n_audio_state, 3, stride=2, padding=1) + projector = MingAudioProjector(config, hidden_size=4) + num_frames, chunk_frames = 3103, 1001 + with torch.inference_mode(): + actual = sum( + projector(encoder_conv(chunk).transpose(1, 2)).shape[1] + for chunk in torch.zeros(1, 1, num_frames).split(chunk_frames, dim=-1) + ) + assert get_audio_token_count(config, num_frames, chunk_frames) == actual + + +@pytest.mark.parametrize("scoped", [False, True]) +def test_chunk_policy_has_separate_cache_entries(asr_processors, scoped): + """The same waveform encoded with different chunks must not share cached data.""" + config, processor, _ = asr_processors + cache = MultiModalProcessorOnlyCache(config) + audio = np.random.RandomState(1).randn(480161).astype(np.float32) + results = [] + for frames in (3000, 6000, 3000): + options: dict[str, object] = {"audio_chunk_frames": frames} + if scoped: + options = {"audio_kwargs": options} + kwargs = dict( + mm_items=processor.info.parse_mm_data({"audio": audio}), + hf_processor_mm_kwargs=options, + ) + result = processor("", **kwargs, cache=cache) + _assert_inputs_equal(result, processor("", **kwargs)) + results.append(result) + assert results[0]["mm_hashes"] != results[1]["mm_hashes"] + _assert_inputs_equal(results[0], results[2]) + assert config.hf_config.audio_config.chunk_frames == 3000 + for result, frames in zip(results, (3000, 6000, 3000)): + data = result["mm_kwargs"].get_data() + assert torch.as_tensor(data["audio_chunk_frames"]).item() == frames + + +@pytest.mark.parametrize( + "options,samples,expected_frames", + [ + ({}, 321, 6000), + ({"audio_chunk_frames": 3000}, 321, 3000), + ({"audio_kwargs": {"audio_chunk_frames": 4500}}, 321, 4500), + ( + {"audio_kwargs": {"audio_chunk_frames": None, "context_chunks": True}}, + 480161, + 6000, + ), + ], +) +def test_processor_defaults_and_request_audio_options( + options, samples, expected_frames +): + """Configured defaults and both request encodings must select the chunk policy.""" + config = ModelConfig( + MODEL, + dtype="bfloat16", + trust_remote_code=True, + mm_processor_kwargs={"audio_kwargs": {"audio_chunk_frames": 6000}}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(config) + actual = processor( + "", + mm_items=processor.info.parse_mm_data( + {"audio": np.zeros(samples, dtype=np.float32)} + ), + hf_processor_mm_kwargs=options, + ) + frames = actual["mm_kwargs"].get_data()["audio_chunk_frames"] + assert torch.as_tensor(frames).item() == expected_frames + + +@pytest.mark.parametrize( + "options,max_model_len", + [({"audio_chunk_frames": 1}, 512), ({"context_chunks": True}, 2048)], +) +def test_dummy_audio_respects_processor_defaults(options, max_model_len): + """Configured chunking must keep profiling within token and duration limits.""" + config = ModelConfig( + MODEL, + dtype="bfloat16", + trust_remote_code=True, + max_model_len=max_model_len, + mm_processor_kwargs={"audio_kwargs": options}, + ) + processor = MULTIMODAL_REGISTRY.create_processor(config) + inputs = processor.get_dummy_mm_inputs({"audio": 1}) + assert len(inputs["prompt_token_ids"]) == max_model_len + + +@pytest.mark.parametrize( + "chunks,expected", + [ + (["Chi", "nese", "\t", " 你好", " 世界 "], "你好 世界"), + (["same ", "way ", "you did"], "same way you did"), + (["Eng", "lish"], ""), + (["Canton", "ese", "\t", "你好"], "你好"), + (["你", "好"], "你好"), + ], +) +def test_streaming_envelope_matches_final_text(chunks, expected): + processor = MingASRStreamingPostProcessor() + deltas = [processor.process_delta(chunk, False) for chunk in chunks] + deltas.append(processor.process_delta("", True)) + assert "".join(deltas) == expected + assert ( + MingV3ASRForConditionalGeneration.post_process_output("".join(chunks)) + == expected + ) diff --git a/tests/models/registry.py b/tests/models/registry.py index c9ff9648cdf5..4bc4831437a2 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1116,6 +1116,17 @@ def check_available_online( "MiMoV2OmniForCausalLM": _HfExamplesInfo( "XiaomiMiMo/MiMo-V2.5-Omni", trust_remote_code=True, is_available_online=False ), + "MingV3ASRForConditionalGeneration": _HfExamplesInfo( + "inclusionAI/Ming-3.0-asr-tiny", + trust_remote_code=True, + use_original_num_layers=True, + hf_overrides={ + "text_config": { + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["full_attention"], + } + }, + ), "MiniCPMO": _HfExamplesInfo( "openbmb/MiniCPM-o-2_6", trust_remote_code=True, diff --git a/vllm/config/speech_to_text.py b/vllm/config/speech_to_text.py index 6d31713c8667..cc56f91b0086 100644 --- a/vllm/config/speech_to_text.py +++ b/vllm/config/speech_to_text.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal from vllm.config.utils import config @@ -50,6 +50,9 @@ class SpeechToTextParams: to_language: str | None = None """Target language for translation (model-dependent).""" + extra_args: dict[str, str | int | float | list[str | int | float]] | None = None + """Request extensions supplied through vllm_xargs.""" + @config class SpeechToTextConfig: @@ -60,6 +63,10 @@ class SpeechToTextConfig: 16kHz audio input. The input audio will be automatically resampled to this rate before processing.""" + audio_resample_method: Literal["pyav", "scipy", "soxr", "torchaudio"] | None = None + """Optional resampler applied after decoding at the native sample rate. + None preserves the audio decoder's default resampling behavior.""" + max_audio_clip_s: int | None = 30 """Maximum duration in seconds for a single audio clip without chunking. Audio longer than this will be split into smaller chunks if diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index 871f69666f51..11a2ee05b09f 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -31,7 +31,7 @@ from vllm.logger import init_logger from vllm.logprobs import FlatLogprobs, Logprob from vllm.model_executor.models import SupportsTranscription -from vllm.multimodal.audio import get_audio_duration, split_audio +from vllm.multimodal.audio import AudioResampler, get_audio_duration, split_audio from vllm.multimodal.media.audio import load_audio from vllm.outputs import RequestOutput from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt @@ -173,14 +173,20 @@ def _decode_and_chunk_speech( # NOTE resample to model SR here for efficiency. This is also a # pre-requisite for chunking, as it assumes Whisper SR. try: + resample_method = self.asr_config.audio_resample_method with io.BytesIO(audio_data) as buf: y, sr = load_audio( buf, - sr=self.asr_config.sample_rate, + sr=None if resample_method else self.asr_config.sample_rate, mono=True, max_duration_s=self.max_audio_decode_duration_s, max_decode_bytes=self.max_audio_decode_bytes, ) + if resample_method: + y = AudioResampler( + target_sr=self.asr_config.sample_rate, method=resample_method + ).resample(y, orig_sr=sr) + sr = self.asr_config.sample_rate except ValueError: raise except Exception as exc: diff --git a/vllm/entrypoints/speech_to_text/transcription/protocol.py b/vllm/entrypoints/speech_to_text/transcription/protocol.py index 7e66c512c720..a9c6f113b93e 100644 --- a/vllm/entrypoints/speech_to_text/transcription/protocol.py +++ b/vllm/entrypoints/speech_to_text/transcription/protocol.py @@ -208,6 +208,7 @@ def build_stt_params( request_prompt=self.prompt, to_language=self.to_language, hotwords=self.hotwords, + extra_args=self.vllm_xargs, ) def to_beam_search_params( diff --git a/vllm/model_executor/models/bailing_moe_v3.py b/vllm/model_executor/models/bailing_moe_v3.py index 5819c89c9cc2..cbb1e8bc8f3f 100644 --- a/vllm/model_executor/models/bailing_moe_v3.py +++ b/vllm/model_executor/models/bailing_moe_v3.py @@ -1100,9 +1100,8 @@ def _forward( assert g1_non_spec is not None and beta_non_spec is not None assert state_indices is not None assert has_initial_state is not None - zero_idx = state_indices[~has_initial_state] - recurrent_state[zero_idx] = 0 initial_state = recurrent_state_active[state_indices].contiguous() + initial_state.masked_fill_(~has_initial_state[:, None, None, None], 0) out, last_state = chunk_kda( q=q, k=k, diff --git a/vllm/model_executor/models/ming_audio.py b/vllm/model_executor/models/ming_audio.py new file mode 100644 index 000000000000..ab29c0a65c2f --- /dev/null +++ b/vllm/model_executor/models/ming_audio.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Adapted from +# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/modeling_ming_v3_asr_encoder.py +# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/modeling_whisper_encoder.py +# Copyright 2026 Ant Group. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Reusable Ming audio tower and checkpoint-compatible projector.""" + +from collections.abc import Iterable + +import torch +from torch import nn +from transformers import WhisperConfig + +from vllm.config import VllmConfig +from vllm.transformers_utils.configs.ming_v3_asr import MingAudioConfig + +from .utils import AutoWeightsLoader, WeightsMapper +from .whisper import WhisperEncoder, _create_fake_bias_for_k_proj + + +def _get_audio_projector_params(config: MingAudioConfig) -> list[tuple[int, int, int]]: + kernels = config.downsampler_kernel_size_list + strides = config.downsampler_stride_list + paddings = config.downsampler_padding_list + if not kernels or not (len(kernels) == len(strides) == len(paddings)): + raise ValueError("Audio convolution lists must have equal positive length") + params = list(zip(kernels, strides, paddings)) + if any( + kernel <= 0 or stride <= 0 or padding < 0 for kernel, stride, padding in params + ): + raise ValueError( + "Audio kernels and strides must be positive and padding nonnegative" + ) + return params + + +def get_audio_token_count( + config: MingAudioConfig, num_frames: int, chunk_frames: int +) -> int: + """Count projected frames with independent encoder positions for each chunk.""" + if num_frames <= 0: + raise ValueError("Audio must contain at least one mel frame") + params = _get_audio_projector_params(config) + + def output_length(length: int) -> int: + length = (length + 1) // 2 + for kernel, stride, padding in params: + length = (length + 2 * padding - kernel) // stride + 1 + if length <= 0: + raise ValueError("Audio chunk is too short for its projector") + return length + + num_chunks, tail = divmod(num_frames, chunk_frames) + return (num_chunks * output_length(chunk_frames) if num_chunks else 0) + ( + output_length(tail) if tail else 0 + ) + + +class MingAudioEncoder(nn.Module): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + "blocks.": "layers.", + ".attn.query.": ".self_attn.q_proj.", + ".attn.key.": ".self_attn.k_proj.", + ".attn.value.": ".self_attn.v_proj.", + ".attn.out.": ".self_attn.out_proj.", + ".attn_ln.": ".self_attn_layer_norm.", + ".mlp_ln.": ".final_layer_norm.", + ".mlp.0.": ".mlp.fc1.", + ".mlp.2.": ".mlp.fc2.", + "ln_post.": "layer_norm.", + }, + orig_to_new_suffix={"positional_embedding": "embed_positions.weight"}, + orig_to_new_stacked={ + ".self_attn.q_proj": (".self_attn.qkv_proj", "q"), + ".self_attn.k_proj": (".self_attn.qkv_proj", "k"), + ".self_attn.v_proj": (".self_attn.qkv_proj", "v"), + }, + ) + + def __init__( + self, *, vllm_config: VllmConfig, config: MingAudioConfig, prefix: str + ) -> None: + super().__init__() + encoder_config = WhisperConfig( + num_mel_bins=config.n_mels, + d_model=config.n_audio_state, + encoder_attention_heads=config.n_audio_head, + encoder_layers=config.n_audio_layer, + encoder_ffn_dim=4 * config.n_audio_state, + max_source_positions=config.n_audio_ctx, + activation_function="gelu", + pos_embed="learned", + ) + self.encoder = WhisperEncoder( + vllm_config=vllm_config.with_hf_config(encoder_config), + prefix=prefix, + enable_pp=False, + init_in_fp32=True, + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + return self.encoder([features.transpose(0, 1)])[0] + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = _create_fake_bias_for_k_proj( + weights, ".attn.key.weight", out_features=self.encoder.conv1.out_channels + ) + loaded = AutoWeightsLoader(self.encoder).load_weights( + weights, mapper=self.hf_to_vllm_mapper + ) + return {f"encoder.{name}" for name in loaded} + + +class Transpose(nn.Module): + def forward(self, value: torch.Tensor) -> torch.Tensor: + return value.transpose(-1, -2) + + +class MingAudioProjector(nn.Module): + def __init__(self, config: MingAudioConfig, hidden_size: int) -> None: + super().__init__() + layers: list[nn.Module] = [] + width = config.n_audio_state + for kernel, stride, padding in _get_audio_projector_params(config): + layers.extend( + [ + nn.Conv1d(width, hidden_size, kernel, stride, padding), + nn.GELU(), + ] + ) + width = hidden_size + layers.extend([Transpose(), nn.Linear(hidden_size, hidden_size)]) + self.proj = nn.Sequential(*layers) + + def forward(self, embeddings: torch.Tensor) -> torch.Tensor: + return self.proj(embeddings.transpose(-1, -2)) diff --git a/vllm/model_executor/models/ming_v3_asr.py b/vllm/model_executor/models/ming_v3_asr.py new file mode 100644 index 000000000000..29e083407980 --- /dev/null +++ b/vllm/model_executor/models/ming_v3_asr.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Adapted from +# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/modeling_ming_v3_asr.py +# https://huggingface.co/inclusionAI/Ming-3.0-asr-tiny/blob/ac95e7b1678c08ea6f4a2581a051e432ed63e9b7/processing_ming_v3_asr.py +# https://github.com/inclusionAI/Ming-3.0-asr/blob/6141e8c552815db724d8df0f72c9d865c547a1ea/examples/hf_runtime.py +# Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved. +# +# The HF model/processor sources retain Apache-2.0 headers. +# ASR helpers from hf_runtime.py are covered by the MIT notice below. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# MIT License +# +# Copyright (c) 2026 inclusionAI +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +"""Inference-only Ming V3 ASR with the native Bailing hybrid decoder.""" + +from collections.abc import Iterable, Mapping, Sequence +from functools import cached_property +from numbers import Integral +from typing import cast + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import BatchFeature + +from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig +from vllm.config.multimodal import MultiModalDummyOptions +from vllm.config.speech_to_text import SpeechToTextParams +from vllm.inputs import MultiModalDataDict, PromptType, TokensPrompt +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, +) +from vllm.sequence import IntermediateTensors +from vllm.tokenizers import cached_tokenizer_from_config +from vllm.transformers_utils.configs.ming_v3_asr import ( + MingAudioConfig, + MingV3ASRConfig, +) +from vllm.transformers_utils.processors.ming_audio import MingASRFeatureExtractor + +from .bailing_moe_v3 import BailingMoeV3ForCausalLM +from .interfaces import ( + HasInnerState, + IsHybrid, + MultiModalEmbeddings, + StreamingTranscriptionPostProcessor, + SupportsMultiModal, + SupportsPP, + SupportsTranscription, +) +from .ming_audio import MingAudioEncoder, MingAudioProjector, get_audio_token_count +from .module_mapping import MultiModelKeys +from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix + +ASR_INSTRUCTION = ( + "Please recognize the language of this speech and transcribe it, " + "placing a `\t` between the identified language and the transcribe." +) +_LANGUAGES = {"zh": "Chinese", "en": "English"} +_OUTPUT_LANGUAGES = ("Chinese", "English", "Canton", "Japanese", "Korean") +_AUDIO_FORMS = ( + (""), + ("<|audio_start|>", "<|audio_pad|>", "<|audio_end|>"), +) + + +def _audio_tokens(tokenizer) -> tuple[str, str, str]: + vocab = tokenizer.get_vocab() + return next(form for form in _AUDIO_FORMS if all(t in vocab for t in form)) + + +def _get_audio_chunk_frames( + config: MingAudioConfig, options: Mapping[str, object], samples: int +) -> int: + chunk_frames = options.get("audio_chunk_frames") + context_chunks = options.get("context_chunks", False) + if context_chunks not in (False, True): + raise ValueError("context_chunks must be a boolean or 0/1") + if context_chunks: + if chunk_frames is not None: + raise ValueError("context_chunks and audio_chunk_frames are exclusive") + if samples > 60 * 16000: + raise ValueError("Context ASR chunk policy supports at most 60 seconds") + chunk_frames = 6000 if samples > 30 * 16000 else 3000 + if chunk_frames is None: + chunk_frames = config.chunk_frames + if ( + isinstance(chunk_frames, bool) + or not isinstance(chunk_frames, Integral) + or not 0 < int(chunk_frames) <= 2 * config.n_audio_ctx + ): + raise ValueError( + "audio_chunk_frames must be a positive integer no larger than " + f"{2 * config.n_audio_ctx}" + ) + return int(chunk_frames) + + +def _transcript(text: str) -> str: + text = text.strip() + if "\t" in text: + return text.split("\t", 1)[1].strip() + for language in _OUTPUT_LANGUAGES: + if text.startswith(language): + return text[len(language) :].strip() + return text + + +class MingASRStreamingPostProcessor(StreamingTranscriptionPostProcessor): + def __init__(self) -> None: + self.raw_text = "" + self.emitted_text = "" + + def process_delta(self, text_delta: str, finished: bool) -> str: + self.raw_text += text_delta + raw = self.raw_text.strip() + if not finished and any( + language.startswith(raw) for language in (*_OUTPUT_LANGUAGES, "Cantonese") + ): + return "" + text = _transcript(self.raw_text) + delta = text[len(self.emitted_text) :] + self.emitted_text = text + return delta + + +class MingASRProcessingInfo(BaseProcessingInfo): + def get_hf_config(self) -> MingV3ASRConfig: + return self.ctx.get_hf_config(MingV3ASRConfig) + + @cached_property + def _feature_extractor(self) -> MingASRFeatureExtractor: + return MingASRFeatureExtractor(self.get_hf_config().audio_config.n_mels) + + def get_feature_extractor(self, **kwargs: object) -> MingASRFeatureExtractor: + return self._feature_extractor + + def get_supported_mm_processor_kwargs(self) -> dict[str, set[str]]: + return {"audio_kwargs": {"audio_chunk_frames", "context_chunks"}} + + def get_audio_options(self, kwargs: Mapping[str, object]) -> Mapping[str, object]: + resolved_kwargs = self._merge_and_resolve_mm_processor_kwargs(kwargs) + audio_options = resolved_kwargs.get("audio_kwargs", {}) + if not isinstance(audio_options, Mapping): + raise ValueError("audio_kwargs must be a mapping") + + return audio_options + + def get_data_parser(self) -> MultiModalDataParser: + return MultiModalDataParser(target_sr=16000, target_channels=1) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + def get_mm_max_tokens_per_item( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> Mapping[str, int]: + tokenizer = self.get_tokenizer() + start, _, end = _audio_tokens(tokenizer) + overhead = len(tokenizer.encode(ASR_INSTRUCTION + start + end)) + return {"audio": max(1, seq_len - overhead)} + + +class MingASRDummyInputsBuilder(BaseDummyInputsBuilder[MingASRProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return "".join(_audio_tokens(self.info.get_tokenizer())) * mm_counts.get( + "audio", 0 + ) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: MultiModalDummyOptions, + ) -> MultiModalDataDict: + audio_options = self.info.get_audio_options({}) + config = self.info.get_hf_config().audio_config + chunk_frames = _get_audio_chunk_frames(config, audio_options, 60 * 16000) + budget = self.info.get_mm_max_tokens_per_item(seq_len, mm_counts)["audio"] + low, high = 1, budget * 4 + chunk_frames + if audio_options.get("context_chunks"): + high = min(high, 6000) + while low < high: + middle = (low + high + 1) // 2 + if get_audio_token_count(config, middle, chunk_frames) <= budget: + low = middle + else: + high = middle - 1 + return { + "audio": self._get_dummy_audios( + length=max(201, low * 160), + num_audios=mm_counts.get("audio", 0), + overrides=mm_options.get("audio"), + ) + } + + +class MingASRMultiModalProcessor(BaseMultiModalProcessor[MingASRProcessingInfo]): + def _call_hf_processor( + self, hf_data: Mapping[str, object], hf_kwargs: Mapping[str, object] + ) -> BatchFeature: + config = self.info.get_hf_config().audio_config + feature_extractor = self.info.get_feature_extractor() + audio_options = self.info.get_audio_options(hf_kwargs) + features, lengths, chunk_frames = [], [], [] + audios = hf_data.get("audio", []) + assert isinstance(audios, list) + for audio in audios: + frames = _get_audio_chunk_frames(config, audio_options, len(audio)) + waveform = torch.as_tensor(audio) + if waveform.ndim != 1 or waveform.numel() <= 200: + raise ValueError( + "Ming audio requires more than 200 mono samples at 16 kHz" + ) + if waveform.is_complex() or not torch.isfinite(waveform).all(): + raise ValueError("Audio must contain finite real samples") + inputs = feature_extractor( + waveform.detach().to(device="cpu", dtype=torch.float32).numpy(), + sampling_rate=feature_extractor.sampling_rate, + padding=False, + truncation=False, + return_tensors="pt", + ) + feature = inputs.input_features[0].T.to(torch.bfloat16) + features.append(feature) + lengths.append(len(feature)) + chunk_frames.append(frames) + return BatchFeature( + { + "audio_feats": features, + "audio_feats_lengths": torch.tensor(lengths), + "audio_chunk_frames": torch.tensor(chunk_frames), + } + ) + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return { + name: MultiModalFieldConfig.batched( + "audio", keep_on_cpu=name != "audio_feats" + ) + for name in ("audio_feats", "audio_feats_lengths", "audio_chunk_frames") + } + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptReplacement]: + config = self.info.get_hf_config() + data = out_mm_kwargs.get_data() + lengths = torch.as_tensor(data.get("audio_feats_lengths", [])).flatten() + frames = torch.as_tensor(data.get("audio_chunk_frames", [])).flatten() + counts = [ + get_audio_token_count(config.audio_config, int(length), int(chunk)) + for length, chunk in zip(lengths, frames, strict=True) + ] + return [ + PromptReplacement( + modality="audio", + target=[config.audio_token_id], + replacement=lambda index: [config.audio_token_id] * counts[index], + ) + ] + + +@MULTIMODAL_REGISTRY.register_processor( + MingASRMultiModalProcessor, + info=MingASRProcessingInfo, + dummy_inputs=MingASRDummyInputsBuilder, +) +class MingV3ASRForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsTranscription, + HasInnerState, + IsHybrid, + SupportsPP, +): + supported_languages = _LANGUAGES + supports_transcription_only = True + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.audio.": "audio_encoder.", + "linear_proj_audio.": "audio_projector.proj.", + "model.visual.": None, + "linear_proj.": None, + "model.": "language_model.model.", + "lm_head.": "language_model.lm_head.", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.config = vllm_config.model_config.hf_config + self.dtype = vllm_config.model_config.dtype + with self._mark_tower_model(vllm_config, "audio"): + self.audio_encoder = MingAudioEncoder( + vllm_config=vllm_config, + config=self.config.audio_config, + prefix=maybe_prefix(prefix, "audio_encoder"), + ) + self.audio_projector = MingAudioProjector( + self.config.audio_config, self.config.text_config.hidden_size + ) + with self._mark_language_model(vllm_config): + self.language_model = BailingMoeV3ForCausalLM( + vllm_config=vllm_config.with_hf_config(self.config.text_config), + prefix=maybe_prefix(prefix, "language_model"), + ) + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str: + if not modality.startswith("audio"): + raise ValueError("Ming V3 ASR only supports audio") + return "" + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + features = cast( + torch.Tensor | list[torch.Tensor] | None, kwargs.get("audio_feats") + ) + if features is None: + return () + lengths = torch.as_tensor(kwargs["audio_feats_lengths"]).flatten().tolist() + chunk_frames = torch.as_tensor(kwargs["audio_chunk_frames"]).flatten().tolist() + outputs = [] + for feature, length, frames in zip( + features, lengths, chunk_frames, strict=True + ): + chunks = [] + feature = feature[:length].to(dtype=self.dtype) + for chunk in feature.split(frames): + encoded = self.audio_encoder(chunk) + projected = self.audio_projector(encoded.unsqueeze(0))[0] + chunks.append(projected) + outputs.append(F.normalize(torch.cat(chunks), dim=-1)) + return outputs + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.language_model.compute_logits(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return AutoWeightsLoader(self).load_weights( + weights, mapper=self.hf_to_vllm_mapper + ) + + def get_mm_mapping(self) -> MultiModelKeys: + return MultiModelKeys.from_string_field( + language_model="language_model", + tower_model=["audio_encoder.", "audio_projector."], + ) + + @classmethod + def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig): + return BailingMoeV3ForCausalLM.get_mamba_state_shape_from_config( + vllm_config.with_hf_config(vllm_config.model_config.hf_config.text_config) + ) + + @classmethod + def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig): + return BailingMoeV3ForCausalLM.get_mamba_state_dtype_from_config( + vllm_config.with_hf_config(vllm_config.model_config.hf_config.text_config) + ) + + @classmethod + def get_mamba_state_copy_func(cls): + return BailingMoeV3ForCausalLM.get_mamba_state_copy_func() + + @classmethod + def validate_language(cls, language: str | None) -> str | None: + if language is not None and language not in cls.supported_languages: + raise ValueError("Ming V3 ASR language hint must be zh or en") + return language + + @classmethod + def get_speech_to_text_config( + cls, model_config: ModelConfig, task_type: str + ) -> SpeechToTextConfig: + return SpeechToTextConfig( + sample_rate=16000, + audio_resample_method="torchaudio", + max_audio_clip_s=None, + min_energy_split_window_size=None, + ) + + @classmethod + def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType: + if stt_params.task_type != "transcribe" or stt_params.to_language is not None: + raise ValueError("Ming V3 ASR does not support translation") + language = cls.validate_language(stt_params.language) + prompt = stt_params.request_prompt + if not prompt.startswith(ASR_INSTRUCTION): + prompt = ASR_INSTRUCTION + prompt + reserved = ("\n", "\r", "