From f53fa26e05c476a43f6db048a9e3b43bcb2b72fb Mon Sep 17 00:00:00 2001 From: Greg Pereira Date: Sun, 5 Apr 2026 10:11:18 -0700 Subject: [PATCH 01/39] [Bugfix] Fix invalid JSON in Gemma 4 streaming tool calls by stripping partial delimiters (#38992) Signed-off-by: greg pereira Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- tests/tool_parsers/test_gemma4_tool_parser.py | 29 +++++++++++++++++++ vllm/tool_parsers/gemma4_tool_parser.py | 7 +++-- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 80cf70d6c7d..26722e68d76 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -502,3 +502,32 @@ class TestStreamingExtraction: results = self._simulate_streaming(parser, mock_request, chunks) name = self._collect_function_name(results) assert name == "get_status" + + def test_streaming_split_delimiter_no_invalid_json(self, parser, mock_request): + """Partial <|"|> delimiter chars must not leak into streamed JSON. + + Reproduces the bug from https://github.com/vllm-project/vllm/issues/38946 + where a token boundary splits the string delimiter, leaving fragments + like '<|' at the end of a parsed value which then corrupt the JSON. + """ + chunks = [ + "<|tool_call>", + "call:todowrite{", + 'content:<|"|>Buy milk<|', + '"|>}', + "", + ] + + results = self._simulate_streaming(parser, mock_request, chunks) + + args_text = self._collect_arguments(results) + assert args_text, "No arguments were streamed" + + # Must be valid JSON — the original bug caused a JSON parse error + parsed_args = json.loads(args_text) + assert parsed_args["content"] == "Buy milk" + + # Ensure no raw delimiter fragments leaked into the JSON + assert "<|" not in args_text, ( + f"Partial delimiter leaked into JSON: {args_text!r}" + ) diff --git a/vllm/tool_parsers/gemma4_tool_parser.py b/vllm/tool_parsers/gemma4_tool_parser.py index 3d0e4e7c4ab..406ba9e7020 100644 --- a/vllm/tool_parsers/gemma4_tool_parser.py +++ b/vllm/tool_parsers/gemma4_tool_parser.py @@ -675,10 +675,11 @@ class Gemma4ToolParser(ToolParser): current_args_json = json.dumps(current_args, ensure_ascii=False) # Withhold trailing closing characters that may shift as more - # tokens arrive. Strip trailing '}', '"', and ']' sequences - # to get the "safe prefix". + # tokens arrive. Strip trailing '}', '"', ']' and partial + # STRING_DELIM fragments ('<', '|', '\\', '>') to get the + # "safe prefix". safe_json = current_args_json - while safe_json and safe_json[-1] in ("}", '"', "]"): + while safe_json and safe_json[-1] in ("}", '"', "]", "<", "|", "\\", ">"): safe_json = safe_json[:-1] prev_streamed = self.streamed_args_for_tool[self.current_tool_id] From 4dd49b06f81af2238ac5a86cfb0b7220083eb125 Mon Sep 17 00:00:00 2001 From: Greg Pereira Date: Sun, 5 Apr 2026 12:11:58 -0700 Subject: [PATCH 02/39] [Bug] Fix Import paths for `encoder_cudagraph` modules (#38997) Signed-off-by: greg pereira Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- tests/v1/cudagraph/test_encoder_cudagraph.py | 8 ++++---- vllm/model_executor/models/interfaces.py | 2 +- vllm/model_executor/models/qwen3_vl.py | 6 +++--- vllm/v1/worker/encoder_cudagraph.py | 2 +- vllm/v1/worker/gpu_model_runner.py | 4 ++-- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 322fcb3caa1..543dfc8bb31 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -14,17 +14,17 @@ from typing import Any import pytest import torch -from vllm.v1.worker.gpu.mm.encoder_cudagraph import ( + +from vllm.platforms import current_platform +from vllm.v1.worker.encoder_cudagraph import ( EncoderCudaGraphManager, ) -from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import ( +from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, EncoderCudaGraphConfig, EncoderCudaGraphReplayBuffers, ) -from vllm.platforms import current_platform - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index df7170aab34..1f1d57493c0 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -46,7 +46,7 @@ if TYPE_CHECKING: from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.multimodal.registry import _ProcessorFactories from vllm.sequence import IntermediateTensors - from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import ( + from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, EncoderCudaGraphConfig, EncoderCudaGraphReplayBuffers, diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index cb48ceb0c77..1aa5dec5390 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -1733,7 +1733,7 @@ class Qwen3VLForConditionalGeneration( # -- SupportsEncoderCudaGraph protocol methods -- def get_encoder_cudagraph_config(self): - from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import ( + from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphConfig, ) @@ -1818,7 +1818,7 @@ class Qwen3VLForConditionalGeneration( device: torch.device, dtype: torch.dtype, ): - from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import ( + from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, ) @@ -1872,7 +1872,7 @@ class Qwen3VLForConditionalGeneration( mm_kwargs: dict[str, Any], max_batch_size: int, ): - from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import ( + from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphReplayBuffers, ) diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index b2930a23474..0fabbc77c07 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -16,7 +16,7 @@ from vllm.distributed import ( from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph from vllm.model_executor.models.vision import get_load_balance_assignment -from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import ( +from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphConfig, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 7a21117fb64..8dfa65da136 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -211,7 +211,7 @@ from .utils import ( if TYPE_CHECKING: from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.spec_decode.ngram_proposer import NgramProposer - from vllm.v1.worker.gpu.mm.encoder_cudagraph import EncoderCudaGraphManager + from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager logger = init_logger(__name__) @@ -5988,7 +5988,7 @@ class GPUModelRunner( SupportsEncoderCudaGraph, supports_encoder_cudagraph, ) - from vllm.v1.worker.gpu.mm.encoder_cudagraph import ( + from vllm.v1.worker.encoder_cudagraph import ( EncoderCudaGraphManager, ) From 56de443db1d50dce2624a21765c11a79dffa8740 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sun, 5 Apr 2026 13:26:11 -0700 Subject: [PATCH 03/39] [ci] Switch some CI jobs to H200 MIG slices (#38956) --- .buildkite/test_areas/basic_correctness.yaml | 1 + .buildkite/test_areas/benchmarks.yaml | 1 + .buildkite/test_areas/cuda.yaml | 1 + .buildkite/test_areas/engine.yaml | 2 ++ .buildkite/test_areas/entrypoints.yaml | 2 ++ .buildkite/test_areas/expert_parallelism.yaml | 1 + .buildkite/test_areas/kernels.yaml | 1 + .buildkite/test_areas/misc.yaml | 2 ++ .buildkite/test_areas/models_basic.yaml | 1 + .buildkite/test_areas/models_language.yaml | 2 ++ .buildkite/test_areas/models_multimodal.yaml | 4 ++++ .buildkite/test_areas/pytorch.yaml | 2 ++ .buildkite/test_areas/ray_compat.yaml | 1 + .buildkite/test_areas/spec_decode.yaml | 4 ++++ 14 files changed, 25 insertions(+) diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index 759d2b53587..042734e8433 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: Basic Correctness timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - vllm/ - tests/basic_correctness/test_basic_correctness diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 72d70a8df2b..4cda6fff144 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: Benchmarks CLI Test timeout_in_minutes: 20 + device: h200_18gb source_file_dependencies: - vllm/ - tests/benchmarks/ diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index b9bb3a2924e..4d1efdb13c8 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: Platform Tests (CUDA) timeout_in_minutes: 15 + device: h200_18gb source_file_dependencies: - vllm/ - tests/cuda diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index ed0df3e4d87..5e4361ec9ad 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: Engine timeout_in_minutes: 15 + device: h200_18gb source_file_dependencies: - vllm/ - tests/engine @@ -25,6 +26,7 @@ steps: - label: e2e Scheduling (1 GPU) timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - vllm/v1/ - tests/v1/e2e/general/ diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index ebe6b9419fc..8c2b529a806 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -61,6 +61,7 @@ steps: - label: Entrypoints Integration (API Server openai - Part 3) timeout_in_minutes: 50 + device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -105,6 +106,7 @@ steps: - label: OpenAI API Correctness timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - csrc/ - vllm/entrypoints/openai/ diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index 90c19701c84..c2adf52a2d5 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: EPLB Algorithm timeout_in_minutes: 15 + device: h200_18gb working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/distributed/eplb diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index da26caf72ef..5fd081699d1 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: vLLM IR Tests timeout_in_minutes: 10 + device: h200_18gb working_dir: "/vllm-workspace/" source_file_dependencies: - vllm/ir diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 5c21e1a7961..b806da88c04 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -19,6 +19,7 @@ steps: - label: V1 Sample + Logits timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - vllm/ - tests/v1/sample @@ -86,6 +87,7 @@ steps: - label: Regression timeout_in_minutes: 20 + device: h200_18gb source_file_dependencies: - vllm/ - tests/test_regression diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index f4e14ff4a94..8ba4484fb02 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: Basic Models Tests (Initialization) timeout_in_minutes: 45 + device: h200_18gb torch_nightly: true source_file_dependencies: - vllm/ diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index a3bd21ccff3..1a7cbc4b6d4 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -67,6 +67,7 @@ steps: - label: Language Models Test (PPL) timeout_in_minutes: 110 + device: h200_18gb optional: true source_file_dependencies: - vllm/ @@ -90,6 +91,7 @@ steps: - label: Language Models Test (MTEB) timeout_in_minutes: 110 + device: h200_18gb optional: true source_file_dependencies: - vllm/ diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index a2bf550dfcd..3bf907bb6c5 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: "Multi-Modal Models (Standard) 1: qwen2" timeout_in_minutes: 45 + device: h200_18gb source_file_dependencies: - vllm/ - tests/models/multimodal @@ -19,6 +20,7 @@ steps: - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma" timeout_in_minutes: 45 + device: h200_18gb source_file_dependencies: - vllm/ - tests/models/multimodal @@ -77,6 +79,7 @@ steps: - label: Multi-Modal Processor # 44min timeout_in_minutes: 60 + device: h200_18gb source_file_dependencies: - vllm/ - tests/models/multimodal @@ -131,6 +134,7 @@ steps: - label: Multi-Modal Models (Extended Pooling) optional: true + device: h200_18gb source_file_dependencies: - vllm/ - tests/models/multimodal/pooling diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index f9968e9a897..ad538e91918 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -49,6 +49,7 @@ steps: - label: PyTorch Fullgraph timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - vllm/ - tests/compile @@ -60,6 +61,7 @@ steps: # if this test fails, it means the nightly torch version is not compatible with some # of the dependencies. Please check the error message and add the package to whitelist # in /vllm/tools/pre_commit/generate_nightly_torch_test.py + device: h200_18gb soft_fail: true source_file_dependencies: - requirements/nightly_torch_test.txt diff --git a/.buildkite/test_areas/ray_compat.yaml b/.buildkite/test_areas/ray_compat.yaml index 7917b0a4ff8..3485e346532 100644 --- a/.buildkite/test_areas/ray_compat.yaml +++ b/.buildkite/test_areas/ray_compat.yaml @@ -7,6 +7,7 @@ steps: # If this fails, it means the PR introduces a dependency that # conflicts with Ray's dependency constraints. # See https://github.com/vllm-project/vllm/issues/33599 + device: h200_18gb soft_fail: true timeout_in_minutes: 10 source_file_dependencies: diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 8dba7a2f8c6..a0b73096867 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: Spec Decode Eagle timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ @@ -13,6 +14,7 @@ steps: - label: Spec Decode Speculators + MTP timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ @@ -23,6 +25,7 @@ steps: - label: Spec Decode Ngram + Suffix timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ @@ -32,6 +35,7 @@ steps: - label: Spec Decode Draft Model timeout_in_minutes: 30 + device: h200_18gb source_file_dependencies: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ From d56e95223917ba21af2dd87a335dcb5c0b347bfe Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Mon, 6 Apr 2026 01:23:45 +0300 Subject: [PATCH 04/39] nano_nemotron_vl: fix tensor device mismatch exception when video profiling (#39029) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- vllm/model_executor/models/nano_nemotron_vl.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 819bb4a3cb2..249b2896910 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -1239,12 +1239,13 @@ class NemotronH_Nano_VL_V2( img_context_token_ids=self._img_context_token_ids, video_temporal_patch_size=video_temporal_patch_size, ) + device = video_embeddings.device # video_repl.full is a list of token IDs - repl_token_ids = torch.tensor(video_repl.full) + repl_token_ids = torch.tensor(video_repl.full, device=device) # Get embedding token IDs for image context (use pre-tokenized version) - embed_token_ids = torch.tensor(self._img_context_token_ids) + embed_token_ids = torch.tensor(self._img_context_token_ids, device=device) # Create mask for video embedding positions is_video_embed = torch.isin(repl_token_ids, embed_token_ids) From 9570654c6d41ffff922cd15d57added9557eb272 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Sun, 5 Apr 2026 20:42:02 -0500 Subject: [PATCH 05/39] [ROCm][CI] Run Kernels Core Operation Test On MI325 and mitigate flakiness (#38184) Signed-off-by: Micah Williamson --- .buildkite/test-amd.yaml | 2 +- tests/kernels/core/test_layernorm.py | 10 +++++++++- vllm/platforms/rocm.py | 5 +++++ 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index f42c495f864..95dc2f688fb 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -751,6 +751,7 @@ steps: timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] agent_pool: mi250_1 + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/ @@ -2035,7 +2036,6 @@ steps: timeout_in_minutes: 38 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] agent_pool: mi325_1 - optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - csrc/ diff --git a/tests/kernels/core/test_layernorm.py b/tests/kernels/core/test_layernorm.py index f8f9660942a..42da24ccb96 100644 --- a/tests/kernels/core/test_layernorm.py +++ b/tests/kernels/core/test_layernorm.py @@ -7,12 +7,20 @@ import torch from tests.kernels.quant_utils import FP8_DTYPE from tests.kernels.utils import opcheck from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx90a + + on_mi250 = on_gfx90a() +else: + on_mi250 = False + DTYPES = [torch.half, torch.bfloat16, torch.float] NUM_TOKENS = [7, 83, 4096] # Arbitrary values for testing HIDDEN_SIZES = [8, 768, 769, 5120, 5125, 8192] # Arbitrary values for testing -ADD_RESIDUAL = [False, True] +ADD_RESIDUAL = [False, True] if not on_mi250 else [True] SEEDS = [0] CUDA_DEVICES = [ f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 72e67bc9eab..7b713536602 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -182,6 +182,7 @@ _ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"]) _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"]) _ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"]) _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) +_ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH _ON_GFX950 = "gfx950" in _GCN_ARCH @@ -273,6 +274,10 @@ def on_gfx9() -> bool: return _ON_GFX9 +def on_gfx90a() -> bool: + return _ON_GFX90A + + def on_gfx942() -> bool: return _ON_GFX942 From 780ba37458362bdc0596c6511e17749d44b145fc Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Sun, 5 Apr 2026 20:42:10 -0500 Subject: [PATCH 06/39] [ROCm][Quantization] Add asymmetric INT8 quantization support to TritonInt8ScaledMMLinearKernel (#38501) Signed-off-by: Andreas Karatzas --- ...lama-4-Maverick-17B-128E-Instruct-FP8.yaml | 3 + .../Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml | 3 + .../Qwen3-235B-A22B-Instruct-2507-FP8.yaml | 3 + .../configs/models-small-rocm.txt | 1 + .../test_lm_eval_correctness.py | 32 +++++++ .buildkite/test-amd.yaml | 18 ++++ .../kernels/linear/scaled_mm/triton.py | 87 ++++++++++++++++--- 7 files changed, 133 insertions(+), 14 deletions(-) diff --git a/.buildkite/lm-eval-harness/configs/Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml b/.buildkite/lm-eval-harness/configs/Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml index 6c0b5540cbb..9a5af854011 100644 --- a/.buildkite/lm-eval-harness/configs/Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml +++ b/.buildkite/lm-eval-harness/configs/Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml @@ -1,6 +1,9 @@ # For hf script, without -t option (tensor parallel size). # bash .buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh -m meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -l 250 -t 8 -f 5 model_name: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8" +required_gpu_arch: + - gfx942 + - gfx950 tasks: - name: "mmlu_pro" metrics: diff --git a/.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml b/.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml index aa4fb9fa03d..ff43fa187b0 100644 --- a/.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml +++ b/.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml @@ -1,6 +1,9 @@ # For vllm script, with -t option (tensor parallel size) # bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic -l 1319 -t 1 model_name: "RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic" +required_gpu_arch: + - gfx942 + - gfx950 tasks: - name: "gsm8k" metrics: diff --git a/.buildkite/lm-eval-harness/configs/Qwen3-235B-A22B-Instruct-2507-FP8.yaml b/.buildkite/lm-eval-harness/configs/Qwen3-235B-A22B-Instruct-2507-FP8.yaml index 514c15d6098..84e4f3fe334 100644 --- a/.buildkite/lm-eval-harness/configs/Qwen3-235B-A22B-Instruct-2507-FP8.yaml +++ b/.buildkite/lm-eval-harness/configs/Qwen3-235B-A22B-Instruct-2507-FP8.yaml @@ -1,4 +1,7 @@ model_name: "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8" +required_gpu_arch: + - gfx942 + - gfx950 tasks: - name: "mmlu_pro" metrics: diff --git a/.buildkite/lm-eval-harness/configs/models-small-rocm.txt b/.buildkite/lm-eval-harness/configs/models-small-rocm.txt index a3bb95e19e2..36e0543879b 100644 --- a/.buildkite/lm-eval-harness/configs/models-small-rocm.txt +++ b/.buildkite/lm-eval-harness/configs/models-small-rocm.txt @@ -1,5 +1,6 @@ Qwen2.5-1.5B-Instruct.yaml Meta-Llama-3.2-1B-Instruct-INT8-compressed-tensors.yaml +Meta-Llama-3-8B-Instruct-INT8-compressed-tensors-asym.yaml Meta-Llama-3-8B-Instruct-nonuniform-compressed-tensors.yaml Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml Qwen1.5-MoE-W4A16-compressed-tensors.yaml diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py index fad5f593be4..d34e603b9e2 100644 --- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py +++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py @@ -13,6 +13,7 @@ import os from contextlib import contextmanager import lm_eval +import pytest import yaml from vllm.platforms import current_platform @@ -89,9 +90,40 @@ def launch_lm_eval(eval_config, tp_size): return results +def _check_rocm_gpu_arch_requirement(eval_config): + """Skip the test if the model requires a ROCm GPU arch not present. + + Model YAML configs can specify:: + + required_gpu_arch: + - gfx942 + - gfx950 + + The check only applies on ROCm. On other platforms (e.g. CUDA) the + field is ignored so that shared config files work for both NVIDIA and + AMD CI pipelines. + """ + required_archs = eval_config.get("required_gpu_arch") + if not required_archs: + return + + if not current_platform.is_rocm(): + return + + from vllm.platforms.rocm import _GCN_ARCH # noqa: E402 + + if not any(arch in _GCN_ARCH for arch in required_archs): + pytest.skip( + f"Model requires GPU arch {required_archs}, " + f"but detected arch is '{_GCN_ARCH}'" + ) + + def test_lm_eval_correctness_param(config_filename, tp_size): eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8")) + _check_rocm_gpu_arch_requirement(eval_config) + results = launch_lm_eval(eval_config, tp_size) rtol = eval_config.get("rtol", DEFAULT_RTOL) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 95dc2f688fb..3e6421a847d 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2690,6 +2690,24 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt +- label: LM Eval Small Models (MI325) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] + agent_pool: mi325_1 + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt + + - label: LM Eval Small Models (B200-MI325) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/triton.py b/vllm/model_executor/kernels/linear/scaled_mm/triton.py index d2d90ed06a7..c68638a6ad9 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/triton.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/triton.py @@ -31,8 +31,6 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel): @classmethod def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: - if not c.input_symmetric: - return False, "supports symmetric input only." return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -62,17 +60,59 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel): # INPUT SCALE if self.config.is_static_input_scheme: assert i_s is not None - replace_parameter( - layer, - i_s_name, - torch.nn.Parameter(i_s.max(), requires_grad=False), - ) - setattr(layer, i_zp_name, None) + + if self.config.input_symmetric: + replace_parameter( + layer, + 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) + + # Reconstruct the ranges to find a single scale and azp + int8_traits = torch.iinfo(torch.int8) + azps = input_zero_point.to(dtype=torch.int32) + range_max = (i_s * (int8_traits.max - azps)).max() + range_min = (i_s * (int8_traits.min - azps)).min() + + scale = (range_max - range_min) / (int8_traits.max - int8_traits.min) + replace_parameter( + layer, + i_s_name, + torch.nn.Parameter(scale, requires_grad=False), + ) + + # AZP loaded as int8 but used as int32 + azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32) + replace_parameter( + layer, + i_zp_name, + torch.nn.Parameter(azp, requires_grad=False), + ) else: setattr(layer, i_s_name, None) setattr(layer, i_zp_name, None) - setattr(layer, azp_adj_name, None) + # azp_adj is the AZP adjustment term, used to account for weights. + # It does not depend on scales or azp, so it is the same for + # static and dynamic quantization. + # See csrc/quantization/w8a8/cutlass/Epilogues.md for the math. + if not self.config.input_symmetric: + weight = getattr(layer, w_q_name) + # weight is already transposed to [K, N], sum over K (dim=0) + azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32) + if self.config.is_static_input_scheme: + # Fold azp into azp_adj for the per-tensor case + azp_adj = getattr(layer, i_zp_name) * azp_adj + setattr( + layer, + azp_adj_name, + torch.nn.Parameter(azp_adj, requires_grad=False), + ) + else: + setattr(layer, azp_adj_name, None) def apply_weights( self, @@ -80,14 +120,33 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - w_q, w_s, i_s, i_zp, _ = self._get_layer_params(layer) + w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer) + symmetric = azp_adj is None x_q, x_s, x_zp = ops.scaled_int8_quant( - x.contiguous(), i_s, i_zp, symmetric=True + x.contiguous(), i_s, i_zp, symmetric=symmetric ) - assert x_zp is None, "Triton kernel only supports symmetric quantization" - - return triton_scaled_mm( + out = triton_scaled_mm( x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias ) + + if azp_adj is not None: + # Asymmetric quantization: subtract the zero-point correction. + # D = scale_a * scale_b * (A_q @ B_q - azp * azp_adj) + bias + # triton_scaled_mm already computed scale_a * scale_b * (A_q @ B_q) + bias + # so we subtract scale_a * scale_b * azp * azp_adj + # + # x_s: [M, 1] or scalar, w_s: [N, 1] or scalar, azp_adj: [1, N] + # Reshape w_s from [N, 1] to [1, N] for proper broadcasting. + w_s_row = w_s.view(1, -1) if w_s.dim() > 0 else w_s + static = i_zp is not None + if not static and x_zp is not None: + # Dynamic per-token: azp is per-token, azp_adj is per-channel + # x_zp: [M, 1], azp_adj: [1, N] + out -= x_s * w_s_row * (x_zp * azp_adj).to(x.dtype) + else: + # Static per-tensor: azp already folded into azp_adj + out -= (x_s * w_s_row * azp_adj).to(x.dtype) + + return out From f6983f01de2bf2e92ab468fa735ebac39cddd670 Mon Sep 17 00:00:00 2001 From: liuchenbing2026 Date: Mon, 6 Apr 2026 10:50:18 +0800 Subject: [PATCH 07/39] MiniMax-M2: add Eagle3 speculative decoding support (#37512) Signed-off-by: liuchenbing Signed-off-by: liucb Co-authored-by: liuchenbing --- tests/models/registry.py | 6 ++++++ vllm/config/speculative.py | 1 + vllm/model_executor/models/minimax_m2.py | 21 ++++++++++++++++----- vllm/model_executor/models/registry.py | 1 + 4 files changed, 24 insertions(+), 5 deletions(-) diff --git a/tests/models/registry.py b/tests/models/registry.py index 895dc457927..f1f80e639d6 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1246,6 +1246,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { use_original_num_layers=True, max_model_len=10240, ), + "Eagle3MiniMaxM2ForCausalLM": _HfExamplesInfo( + "MiniMaxAI/MiniMax-M2", + trust_remote_code=True, + speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", + tokenizer="MiniMaxAI/MiniMax-M2", + ), "EagleMistralLarge3ForCausalLM": _HfExamplesInfo( "mistralai/Mistral-Large-3-675B-Instruct-2512", speculative_model="mistralai/Mistral-Large-3-675B-Instruct-2512-Eagle", diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index f1fda9afd31..0e74501dd9a 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -817,6 +817,7 @@ class SpeculativeConfig: "deepseek_v3", "kimi_k2", "kimi_k25", + "minimax_m2", ] if ( self.method in ("eagle3", "extract_hidden_states", "dflash") diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index 0f43bc0cdce..f10452c5738 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -24,6 +24,7 @@ """Inference-only MiniMaxM2 model.""" from collections.abc import Iterable +from itertools import islice from typing import Any import torch @@ -59,7 +60,7 @@ from vllm.model_executor.model_loader.weight_utils import ( ) from vllm.sequence import IntermediateTensors -from .interfaces import SupportsLoRA, SupportsPP +from .interfaces import EagleModelMixin, SupportsEagle3, SupportsLoRA, SupportsPP from .utils import ( AutoWeightsLoader, PPMissingLayer, @@ -313,7 +314,7 @@ class MiniMaxM2DecoderLayer(nn.Module): @support_torch_compile -class MiniMaxM2Model(nn.Module): +class MiniMaxM2Model(nn.Module, EagleModelMixin): fall_back_to_pt_during_load = False def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -366,7 +367,7 @@ class MiniMaxM2Model(nn.Module): positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: + ) -> 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 @@ -378,14 +379,24 @@ class MiniMaxM2Model(nn.Module): hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] - for layer in self.layers[self.start_layer : self.end_layer]: + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer) + ): hidden_states, residual = layer(positions, hidden_states, residual) + self._maybe_add_hidden_state( + aux_hidden_states, idx + 1, hidden_states, residual + ) if not get_pp_group().is_last_rank: return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} ) hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: @@ -496,7 +507,7 @@ class MiniMaxM2Model(nn.Module): return loaded_params -class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): +class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3): packed_modules_mapping = { "qkv_proj": [ "q_proj", diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 1901381cbd3..4b354add384 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -554,6 +554,7 @@ _SPECULATIVE_DECODING_MODELS = { "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "Eagle3MiniMaxM2ForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), From c5e3454e5adf063d5af75140c36a7f900a9e4c4c Mon Sep 17 00:00:00 2001 From: bhargav-patel-29 Date: Mon, 6 Apr 2026 13:49:56 +0530 Subject: [PATCH 08/39] [Model] Add support for BharatGen's Param2MoE model (#38000) Signed-off-by: bhargav-patel-29 Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- docs/models/supported_models.md | 1 + tests/models/registry.py | 4 + vllm/model_executor/models/param2moe.py | 900 ++++++++++++++++++++++++ vllm/model_executor/models/registry.py | 1 + 4 files changed, 906 insertions(+) create mode 100644 vllm/model_executor/models/param2moe.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index c987acfa3f9..b86701c3ceb 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -457,6 +457,7 @@ th { | `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ | | `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ | | `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ | +| `Param2MoEForCausalLM` | param2moe | `bharatgenai/Param2-17B-A2.4B-Thinking`, etc. | ✅︎ | ✅︎ | | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | diff --git a/tests/models/registry.py b/tests/models/registry.py index f1f80e639d6..61f8958d11b 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -461,6 +461,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { trust_remote_code=True, is_available_online=False, ), + "Param2MoEForCausalLM": _HfExamplesInfo( + "bharatgenai/Param2-17B-A2.4B-Thinking", + trust_remote_code=True, + ), "PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"), "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"), "Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"), diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py new file mode 100644 index 00000000000..6812b1812e5 --- /dev/null +++ b/vllm/model_executor/models/param2moe.py @@ -0,0 +1,900 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright 2026 BharatGen AI team. All rights reserved. +# +# This code has been modified to accommodate Param2MoE's GQA-based MoE architecture. +# +# 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. +# limitations under the License. +from __future__ import annotations + +from collections.abc import Iterable, Iterator +from itertools import islice + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import SharedFusedMoE +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +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 default_weight_loader +from vllm.sequence import IntermediateTensors + +from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP +from .utils import ( + AutoWeightsLoader, + PPMissingLayer, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) + + +def _is_expert_bias_name(name: str) -> bool: + """True when the weight is the MoE router's per-expert score bias.""" + return name.endswith(".mlp.gate.expert_bias") + + +def _zero_mean_tensor(t: torch.Tensor) -> torch.Tensor: + if t.numel() == 0: + return t + return t - t.mean() + + +def _rename_and_normalize_weights( + weights: Iterable[tuple[str, torch.Tensor]], +) -> Iterator[tuple[str, torch.Tensor]]: + """ + Translate HuggingFace Param2MoE weight names to vLLM internal names + and zero-mean the expert-bias tensor so the router stays balanced. + + Mapping table (HF → vLLM): + model.word_embeddings.* → model.embed_tokens.* + *.attention.query_key_value.* → *.self_attn.qkv_proj.* + *.attention.dense.* → *.self_attn.o_proj.* + *.attention.query_layernorm.* → *.self_attn.q_layernorm.* + *.attention.key_layernorm.* → *.self_attn.k_layernorm.* + *.mlp.gate.expert_bias → *.mlp.gate.e_score_correction_bias + (also zero-meant for load balance) + """ + for name, w in weights: + # Embedding table + name = name.replace("model.word_embeddings.", "model.embed_tokens.") + # Fused QKV projection (HF: query_key_value → vLLM: qkv_proj) + name = name.replace(".attention.query_key_value.", ".self_attn.qkv_proj.") + # Output projection (HF: dense → vLLM: o_proj) + name = name.replace(".attention.dense.", ".self_attn.o_proj.") + # Per-head query norm + name = name.replace(".attention.query_layernorm.", ".self_attn.q_layernorm.") + # Per-head key norm + name = name.replace(".attention.key_layernorm.", ".self_attn.k_layernorm.") + # Catch any remaining .attention. → .self_attn. prefixes + # (e.g. future bias params on the projection layers) + name = name.replace(".attention.", ".self_attn.") + + # Expert-score bias: rename + zero-mean + if name.endswith(".mlp.gate.expert_bias"): + name = name.replace( + ".mlp.gate.expert_bias", + ".mlp.gate.e_score_correction_bias", + ) + w = _zero_mean_tensor(w) + + yield name, w + + +class Param2MoEAttention(nn.Module): + """ + Grouped-Query Attention (GQA) for Param2MoE. + + Notable differences from a vanilla GQA layer: + * The checkpoint fuses Q, K, V into a single ``query_key_value`` weight. + vLLM receives it already renamed to ``qkv_proj`` by the weight-name + translator and splits it during ``load_weights``. + * Optional per-head RMS norms on Q and K (``use_qk_norm=True``). + """ + + def __init__( + self, + config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim or (self.hidden_size // self.num_heads) + self.use_qk_norm: bool = getattr(config, "use_qk_norm", False) + + tp_size = get_tensor_model_parallel_world_size() + assert self.num_heads % tp_size == 0, ( + f"num_attention_heads ({self.num_heads}) must be divisible " + f"by tensor-parallel world size ({tp_size})." + ) + assert self.num_kv_heads % tp_size == 0, ( + f"num_key_value_heads ({self.num_kv_heads}) must be divisible " + f"by tensor-parallel world size ({tp_size})." + ) + self.num_local_heads = self.num_heads // tp_size + self.num_local_kv_heads = self.num_kv_heads // tp_size + + # Sizes after TP split (used in forward to split qkv output) + self.q_size_local = self.num_local_heads * self.head_dim + self.kv_size_local = self.num_local_kv_heads * self.head_dim + + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + hidden_size=self.hidden_size, + head_size=self.head_dim, + total_num_heads=self.num_heads, + total_num_kv_heads=self.num_kv_heads, + bias=getattr(config, "use_qkv_bias", False), + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.num_heads * self.head_dim, + output_size=self.hidden_size, + bias=getattr(config, "use_bias", False), + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + if self.use_qk_norm: + self.q_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + # `partial_rotary_factor` defaults to 1.0 (full RoPE) if not in config + partial_rotary_factor: float = getattr(config, "partial_rotary_factor", 1.0) + rope_dim = int(self.head_dim * partial_rotary_factor) + + rope_parameters: dict = { + "rope_type": "default", + "base": config.rope_theta, + } + if config.rope_scaling is not None: + rope_parameters.update(config.rope_scaling) + # Normalise key: some checkpoints use "type", vLLM wants "rope_type" + if "type" in rope_parameters and "rope_type" not in rope_parameters: + rope_parameters["rope_type"] = rope_parameters.pop("type") + + self.rotary_emb = get_rope( + rope_dim, + max_position=config.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=True, + ) + + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + scale=self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # 1. Fused QKV projection → split into local Q / K / V + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split( + [self.q_size_local, self.kv_size_local, self.kv_size_local], + dim=-1, + ) + + # 2. Optional per-head QK norms + # Reshape to (T, num_local_heads, head_dim), norm, reshape back. + if self.use_qk_norm: + T = q.shape[0] + q = self.q_layernorm(q.view(T, self.num_local_heads, self.head_dim)).view( + T, self.q_size_local + ) + k = self.k_layernorm( + k.view(T, self.num_local_kv_heads, self.head_dim) + ).view(T, self.kv_size_local) + + # 3. Rotary position embeddings + q, k = self.rotary_emb(positions, q, k) + + # 4. Paged attention + attn_output = self.attn(q, k, v) + + # 5. Output projection + output, _ = self.o_proj(attn_output) + return output + + +class Param2MoEMLP(nn.Module): + """SwiGLU feed-forward block used for dense layers.""" + + def __init__( + self, + intermediate_size: int, + config, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[intermediate_size, intermediate_size], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=config.hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Param2MoEMoEBlock(nn.Module): + """ + Mixture-of-Experts block for Param2MoE. + + Routing: + * Sigmoid scoring (config.score_function = "sigmoid") + * Grouped top-k (n_group, topk_group) + * Per-expert bias (gate.expert_bias → e_score_correction_bias) + * routed_scaling_factor normalisation + + One set of shared (always-active) experts is added on top. + """ + + def __init__( + self, + config, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.tp_size = get_tensor_model_parallel_world_size() + self.hidden_size = config.hidden_size + + self.num_experts: int = config.num_experts + self.top_k: int = config.num_experts_per_tok + self.routed_scaling_factor: float = getattr( + config, "routed_scaling_factor", 1.0 + ) + + self.n_group: int | None = getattr(config, "n_group", None) + self.topk_group: int | None = getattr(config, "topk_group", None) + self.use_grouped_topk: bool = ( + self.n_group is not None and self.topk_group is not None + ) + + self.norm_expert_prob: bool = getattr(config, "norm_topk_prob", True) + self.score_function: str = getattr(config, "score_function", "sigmoid") + + self.gate = nn.Linear( + self.hidden_size, + self.num_experts, + bias=False, + ) + + if getattr(config, "moe_router_enable_expert_bias", True): + self.gate.e_score_correction_bias = nn.Parameter( + torch.zeros(self.num_experts, dtype=torch.float32) + ) + else: + self.gate.e_score_correction_bias = None # type: ignore[assignment] + + self.num_shared_experts: int = getattr(config, "num_shared_experts", 1) + if self.num_shared_experts > 0: + # If moe_shared_expert_intermediate_size is present in the config + # it already encodes the TOTAL intermediate size across all shared + # experts (i.e. it equals moe_intermediate_size * num_shared_experts). + # Do NOT multiply again. Fall back to computing the product only + # when the dedicated field is absent. + if ( + hasattr(config, "moe_shared_expert_intermediate_size") + and config.moe_shared_expert_intermediate_size is not None + ): + shared_int: int = config.moe_shared_expert_intermediate_size + else: + shared_int = config.moe_intermediate_size * self.num_shared_experts + self.shared_experts = Param2MoEMLP( + intermediate_size=shared_int, + config=config, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + else: + self.shared_experts = None # type: ignore[assignment] + + self.experts = SharedFusedMoE( + shared_experts=self.shared_experts, + num_experts=self.num_experts, + top_k=self.top_k, + hidden_size=self.hidden_size, + intermediate_size=config.moe_intermediate_size, + reduce_results=False, + renormalize=self.norm_expert_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + scoring_func=self.score_function, + e_score_correction_bias=self.gate.e_score_correction_bias, + num_expert_group=self.n_group, + topk_group=self.topk_group, + use_grouped_topk=self.use_grouped_topk, + routed_scaling_factor=self.routed_scaling_factor, + ) + + def maybe_get_fused_moe(self) -> SharedFusedMoE: + return self.experts + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + # Router: both input and weight must be float32 for numerical + # stability (mirrors the original Param2MoEGate behaviour). + # The gate nn.Linear weight lives in the model dtype (bfloat16), + # so we must cast both explicitly via F.linear instead of calling + # self.gate() which would hit a dtype mismatch. + router_logits = F.linear( + hidden_states.float(), + self.gate.weight.float(), + ).to(hidden_states.dtype) + + final_hidden = self.experts( + hidden_states=hidden_states, + router_logits=router_logits, + ) + + if self.shared_experts is not None: + shared_output, expert_output = final_hidden + else: + shared_output, expert_output = None, final_hidden + + if shared_output is not None: + expert_output = expert_output + shared_output + + if self.tp_size > 1: + expert_output = self.experts.maybe_all_reduce_tensor_model_parallel( + expert_output + ) + + return expert_output.view(num_tokens, hidden_dim) + + +class Param2MoEDecoderLayer(nn.Module): + """ + Single transformer decoder block. + + Dense for the first ``first_k_dense_replace`` layers; MoE thereafter. + """ + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + + config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + hidden_size = config.hidden_size + # Derive the layer index from the prefix (e.g. "model.layers.3") + layer_idx = int(prefix.split(".")[-1]) + + self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps) + self.self_attn = Param2MoEAttention( + config=config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps) + + first_k_dense: int = getattr(config, "first_k_dense_replace", 1) + is_moe_layer = config.num_experts is not None and layer_idx >= first_k_dense + + if is_moe_layer: + self.mlp = Param2MoEMoEBlock( + config=config, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = Param2MoEMLP( # type: ignore[assignment] + intermediate_size=config.intermediate_size, + config=config, + quant_config=quant_config, + reduce_results=True, + prefix=f"{prefix}.mlp", + ) + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Pre-norm + attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + # Pre-norm + MLP + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class Param2MoEModel(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + ) -> None: + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + self.config = config + self.vocab_size = config.vocab_size + self.embed_dim = config.hidden_size + self.tie_word_embeddings: bool = getattr(config, "tie_word_embeddings", False) + + # Embedding (HF name: word_embeddings → vLLM name: embed_tokens) + if get_pp_group().is_first_rank or ( + self.tie_word_embeddings and get_pp_group().is_last_rank + ): + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + self.embed_dim, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: Param2MoEDecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + ), + prefix=f"{prefix}.layers", + ) + + 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 = RMSNorm(self.embed_dim, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + 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, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + 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"] + + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states, residual = layer(hidden_states, positions, residual) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + """ + Custom weight loader for the inner Param2MoEModel. + + Receives weights that have already been renamed/normalised by the + outer model and whose ``model.`` prefix has been stripped by + ``AutoWeightsLoader``. Handles: + 1. Fused QKV split (query_key_value → qkv_proj q/k/v shards). + 2. gate_proj + up_proj → gate_up_proj stacking (dense + shared-exp). + 3. Routed-expert weights via the fused-MoE mapping. + 4. All remaining weights via their default loader. + """ + config = self.config + num_heads: int = config.num_attention_heads + num_kv_heads: int = config.num_key_value_heads + head_dim: int = config.head_dim or (config.hidden_size // num_heads) + q_split = num_heads * head_dim + kv_split = num_kv_heads * head_dim + + stacked_params_mapping = [ + # (vllm_param_name, ckpt_weight_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + expert_params_mapping = self.get_expert_mapping() + + for name, loaded_weight in weights: + # ------------------------------------------------------------------ + # 1. Fused QKV: split into q / k / v shards for QKVParallelLinear + # ------------------------------------------------------------------ + if name.endswith(".self_attn.qkv_proj.weight"): + if name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + q_w = loaded_weight[:q_split, :] + k_w = loaded_weight[q_split : q_split + kv_split, :] + v_w = loaded_weight[q_split + kv_split :, :] + weight_loader(param, q_w, "q") + weight_loader(param, k_w, "k") + weight_loader(param, v_w, "v") + loaded_params.add(name) + continue + + # ------------------------------------------------------------------ + # 2. gate_proj / up_proj → gate_up_proj (dense MLP + shared-exp.) + # ------------------------------------------------------------------ + matched_stacked = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: # routed experts handled below + continue + new_name = name.replace(weight_name, param_name) + if new_name.endswith(".bias") and new_name not in params_dict: + continue + if new_name not in params_dict: + continue + if is_pp_missing_parameter(new_name, self): + continue + + param = params_dict[new_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(new_name) + matched_stacked = True + break + + if matched_stacked: + continue + + # ------------------------------------------------------------------ + # 3. Routed expert weights → fused-MoE kernel layout + # ------------------------------------------------------------------ + matched_expert = False + for ( + param_name, + weight_name, + expert_id, + shard_id, + ) in expert_params_mapping: + if weight_name not in name: + continue + new_name = name.replace(weight_name, param_name) + if is_pp_missing_parameter(new_name, self): + continue + if new_name not in params_dict: + continue + + param = params_dict[new_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + loaded_params.add(new_name) + matched_expert = True + break + + if matched_expert: + continue + + # ------------------------------------------------------------------ + # 4. All other weights: direct load (layernorms, embed_tokens, …) + # ------------------------------------------------------------------ + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + try: + weight_loader(param, loaded_weight) + except Exception as e: + raise RuntimeError( + f"[param2moe] Failed to load weight '{name}' " + f"with shape {tuple(loaded_weight.shape)} " + f"into param type {type(param).__name__}: {e}" + ) from e + loaded_params.add(name) + + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return SharedFusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + + +class Param2MoEMixtureOfExperts(MixtureOfExperts): + """Implements the vLLM MixtureOfExperts protocol for Param2MoE.""" + + expert_weights: list[torch.Tensor] + + def extract_moe_parameters(self, example_moe: Param2MoEMoEBlock | None) -> None: + if example_moe is None: + raise RuntimeError( + "No Param2MoEMoEBlock found in model.layers. " + "Check first_k_dense_replace and num_experts in config." + ) + self.num_logical_experts = example_moe.num_experts + self.num_routed_experts = example_moe.num_experts + self.num_shared_experts = example_moe.num_shared_experts + + self.num_physical_experts = self.num_logical_experts + self.num_local_physical_experts = self.num_logical_experts + self.num_redundant_experts = 0 + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + 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 moe in self.moe_mlp_layers: + moe.n_physical_experts = num_physical_experts + moe.n_local_physical_experts = num_local_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + + fused = moe.experts + if hasattr(fused, "n_local_physical_experts"): + fused.n_local_physical_experts = num_local_physical_experts + if hasattr(fused, "n_physical_experts"): + fused.n_physical_experts = num_physical_experts + if hasattr(fused, "n_redundant_experts"): + fused.n_redundant_experts = self.num_redundant_experts + if hasattr(fused, "update_expert_map"): + fused.update_expert_map() + + def set_eplb_state( + self, + expert_load_view: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + ) -> None: + self.expert_weights.clear() + for layer_idx, layer in enumerate(self.moe_layers): + if hasattr(layer, "get_expert_weights"): + self.expert_weights.append(layer.get_expert_weights()) + if hasattr(layer, "set_eplb_state"): + layer.set_eplb_state( + moe_layer_idx=layer_idx, + expert_load_view=expert_load_view, + logical_to_physical_map=logical_to_physical_map, + logical_replica_count=logical_replica_count, + ) + + +class Param2MoEForCausalLM( + nn.Module, SupportsPP, SupportsLoRA, Param2MoEMixtureOfExperts +): + """ + vLLM-native Param2MoE CausalLM. + + Uses Grouped-Query Attention (GQA) with a Sigmoid-scored, + grouped-topk Mixture-of-Experts MLP. + """ + + # LoRA packed-module mapping. The fused gate_up_proj handles + # gate_proj and up_proj from the checkpoint. + packed_modules_mapping = { + "qkv_proj": ["query_key_value"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + # Modules eligible for LoRA adaptation. + supported_lora_modules = [ + "qkv_proj", + "o_proj", + "gate_up_proj", + "down_proj", + ] + + # Embedding layers and their weight-tying counterparts. + embedding_modules = { + "embed_tokens": "input_embeddings", + "lm_head": "output_embeddings", + } + + # Modules that need vocab-size padding for LoRA. + embedding_padding_modules = ["lm_head"] + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + self.config = config + self.quant_config = quant_config + + self.model = Param2MoEModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + self.tie_word_embeddings: bool = getattr(config, "tie_word_embeddings", False) + if get_pp_group().is_last_rank: + if self.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = None # type: ignore[assignment] + + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + self.expert_weights: list[torch.Tensor] = [] + self.num_moe_layers: int = 0 + self.moe_layers: list = [] + self.moe_mlp_layers: list = [] + + example_moe: Param2MoEMoEBlock | None = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if isinstance(layer.mlp, Param2MoEMoEBlock): + example_moe = layer.mlp + self.moe_mlp_layers.append(layer.mlp) + self.moe_layers.append(layer.mlp.experts) + self.num_moe_layers += 1 + + if self.config.num_experts is not None: + self.extract_moe_parameters(example_moe) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + if not get_pp_group().is_last_rank: + return None + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(_rename_and_normalize_weights(weights)) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 4b354add384..fa129bfb42c 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -182,6 +182,7 @@ _TEXT_GENERATION_MODELS = { "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"), "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"), "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"), + "Param2MoEForCausalLM": ("param2moe", "Param2MoEForCausalLM"), "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"), "PhiForCausalLM": ("phi", "PhiForCausalLM"), "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), From fef56c18555e881c671acf654630732b7271c14f Mon Sep 17 00:00:00 2001 From: Julien Denize <40604584+juliendenize@users.noreply.github.com> Date: Mon, 6 Apr 2026 16:28:51 +0200 Subject: [PATCH 09/39] [Mistral Grammar] Support Grammar Factory (#38150) Signed-off-by: juliendenize --- requirements/common.txt | 2 +- requirements/rocm-test.txt | 2 +- requirements/test.txt | 2 +- tests/tokenizers_/test_mistral.py | 28 ++ .../tool_parsers/test_mistral_tool_parser.py | 347 +++++++++++++++++- .../test_backend_guidance.py | 44 +++ vllm/sampling_params.py | 28 +- vllm/tokenizers/mistral.py | 25 ++ vllm/tool_parsers/mistral_tool_parser.py | 142 ++++++- vllm/v1/structured_output/backend_guidance.py | 10 +- 10 files changed, 601 insertions(+), 29 deletions(-) diff --git a/requirements/common.txt b/requirements/common.txt index 05666c5d14b..b610fd67868 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -31,7 +31,7 @@ partial-json-parser # used for parsing partial JSON outputs pyzmq >= 25.0.0 msgspec gguf >= 0.17.0 -mistral_common[image] >= 1.10.0 +mistral_common[image] >= 1.11.0 opencv-python-headless >= 4.13.0 # required for video IO pyyaml six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12 diff --git a/requirements/rocm-test.txt b/requirements/rocm-test.txt index d5afde3c838..a441bfef04d 100644 --- a/requirements/rocm-test.txt +++ b/requirements/rocm-test.txt @@ -604,7 +604,7 @@ mcp==1.27.0 # via -r requirements/common.txt mdurl==0.1.2 # via markdown-it-py -mistral-common==1.10.0 +mistral-common==1.11.0 # via # -c requirements/common.txt # -r requirements/common.txt diff --git a/requirements/test.txt b/requirements/test.txt index 642e589a6a2..c8ff5fcabb2 100644 --- a/requirements/test.txt +++ b/requirements/test.txt @@ -508,7 +508,7 @@ mbstrdecoder==1.1.3 # typepy mdurl==0.1.2 # via markdown-it-py -mistral-common==1.10.0 +mistral-common==1.11.0 # via # -c requirements/common.txt # -r requirements/test.in diff --git a/tests/tokenizers_/test_mistral.py b/tests/tokenizers_/test_mistral.py index faff6115026..2b101e8f98d 100644 --- a/tests/tokenizers_/test_mistral.py +++ b/tests/tokenizers_/test_mistral.py @@ -3,8 +3,10 @@ from typing import Any +import llguidance import pytest from mistral_common.exceptions import InvalidMessageStructureException +from mistral_common.guidance.grammar_factory import GrammarFactory from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy from vllm.tokenizers.mistral import ( @@ -2407,3 +2409,29 @@ class TestMistralTokenizer: assert actual_tokens == expected_tokens assert mistral_tokenizer.convert_ids_to_tokens([]) == [] + + def test_grammar_factory(self, mistral_tokenizer: MistralTokenizer) -> None: + # works in this case cause Mistral 7B is < v11 and SPM + if not mistral_tokenizer.is_tekken: + with pytest.raises(AttributeError): + mistral_tokenizer.grammar_factory # noqa: B018 + return + factory = mistral_tokenizer.grammar_factory + assert isinstance(factory, GrammarFactory) + + # Test caching + factory_2 = mistral_tokenizer.grammar_factory + assert factory is factory_2 + + def test_llg_tokenizer(self, mistral_tokenizer: MistralTokenizer) -> None: + if not mistral_tokenizer.is_tekken: + with pytest.raises(ValueError): + mistral_tokenizer.llg_tokenizer # noqa: B018 + return + + llg_tokenizer = mistral_tokenizer.llg_tokenizer + assert isinstance(llg_tokenizer, llguidance.LLTokenizer) + + # Test caching + llg_tokenizer_2 = mistral_tokenizer.llg_tokenizer + assert llg_tokenizer is llg_tokenizer_2 diff --git a/tests/tool_parsers/test_mistral_tool_parser.py b/tests/tool_parsers/test_mistral_tool_parser.py index 4be5646669b..064ccb39ef4 100644 --- a/tests/tool_parsers/test_mistral_tool_parser.py +++ b/tests/tool_parsers/test_mistral_tool_parser.py @@ -3,19 +3,43 @@ import json from collections.abc import Generator +from unittest.mock import MagicMock, patch import partial_json_parser import pytest from mistral_common.protocol.instruct.messages import AssistantMessage from mistral_common.protocol.instruct.request import InstructRequest -from mistral_common.protocol.instruct.tool_calls import FunctionCall, ToolCall +from mistral_common.protocol.instruct.tool_calls import ( + FunctionCall, + ToolCall, +) +from mistral_common.protocol.instruct.tool_calls import ( + NamedToolChoice as MistralNamedToolChoice, +) +from mistral_common.protocol.instruct.tool_calls import ( + ToolChoice as MistralToolChoice, +) +from mistral_common.protocol.instruct.tool_calls import ( + ToolChoiceEnum as MistralToolChoiceEnum, +) from partial_json_parser.core.options import Allow -from vllm.entrypoints.openai.engine.protocol import DeltaMessage, DeltaToolCall +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, + DeltaToolCall, + StructuralTagResponseFormat, +) +from vllm.sampling_params import StructuredOutputsParams from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.tokenizers.detokenizer_utils import detokenize_incrementally from vllm.tokenizers.mistral import MistralTokenizer -from vllm.tool_parsers.mistral_tool_parser import MistralToolParser +from vllm.tool_parsers.mistral_tool_parser import ( + _DEFAULT_JSON_SCHEMA, + MistralToolParser, +) @pytest.fixture(scope="module") @@ -40,6 +64,13 @@ def mistral_tool_parser(mistral_tokenizer): return MistralToolParser(mistral_tokenizer) +@pytest.fixture +def non_mistral_parser() -> MistralToolParser: + mock_tokenizer = MagicMock() + mock_tokenizer.get_vocab.return_value = {"[TOOL_CALLS]": 1} + return MistralToolParser(mock_tokenizer) + + def assert_tool_calls( actual_tool_calls: list[ToolCall] | list[DeltaToolCall], expected_tool_calls: list[ToolCall], @@ -951,3 +982,313 @@ def test_fast_detokenization_text_detection_pre_v11( assert len(delta_message.tool_calls) > 0 assert delta_message.tool_calls[0].function is not None assert delta_message.tool_calls[0].function.name == "add" + + +SAMPLE_TOOLS_DICTS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "add", + "description": "Add two numbers", + "parameters": { + "type": "object", + "properties": { + "a": {"type": "number"}, + "b": {"type": "number"}, + }, + "required": ["a", "b"], + }, + }, + }, +] + + +def _make_request(**kwargs) -> ChatCompletionRequest: + defaults: dict = { + "messages": [], + "model": "mistralai/Mistral-Small-3.2-24B-Instruct-2506", + "tools": SAMPLE_TOOLS_DICTS, + "tool_choice": "auto", + } + defaults.update(kwargs) + return ChatCompletionRequest(**defaults) + + +@pytest.mark.parametrize( + "request_kwargs,expected_mode,expected_parallel", + [ + ({"tool_choice": "auto"}, MistralToolChoiceEnum.auto, True), + ({"tool_choice": "none"}, MistralToolChoiceEnum.none, True), + ({"tool_choice": "required"}, MistralToolChoiceEnum.required, True), + ({"tool_choice": None, "tools": None}, MistralToolChoiceEnum.auto, True), + ( + { + "tool_choice": { + "type": "function", + "function": {"name": "get_weather"}, + } + }, + MistralNamedToolChoice.model_validate( + {"type": "function", "function": {"name": "get_weather"}} + ), + True, + ), + ( + {"tool_choice": "auto", "parallel_tool_calls": False}, + MistralToolChoiceEnum.auto, + False, + ), + ( + {"tool_choice": "auto", "response_format": {"type": "text"}}, + MistralToolChoiceEnum.auto, + True, + ), + ], + ids=[ + "auto", + "none", + "required", + "null_tool_choice", + "named_tool_choice", + "parallel_false", + "response_format_text", + ], +) +def test_adjust_request_grammar_factory( + mistral_tool_parser: MistralToolParser, + request_kwargs: dict, + expected_mode: MistralToolChoice, + expected_parallel: bool, +) -> None: + request = _make_request(**request_kwargs) + factory = mistral_tool_parser.model_tokenizer.grammar_factory + + with patch.object( + factory, + "get_lark_from_jinja", + wraps=factory.get_lark_from_jinja, + ) as mock_get_lark: + result = mistral_tool_parser.adjust_request(request) + + mock_get_lark.assert_called_once() + call_kwargs = mock_get_lark.call_args + + assert call_kwargs.kwargs["mode"] == expected_mode + assert call_kwargs.kwargs["json_schema"] is None + assert call_kwargs.kwargs["parallel_tool_calls"] == expected_parallel + + assert result.structured_outputs is not None + assert isinstance(result.structured_outputs.grammar, str) + assert len(result.structured_outputs.grammar) > 0 + + +def test_adjust_request_unsupported_grammar_for_tokenizer(mistral_tokenizer) -> None: + with patch.object( + type(mistral_tokenizer), + "supports_grammar", + new_callable=lambda: property(lambda self: False), + ): + parser = MistralToolParser(mistral_tokenizer) + request = _make_request() + result = parser.adjust_request(request) + + assert result.structured_outputs is None + + +@pytest.mark.parametrize( + "tool_choice,expected_skip", + [("auto", False), ("none", True)], + ids=["auto_skip_false", "none_skip_true"], +) +def test_adjust_request_non_mistral_tokenizer( + non_mistral_parser: MistralToolParser, + tool_choice: str, + expected_skip: bool, +) -> None: + request = _make_request(tool_choice=tool_choice) + result = non_mistral_parser.adjust_request(request) + + assert result.skip_special_tokens is expected_skip + + +@pytest.mark.parametrize( + "so_kwargs", + [ + {"regex": r"\d+"}, + {"choice": ["a", "b"]}, + {"structural_tag": '{"key": "value"}'}, + {"grammar": "start: 'hello'"}, + ], + ids=["regex", "choice", "structural_tag", "grammar"], +) +def test_adjust_request_unsupported_structured_outputs( + mistral_tool_parser: MistralToolParser, + so_kwargs: dict, +) -> None: + request = _make_request( + structured_outputs=StructuredOutputsParams(**so_kwargs), + ) + result = mistral_tool_parser.adjust_request(request) + + assert result.structured_outputs == request.structured_outputs + + +def test_adjust_request_unsupported_response_format( + mistral_tool_parser: MistralToolParser, +) -> None: + request = _make_request( + response_format=StructuralTagResponseFormat( + type="structural_tag", format={"some": "config"} + ), + ) + result = mistral_tool_parser.adjust_request(request) + assert result.structured_outputs is None + assert result.response_format == request.response_format + + +@pytest.mark.parametrize( + "so_kwargs,expected_json_schema", + [ + ({"json_object": True}, _DEFAULT_JSON_SCHEMA), + ({"json": '{"type": "object"}'}, {"type": "object"}), + ( + {"json": {"type": "object", "properties": {"x": {"type": "integer"}}}}, + {"type": "object", "properties": {"x": {"type": "integer"}}}, + ), + ], + ids=["json_object", "json_str", "json_dict"], +) +def test_adjust_request_structured_outputs_generates_grammar( + mistral_tool_parser: MistralToolParser, + so_kwargs: dict, + expected_json_schema: str, +) -> None: + request = _make_request( + structured_outputs=StructuredOutputsParams(**so_kwargs), + ) + factory = mistral_tool_parser.model_tokenizer.grammar_factory + + with patch.object( + factory, + "get_lark_from_jinja", + wraps=factory.get_lark_from_jinja, + ) as mock_get_lark: + result = mistral_tool_parser.adjust_request(request) + + mock_get_lark.assert_called_once() + assert mock_get_lark.call_args.kwargs["json_schema"] == expected_json_schema + + assert result.structured_outputs is not None + assert isinstance(result.structured_outputs.grammar, str) + assert len(result.structured_outputs.grammar) > 0 + + +@pytest.mark.parametrize( + "response_format_kwargs,expected_json_schema", + [ + ({"type": "json_object"}, _DEFAULT_JSON_SCHEMA), + ( + { + "type": "json_schema", + "json_schema": { + "name": "my_schema", + "schema": { + "type": "object", + "properties": {"x": {"type": "integer"}}, + }, + }, + }, + {"type": "object", "properties": {"x": {"type": "integer"}}}, + ), + ], + ids=["json_object", "json_schema_with_schema"], +) +def test_adjust_request_response_format_generates_grammar( + mistral_tool_parser: MistralToolParser, + response_format_kwargs: dict, + expected_json_schema: str, +) -> None: + request = _make_request(response_format=response_format_kwargs) + factory = mistral_tool_parser.model_tokenizer.grammar_factory + + with patch.object( + factory, + "get_lark_from_jinja", + wraps=factory.get_lark_from_jinja, + ) as mock_get_lark: + result = mistral_tool_parser.adjust_request(request) + + mock_get_lark.assert_called_once() + assert mock_get_lark.call_args.kwargs["json_schema"] == expected_json_schema + + assert result.structured_outputs is not None + assert isinstance(result.structured_outputs.grammar, str) + assert len(result.structured_outputs.grammar) > 0 + + +def test_adjust_request_tool_choice_none_with_json_schema_uses_json_schema_factory( + mistral_tool_parser: MistralToolParser, +) -> None: + request = _make_request( + tool_choice="none", + structured_outputs=StructuredOutputsParams(json='{"type": "object"}'), + ) + factory = mistral_tool_parser.model_tokenizer.grammar_factory + + with patch.object( + factory, + "get_lark_for_json_schema", + wraps=factory.get_lark_for_json_schema, + ) as mock_json_schema: + result = mistral_tool_parser.adjust_request(request) + + mock_json_schema.assert_called_once() + assert mock_json_schema.call_args.kwargs["json_schema"] == {"type": "object"} + + assert result.structured_outputs is not None + assert isinstance(result.structured_outputs.grammar, str) + assert len(result.structured_outputs.grammar) > 0 + + +def test_adjust_request_tool_choice_auto_with_json_schema_uses_jinja_factory( + mistral_tool_parser: MistralToolParser, +) -> None: + request = _make_request( + tool_choice="auto", + structured_outputs=StructuredOutputsParams(json='{"type": "object"}'), + ) + factory = mistral_tool_parser.model_tokenizer.grammar_factory + + with ( + patch.object( + factory, + "get_lark_for_json_schema", + wraps=factory.get_lark_for_json_schema, + ) as mock_json_schema, + patch.object( + factory, + "get_lark_from_jinja", + wraps=factory.get_lark_from_jinja, + ) as mock_jinja, + ): + result = mistral_tool_parser.adjust_request(request) + + mock_jinja.assert_called_once() + assert mock_jinja.call_args.kwargs["json_schema"] == {"type": "object"} + mock_json_schema.assert_not_called() + + assert result.structured_outputs is not None + assert isinstance(result.structured_outputs.grammar, str) + assert len(result.structured_outputs.grammar) > 0 diff --git a/tests/v1/structured_output/test_backend_guidance.py b/tests/v1/structured_output/test_backend_guidance.py index 704ed8b9c9e..ca8c9b0d785 100644 --- a/tests/v1/structured_output/test_backend_guidance.py +++ b/tests/v1/structured_output/test_backend_guidance.py @@ -11,6 +11,7 @@ from vllm.config.model import ModelConfig from vllm.config.parallel import ParallelConfig from vllm.config.speculative import SpeculativeConfig from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.tokenizers import get_tokenizer from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.structured_output.backend_guidance import GuidanceBackend @@ -19,6 +20,14 @@ from vllm.v1.structured_output.backend_types import StructuredOutputOptions TOKENIZER = "gpt2" +@pytest.fixture(scope="module") +def mistral_tokenizer(): + return get_tokenizer( + tokenizer_name="mistralai/Mistral-Small-3.2-24B-Instruct-2506", + tokenizer_mode="mistral", + ) + + def test_backend_guidance_rollback_terminated(): # Test that the backend guidance successfully rollbacks from a # terminated state. This can happen with speculative decoding, @@ -187,3 +196,38 @@ def test_grammar_init_async_and_sync(async_grammar): # Verify the grammar can accept valid tokens assert grammar.accept_tokens(request.request_id, prompt) + + +@pytest.mark.parametrize( + "request_type,grammar_spec", + [ + pytest.param( + StructuredOutputOptions.JSON, + '{"type": "object"}', + id="json", + ), + pytest.param( + StructuredOutputOptions.GRAMMAR, + 'start: "hello" | "world"', + id="lark", + ), + ], +) +def test_mistral_tokenizer_compile_grammar( + mistral_tokenizer, + request_type: StructuredOutputOptions, + grammar_spec: str, +) -> None: + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig(backend="guidance"), + ) + backend = GuidanceBackend( + vllm_config, + tokenizer=mistral_tokenizer, + vocab_size=mistral_tokenizer.vocab_size, + ) + assert backend.ll_tokenizer is mistral_tokenizer.llg_tokenizer + + grammar = backend.compile_grammar(request_type, grammar_spec) + assert grammar is not None + assert not grammar.is_terminated() diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 97976b83209..9bcc669591e 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -153,6 +153,10 @@ class RequestOutputKind(Enum): FINAL_ONLY = 2 +def _is_non_tekken_mistral(tokenizer: TokenizerLike) -> bool: + return is_mistral_tokenizer(tokenizer) and not tokenizer.is_tekken + + class SamplingParams( PydanticMsgspecMixin, msgspec.Struct, @@ -801,16 +805,17 @@ class SamplingParams( # xgrammar with no fallback validate_xgrammar_grammar(self) elif backend.startswith("guidance"): + if _is_non_tekken_mistral(tokenizer=tokenizer): + raise ValueError( + "Non-tekken Mistral tokenizers are not supported for the 'guidance'" + " structured output backend. Please either use a more recent " + "Mistral model, the ['xgrammar', 'outlines'] " + "backends or tokenizer_mode='hf' instead." + ) # TODO: ideally we would have the LLTokenizer here as Lark syntax # allows <|special_token|> and similar, see # https://github.com/guidance-ai/llguidance/blob/main/docs/syntax.md#special-tokens # Without tokenizer these are disallowed in grammars. - if is_mistral_tokenizer(tokenizer): - raise ValueError( - "Mistral tokenizer is not supported for the 'guidance' " - "structured output backend. Please use ['xgrammar', 'outlines'] " - "backends or tokenizer_mode='hf' instead." - ) validate_guidance_grammar(self, tokenizer=None) elif backend == "outlines": # outlines backend @@ -839,19 +844,20 @@ class SamplingParams( # or includes some jsonschema feature(s) that # are not supported in xgrammar. + skip_guidance = _is_non_tekken_mistral(tokenizer) + # Check if schema has features unsupported by guidance so_params = self.structured_outputs - skip_guidance = False - if so_params.json: + if not skip_guidance and so_params.json: if isinstance(so_params.json, str): schema = json_mod.loads(so_params.json) else: schema = so_params.json skip_guidance = has_guidance_unsupported_json_features(schema) - if is_mistral_tokenizer(tokenizer) or skip_guidance: - # Fall back to outlines if the tokenizer is Mistral - # or if schema contains features unsupported by guidance + if skip_guidance: + # Fall back to outlines if the tokenizer is non-tekken Mistral or + # the schema contains features unsupported by guidance validate_structured_output_request_outlines(self) self.structured_outputs._backend = "outlines" else: diff --git a/vllm/tokenizers/mistral.py b/vllm/tokenizers/mistral.py index e20f1edd472..147dca88877 100644 --- a/vllm/tokenizers/mistral.py +++ b/vllm/tokenizers/mistral.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence +from functools import cached_property from pathlib import Path from typing import TYPE_CHECKING, Any, cast, overload +from mistral_common.guidance.grammar_factory import GrammarFactory +from mistral_common.guidance.tokenizer import from_mistral_tokenizer from mistral_common.protocol.instruct.request import ( ChatCompletionRequest as MistralChatCompletionRequest, ) @@ -45,6 +48,7 @@ except ImportError: ) if TYPE_CHECKING: + import llguidance from transformers import BatchEncoding logger = init_logger(__name__) @@ -574,3 +578,24 @@ class MistralTokenizer(TokenizerLike): ] return tokens + + @property + def supports_grammar(self) -> bool: + return GrammarFactory.is_supported(self.mistral) + + @cached_property + def grammar_factory(self) -> GrammarFactory: + if not self.supports_grammar: + raise AttributeError( + "This tokenizer does not support `grammar_factory`. " + "This is only supported for tekken tokenizers with " + "version >= 11." + ) + # Cache grammar factory to avoid creating a llguidance tokenizer at every usage. + return GrammarFactory(self.mistral) + + @cached_property + def llg_tokenizer(self) -> "llguidance.LLTokenizer": + if not self.is_tekken: + raise ValueError("`llg_tokenizer` is only supported for Tekkenizers.") + return from_mistral_tokenizer(self.mistral) diff --git a/vllm/tool_parsers/mistral_tool_parser.py b/vllm/tool_parsers/mistral_tool_parser.py index dc92522a052..4d1aaffedd0 100644 --- a/vllm/tool_parsers/mistral_tool_parser.py +++ b/vllm/tool_parsers/mistral_tool_parser.py @@ -10,6 +10,18 @@ from typing import Any import ijson import regex as re +from mistral_common.protocol.instruct.tool_calls import ( + NamedToolChoice as MistralNamedToolChoice, +) +from mistral_common.protocol.instruct.tool_calls import ( + Tool as MistralTool, +) +from mistral_common.protocol.instruct.tool_calls import ( + ToolChoice as MistralToolChoice, +) +from mistral_common.protocol.instruct.tool_calls import ( + ToolChoiceEnum as MistralToolChoiceEnum, +) from pydantic import Field from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -25,6 +37,7 @@ from vllm.entrypoints.openai.engine.protocol import ( ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger +from vllm.sampling_params import StructuredOutputsParams from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import ( Tool, @@ -36,6 +49,8 @@ logger = init_logger(__name__) ALPHANUMERIC = ascii_letters + digits +_DEFAULT_JSON_SCHEMA = {"anyOf": [{"type": "object"}, {"type": "array"}]} + class StreamingState(Enum): """Enum for tracking the current streaming parsing state.""" @@ -80,6 +95,9 @@ class MistralToolParser(ToolParser): Used when --enable-auto-tool-choice --tool-call-parser mistral are all set """ + # Used to generate correct grammar in `adjust_request` + model_can_reason: bool = False + def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): super().__init__(tokenizer, tools) @@ -115,18 +133,124 @@ class MistralToolParser(ToolParser): def adjust_request( self, request: ChatCompletionRequest | ResponsesRequest ) -> ChatCompletionRequest | ResponsesRequest: - request = super().adjust_request(request) + so_non_supported_attributes = [ + "regex", + "choice", + "grammar", + # whitespace_pattern is not a constraint type but an option; + # Mistral grammar factory does not support it. + "whitespace_pattern", + "structural_tag", + ] + any_so_non_supported_active = request.structured_outputs is not None and any( + getattr(request.structured_outputs, attribute) is not None + for attribute in so_non_supported_attributes + ) + response_format_non_supported_active = ( + isinstance(request, ResponsesRequest) + or request.response_format is not None + and request.response_format.type == "structural_tag" + ) + if ( not is_mistral_tokenizer(self.model_tokenizer) - and request.tools - and request.tool_choice != "none" + or isinstance(request, ResponsesRequest) + or not self.model_tokenizer.supports_grammar + or any_so_non_supported_active + or response_format_non_supported_active ): - # Do not skip special tokens when using chat template - # with Mistral parser as TOOL_CALL token is needed - # for tool detection. - # Note: we don't want skip_special_tokens=False - # with MistralTokenizer as it is incompatible - request.skip_special_tokens = False + request = super().adjust_request(request) + if request.tools and request.tool_choice != "none": + # Do not skip special tokens when using chat template + # with Mistral parser as TOOL_CALL token is needed + # for tool detection. + # Note: we don't want skip_special_tokens=False + # with MistralTokenizer as it is incompatible + request.skip_special_tokens = False + return request + + json_schema: dict[str, Any] | None = None + if request.structured_outputs is not None: + if request.structured_outputs.json_object is not None: + json_schema = _DEFAULT_JSON_SCHEMA + elif request.structured_outputs.json is not None: + if isinstance(request.structured_outputs.json, str): + json_schema = json.loads(request.structured_outputs.json) + else: + json_schema = request.structured_outputs.json + else: + raise ValueError( + "Unsupported request.structured_outputs for MistralToolParser. " + "Only `json` and `json_object` are supported." + ) + elif ( + request.response_format is not None + and request.response_format.type != "text" + ): + if request.response_format.type == "json_object": + json_schema = _DEFAULT_JSON_SCHEMA + elif request.response_format.type == "json_schema": + if request.response_format.json_schema is not None: + json_schema = request.response_format.json_schema.json_schema + else: + json_schema = _DEFAULT_JSON_SCHEMA + else: + raise ValueError( + "MistralToolParser only accepts `text`, `json_object` or " + f"`json_schema`, got {request.response_format=}" + ) + # Structured Outputs will be defined. + request.response_format = None + + grammar_factory = self.model_tokenizer.grammar_factory + + # TODO: Once unified parser, improve this. + # The issue is figuring out when a model is a reasoning one or not. + template = grammar_factory.select_jinja_template( + reasoning=self.model_can_reason + ) + + tools = ( + [ + MistralTool.from_openai(openai_tool=tool.model_dump()) + for tool in request.tools + ] + if request.tools is not None + else None + ) + + tool_choice: MistralToolChoice + match request.tool_choice: + case "none" | "auto" | "required": + tool_choice = MistralToolChoiceEnum(request.tool_choice) + case None: + tool_choice = MistralToolChoiceEnum.auto + # _ == Named tool choice + case _: + tool_choice = MistralNamedToolChoice.model_validate( + { + "type": "function", + "function": {"name": request.tool_choice.function.name}, + } + ) + + # Rendering grammar is cached in mistral-common given tools, template and mode. + match tool_choice, json_schema is not None: + case MistralToolChoiceEnum.none, True: + lark_grammar = grammar_factory.get_lark_for_json_schema( + template=template, json_schema=json_schema + ) + case _, _: + lark_grammar = grammar_factory.get_lark_from_jinja( + template=template, + mode=tool_choice, + tools=tools, + json_schema=json_schema, + parallel_tool_calls=request.parallel_tool_calls, + json_only=False, + ) + + request.structured_outputs = StructuredOutputsParams(grammar=lark_grammar) return request def extract_tool_calls( diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 6063a2dc2a6..31178e9f246 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -12,6 +12,7 @@ import torch from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader +from vllm.utils.mistral import is_mistral_tokenizer from vllm.v1.structured_output.backend_types import ( StructuredOutputBackend, StructuredOutputGrammar, @@ -92,9 +93,12 @@ class GuidanceBackend(StructuredOutputBackend): self.vllm_config.structured_outputs_config.disable_additional_properties ) - self.ll_tokenizer = llguidance_hf.from_tokenizer( - self.tokenizer, max(self.vocab_size, len(self.tokenizer)) - ) + if is_mistral_tokenizer(self.tokenizer): + self.ll_tokenizer = self.tokenizer.llg_tokenizer + else: + self.ll_tokenizer = llguidance_hf.from_tokenizer( + self.tokenizer, max(self.vocab_size, len(self.tokenizer)) + ) def compile_grammar( self, request_type: StructuredOutputOptions, grammar_spec: str From e69a265135ef48312d78130f64b7bfce4cd81a37 Mon Sep 17 00:00:00 2001 From: Walter Beller-Morales Date: Mon, 6 Apr 2026 11:00:16 -0400 Subject: [PATCH 10/39] [Feat][Core] safely abort requests when FSM fails to advance (#38663) Signed-off-by: walterbm --- tests/v1/core/test_async_scheduler.py | 67 +++++++++++++++++++++- tests/v1/core/test_scheduler.py | 81 +++++++++++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 29 ++++++---- 3 files changed, 164 insertions(+), 13 deletions(-) diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index a77ae81bae5..e821e47172c 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import deque +from unittest.mock import Mock import pytest -from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.async_scheduler import AsyncScheduler +from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import RequestStatus from vllm.v1.utils import ConstantList @@ -247,3 +249,66 @@ def test_prefix_caching_for_multi_turn(): # requests. for req in next_turn_requests: assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE + + +def test_abort_request_when_structured_output_fsm_cannot_advance(): + scheduler = object.__new__(AsyncScheduler) + request = create_requests(num_requests=1, num_tokens=1)[0] + request.structured_output_request = Mock() + request.structured_output_request.grammar = Mock() + request.structured_output_request.grammar.accept_tokens.return_value = False + request.status = RequestStatus.RUNNING + request.num_computed_tokens = request.num_tokens + request.num_output_placeholders = 1 + + scheduler.perf_metrics = None + scheduler.connector = None + scheduler.structured_output_manager = Mock() + scheduler.structured_output_manager.should_advance.return_value = True + scheduler.requests = {request.request_id: request} + scheduler.running = [request] + scheduler.waiting = Mock() + scheduler.kv_cache_manager = Mock() + scheduler.kv_cache_manager.take_events.return_value = None + scheduler.kv_event_publisher = Mock() + scheduler.finished_req_ids = set() + scheduler.finished_req_ids_dict = None + scheduler.vllm_config = Mock() + scheduler.vllm_config.model_config.enable_return_routed_experts = False + scheduler.recompute_kv_load_failures = False + scheduler.make_stats = Mock(return_value=None) + scheduler.max_model_len = 128 + + def free_request(req, delay_free_blocks=False): + scheduler.finished_req_ids.add(req.request_id) + scheduler.requests.pop(req.request_id, None) + return None + + scheduler._free_request = Mock(side_effect=free_request) + + output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 1}, + total_num_scheduled_tokens=1, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_runner_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[123]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + scheduler.update_from_output(output, model_runner_output) + + assert request.resumable is False + assert request.status == RequestStatus.FINISHED_ERROR + assert request.request_id not in scheduler.requests + assert not scheduler.running diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 2304bf7ecae..8fd2309f7ec 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -26,6 +26,7 @@ from vllm.v1.core.encoder_cache_manager import EncoderCacheManager from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.engine import FinishReason from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, @@ -2463,6 +2464,86 @@ def test_schedule_skip_tokenizer_init_structured_output_request(): assert len(scheduler.skipped_waiting) == 1 +def test_abort_request_when_structured_output_fsm_cannot_advance(): + scheduler = object.__new__(Scheduler) + sampling_params = SamplingParams(ignore_eos=True, max_tokens=4) + sampling_params.update_from_generation_config({}, EOS_TOKEN_ID) + + request = Request( + request_id="0", + prompt_token_ids=[0, 1], + mm_features=None, + sampling_params=sampling_params, + pooling_params=None, + ) + request.structured_output_request = Mock() + request.structured_output_request.grammar = Mock() + request.structured_output_request.grammar.accept_tokens.return_value = False + request.status = RequestStatus.RUNNING + request.num_computed_tokens = request.num_tokens + + scheduler.perf_metrics = None + scheduler.connector = None + scheduler.structured_output_manager = Mock() + scheduler.structured_output_manager.should_advance.return_value = True + scheduler.requests = {request.request_id: request} + scheduler.running = [request] + scheduler.waiting = Mock() + scheduler.kv_cache_manager = Mock() + scheduler.kv_cache_manager.take_events.return_value = None + scheduler.kv_event_publisher = Mock() + scheduler.finished_req_ids = set() + scheduler.finished_req_ids_dict = None + scheduler.vllm_config = Mock() + scheduler.vllm_config.model_config.enable_return_routed_experts = False + scheduler.recompute_kv_load_failures = False + scheduler.make_stats = Mock(return_value=None) + scheduler.max_model_len = 128 + + def free_request(req: Request, delay_free_blocks: bool = False): + scheduler.finished_req_ids.add(req.request_id) + scheduler.requests.pop(req.request_id, None) + return None + + scheduler._free_request = Mock(side_effect=free_request) + + output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 1}, + total_num_scheduled_tokens=1, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + + model_runner_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[123]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + engine_core_outputs = scheduler.update_from_output(output, model_runner_output) + + request.structured_output_request.grammar.accept_tokens.assert_called_once_with( + request.request_id, [123] + ) + assert request.resumable is False + assert request.status == RequestStatus.FINISHED_ERROR + assert request.request_id not in scheduler.requests + assert not scheduler.running + scheduler._free_request.assert_called_once_with(request) + assert len(engine_core_outputs[0].outputs) == 1 + engine_core_output = engine_core_outputs[0].outputs[0] + assert engine_core_output.request_id == request.request_id + assert engine_core_output.new_token_ids == [123] + assert engine_core_output.finish_reason == FinishReason.ERROR + + @pytest.mark.parametrize( "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")] ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index fe524ccace1..2a0bf463cf0 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1406,6 +1406,23 @@ class Scheduler(SchedulerInterface): request.status = RequestStatus.FINISHED_STOPPED stopped = True + if new_token_ids and self.structured_output_manager.should_advance(request): + struct_output_request = request.structured_output_request + assert struct_output_request is not None + assert struct_output_request.grammar is not None + if not struct_output_request.grammar.accept_tokens( # type: ignore[union-attr] + req_id, new_token_ids + ): + logger.error( + "Unexpected: grammar rejected tokens %s for request %s. " + "Terminating request.", + new_token_ids, + req_id, + ) + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True + routed_experts = None finish_reason = None if stopped: @@ -1431,18 +1448,6 @@ class Scheduler(SchedulerInterface): ): new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) - if new_token_ids and self.structured_output_manager.should_advance(request): - struct_output_request = request.structured_output_request - assert struct_output_request is not None - assert struct_output_request.grammar is not None - ok = struct_output_request.grammar.accept_tokens(req_id, new_token_ids) - if not ok: - logger.warning( - "Unexpected: grammar rejected tokens %s for request %s.", - new_token_ids, - req_id, - ) - if num_nans_in_logits is not None and req_id in num_nans_in_logits: request.num_nans_in_logits = num_nans_in_logits[req_id] From 47e605092b7fce3d64264b34250b1a286f344633 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Mon, 6 Apr 2026 11:19:39 -0400 Subject: [PATCH 11/39] [Gemma4] Enable Fast Prefill Optimization (#38879) Signed-off-by: Lucas Wilkinson --- vllm/model_executor/models/gemma4.py | 420 +++++++++++++++++++++++---- 1 file changed, 371 insertions(+), 49 deletions(-) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index edb53313499..2e9fc681903 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -19,6 +19,7 @@ """Gemma 4 model implementation for vLLM.""" from collections.abc import Iterable +from dataclasses import replace from itertools import islice import regex as re @@ -32,6 +33,7 @@ from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.activation import GeluAndMul from vllm.model_executor.layers.attention import Attention @@ -56,6 +58,7 @@ from vllm.model_executor.model_loader.weight_utils import ( maybe_remap_kv_scale_name, ) from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP from .utils import ( @@ -636,7 +639,205 @@ class Gemma4DecoderLayer(nn.Module): return hidden_states, None -@support_torch_compile +def _run_decoder_layers( + decoder_layers: list[Gemma4DecoderLayer], + layer_idx_start: int, + positions: torch.Tensor, + hidden_states: torch.Tensor, + per_layer_inputs: torch.Tensor | None = None, + **kwargs, +) -> torch.Tensor: + """Run a slice of decoder layers with PLE extraction.""" + residual = None + for idx, layer in enumerate(decoder_layers): + layer_idx = idx + layer_idx_start + layer_per_input = ( + per_layer_inputs[:, layer_idx, :] if per_layer_inputs is not None else None + ) + hidden_states, residual = layer( + positions, + hidden_states, + residual, + per_layer_input=layer_per_input, + **kwargs, + ) + return hidden_states + + +@support_torch_compile( + enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill +) +class Gemma4SelfDecoderLayers(nn.Module): + """Compiled wrapper: embedding + non-KV-shared layers (YOCO first half). + + Owns the embedding and PLE modules so they are inside the compiled + graph. Gemma4Model delegates embedding methods here. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layers: list[Gemma4DecoderLayer], + layer_idx_start: int, + embed_tokens: VocabParallelEmbedding, + normalizer: torch.Tensor, + embed_tokens_per_layer: VocabParallelEmbedding | None, + embed_scale_per_layer: torch.Tensor | None, + per_layer_model_projection: ColumnParallelLinear | None, + per_layer_projection_norm: RMSNorm | None, + per_layer_input_scale: torch.Tensor | None, + per_layer_projection_scale: torch.Tensor | None, + ): + super().__init__() + self.decoder_layers = decoder_layers + self.layer_idx_start = layer_idx_start + + config = _get_text_config(vllm_config.model_config.hf_config) + self.config = config + self.hidden_size_per_layer_input = getattr( + config, "hidden_size_per_layer_input", 0 + ) + self.vocab_size_per_layer_input = getattr( + config, "vocab_size_per_layer_input", config.vocab_size + ) + + # Shared references to modules owned by Gemma4Model — must be + # inside this nn.Module so torch.compile captures them. + self.embed_tokens = embed_tokens + self.normalizer = normalizer + self.embed_tokens_per_layer = embed_tokens_per_layer + self.embed_scale_per_layer = embed_scale_per_layer + self.per_layer_model_projection = per_layer_model_projection + self.per_layer_projection_norm = per_layer_projection_norm + self.per_layer_input_scale = per_layer_input_scale + self.per_layer_projection_scale = per_layer_projection_scale + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) * self.normalizer + + def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None: + """Get per-layer embeddings from embed_tokens_per_layer. + + Returns: + Per-layer embeddings (num_tokens, num_layers, + hidden_size_per_layer_input) + """ + if self.embed_tokens_per_layer is None: + return None + per_layer_inputs_mask = torch.logical_and( + input_ids >= 0, + input_ids < self.vocab_size_per_layer_input, + ) + per_layer_inputs_tokens = torch.where( + per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids) + ) + per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens) + per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer + return per_layer_embeds.reshape( + *input_ids.shape, + self.config.num_hidden_layers, + self.hidden_size_per_layer_input, + ) + + def project_per_layer_inputs( + self, + inputs_embeds: torch.Tensor, + per_layer_inputs: torch.Tensor | None, + ) -> torch.Tensor | None: + """Project inputs_embeds and combine with per_layer_inputs. + + Steps: + 1. Project inputs_embeds: hidden_size → total_ple_dim + 2. Scale by hidden_size^{-0.5} + 3. Reshape to (num_tokens, num_layers, per_layer_dim) + 4. Normalize with per_layer_projection_norm + 5. Combine: (projection + per_layer_inputs) * 1/sqrt(2) + """ + if self.per_layer_model_projection is None: + return None + per_layer_projection = self.per_layer_model_projection(inputs_embeds) + per_layer_projection = per_layer_projection * self.per_layer_projection_scale + per_layer_projection = per_layer_projection.reshape( + *inputs_embeds.shape[:-1], + self.config.num_hidden_layers, + self.hidden_size_per_layer_input, + ) + per_layer_projection = self.per_layer_projection_norm(per_layer_projection) + if per_layer_inputs is None: + return per_layer_projection + return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + per_layer_inputs: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if inputs_embeds is not None: + hidden_states = inputs_embeds + per_layer_inputs = self.project_per_layer_inputs( + hidden_states, per_layer_inputs + ) + else: + hidden_states = self.embed_input_ids(input_ids) + per_layer_embeds = self.get_per_layer_inputs(input_ids) + per_layer_inputs = self.project_per_layer_inputs( + hidden_states, per_layer_embeds + ) + + hidden_states = _run_decoder_layers( + self.decoder_layers, + self.layer_idx_start, + positions, + hidden_states, + per_layer_inputs, + **kwargs, + ) + return hidden_states, per_layer_inputs + + +@support_torch_compile( + enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill +) +class Gemma4CrossDecoderLayers(nn.Module): + """Cross-decoder layers (YOCO second half, KV-shared).""" + + def __init__( + self, + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layers: list[Gemma4DecoderLayer], + layer_idx_start: int, + ): + super().__init__() + self.decoder_layers = decoder_layers + self.layer_idx_start = layer_idx_start + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + per_layer_inputs: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return _run_decoder_layers( + self.decoder_layers, + self.layer_idx_start, + positions, + hidden_states, + per_layer_inputs, + **kwargs, + ) + + +@support_torch_compile( + enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill +) class Gemma4Model(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -740,6 +941,75 @@ class Gemma4Model(nn.Module): torch.tensor(config.hidden_size**0.5), persistent=False, ) + + # --- You Only Cache Once (YOCO) split for fast prefill --- + first_kv_shared_layer_idx = config.num_hidden_layers - getattr( + config, "num_kv_shared_layers", 0 + ) + + from vllm.compilation.backends import set_model_tag + + # Layers 0..(K-1) are self-decoder layers in YOCO + with set_model_tag("self_decoder"): + self.self_decoder = Gemma4SelfDecoderLayers( + vllm_config=vllm_config, + prefix=f"{prefix}.self_decoder", + decoder_layers=self.layers[:first_kv_shared_layer_idx], + layer_idx_start=0, + embed_tokens=self.embed_tokens, + normalizer=self.normalizer, + embed_tokens_per_layer=getattr(self, "embed_tokens_per_layer", None), + embed_scale_per_layer=getattr(self, "embed_scale_per_layer", None), + per_layer_model_projection=getattr( + self, "per_layer_model_projection", None + ), + per_layer_projection_norm=getattr( + self, "per_layer_projection_norm", None + ), + per_layer_input_scale=getattr(self, "per_layer_input_scale", None), + per_layer_projection_scale=getattr( + self, "per_layer_projection_scale", None + ), + ) + # Layers K..(N-1) are cross-decoder layers in YOCO + with set_model_tag("cross_decoder"): + self.cross_decoder = Gemma4CrossDecoderLayers( + vllm_config=vllm_config, + prefix=f"{prefix}.cross_decoder", + decoder_layers=self.layers[first_kv_shared_layer_idx:], + layer_idx_start=first_kv_shared_layer_idx, + ) + + self.fast_prefill_enabled = cache_config.kv_sharing_fast_prefill + + if self.fast_prefill_enabled: + # Allocate static buffers for CUDAGraph + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + device = next(self.parameters()).device + self.positions = torch.zeros( + max_num_tokens, dtype=torch.int64, device=device + ) + self.hidden_states = torch.zeros( + (max_num_tokens, config.hidden_size), + dtype=self.embed_tokens.weight.dtype, + device=device, + ) + if ( + self.hidden_size_per_layer_input + and self.hidden_size_per_layer_input > 0 + ): + self.per_layer_inputs = torch.zeros( + ( + max_num_tokens, + config.num_hidden_layers, + self.hidden_size_per_layer_input, + ), + dtype=self.embed_tokens.weight.dtype, + device=device, + ) + else: + self.per_layer_inputs = None + # Custom factory that includes per_layer_inputs for PLE-enabled PP. # per_layer_inputs has shape (batch, num_layers, per_layer_dim), # which differs from the standard (batch, hidden_size) shape, @@ -776,47 +1046,22 @@ class Gemma4Model(nn.Module): self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) * self.normalizer + return self.self_decoder.embed_input_ids(input_ids) - def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor: + def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None: """Get per-layer embeddings from embed_tokens_per_layer. Returns: Per-layer embeddings (num_tokens, num_layers, hidden_size_per_layer_input) """ - if self.embed_tokens_per_layer is None: - return None - - # Handle out-of-vocab tokens for PLE (vocab_size_per_layer_input may - # be smaller than the main vocab_size). - per_layer_inputs_mask = torch.logical_and( - input_ids >= 0, - input_ids < self.vocab_size_per_layer_input, - ) - per_layer_inputs_tokens = torch.where( - per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids) - ) - - # Get packed per-layer embeddings: (num_tokens, total_ple_dim) - per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens) - - # Apply embed_scale (sqrt of per-layer hidden dim) - per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer - - # Reshape to (num_tokens, num_layers, hidden_size_per_layer_input) - per_layer_embeds = per_layer_embeds.reshape( - *input_ids.shape, - self.config.num_hidden_layers, - self.hidden_size_per_layer_input, - ) - return per_layer_embeds + return self.self_decoder.get_per_layer_inputs(input_ids) def project_per_layer_inputs( self, inputs_embeds: torch.Tensor, per_layer_inputs: torch.Tensor | None, - ) -> torch.Tensor: + ) -> torch.Tensor | None: """Project inputs_embeds and combine with per_layer_inputs. Steps: @@ -826,29 +1071,94 @@ class Gemma4Model(nn.Module): 4. Normalize with per_layer_projection_norm 5. Combine: (projection + per_layer_inputs) * 1/sqrt(2) """ - if self.per_layer_model_projection is None: - return None - - # Project from hidden_size to total_ple_dim - # Scaled projection: output = linear(input, weight) * scale - per_layer_projection = self.per_layer_model_projection(inputs_embeds) - per_layer_projection = per_layer_projection * self.per_layer_projection_scale - - # Reshape to (num_tokens, num_layers, hidden_size_per_layer_input) - per_layer_projection = per_layer_projection.reshape( - *inputs_embeds.shape[:-1], - self.config.num_hidden_layers, - self.hidden_size_per_layer_input, + return self.self_decoder.project_per_layer_inputs( + inputs_embeds, per_layer_inputs ) - # Normalize - per_layer_projection = self.per_layer_projection_norm(per_layer_projection) + def fast_prefill_forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + per_layer_inputs: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + logits_indices_padded, num_logits_indices = None, None + attn_metadata = get_forward_context().attn_metadata - if per_layer_inputs is None: - return per_layer_projection + if attn_metadata is not None: + assert isinstance(attn_metadata, dict) + layer_attn_metadata = attn_metadata[ + self.layers[-1].self_attn.attn.layer_name + ] + if isinstance(layer_attn_metadata, KVSharingFastPrefillMetadata): + logits_indices_padded = layer_attn_metadata.logits_indices_padded + num_logits_indices = layer_attn_metadata.num_logits_indices - # Combine: (projection + per_layer_inputs) * scale - return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale + batch_size = positions.size(0) + self.positions[:batch_size].copy_(positions) + self_decoder_hidden_states, per_layer_inputs = self.self_decoder( + input_ids=input_ids, + positions=self.positions[:batch_size], + inputs_embeds=inputs_embeds, + per_layer_inputs=per_layer_inputs, + **kwargs, + ) + + if logits_indices_padded is None: + logits_indices_padded = torch.arange( + batch_size, + dtype=positions.dtype, + device=positions.device, + ) + + # NOTE: Keep .clone() until fix in + # https://github.com/vllm-project/vllm/pull/22282 + hidden_states = self_decoder_hidden_states.clone() + + num_padded = logits_indices_padded.size(0) + self.positions[:num_padded].copy_(positions[logits_indices_padded]) + self.hidden_states[:num_padded].copy_( + self_decoder_hidden_states[logits_indices_padded] + ) + if self.per_layer_inputs is not None and per_layer_inputs is not None: + self.per_layer_inputs[:num_padded].copy_( + per_layer_inputs[logits_indices_padded] + ) + + # Update batch_descriptor so the cross-decoder's piecewise + # CUDAGraphWrapper dispatches to the correct (reduced) batch size. + forward_context = get_forward_context() + orig_batch_desc = forward_context.batch_descriptor + if orig_batch_desc is not None: + forward_context.batch_descriptor = replace( + orig_batch_desc, num_tokens=num_padded + ) + + cross_per_layer = ( + self.per_layer_inputs[:num_padded] + if self.per_layer_inputs is not None + else None + ) + cross_hidden_states = self.cross_decoder( + self.positions[:num_padded], + self.hidden_states[:num_padded], + cross_per_layer, + **kwargs, + ) + + # Restore the original batch_descriptor + forward_context.batch_descriptor = orig_batch_desc + + if num_logits_indices is not None: + assert num_logits_indices > 0 + hidden_states[logits_indices_padded[:num_logits_indices]] = ( + cross_hidden_states[:num_logits_indices] + ) + else: + hidden_states = cross_hidden_states + + return hidden_states def forward( self, @@ -859,6 +1169,18 @@ class Gemma4Model(nn.Module): per_layer_inputs: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor | IntermediateTensors: + if self.fast_prefill_enabled: + hidden_states = self.fast_prefill_forward( + input_ids, + positions, + inputs_embeds, + per_layer_inputs, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + return hidden_states + + # Normal (non-fast-prefill) path with PP support if get_pp_group().is_first_rank: if inputs_embeds is not None: hidden_states = inputs_embeds From f40d9879f2dfe4d878b77768ad30935ea4e42b1f Mon Sep 17 00:00:00 2001 From: Lukas Geiger Date: Mon, 6 Apr 2026 16:39:37 +0100 Subject: [PATCH 12/39] [Models][GDN] Remove GPU/CPU syncs in `GDNAttentionMetadata.build` during speculative decoding (#38047) Signed-off-by: Lukas Geiger --- vllm/v1/attention/backends/gdn_attn.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 5ebf040be7a..85715e91ab4 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -253,7 +253,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] ) # Filter by spec_sequence_masks to exclude padded sequences spec_state_indices_tensor = block_table_tensor[ - spec_sequence_masks, : self.num_spec + 1 + spec_sequence_masks_cpu, : self.num_spec + 1 ] non_spec_state_indices_tensor = None # Padded sequences are always at the back, so the first @@ -264,7 +264,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] non_spec_query_start_loc_cpu = None else: spec_token_masks = torch.repeat_interleave( - spec_sequence_masks, query_lens + spec_sequence_masks, + query_lens, + output_size=query_start_loc_cpu[-1].item(), ) index = torch.argsort(spec_token_masks, stable=True) num_non_spec_tokens = num_prefill_tokens + num_decode_tokens @@ -272,10 +274,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] spec_token_indx = index[num_non_spec_tokens:] spec_state_indices_tensor = block_table_tensor[ - spec_sequence_masks, : self.num_spec + 1 + spec_sequence_masks_cpu, : self.num_spec + 1 ] non_spec_state_indices_tensor = block_table_tensor[ - ~spec_sequence_masks, 0 + ~spec_sequence_masks_cpu, 0 ] spec_query_start_loc = torch.zeros( @@ -284,7 +286,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] device=query_start_loc.device, ) torch.cumsum( - query_lens[spec_sequence_masks], dim=0, out=spec_query_start_loc[1:] + query_lens[spec_sequence_masks_cpu], + dim=0, + out=spec_query_start_loc[1:], ) non_spec_query_start_loc = torch.zeros( query_lens.size(0) - num_spec_decodes + 1, @@ -292,7 +296,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] device=query_start_loc.device, ) torch.cumsum( - query_lens[~spec_sequence_masks], + query_lens[~spec_sequence_masks_cpu], dim=0, out=non_spec_query_start_loc[1:], ) @@ -307,7 +311,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] ) assert num_accepted_tokens is not None - num_accepted_tokens = num_accepted_tokens[spec_sequence_masks] + num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] chunk_indices: torch.Tensor | None = None chunk_offsets: torch.Tensor | None = None @@ -331,8 +335,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata] if num_prefills > 0: has_initial_state = context_lens_tensor > 0 - if spec_sequence_masks is not None: - has_initial_state = has_initial_state[~spec_sequence_masks] + if spec_sequence_masks_cpu is not None: + has_initial_state = has_initial_state[~spec_sequence_masks_cpu] assert non_spec_query_start_loc_cpu is not None nums_dict, batch_ptr, token_chunk_offset_ptr = ( compute_causal_conv1d_metadata( From 4ae218c122f768ed9e2ff454b8567fd8c9373f6d Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 6 Apr 2026 11:52:05 -0400 Subject: [PATCH 13/39] [Refactor] Remove unused dead code (#38842) Signed-off-by: yewentao256 --- vllm/model_executor/models/mlp_speculator.py | 53 -------------------- vllm/v1/attention/ops/flashmla.py | 13 ----- vllm/v1/executor/ray_distributed_executor.py | 8 --- 3 files changed, 74 deletions(-) delete mode 100644 vllm/v1/executor/ray_distributed_executor.py diff --git a/vllm/model_executor/models/mlp_speculator.py b/vllm/model_executor/models/mlp_speculator.py index 48604d8e510..612baba8eaa 100644 --- a/vllm/model_executor/models/mlp_speculator.py +++ b/vllm/model_executor/models/mlp_speculator.py @@ -17,8 +17,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from .utils import maybe_prefix -SQRT2 = 2**0.5 - class MLPSpeculatorLayerNorm(nn.Module): """ @@ -171,57 +169,6 @@ class MLPSpeculator(nn.Module): config.vocab_size, config.vocab_size, 1.0 ) - # NOTE(woosuk): This method is commented out because it is old code - # using V0. We should either port it to V1 or remove it. - - # def generate_proposals( - # self, - # input_ids: torch.Tensor, - # previous_hidden_states: torch.Tensor, - # num_predict_tokens: int, - # sampling_metadata: SamplingMetadata, - # ) -> list[SamplerOutput]: - # if num_predict_tokens > self.max_speculative_tokens: - # raise ValueError(f"Max speculative tokens for model is " - # f"{self.max_speculative_tokens}, but " - # f"{num_predict_tokens} were requested") - - # # b x 1 x d - # previous_hidden_states = previous_hidden_states.unsqueeze(1) - - # if self.scale_input: - # previous_hidden_states = self.ln0(previous_hidden_states) / SQRT2 - - # # b x 1 - # last_tokens = input_ids.unsqueeze(1) - - # next_tokens = [] - - # for head_index in range(num_predict_tokens): - - # # Project and predict - # z = self.emb[head_index](last_tokens) # b k d - # states = self.proj[head_index](previous_hidden_states) - - # # Weighted add of state_weight*state and emb_weight*z - # # Let subsequent LN take care of denominator - # # state_weight is close to 1, so shouldn't be any precision issues - # states.add_(z, alpha=self.emb_weight / self.state_weight) - - # states = self.activation(self.ln[head_index](states)) # b k d - # previous_hidden_states = states - # # TODO: not yet supporting top_k_tokens_per_head - # states = states.flatten(0, 1) - - # logits = self.logits_processor(self.head[head_index], states, - # sampling_metadata) - - # output = self.sampler(logits, sampling_metadata) - # last_tokens = output.sampled_token_ids - # next_tokens.append(output) - - # return next_tokens - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() diff --git a/vllm/v1/attention/ops/flashmla.py b/vllm/v1/attention/ops/flashmla.py index aa667570a82..df04f5bf228 100644 --- a/vllm/v1/attention/ops/flashmla.py +++ b/vllm/v1/attention/ops/flashmla.py @@ -151,16 +151,3 @@ def flash_mla_with_kvcache_fp8( descale_k, ) return out, softmax_lse - - -# -# TODO: Add fake functions -# -# @register_fake("_flashmla_C::get_mla_metadata") -# def _get_mla_metadata_fake(....) -> Tuple[torch.Tensor, torch.Tensor]: -# return .... -# -# @register_fake("_flashmla_C::fwd_kvcache_mla") -# def _fwd_kvcache_mla_fake(....) -> Tuple[torch.Tensor, torch.Tensor]: -# return .... -# diff --git a/vllm/v1/executor/ray_distributed_executor.py b/vllm/v1/executor/ray_distributed_executor.py deleted file mode 100644 index 9a56c093ad6..00000000000 --- a/vllm/v1/executor/ray_distributed_executor.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.v1.executor.ray_executor import ( - RayDistributedExecutor as _RayDistributedExecutor, -) - -# For backwards compatibility. -RayDistributedExecutor = _RayDistributedExecutor From 608914de30380e3505810c4e01187da2f71e356f Mon Sep 17 00:00:00 2001 From: Frederik Gossen Date: Mon, 6 Apr 2026 12:37:13 -0400 Subject: [PATCH 14/39] [Core] Re-enable Inductor pre-grad passes in standalone compile (torch>=2.12) (#38944) Signed-off-by: Frederik Gossen --- vllm/compilation/compiler_interface.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index bddacfbbc29..5c34d3a1b38 100644 --- a/vllm/compilation/compiler_interface.py +++ b/vllm/compilation/compiler_interface.py @@ -345,9 +345,9 @@ class InductorStandaloneAdaptor(CompilerInterface): # Inductor's pre-grad passes don't do anything for vLLM. # The pre-grad passes get run even on cache-hit and negatively impact # vllm cold compile times by O(1s) - # Can remove this after the following issue gets fixed + # Fixed upstream in PyTorch 2.12: # https://github.com/pytorch/pytorch/issues/174502 - if envs.VLLM_ENABLE_PREGRAD_PASSES: + if is_torch_equal_or_newer("2.12.0.dev") or envs.VLLM_ENABLE_PREGRAD_PASSES: pregrad_ctx: Any = contextlib.nullcontext() else: pregrad_ctx = patch( From 93bada494f78b274867772ff337a3e3fb15976d6 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Mon, 6 Apr 2026 12:41:59 -0400 Subject: [PATCH 15/39] [MoE Refactor] Split of DefaultMoERunner class (#35326) Signed-off-by: Bill Nell Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- vllm/model_executor/layers/fused_moe/layer.py | 8 +- .../fused_moe/runner/chunking_moe_runner.py | 243 ++++++ .../fused_moe/runner/default_moe_runner.py | 704 +----------------- .../layers/fused_moe/runner/moe_runner.py | 17 + .../fused_moe/runner/moe_runner_base.py | 527 +++++++++++++ .../fused_moe/runner/moe_runner_factory.py | 51 ++ .../layers/fused_moe/runner/shared_experts.py | 22 +- vllm/model_executor/models/nemotron_h.py | 1 + 8 files changed, 868 insertions(+), 705 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py create mode 100644 vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 2fb61615b02..c4fc1fd2557 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -39,8 +39,8 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) -from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( - DefaultMoERunner, +from vllm.model_executor.layers.fused_moe.runner.moe_runner_factory import ( + create_moe_runner, ) from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, @@ -572,8 +572,8 @@ class FusedMoE(CustomOp): # Storing the runner in the FusedMoE is an intermediate state, eventually # the runner will own the FusedMoE layer and provide the execution interface # for MoE ops. - self.runner = DefaultMoERunner( - layer=self, + self.runner = create_moe_runner( + layer_name=self.layer_name, moe_config=self.moe_config, router=self.router, routed_input_transform=self._routed_input_transform, diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py new file mode 100644 index 00000000000..a8c75486d71 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.forward_context import ( + get_forward_context, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) +from vllm.utils.math_utils import cdiv +from vllm.v1.worker.ubatching import dbo_current_ubatch_id +from vllm.v1.worker.workspace import current_workspace_manager + + +class ChunkingMoERunner(MoERunnerBase): + """ + MoE runner wrapper that adds chunked processing to any MoERunnerBase. + + This runner wraps an inner MoERunnerBase and overrides _forward_impl to + process large batches by breaking them into smaller chunks. Each chunk + is delegated to the inner runner's _forward_impl, making chunking + composable with any runner implementation. + + All MoERunnerBase state (moe_config, router, quant_method, etc.) is + transparently delegated to the inner runner via __getattr__. + ChunkingMoERunner only owns chunking-specific state: the pre-allocated + workspace buffers and the reduce_results override. + + Key behaviors: + - Pre-allocates workspace tensors for CUDA graph compatibility + - Processes chunks via inner._forward_impl per chunk + - Never reduces results (reduce_results always returns False) + """ + + def __init__(self, inner: MoERunnerBase): + # Assert that _maybe_dispatch/_maybe_combine will be nops. + assert inner.moe_config.pcp_size == 1 + + # Skip MoERunnerBase.__init__ — all state is delegated to inner + # via __getattr__. Only chunking-specific state lives here. + self._inner = inner + + # Pre-allocated staging buffers. These need to exist ahead of time + # due to CUDA graph construction needing fixed buffer addresses. + self.batched_hidden_states, self.batched_router_logits = ( + self._init_dp_chunking() + ) + + def __getattr__(self, name): + # Delegate attribute access to the inner runner. This is only + # called when normal lookup (instance __dict__, class MRO) fails, + # so ChunkingMoERunner's own attributes and methods take priority. + return getattr(self._inner, name) + + @property + def shared_experts(self) -> SharedExperts | None: + return self._inner.shared_experts + + # TODO(bnell): temporary hack, do not call this method. + def _replace_quant_method(self, quant_method: FusedMoEMethodBase): + self._inner._replace_quant_method(quant_method) + self.quant_method = quant_method + + def is_internal_router(self) -> bool: + return self._inner.gate is not None + + # Reducing results when chunking is handled by the MK finalize operations + # when DP chunking is enabled.. + # This will be removed by #35949 + @property + def reduce_results(self) -> bool: + return False + + def _init_dp_chunking(self) -> list[torch.Tensor]: + states_shape: tuple[int, ...] + logits_shape: tuple[int, ...] + + moe = self.moe_config + + if self.enable_dbo: + states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim) + logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts) + else: + states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim) + logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts) + + # Does this need some kind of profiling run check like modular_kernel.py? + return current_workspace_manager().get_simultaneous( + (states_shape, moe.in_dtype), + (logits_shape, moe.router_logits_dtype), + ) + + def _allocate_dp_chunking_outputs( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + # Assert the inputs are of the proper type and shape. + assert self.batched_hidden_states is not None + assert self.batched_router_logits is not None + + assert self.batched_hidden_states.dtype == hidden_states.dtype, ( + f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}" + ) + assert self.batched_router_logits.dtype == router_logits.dtype, ( + f"{self.batched_router_logits.dtype} == {router_logits.dtype}" + ) + + # Check size compatibility. + assert self.batched_hidden_states.size(-1) == hidden_states.size(-1) + assert self.batched_router_logits.size(-1) == router_logits.size(-1) + + final_fused_hidden_states = torch.empty_like(hidden_states) + if self.shared_experts is not None: + if shared_experts_input is not None: + final_shared_hidden_states = torch.empty_like(shared_experts_input) + else: + final_shared_hidden_states = torch.empty_like(hidden_states) + else: + final_shared_hidden_states = None + + return final_shared_hidden_states, final_fused_hidden_states + + def _slice_and_copy_input( + self, + out_slice: torch.Tensor, + orig: torch.Tensor | None, + start: int, + end: int, + ) -> torch.Tensor: + assert orig is not None + slice_size = end - start + orig_slice = orig[start:end, :] + if self.enable_dbo: + assert out_slice.dim() == 3 + batch_buffer_idx = dbo_current_ubatch_id() + out_slice = out_slice[batch_buffer_idx, :] + + assert out_slice.size(0) >= slice_size + out_slice = out_slice[:slice_size, :] + out_slice.copy_(orig_slice, non_blocking=True) + return out_slice + + def _forward_impl( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + final_shared_hidden_states, final_fused_hidden_states = ( + self._allocate_dp_chunking_outputs( + hidden_states, router_logits, shared_experts_input + ) + ) + + ctx = get_forward_context() + # flashinfer_cutlass_kernels can handle: optional DP + TP/EP + max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu + moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens + + # If the input to the MoE is sequence parallel then divide by sp_size + # to find the maximum number of tokens for any individual dispatcher. + if self.moe_config.is_sequence_parallel: + max_tokens_across_dispatchers = cdiv( + max_tokens_across_dispatchers, self.moe_config.sp_size + ) + + num_tokens = hidden_states.size(0) + for chunk_idx, chunk_start_ in enumerate( + range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank) + ): + chunk_start = chunk_start_ + chunk_end = min( + chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers + ) + # clamp start and end + chunk_start = min(chunk_start, num_tokens - 1) + chunk_end = min(chunk_end, num_tokens) + chunk_sizes = ctx.dp_metadata.chunked_sizes( + self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx + ) + with chunk_sizes: + hidden_states_chunk = self._slice_and_copy_input( + self.batched_hidden_states, + hidden_states, + chunk_start, + chunk_end, + ) + + router_logits_chunk = self._slice_and_copy_input( + self.batched_router_logits, + router_logits, + chunk_start, + chunk_end, + ) + + shared_experts_input_chunk = ( + shared_experts_input[chunk_start:chunk_end, :] + if shared_experts_input is not None + else None + ) + + # Delegate per-chunk computation to the inner runner. + chunk_result = self._inner._forward_impl( + layer=layer, + hidden_states=hidden_states_chunk, + router_logits=router_logits_chunk, + shared_experts_input=shared_experts_input_chunk, + ) + + # Store outputs + # TODO(bnell): document when chunk_start >= num_tokens + if chunk_start < num_tokens: + if self.shared_experts is not None: + assert isinstance(chunk_result, tuple) + shared_output_chunk, hidden_states_chunk = chunk_result + final_fused_hidden_states[chunk_start:chunk_end, :].copy_( + hidden_states_chunk, non_blocking=True + ) + assert shared_output_chunk is not None + assert final_shared_hidden_states is not None + final_shared_hidden_states[chunk_start:chunk_end, :].copy_( + shared_output_chunk, non_blocking=True + ) + else: + assert isinstance(chunk_result, torch.Tensor) + final_fused_hidden_states[chunk_start:chunk_end, :].copy_( + chunk_result, non_blocking=True + ) + + if self.shared_experts is None: + return final_fused_hidden_states + else: + assert final_shared_hidden_states is not None + return (final_shared_hidden_states, final_fused_hidden_states) diff --git a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py index 4f9409e2cec..85c6563c084 100644 --- a/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py @@ -1,516 +1,45 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable -from contextlib import nullcontext -from typing import TYPE_CHECKING import torch -import torch.nn.functional as F from vllm.distributed import ( get_ep_group, get_pcp_group, - tensor_model_parallel_all_reduce, ) -from vllm.forward_context import ( - ForwardContext, - get_forward_context, - is_forward_context_available, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, -) -from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( - FusedMoEMethodBase, -) -from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( - FusedMoERouter, -) -from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner -from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( - SharedExperts, - SharedExpertsOrder, -) -from vllm.platforms import current_platform -from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import ( - HAS_OPAQUE_TYPE, - ModuleName, - direct_register_custom_op, -) -from vllm.v1.worker.ubatching import dbo_current_ubatch_id - -logger = init_logger(__name__) +from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase -def get_layer_from_name(layer_name: str) -> torch.nn.Module: - forward_context: ForwardContext = get_forward_context() - if layer_name == "from_forward_context": - all_moe_layers = forward_context.all_moe_layers - assert all_moe_layers is not None - moe_layer_index = forward_context.moe_layer_index - if moe_layer_index >= len(all_moe_layers): - raise AssertionError( - "We expected the number of MOE layers in `all_moe_layers` " - "to be equal to the number of " - "{vllm.moe_forward, vllm.moe_forward_shared} calls." - ) - layer_name = all_moe_layers[moe_layer_index] - forward_context.moe_layer_index += 1 - return forward_context.no_compile_layers[layer_name] - - -# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object; -# on older versions it remains a plain str. -if TYPE_CHECKING: - from typing import TypeAlias - - _layer_name_type: TypeAlias = str | ModuleName -else: - _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str - - -def _resolve_layer_name(layer_name: str | ModuleName) -> str: - return layer_name.value if isinstance(layer_name, ModuleName) else layer_name - - -# Note: _moe_forward and _moe_forward_shared should not contain any -# implementation details, They should merely pass along control to -# the runner's 'forward_dispatch' method. -def _moe_forward( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> torch.Tensor: - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_fake( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> torch.Tensor: - return torch.empty_like(hidden_states) - - -def _moe_forward_shared( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> tuple[torch.Tensor, torch.Tensor]: - layer = get_layer_from_name(_resolve_layer_name(layer_name)) - return layer.runner.forward_dispatch( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -def _moe_forward_shared_fake( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - layer_name: _layer_name_type, -) -> tuple[torch.Tensor, torch.Tensor]: - # Output shapes: - # - fused_out: same as hidden_states (routed experts use transformed size) - # - shared_out: same as shared_experts_input if provided, else same as - # hidden_states - # (For latent MoE: shared experts use original hidden_size, not latent size) - fused_out = torch.empty_like(hidden_states) - if shared_experts_input is not None: - shared_out = torch.empty_like(shared_experts_input) - else: - shared_out = torch.empty_like(hidden_states) - return shared_out, fused_out - - -direct_register_custom_op( - op_name="moe_forward", - op_func=_moe_forward, - mutates_args=["hidden_states"], # is this still true? - fake_impl=_moe_forward_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -direct_register_custom_op( - op_name="moe_forward_shared", - op_func=_moe_forward_shared, - fake_impl=_moe_forward_shared_fake, - tags=(torch.Tag.needs_fixed_stride_order,), -) - - -class DefaultMoERunner(MoERunner): +class DefaultMoERunner(MoERunnerBase): """ - Default implementation of the MoE runner for executing Mixture of Experts layers. + Standard MoE runner implementation for executing Mixture of Experts layers. - This class provides a comprehensive implementation for running MoE computations - with support for: - - Expert routing and token dispatching + This is the primary concrete implementation of MoE execution logic, providing + comprehensive support for standard MoE operations. It handles: + - Expert routing and token dispatching using various routing strategies - Shared experts computation with optional parallel execution using CUDA streams - - Data parallel (DP) chunking for large batch processing - Tensor model parallel and expert parallel operations - - Various quantization methods and custom operators + - Multiple quantization methods and optimized kernel selection - Both monolithic and decomposed expert execution paths + - Integration with various parallel execution modes (TP, EP, DP) - The runner handles the complete MoE forward pass including routing tokens to - experts, executing expert computations, and combining results. It supports - advanced features like overlapped execution of shared experts and optimized - kernels for different parallel execution modes. + The runner orchestrates the complete MoE forward pass including routing tokens + to experts, executing expert computations in parallel, and combining results. + It supports advanced features like overlapped execution of shared experts, + optimized kernels for different parallel configurations, and seamless + integration with vLLM's distributed execution framework. - Eventually, this class will be split up and specialized for different - configurations, e.g. the presence or absence of shared experts, a gate, etc. + This implementation is suitable for most standard MoE use cases. For specialized + scenarios like large batch chunking, alternative runners like ChunkingMoERunner + may be more appropriate. + + Eventually, this class may be split into more specialized implementations + for different configurations (e.g., with/without shared experts, gates, etc.). """ - def __init__( - self, - layer: torch.nn.Module, - moe_config: FusedMoEConfig, - router: FusedMoERouter, - routed_input_transform: torch.nn.Module | None, - gate: torch.nn.Module | None, - shared_experts: torch.nn.Module | None, - quant_method: FusedMoEMethodBase, - reduce_results: bool, - enable_dbo: bool, - ): - super().__init__() - self.moe_config = moe_config - self.router = router - self.routed_input_transform = routed_input_transform - self.gate = gate - self.quant_method = quant_method - self.reduce_results = reduce_results - self.enable_dbo = enable_dbo - - self.shared_experts: SharedExperts | None = None - if shared_experts is not None: - self.shared_experts = SharedExperts( - shared_experts, - moe_config=moe_config, - # Note: For now we must pass quant_method along to SharedExperts so it - # can property determine where the shared experts are supposed to be - # called, i.e. by a MK or by the MoERunner. - # Once the MK can be created upfront, we can just pass in the proper - # flags derived from the quant_method's MK. - reduce_results=reduce_results, - quant_method=quant_method, - enable_dbo=enable_dbo, - ) - - # Chunked all2all staging tensor - # These need to exist ahead of time due to CUDAgraph construction - # needing a fixed buffer address. - self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking - self.batched_hidden_states: torch.Tensor | None = None - self.batched_router_logits: torch.Tensor | None = None - self._maybe_init_dp_chunking() - - # Needed for string -> FusedMoE layer lookup in custom ops. - self.layer_name = layer.layer_name - - self.forward_entry, self.forward_impl = self._select_forward(layer) - - def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]: - # Select implementation based on presence of DP chunking. - forward_impl_fn = ( - self._forward_impl_chunked if self.use_dp_chunking else self._forward_impl - ) - - if current_platform.is_tpu() or current_platform.is_cpu(): - # TODO: Once the OOM issue for the TPU backend is resolved, we - # will switch to using the moe_forward custom op. - # Note: CPU doesn't require wrapped forward_impl. - return ( - _moe_forward if self.shared_experts is None else _moe_forward_shared, - forward_impl_fn, - ) - - return ( - torch.ops.vllm.moe_forward - if self.shared_experts is None - else torch.ops.vllm.moe_forward_shared, - forward_impl_fn, - ) - - # TODO(bnell): temporary hack, do not call this method. - def _replace_quant_method(self, quant_method: FusedMoEMethodBase): - if self.shared_experts is not None: - self.shared_experts._quant_method = quant_method - self.quant_method = quant_method - - def is_internal_router(self) -> bool: - return self.gate is not None - - def _maybe_init_dp_chunking(self): - if not self.use_dp_chunking: - return - - assert self.batched_hidden_states is None - states_shape: tuple[int, ...] - logits_shape: tuple[int, ...] - - moe = self.moe_config - - if self.enable_dbo: - states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim) - logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts) - else: - states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim) - logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts) - - device = torch.accelerator.current_device_index() - self.batched_hidden_states = torch.zeros( - states_shape, - dtype=moe.in_dtype, - device=device, - ) - - self.batched_router_logits = torch.zeros( - logits_shape, - dtype=moe.router_logits_dtype, - device=device, - ) - - def must_reduce_shared_expert_outputs(self) -> bool: - """ - The shared_experts are typically computed using the RowParallelLinear - layer. The result of this function is typically used as - the reduce_results argument to the module. - When just tensor-parallel is used, it is not required to reduce - the shared_experts results immediately. Instead we reduce at the - once at the end of the MoE op. (Refer to DeepSeekV2MoE module) - With EP and all2all kernels - this is no longer viable as all - GPU ranks in DP, produce the complete set of hidden_states. - Therefore it is required that we reduce the shared_experts output - early. - """ - return ( - self.quant_method.moe_kernel is not None - and self.quant_method.moe_kernel.output_is_reduced() - ) - - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """ - Some combine kernels reduce across GPU ranks by default. - """ - if self.must_reduce_shared_expert_outputs(): - return final_hidden_states - else: - return tensor_model_parallel_all_reduce(final_hidden_states) - - def apply_routed_input_transform( - self, hidden_states: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Apply transform for routed experts (e.g., latent projection). - - This is called by FusedMoE.forward_native. The original hidden_states - is saved separately so shared experts get [S, hidden_size] while - routed experts get the transformed [S, moe_latent_size]. - - TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be - moved inside SharedFusedMoE to all-reduce on the smaller latent - dimension. - - Returns (possibly transformed) hidden states and the input for shared - experts (or None if there are no shared experts). - """ - if self.routed_input_transform is not None: - result = self.routed_input_transform(hidden_states) - # ReplicatedLinear returns (output, extra_bias) tuple. - # We only need the output tensor; extra_bias is not used here. - if isinstance(result, tuple): - return result[0], hidden_states - return result, hidden_states - - return ( - hidden_states, - hidden_states if self.shared_experts is not None else None, - ) - - def _maybe_reduce_output( - self, - states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - trunc_sizes: list[int], - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return x[..., :trunc_size] - - def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: - return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) - - if ( - not self.moe_config.is_sequence_parallel - and not self.use_dp_chunking - and self.reduce_results - and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) - ): - func = reduce_and_trunc - else: - func = trunc - - if isinstance(states, tuple): - return tuple( - [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] - ) - else: - assert len(trunc_sizes) == 1 - return func(states, trunc_sizes[0]) - - def _encode_layer_name(self) -> str | ModuleName: - if HAS_OPAQUE_TYPE: - return ModuleName(self.layer_name) - # Can be unavailable or None in unittests - if ( - is_forward_context_available() - and get_forward_context().all_moe_layers is not None - ): - return "from_forward_context" - return self.layer_name - - def _maybe_pad_hidden_states( - self, - shared_experts_input: torch.Tensor | None, - hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, list[int]]: - shared_experts_hidden_dim = ( - shared_experts_input.shape[-1] if shared_experts_input is not None else 0 - ) - transformed_hidden_dim = hidden_states.shape[-1] - if ( - not self.quant_method.skip_forward_padding - and self.moe_config.hidden_dim != transformed_hidden_dim - ): - hidden_states = F.pad( - hidden_states, - (0, self.moe_config.hidden_dim - transformed_hidden_dim), - mode="constant", - value=0.0, - ) - - if self.shared_experts is not None: - orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] - else: - orig_hidden_dims = [transformed_hidden_dim] - - return hidden_states, orig_hidden_dims - - def _maybe_apply_shared_experts( - self, - shared_experts_input: torch.Tensor | None, - order: SharedExpertsOrder, - ): - if self.shared_experts is not None: - assert shared_experts_input is not None - self.shared_experts.apply(shared_experts_input, order) - - def _apply_quant_method( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - # Run this before quant_method to avoid inplace issues. - # TODO(bnell): probably not needed anymore since inplace is - # disabled when shared experts are present. - self._maybe_apply_shared_experts( - shared_experts_input, SharedExpertsOrder.NO_OVERLAP - ) - - if self.quant_method.is_monolithic: - fused_out = self.quant_method.apply_monolithic( - layer=layer, - x=hidden_states, - router_logits=router_logits, - ) - else: - topk_weights, topk_ids = self.router.select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - ) - - # Passing shared_experts_input in case SharedExpertsOrder is - # NO_OVERLAP or MK_INTERNAL_OVERLAPPED. - fused_out = self.quant_method.apply( - layer=layer, - x=hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_experts_input=shared_experts_input, - ) - - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.MULTI_STREAM_OVERLAPPED, - ) - - return ( - self.shared_experts.output if self.shared_experts is not None else None, - fused_out, - ) - - def _sequence_parallel_context(self): - ctx = get_forward_context() - return ( - ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) - if ctx.dp_metadata - else nullcontext() - ) - - def _allocate_dp_chunking_outputs( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - assert self.use_dp_chunking - - # Assert the inputs are of the proper type and shape. - assert self.batched_hidden_states is not None - assert self.batched_router_logits is not None - - assert self.batched_hidden_states.dtype == hidden_states.dtype, ( - f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}" - ) - assert self.batched_router_logits.dtype == router_logits.dtype, ( - f"{self.batched_router_logits.dtype} == {router_logits.dtype}" - ) - - # Check size compatibility. - assert self.batched_hidden_states.size(-1) == hidden_states.size(-1) - assert self.batched_router_logits.size(-1) == router_logits.size(-1) - - final_fused_hidden_states = torch.empty_like(hidden_states) - if self.shared_experts is not None: - final_shared_hidden_states = torch.empty_like(hidden_states) - else: - final_shared_hidden_states = None - - return final_shared_hidden_states, final_fused_hidden_states - - def _maybe_sync_shared_experts_stream( - self, - shared_experts_input: torch.Tensor | None, - ): - # If router/gate provided, then apply it here. - # (Note: This code runs only when "overlapped mode" is on to allow - # parallel execution of shared experts with the FusedMoE via - # separate cuda stream) - if self.shared_experts is not None: - self.shared_experts.maybe_sync_shared_experts_stream(shared_experts_input) + @property + def reduce_results(self) -> bool: + return self._reduce_results @property def do_naive_dispatch_combine(self) -> bool: @@ -572,195 +101,6 @@ class DefaultMoERunner(MoERunner): else: return hidden_states - def forward( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """Invoke the fused moe layer. - - Input: - - hidden_states - - router_logits - - Output: - - The new hidden_states. - or - - A tuple of (shared experts output, new hidden_states). - - Calling sequence - - forward - - self.forward_entry (_moe_forward or _moe_forward_shared custom op) - - forward_dispatch - - forward_impl (_forward_impl or _forward_impl_chunked) - - Note: The existence of _moe_forward and _moe_forward_shared custom ops are due - to the following reasons: - 1. the chunking loop in _forward_impl_chunked cannot be compiled by - torch.compile - 2. pytorch cannot handle union types in custom op signatures so _moe_forward - and _moe_forward_shared must be split. - - If _forward_impl_chunked can be implemented via torch.scan we can potentially - get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence - into the 'forward' method. - """ - - # Apply transform for routed experts (e.g., latent projection for latent MoE) - hidden_states, shared_experts_input = self.apply_routed_input_transform( - hidden_states - ) - - hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( - shared_experts_input, - hidden_states, - ) - - fused_output = self.forward_entry( - hidden_states, - router_logits, - shared_experts_input, - self._encode_layer_name(), - ) - - return self._maybe_reduce_output(fused_output, og_hidden_dims) - - def forward_dispatch( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - # TODO(bnell): this can be removed after MK migration is complete. - layer.ensure_moe_quant_config_init() - - # Sync aux and main stream for shared expert multi-stream overlap. - self._maybe_sync_shared_experts_stream(shared_experts_input) - - # If the Runner holds the gate, apply it after the stream sync, - # so it can run overlapped with the - # NOTE: in future PR, MoE runner will always hold the gate. - if self.gate is not None: - router_logits, _ = self.gate(hidden_states) - - self._maybe_apply_shared_experts( - shared_experts_input, - SharedExpertsOrder.EXTERNAL, - ) - - with self._sequence_parallel_context(): - return self.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - def _slice_and_copy_input( - self, - out_slice: torch.Tensor, - orig: torch.Tensor | None, - start: int, - end: int, - ) -> torch.Tensor: - assert orig is not None - slice_size = end - start - orig_slice = orig[start:end, :] - if self.enable_dbo: - assert out_slice.dim() == 3 - batch_buffer_idx = dbo_current_ubatch_id() - out_slice = out_slice[batch_buffer_idx, :] - - assert out_slice.size(0) >= slice_size - out_slice = out_slice[:slice_size, :] - out_slice.copy_(orig_slice, non_blocking=True) - return out_slice - - def _forward_impl_chunked( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - final_shared_hidden_states, final_fused_hidden_states = ( - self._allocate_dp_chunking_outputs(hidden_states, router_logits) - ) - - ctx = get_forward_context() - # flashinfer_cutlass_kernels can handle: optional DP + TP/EP - max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu - moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens - - # If the input to the MoE is sequence parallel then divide by sp_size - # to find the maximum number of tokens for any individual dispatcher. - if self.moe_config.is_sequence_parallel: - max_tokens_across_dispatchers = cdiv( - max_tokens_across_dispatchers, self.moe_config.sp_size - ) - - num_tokens = hidden_states.size(0) - for chunk_idx, chunk_start_ in enumerate( - range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank) - ): - chunk_start = chunk_start_ - chunk_end = min( - chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers - ) - # clamp start and end - chunk_start = min(chunk_start, num_tokens - 1) - chunk_end = min(chunk_end, num_tokens) - chunk_sizes = ctx.dp_metadata.chunked_sizes( - self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx - ) - with chunk_sizes: - hidden_states_chunk = self._slice_and_copy_input( - self.batched_hidden_states, - hidden_states, - chunk_start, - chunk_end, - ) - - router_logits_chunk = self._slice_and_copy_input( - self.batched_router_logits, - router_logits, - chunk_start, - chunk_end, - ) - - shared_experts_input_chunk = ( - shared_experts_input[chunk_start:chunk_end, :] - if shared_experts_input is not None - else None - ) - - shared_output_chunk, hidden_states_chunk = self._apply_quant_method( - layer=layer, - hidden_states=hidden_states_chunk, - router_logits=router_logits_chunk, - shared_experts_input=shared_experts_input_chunk, - ) - - # Store outputs - # TODO(bnell): document when chunk_start >= num_tokens - if chunk_start < num_tokens: - final_fused_hidden_states[chunk_start:chunk_end, :].copy_( - hidden_states_chunk, non_blocking=True - ) - if self.shared_experts is not None: - assert shared_output_chunk is not None - assert final_shared_hidden_states is not None - final_shared_hidden_states[chunk_start:chunk_end, :].copy_( - shared_output_chunk, non_blocking=True - ) - - if self.shared_experts is None: - return final_fused_hidden_states - else: - assert final_shared_hidden_states is not None - return (final_shared_hidden_states, final_fused_hidden_states) - def _forward_impl( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 720e997cda3..9ffbf3108f8 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -4,6 +4,13 @@ from abc import ABC, abstractmethod import torch +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) + class MoERunner(ABC): """ @@ -36,3 +43,13 @@ class MoERunner(ABC): @abstractmethod def is_internal_router(self) -> bool: raise NotImplementedError + + @property + @abstractmethod + def shared_experts(self) -> SharedExperts | None: + raise NotImplementedError + + # TODO(bnell): temporary hack, do not call this method. + @abstractmethod + def _replace_quant_method(self, quant_method: FusedMoEMethodBase): + raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py new file mode 100644 index 00000000000..d8788d47d18 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from abc import abstractmethod +from collections.abc import Callable +from contextlib import nullcontext +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F + +from vllm.distributed import ( + tensor_model_parallel_all_reduce, +) +from vllm.forward_context import ( + ForwardContext, + get_forward_context, + is_forward_context_available, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, + SharedExpertsOrder, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + HAS_OPAQUE_TYPE, + ModuleName, + direct_register_custom_op, +) + + +def get_layer_from_name(layer_name: str) -> torch.nn.Module: + forward_context: ForwardContext = get_forward_context() + if layer_name == "from_forward_context": + all_moe_layers = forward_context.all_moe_layers + assert all_moe_layers is not None + moe_layer_index = forward_context.moe_layer_index + if moe_layer_index >= len(all_moe_layers): + raise AssertionError( + "We expected the number of MOE layers in `all_moe_layers` " + "to be equal to the number of " + "{vllm.moe_forward, vllm.moe_forward_shared} calls." + ) + layer_name = all_moe_layers[moe_layer_index] + forward_context.moe_layer_index += 1 + return forward_context.no_compile_layers[layer_name] + + +# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object; +# on older versions it remains a plain str. +if TYPE_CHECKING: + from typing import TypeAlias + + _layer_name_type: TypeAlias = str | ModuleName +else: + _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str + + +def _resolve_layer_name(layer_name: str | ModuleName) -> str: + return layer_name.value if isinstance(layer_name, ModuleName) else layer_name + + +# Note: _moe_forward and _moe_forward_shared should not contain any +# implementation details, They should merely pass along control to +# the runner's 'forward_dispatch' method. +def _moe_forward( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: _layer_name_type, +) -> torch.Tensor: + layer = get_layer_from_name(_resolve_layer_name(layer_name)) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + +def _moe_forward_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: _layer_name_type, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +def _moe_forward_shared( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: _layer_name_type, +) -> tuple[torch.Tensor, torch.Tensor]: + layer = get_layer_from_name(_resolve_layer_name(layer_name)) + return layer.runner.forward_dispatch( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + +def _moe_forward_shared_fake( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + layer_name: _layer_name_type, +) -> tuple[torch.Tensor, torch.Tensor]: + # Output shapes: + # - fused_out: same as hidden_states (routed experts use transformed size) + # - shared_out: same as shared_experts_input if provided, else same as + # hidden_states + # (For latent MoE: shared experts use original hidden_size, not latent size) + fused_out = torch.empty_like(hidden_states) + if shared_experts_input is not None: + shared_out = torch.empty_like(shared_experts_input) + else: + shared_out = torch.empty_like(hidden_states) + return shared_out, fused_out + + +direct_register_custom_op( + op_name="moe_forward", + op_func=_moe_forward, + mutates_args=["hidden_states"], # is this still true? + fake_impl=_moe_forward_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +direct_register_custom_op( + op_name="moe_forward_shared", + op_func=_moe_forward_shared, + fake_impl=_moe_forward_shared_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +class MoERunnerBase(MoERunner): + """ + Abstract base class providing common functionality for MoE runner implementations. + + This class serves as the foundation for concrete MoE runner implementations by + providing shared state management and common utilities. It handles: + - Common initialization and configuration management + - Shared expert output reduction logic for tensor parallel scenarios + - Base methods for tensor model parallel reductions + - Common properties and utility functions used across different runner types + + Concrete subclasses must implement the abstract methods to define their specific + execution strategies, such as standard execution, chunked processing, or other + specialized approaches. The base class provides the infrastructure while + allowing flexibility in the actual MoE computation implementation. + + Key abstract methods that subclasses must implement: + - reduce_results: Determines whether results should be reduced across ranks + - _forward_impl: The core MoE computation logic specific to each runner type + """ + + def __init__( + self, + layer_name: str, + moe_config: FusedMoEConfig, + router: FusedMoERouter, + routed_input_transform: torch.nn.Module | None, + gate: torch.nn.Module | None, + shared_experts: torch.nn.Module | None, + quant_method: FusedMoEMethodBase, + reduce_results: bool, + enable_dbo: bool, + ): + super().__init__() + self.moe_config = moe_config + self.router = router + self.routed_input_transform = routed_input_transform + self.gate = gate + self.quant_method = quant_method + self._reduce_results = reduce_results + self.enable_dbo = enable_dbo + + self._shared_experts: SharedExperts | None = None + if shared_experts is not None: + self._shared_experts = SharedExperts( + shared_experts, + moe_config=moe_config, + # Note: For now we must pass quant_method along to SharedExperts so it + # can property determine where the shared experts are supposed to be + # called, i.e. by a MK or by the MoERunner. + # Once the MK can be created upfront, we can just pass in the proper + # flags derived from the quant_method's MK. + reduce_results=reduce_results, + quant_method=quant_method, + enable_dbo=enable_dbo, + ) + + # Needed for string -> FusedMoE layer lookup in custom ops. + self.layer_name = layer_name + + self.forward_entry = self._select_forward() + + def _select_forward(self) -> Callable: + if current_platform.is_tpu() or current_platform.is_cpu(): + # TODO: Once the OOM issue for the TPU backend is resolved, we + # will switch to using the moe_forward custom op. + # Note: CPU doesn't require wrapped _forward_impl. + return _moe_forward if self._shared_experts is None else _moe_forward_shared + + return ( + torch.ops.vllm.moe_forward + if self._shared_experts is None + else torch.ops.vllm.moe_forward_shared + ) + + @property + def shared_experts(self) -> SharedExperts | None: + return self._shared_experts + + # TODO(bnell): temporary hack, do not call this method. + def _replace_quant_method(self, quant_method: FusedMoEMethodBase): + if self._shared_experts is not None: + self._shared_experts._quant_method = quant_method + self.quant_method = quant_method + + def is_internal_router(self) -> bool: + return self.gate is not None + + @property + @abstractmethod + def reduce_results(self) -> bool: + raise NotImplementedError + + def must_reduce_shared_expert_outputs(self) -> bool: + """ + The shared_experts are typically computed using the RowParallelLinear + layer. The result of this function is typically used as + the reduce_results argument to the module. + When just tensor-parallel is used, it is not required to reduce + the shared_experts results immediately. Instead we reduce at the + once at the end of the MoE op. (Refer to DeepSeekV2MoE module) + With EP and all2all kernels - this is no longer viable as all + GPU ranks in DP, produce the complete set of hidden_states. + Therefore it is required that we reduce the shared_experts output + early. + """ + return ( + self.quant_method.moe_kernel is not None + and self.quant_method.moe_kernel.output_is_reduced() + ) + + def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): + """ + Some combine kernels reduce across GPU ranks by default. + """ + if self.must_reduce_shared_expert_outputs(): + return final_hidden_states + else: + return tensor_model_parallel_all_reduce(final_hidden_states) + + def apply_routed_input_transform( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Apply transform for routed experts (e.g., latent projection). + + This is called by FusedMoE.forward_native. The original hidden_states + is saved separately so shared experts get [S, hidden_size] while + routed experts get the transformed [S, moe_latent_size]. + + TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be + moved inside SharedFusedMoE to all-reduce on the smaller latent + dimension. + + Returns (possibly transformed) hidden states and the input for shared + experts (or None if there are no shared experts). + """ + if self.routed_input_transform is not None: + result = self.routed_input_transform(hidden_states) + # ReplicatedLinear returns (output, extra_bias) tuple. + # We only need the output tensor; extra_bias is not used here. + if isinstance(result, tuple): + return result[0], hidden_states + return result, hidden_states + + return ( + hidden_states, + hidden_states if self._shared_experts is not None else None, + ) + + def _maybe_reduce_output( + self, + states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + trunc_sizes: list[int], + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: + return x[..., :trunc_size] + + def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor: + return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size) + + if ( + not self.moe_config.is_sequence_parallel + and self.reduce_results + and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) + ): + func = reduce_and_trunc + else: + func = trunc + + if isinstance(states, tuple): + return tuple( + [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)] + ) + else: + assert len(trunc_sizes) == 1 + return func(states, trunc_sizes[0]) + + def _encode_layer_name(self) -> str | ModuleName: + if HAS_OPAQUE_TYPE: + return ModuleName(self.layer_name) + # Can be unavailable or None in unittests + if ( + is_forward_context_available() + and get_forward_context().all_moe_layers is not None + ): + return "from_forward_context" + return self.layer_name + + def _maybe_pad_hidden_states( + self, + shared_experts_input: torch.Tensor | None, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, list[int]]: + shared_experts_hidden_dim = ( + shared_experts_input.shape[-1] if shared_experts_input is not None else 0 + ) + transformed_hidden_dim = hidden_states.shape[-1] + if ( + not self.quant_method.skip_forward_padding + and self.moe_config.hidden_dim != transformed_hidden_dim + ): + hidden_states = F.pad( + hidden_states, + (0, self.moe_config.hidden_dim - transformed_hidden_dim), + mode="constant", + value=0.0, + ) + + if self._shared_experts is not None: + orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim] + else: + orig_hidden_dims = [transformed_hidden_dim] + + return hidden_states, orig_hidden_dims + + def _maybe_apply_shared_experts( + self, + shared_experts_input: torch.Tensor | None, + order: SharedExpertsOrder, + ): + if self._shared_experts is not None: + assert shared_experts_input is not None + self._shared_experts.apply(shared_experts_input, order) + + def _apply_quant_method( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + # Run this before quant_method to avoid inplace issues. + # TODO(bnell): probably not needed anymore since inplace is + # disabled when shared experts are present. + self._maybe_apply_shared_experts( + shared_experts_input, SharedExpertsOrder.NO_OVERLAP + ) + + if self.quant_method.is_monolithic: + fused_out = self.quant_method.apply_monolithic( + layer=layer, + x=hidden_states, + router_logits=router_logits, + ) + else: + topk_weights, topk_ids = self.router.select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + ) + + # Passing shared_experts_input in case SharedExpertsOrder is + # NO_OVERLAP or MK_INTERNAL_OVERLAPPED. + fused_out = self.quant_method.apply( + layer=layer, + x=hidden_states, + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts_input=shared_experts_input, + ) + + self._maybe_apply_shared_experts( + shared_experts_input, + SharedExpertsOrder.MULTI_STREAM_OVERLAPPED, + ) + + return ( + self._shared_experts.output if self._shared_experts is not None else None, + fused_out, + ) + + def _sequence_parallel_context(self): + ctx = get_forward_context() + return ( + ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size) + if ctx.dp_metadata + else nullcontext() + ) + + def _maybe_sync_shared_experts_stream( + self, + shared_experts_input: torch.Tensor | None, + ): + # If router/gate provided, then apply it here. + # (Note: This code runs only when "overlapped mode" is on to allow + # parallel execution of shared experts with the FusedMoE via + # separate cuda stream) + if self._shared_experts is not None: + self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input) + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Invoke the fused moe layer. + + Input: + - hidden_states + - router_logits + + Output: + - The new hidden_states. + or + - A tuple of (shared experts output, new hidden_states). + + Calling sequence + - forward + - self.forward_entry (_moe_forward or _moe_forward_shared custom op) + - forward_dispatch + - _forward_impl + + Note: The existence of _moe_forward and _moe_forward_shared custom ops are due + to the following reasons: + 1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by + torch.compile + 2. pytorch cannot handle union types in custom op signatures so _moe_forward + and _moe_forward_shared must be split. + + If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can + potentially get rid of _moe_forward and _moe_forward_shared and collapse the + whole sequence into the 'forward' method. + """ + + # Apply transform for routed experts (e.g., latent projection for latent MoE) + hidden_states, shared_experts_input = self.apply_routed_input_transform( + hidden_states + ) + + hidden_states, og_hidden_dims = self._maybe_pad_hidden_states( + shared_experts_input, + hidden_states, + ) + + fused_output = self.forward_entry( + hidden_states, + router_logits, + shared_experts_input, + self._encode_layer_name(), + ) + + return self._maybe_reduce_output(fused_output, og_hidden_dims) + + def forward_dispatch( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # TODO(bnell): this can be removed after MK migration is complete. + layer.ensure_moe_quant_config_init() + + # Sync aux and main stream for shared expert multi-stream overlap. + self._maybe_sync_shared_experts_stream(shared_experts_input) + + # If the Runner holds the gate, apply it after the stream sync, + # so it can run overlapped with the + # NOTE: in future PR, MoE runner will always hold the gate. + if self.gate is not None: + router_logits, _ = self.gate(hidden_states) + + with self._sequence_parallel_context(): + return self._forward_impl( + layer, + hidden_states, + router_logits, + shared_experts_input, + ) + + @abstractmethod + def _forward_impl( + self, + layer: torch.nn.Module, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py new file mode 100644 index 00000000000..da5068fa091 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) +from vllm.model_executor.layers.fused_moe.runner.chunking_moe_runner import ( + ChunkingMoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( + DefaultMoERunner, +) +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) + + +def create_moe_runner( + layer_name: str, + moe_config: FusedMoEConfig, + router: FusedMoERouter, + routed_input_transform: torch.nn.Module | None, + gate: torch.nn.Module | None, + shared_experts: SharedExperts | None, + quant_method: FusedMoEMethodBase, + reduce_results: bool, + enable_dbo: bool, +) -> MoERunner: + runner = DefaultMoERunner( + layer_name, + moe_config, + router, + routed_input_transform, + gate, + shared_experts, + quant_method, + reduce_results, + enable_dbo, + ) + if moe_config.moe_parallel_config.use_dp_chunking: + return ChunkingMoERunner(runner) + return runner diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index bb8645c8f61..f5b07a6a51a 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -32,19 +32,14 @@ class SharedExpertsOrder(IntEnum): # No shared experts. NONE = (0,) - # Get rid of this one? combine with BEFORE? - # Note: this might be important for torch.compile reasons. Can - # get rid of it after _moe_forward is undone. - EXTERNAL = (1,) - # No overlap - defensively called before MK. - NO_OVERLAP = (2,) + NO_OVERLAP = (1,) # Overlapped with dispatch/combine in DP/EP - called by the MK. - MK_INTERNAL_OVERLAPPED = (3,) + MK_INTERNAL_OVERLAPPED = (2,) # Overlapped with the gate, router, experts in aux stream. - MULTI_STREAM_OVERLAPPED = (4,) + MULTI_STREAM_OVERLAPPED = (3,) class SharedExperts: @@ -110,9 +105,6 @@ class SharedExperts: self, hidden_states: torch.Tensor, ) -> SharedExpertsOrder: - if self._use_external_experts: - return SharedExpertsOrder.EXTERNAL - if self._quant_method.mk_owns_shared_expert: return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED @@ -205,12 +197,4 @@ class SharedExperts: else: self._output[self._output_idx] = self._layer(shared_experts_input) - if order == SharedExpertsOrder.EXTERNAL: - # TODO: figure out how to combine this with maybe_reduce_output? - # or get rid of it completely. - assert self._output[self._output_idx] is not None - self._output[self._output_idx] = self._maybe_reduce_shared_out( - self._output[self._output_idx] - ) - assert self._output[self._output_idx] is not None diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 4ec794eccf7..8abbc808cce 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -231,6 +231,7 @@ class NemotronHMoE(nn.Module): num_redundant_experts=self.n_redundant_experts, is_sequence_parallel=self.is_sequence_parallel, routed_input_transform=self.fc1_latent_proj, + router_logits_dtype=self.gate.out_dtype, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: From bfdc0a3a996f82f592f43ff803adad173d3fbe11 Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Mon, 6 Apr 2026 13:07:02 -0400 Subject: [PATCH 16/39] [NIXL][Mamba][3/N] Heterogeneous TP: 3-read conv state transfer (#37635) --- .../config_sweep_accuracy_test.sh | 4 +- .../unit/test_nixl_connector_hma.py | 45 +- .../kv_transfer/kv_connector/utils.py | 381 ++++++++++++++- .../kv_connector/v1/nixl_connector.py | 451 +++++++++++++++--- .../v1/ssm_conv_transfer_utils.py | 164 +++++++ 5 files changed, 970 insertions(+), 75 deletions(-) create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index fe79a99fced..b0794bfa38a 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -19,9 +19,9 @@ dp_ep_configs=( "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1) ) hybrid_ssm_configs=( - "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" + "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models. - "ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling" + "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling" ) sw_attn_configs=( "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 898f8e4b35b..adb0acae1cb 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -224,6 +224,8 @@ def test_get_block_descs_ids_hybrid_ssm(): worker._has_mamba = True worker._is_mamba_group = [False, True] worker._physical_blocks_per_logical_kv_block = 1 + worker._mamba_phys_ratio = {engine_id: 1} + worker.block_len_per_layer = [100] # num_descs = num_regions * num_blocks (no blocks_first doubling) worker.num_descs = 2 * num_blocks @@ -234,9 +236,10 @@ def test_get_block_descs_ids_hybrid_ssm(): # FA group: stride=num_blocks=100, offset=0 # region0: [3, 5], region1: [103, 105] # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1), - # offset=num_descs=200 - # region0: [201, 202], region1: [301, 302] - expected = [3, 5, 103, 105, 201, 202, 301, 302] + # offset=num_fa_descs=200, 4 regions per Mamba layer (x, B, C, ssm) + # region0: [201, 202], region1: [301, 302], + # region2: [401, 402], region3: [501, 502] + expected = [3, 5, 103, 105, 201, 202, 301, 302, 401, 402, 501, 502] assert list(result) == expected, f"Expected {expected}, got {list(result)}" @@ -259,6 +262,8 @@ def test_get_block_descs_ids_kernel_block_mismatch(): worker._has_mamba = True worker._is_mamba_group = [False, True] worker._physical_blocks_per_logical_kv_block = ratio + worker._mamba_phys_ratio = {engine_id: ratio} + worker.block_len_per_layer = [100] worker.num_descs = 2 * num_blocks # 800 fa_blocks = [3, 7] # kernel-level block IDs @@ -267,9 +272,11 @@ def test_get_block_descs_ids_kernel_block_mismatch(): # FA group: stride=num_blocks=400, offset=0 # region0: [3, 7], region1: [403, 407] - # SSM group: stride=logical_blocks=400//4=100, offset=num_descs=800 - # region0: [801, 802], region1: [901, 902] - expected = [3, 7, 403, 407, 801, 802, 901, 902] + # SSM group: stride=logical_blocks=400//4=100, offset=num_fa_descs=800, + # 4 regions per Mamba layer (x, B, C, ssm) + # region0: [801, 802], region1: [901, 902], + # region2: [1001, 1002], region3: [1101, 1102] + expected = [3, 7, 403, 407, 801, 802, 901, 902, 1001, 1002, 1101, 1102] assert list(result) == expected, f"Expected {expected}, got {list(result)}" @@ -418,3 +425,29 @@ def test_has_mamba_init( ) assert scheduler._has_mamba is expected_has_mamba assert scheduler._is_hma_required is expected_is_hma + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "ssm_sizes,block_len,expected_ratio", + [ + # Nemotron 30B TP=1: ceil((36864 + 2097152) / 8192) = 261 + ((36864, 2097152), 8192, 261), + # Nemotron 30B TP=2: ceil((18432 + 1048576) / 4096) = 261 + ((18432, 1048576), 4096, 261), + # Nemotron 30B TP=4: ceil((9216 + 524288) / 4096) = 131 + ((9216, 524288), 4096, 131), + ], +) +def test_compute_mamba_phys_ratio(ssm_sizes, block_len, expected_ratio): + """Verify that compute_mamba_phys_ratio is TP-dependent. + + With dimension-sharded Mamba state, the ratio differs across TP sizes + (e.g. TP=1 → 261, TP=4 → 131 for Nemotron 30B). This is why + _mamba_phys_ratio must be stored per-engine. + """ + from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( + compute_mamba_phys_ratio, + ) + + assert compute_mamba_phys_ratio(ssm_sizes, block_len) == expected_ratio diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 72980a85ab2..8e66fce4c64 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -5,7 +5,7 @@ KV cache helper for store. """ from collections.abc import Iterator -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast import torch @@ -516,6 +516,338 @@ class TpKVTopology: return cache if self.split_k_and_v else [cache] +# ---- Mamba-HMA hetero-TP transfer config ---- +# +# Key insight: with hetero-TP (P_TP > D_TP), FA KV cache may be +# replicated across P ranks (when P_TP > num_kv_heads), but Mamba +# conv/SSM state is almost always uniquely sharded per P rank. So the +# number of P ranks D must read from can differ between FA and Mamba, +# and they must be handled separately. + + +def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range: + """Physical KV head range stored in a rank's KV cache tensor. + + When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank. + When ``tp_size > num_heads``: 1 physical head per rank. Heads are + distributed **contiguously** (matching vLLM's GQA weight partitioning): + consecutive ranks share a head before moving to the next one. + """ + if tp_size <= num_heads: + assert num_heads % tp_size == 0 + per_rank = num_heads // tp_size + return range(rank * per_rank, (rank + 1) * per_rank) + else: + h = rank * num_heads // tp_size + return range(h, h + 1) + + +def _range_overlap(a: range, b: range) -> range: + start = max(a.start, b.start) + stop = min(a.stop, b.stop) + return range(start, max(start, stop)) + + +@dataclass +class HeteroTPTransferConfig: + """Precomputed transfer plan for one (D rank, P engine) pair. + + Currently only instantiated for Mamba-HMA (hybrid SSM+Attention) models + where FA and mamba require different splitting factors. Could be extended + to other model types that need non-uniform hetero-TP transfer sizing. + + All descriptor sizes are computed here. The guarantee is: + local_entry_size == remote_entry_size (for NIXL) + + Attributes that start with ``fa_`` concern FlashAttention KV cache. + Attributes that start with ``mamba_`` concern Mamba conv/SSM state. + """ + + # ---- Input parameters (from handshake) ---- + tp_ratio: int + K: int # total_num_kv_heads (before TP sharding) + d_tp: int # D engine's tensor_parallel_size + p_tp: int # P engine's tensor_parallel_size + d_rank: int # this D worker's TP rank + use_mla: bool + + # Per-layer block lengths (bytes, K+V combined for blocks_first). + # Uniform across layers for current models. + d_block_len: int # D's block_len_per_layer (representative) + p_block_len: int # P's block_len_per_layer (from handshake) + is_blocks_first: bool # kv_topo.is_kv_layout_blocks_first + + # ---- Derived: computed in __post_init__ ---- + # + # Physical heads per rank (what the KV tensor actually stores) + d_physical_heads: int = field(init=False) + p_physical_heads: int = field(init=False) + + # How many distinct P ranks D needs for FA data + physical_fa_num_reads: int = field(init=False) + + # Which P ranks contribute unique FA heads (ordered by head index) + fa_read_targets: list[int] = field(init=False) + + # All P ranks needed for mamba (always abs_tp for tp_ratio < 0) + mamba_num_reads: int = field(init=False) + + # All P ranks this D rank communicates with (FA ∪ mamba) + transfer_targets: list[int] = field(init=False) + + # FA descriptor entry size (K or V side, for blocks_first layout) + # Guaranteed: fa_entry_size is the SAME for local handle AND remote desc. + fa_entry_size: int = field(init=False) + + # Replication flags + is_d_replicated: bool = field(init=False) + is_p_replicated: bool = field(init=False) + + # Pre-built set for fast lookup + _fa_target_set: frozenset[int] = field(init=False, repr=False) + # Map: P rank → index in fa_read_targets (for head slot offset) + _fa_target_index: dict[int, int] = field(init=False, repr=False) + + def __post_init__(self) -> None: + K = self.K + self.is_d_replicated = self.d_tp > K + self.is_p_replicated = self.p_tp > K + + self.d_physical_heads = max(1, K // self.d_tp) + self.p_physical_heads = max(1, K // self.p_tp) + + abs_tp = -self.tp_ratio if self.tp_ratio < 0 else 1 + + # ---- Mamba range (computed first so FA can prefer ranks in it) ---- + mamba_range: range | None = None + if self.tp_ratio < 0: + mamba_range = range(self.d_rank * abs_tp, (self.d_rank + 1) * abs_tp) + + # ---- FA read targets ---- + if self.use_mla or self.tp_ratio >= 0: + self.physical_fa_num_reads = 1 + self.fa_read_targets = ( + [0] + if self.use_mla + # Must match kv_topo.get_target_remote_ranks (d_rank // tp_ratio). + else [ + self.d_rank // self.tp_ratio if self.tp_ratio > 0 else self.d_rank + ] + ) + else: + d_needs = _physical_head_range(self.d_tp, K, self.d_rank) + # When mamba range exists, prefer P ranks within it so that + # FA targets are a subset of mamba transfer_targets (avoids + # orphaned FA targets outside the transfer loop). + search_range = mamba_range if mamba_range is not None else range(self.p_tp) + seen: set[tuple[int, int]] = set() + targets: list[int] = [] + for p in search_range: + p_has = _physical_head_range(self.p_tp, K, p) + ov = _range_overlap(d_needs, p_has) + if len(ov) > 0: + key = (ov.start, ov.stop) + if key not in seen: + seen.add(key) + targets.append(p) + if not targets: + # Fallback: search globally (should not happen in practice) + for p in range(self.p_tp): + p_has = _physical_head_range(self.p_tp, K, p) + ov = _range_overlap(d_needs, p_has) + if len(ov) > 0: + key = (ov.start, ov.stop) + if key not in seen: + seen.add(key) + targets.append(p) + self.fa_read_targets = targets + self.physical_fa_num_reads = len(targets) + + self._fa_target_set = frozenset(self.fa_read_targets) + self._fa_target_index = {r: i for i, r in enumerate(self.fa_read_targets)} + + # ---- Mamba targets ---- + if mamba_range is not None and abs_tp > self.physical_fa_num_reads: + self.mamba_num_reads = abs_tp + self.transfer_targets = list(mamba_range) + else: + self.mamba_num_reads = self.physical_fa_num_reads + self.transfer_targets = list(self.fa_read_targets) + + # ---- FA entry size ---- + # For blocks_first: block_len_per_layer includes K+V; // 2 gives K (or V). + # Use min(D, P) because D indexes into P when tp_ratio > 0, + # and P is the natural unit when tp_ratio < 0. + effective_block_len = min(self.d_block_len, self.p_block_len) + if self.is_blocks_first: + self.fa_entry_size = effective_block_len // 2 + else: + self.fa_entry_size = effective_block_len + + self._validate() + + def _validate(self) -> None: + """Cross-check internal consistency.""" + if self.is_d_replicated and self.is_p_replicated and self.tp_ratio > 0: + logger.info( + "Both-replicated hetero-TP: D_TP=%d > P_TP=%d > K=%d. " + "Using d_rank // tp_ratio routing with relative head offset.", + self.d_tp, + self.p_tp, + self.K, + ) + + # FA targets must be a subset of transfer_targets + tt_set = set(self.transfer_targets) + for t in self.fa_read_targets: + if t not in tt_set: + logger.error( + "FA target P rank %d is NOT in transfer_targets %s. " + "This will cause missed FA reads!", + t, + self.transfer_targets, + ) + + # For tp_ratio < 0 with blocks_first: D_K_half / reads should == P_K_half + if ( + self.is_blocks_first + and self.tp_ratio < 0 + and self.physical_fa_num_reads > 0 + ): + d_k_half = self.d_block_len // 2 + p_k_half = self.p_block_len // 2 + expected_local = d_k_half // self.physical_fa_num_reads + if expected_local != p_k_half: + logger.warning( + "FA size mismatch: D_K_half=%d / reads=%d = %d, " + "but P_K_half=%d. This may indicate a head count or " + "Mamba-HMA inflation inconsistency.", + d_k_half, + self.physical_fa_num_reads, + expected_local, + p_k_half, + ) + + # ---- Query methods ---- + + def should_skip_fa(self, p_rank: int) -> bool: + """Whether to skip FA groups for this P rank (mamba-only transfer).""" + return p_rank not in self._fa_target_set + + def fa_head_slot(self, p_rank: int) -> int: + """Index into D's FA block for this P rank's head data. + + For P ranks in fa_read_targets, returns 0, 1, ..., reads-1. + For P ranks NOT in fa_read_targets (replicated duplicates), + returns the slot of the matching FA target with the same head. + """ + if p_rank in self._fa_target_index: + return self._fa_target_index[p_rank] + # Duplicate head: find which fa_target has the same physical head + p_head = _physical_head_range(self.p_tp, self.K, p_rank) + for target in self.fa_read_targets: + t_head = _physical_head_range(self.p_tp, self.K, target) + if _range_overlap(p_head, t_head): + return self._fa_target_index[target] + return 0 # fallback + + def fa_rank_offset(self, remote_kv_block_len: int) -> int: + """Byte offset into P's FA block for this D rank. + + When D is replicated (D_TP > K), multiple D ranks share a head. + Computes offset *relative to the target P rank's first head* + so it works regardless of how many heads P has. + When neither side replicates, falls back to tp_rank % tp_ratio. + Returns 0 when D does not index into P's block. + """ + if self.use_mla or self.tp_ratio <= 0: + return 0 + if self.is_d_replicated: + d_head = self.d_rank * self.K // self.d_tp + p_rank = self.fa_read_targets[0] + p_start = p_rank * self.K // self.p_tp + return (d_head - p_start) * remote_kv_block_len + return self.d_rank % self.tp_ratio * remote_kv_block_len + + @property + def needs_split_handles(self) -> bool: + """Whether per-P-rank split handles are needed. + + True when FA and mamba have different read counts, requiring + different splitting factors in the local handle. + """ + return self.tp_ratio < 0 and not self.use_mla and len(self.transfer_targets) > 1 + + def compute_split_handle_data( + self, + src_blocks_data: list[tuple[int, int, int]], + num_fa_descs: int, + abs_tp: int, + ) -> list[list[tuple[int, int, int]]]: + """Compute per-P-rank (addr, len, tp) triples for Mamba-HMA split handles. + + FA descriptors (indices < num_fa_descs) are sliced by + ``physical_fa_num_reads``; mamba descriptors are sliced uniformly + by ``abs_tp``. + + Returns one list of triples per transfer target. + """ + all_handle_data: list[list[tuple[int, int, int]]] = [] + for p_idx, p_rank in enumerate(self.transfer_targets): + handle_data: list[tuple[int, int, int]] = [] + skip_fa = self.should_skip_fa(p_rank) + fa_slot = self.fa_head_slot(p_rank) if not skip_fa else 0 + + for j, (addr, local_len, tp) in enumerate(src_blocks_data): + if j < num_fa_descs: + assert self.physical_fa_num_reads >= 1 + fa_chunk = local_len // self.physical_fa_num_reads + handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, tp)) + else: + mamba_chunk = local_len // abs_tp + handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, tp)) + all_handle_data.append(handle_data) + return all_handle_data + + def filter_block_ids_for_rank( + self, + remote_rank: int, + local_ids: BlockIds, + remote_ids: BlockIds, + is_mamba_group: list[bool], + ) -> tuple[BlockIds, BlockIds]: + """Zero out FA groups for P ranks outside fa_read_targets. + + Returns (filtered_local_ids, filtered_remote_ids). When the + remote rank carries FA data for this D rank, returns the inputs + unchanged. + """ + if not self.should_skip_fa(remote_rank): + return local_ids, remote_ids + num_groups = len(local_ids) + filtered_local: list[list[int]] = [ + [] if not is_mamba_group[g] else local_ids[g] for g in range(num_groups) + ] + filtered_remote: list[list[int]] = [ + [] if not is_mamba_group[g] else remote_ids[g] for g in range(num_groups) + ] + return filtered_local, filtered_remote + + def describe(self) -> str: + """One-line summary for logging.""" + return ( + f"HeteroTPTransferConfig(" + f"tp_ratio={self.tp_ratio}, K={self.K}, " + f"d_tp={self.d_tp}, p_tp={self.p_tp}, d_rank={self.d_rank}, " + f"physical_fa_reads={self.physical_fa_num_reads}, " + f"mamba_reads={self.mamba_num_reads}, " + f"fa_targets={self.fa_read_targets}, " + f"transfer_targets={self.transfer_targets}, " + f"fa_entry_size={self.fa_entry_size}, " + f"d_block_len={self.d_block_len}, p_block_len={self.p_block_len})" + ) + + def get_current_attn_backends( vllm_config: VllmConfig, layer_names: list[str] | None = None ) -> list[type[AttentionBackend]]: @@ -559,3 +891,50 @@ def get_current_attn_backend( ) -> type[AttentionBackend]: """Get the first attention backend for the given layers.""" return get_current_attn_backends(vllm_config, layer_names)[0] + + +# TODO (ZhanqiuHu): Consolidate TpKVTopology and HeteroTPTransferConfig +# into a single engine-agnostic TransferTopology class. +# 6 of 9 HeteroTPTransferConfig init fields duplicate TpKVTopology data. +# +# @dataclass +# class EngineTransferInfo: +# """Per-remote-engine transfer state, computed at handshake.""" +# p_tp: int +# tp_ratio: int +# p_block_len: int +# block_size: int +# # Mamba-specific (None for non-mamba models) +# fa_read_targets: list[int] | None = None +# transfer_targets: list[int] | None = None +# physical_fa_num_reads: int | None = None +# mamba_num_reads: int | None = None +# fa_entry_size: int | None = None +# +# class TransferTopology: +# """Single source of truth for TP topology + transfer sizing.""" +# # Shared (set once at init, replaces duplicate fields) +# tp_rank: int # == TpKVTopology.tp_rank == HeteroTP.d_rank +# tp_size: int # == TpKVTopology.tp_size == HeteroTP.d_tp +# total_num_kv_heads: int # == HeteroTP.K +# is_mla: bool # == HeteroTP.use_mla +# is_mamba: bool +# is_blocks_first: bool # == HeteroTP.is_blocks_first +# d_block_len: int +# +# # Per-engine (populated via register_engine() at handshake) +# _engines: dict[EngineId, EngineTransferInfo] +# +# def register_engine(self, engine_id, p_tp, p_block_len, ...): ... +# +# # General (from TpKVTopology) +# def tp_ratio(self, engine_id) -> int: ... +# def target_remote_ranks(self, engine_id) -> list[int]: ... +# def is_kv_replicated(self, engine_id) -> bool: ... +# +# # Mamba-specific (from HeteroTPTransferConfig, gated by is_mamba) +# def fa_rank_offset(self, engine_id, block_len) -> int: ... +# def physical_fa_num_reads(self, engine_id) -> int: ... +# def transfer_targets(self, engine_id) -> list[int]: ... +# def should_skip_fa(self, engine_id, p_rank) -> bool: ... +# def filter_block_ids_for_rank(self, engine_id, ...) -> ...: ... diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index 0aaf3b6e938..c575043fb34 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -25,6 +25,7 @@ from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.utils import ( BlockIds, EngineId, + HeteroTPTransferConfig, TpKVTopology, get_current_attn_backend, get_current_attn_backends, @@ -47,12 +48,18 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( PromMetric, PromMetricT, ) +from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( + MambaConvSplitInfo, + compute_mamba_phys_ratio, + derive_mamba_conv_split, +) from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) from vllm.forward_context import ForwardContext from vllm.logger import init_logger +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first from vllm.platforms import current_platform from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import make_zmq_path, make_zmq_socket @@ -1038,7 +1045,7 @@ class NixlConnectorWorker: } self.hma_group_size = len(kv_cache_config.kv_cache_tensors) - # Mamba metadata + # ---- Mamba model state (derived from model config) ---- self._is_mamba_group = [ isinstance(group.kv_cache_spec, MambaSpec) for group in kv_cache_config.kv_cache_groups @@ -1065,6 +1072,17 @@ class NixlConnectorWorker: ssm_shape.numel() * ssm_nbytes, ) self._mamba_ssm_size = mamba_ssm_size + # Conv state sub-projection decomposition (None when no Mamba). + # The 3-read transfer requires DS (dim, state_len) conv layout so + # that x/B/C sub-projections are contiguous in memory. + self._conv_decomp: MambaConvSplitInfo | None = None + if self._has_mamba: + assert is_conv_state_dim_first(), ( + "3-read Mamba conv transfer requires DS conv state layout. " + "Set VLLM_SSM_CONV_STATE_LAYOUT=DS" + ) + local_tp = vllm_config.parallel_config.tensor_parallel_size + self._conv_decomp = derive_mamba_conv_split(mamba_spec, local_tp) # Agent. non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"] @@ -1175,6 +1193,16 @@ class NixlConnectorWorker: self.dst_num_blocks: dict[EngineId, int] = {} self._registered_descs: list[Any] = [] + # ---- Mamba-HMA per-engine state (only used when self._has_mamba) ---- + # Per-engine transfer config (source of truth for FA/mamba sizing). + self._transfer_configs: dict[str, HeteroTPTransferConfig] = {} + # NOTE (ZhanqiuHu): _mamba_phys_ratio MUST be per-engine. + # compute_mamba_phys_ratio = ceil((conv_bytes + ssm_bytes) / block_len) + # where conv/ssm bytes are per-TP-rank (dimension-sharded). With + # heterogeneous TP the per-rank sizes differ, so the ratio differs: + # e.g. Nemotron 30B: P(TP=4) → 131, D(TP=1) → 261. + self._mamba_phys_ratio: dict[EngineId, int] = {} + # In progress transfers. # [req_id -> list[handle]] self._recving_metadata: dict[ReqId, ReqMeta] = {} @@ -1701,8 +1729,7 @@ class NixlConnectorWorker: # then duplicate it logically to be able to index SSM/Conv separately. self.num_regions *= 2 - # TODO (NickLucche) Adapt to different descs views (engine_id->tp_rank) to - # support heterogeneous TP. + # Total local FA descriptors (boundary between FA and mamba descs). self.num_descs = self.num_regions * self.num_blocks descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type) @@ -1715,6 +1742,9 @@ class NixlConnectorWorker: self.dst_num_blocks[self.engine_id] = self.num_blocks if self._has_mamba: + self._mamba_phys_ratio[self.engine_id] = ( + self._physical_blocks_per_logical_kv_block + ) logger.info( "Hybrid SSM registration: num_blocks=%s, " "logical_num_blocks=%s, ratio=%s, num_regions=%s, " @@ -1755,6 +1785,149 @@ class NixlConnectorWorker: agent_metadata_bytes=encoder.encode(agent_metadata), ) + def _build_mamba_local( + self, + base_addresses: list[int], + block_size_ratio: int, + ) -> list[tuple[int, int, int]]: + """Build 4 desc regions (x, B, C, ssm) per layer for local mamba + blocks, enabling the 3-read transfer with DS conv layout.""" + assert block_size_ratio == 1, ( + "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " + f"Got block_size_ratio={block_size_ratio}." + ) + assert self._conv_decomp is not None + conv_offsets = self._conv_decomp.local_conv_offsets + conv_size, ssm_size = self._mamba_ssm_size + num_blocks = self._logical_num_blocks * block_size_ratio + phys_ratio = self._physical_blocks_per_logical_kv_block + + result: list[tuple[int, int, int]] = [] + for i, base_addr in enumerate(base_addresses): + page_stride = self.block_len_per_layer[i] // block_size_ratio * phys_ratio + for off, sz in conv_offsets: + for blk in range(num_blocks): + result.append( + (base_addr + blk * page_stride + off, sz, self.device_id) + ) + # SSM temporal state follows the conv state. + for blk in range(num_blocks): + result.append( + ( + base_addr + blk * page_stride + conv_size, + ssm_size, + self.device_id, + ) + ) + return result + + def _build_fa_remote_for_mamba( + self, + nixl_agent_meta: NixlAgentMetadata, + transfer_cfg: HeteroTPTransferConfig, + block_size_ratio: int, + kv_topo: TpKVTopology, + ) -> list[tuple[int, int, int]]: + """Build remote FA descriptors for mamba models. + + Uses transfer_cfg for GQA-aware FA divisor and head-based rank offset + instead of the standard uniform tp_ratio split. + """ + assert block_size_ratio == 1, ( + "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " + f"Got block_size_ratio={block_size_ratio}." + ) + # TODO (ZhanqiuHu): unify with register_remote_blocks when Mamba-HMA + # hetero-TP logic stabilizes. + tp_ratio = transfer_cfg.tp_ratio + result: list[tuple[int, int, int]] = [] + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): + local_block_len = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=True, mamba_view=False + ) + remote_kv_block_len = local_block_len // block_size_ratio + if block_size_ratio > 1: + local_block_len = remote_kv_block_len + + if tp_ratio < 0 and not self.use_mla: + local_block_len = local_block_len // transfer_cfg.physical_fa_num_reads + + rank_offset = transfer_cfg.fa_rank_offset(remote_kv_block_len) + + num_blocks = nixl_agent_meta.num_blocks + page_size = nixl_agent_meta.block_lens[i] + for block_id in range(num_blocks): + block_offset = block_id * page_size + addr = base_addr + block_offset + rank_offset + result.append((addr, local_block_len, nixl_agent_meta.device_id)) + + if kv_topo.is_kv_layout_blocks_first: + second_split = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=False, mamba_view=False + ) + if tp_ratio < 0 and not self.use_mla: + second_split = second_split // transfer_cfg.physical_fa_num_reads + for block_id in range(num_blocks): + block_offset = block_id * page_size + addr = base_addr + block_offset + rank_offset + v_addr = addr + nixl_agent_meta.block_lens[i] // 2 + result.append((v_addr, second_split, nixl_agent_meta.device_id)) + return result + + def _build_mamba_remote( + self, + nixl_agent_meta: NixlAgentMetadata, + tp_ratio: int, + ) -> list[tuple[int, int, int]]: + """Build 4 remote desc regions (x, B, C, ssm) per layer for + the 3-read transfer. For hetero-TP, each D rank reads only its + sub-projection slice from the P rank.""" + assert self._conv_decomp is not None + effective_ratio = max(tp_ratio, 1) + # Mamba conv state is always TP-sharded, even when attention KV + # is replicated (num_kv_heads < tp_size). + local_offset = self.tp_rank % effective_ratio + conv_size_remote = nixl_agent_meta.ssm_sizes[0] + + if tp_ratio >= 1: + # D_TP >= P_TP: P page is larger, D reads its slice. + conv_offsets = self._conv_decomp.remote_conv_offsets( + local_offset, effective_ratio + ) + ssm_read_size = self._mamba_ssm_size[1] + else: + # NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages + # are smaller than D's. self._conv_decomp has D-sized dimensions, + # but we need P-sized offsets. Scale down by |tp_ratio|. + abs_ratio = -tp_ratio + xb_p = self._conv_decomp.x_bytes // abs_ratio + bb_p = self._conv_decomp.b_bytes // abs_ratio + conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)] + ssm_read_size = nixl_agent_meta.ssm_sizes[1] + + remote_ratio = self._mamba_phys_ratio[nixl_agent_meta.engine_id] + num_blocks = nixl_agent_meta.num_blocks // remote_ratio + device_id = nixl_agent_meta.device_id + + result: list[tuple[int, int, int]] = [] + # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case + # block lengths vary across layers (e.g. MLA). + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): + page_stride = nixl_agent_meta.block_lens[i] * remote_ratio + for off, sz in conv_offsets: + for blk in range(num_blocks): + result.append((base_addr + blk * page_stride + off, sz, device_id)) + # SSM temporal state is also TP-sharded on the heads dimension. + for blk in range(num_blocks): + ssm_addr = ( + base_addr + + blk * page_stride + + conv_size_remote + + local_offset * ssm_read_size + ) + result.append((ssm_addr, ssm_read_size, device_id)) + return result + def register_local_xfer_handler( self, block_size: int, @@ -1823,13 +1996,22 @@ class NixlConnectorWorker: self.device_id, ) + # NOTE (ZhanqiuHu): mamba=True path in register_blocks is not used + # right now — we use _build_mamba_local instead for the 3-read + # approach. However, we might still need this as a fallback for homogeneous TP. register_blocks(blocks_data, mamba=False) if self._has_mamba: assert self.num_descs == len(blocks_data) - logger.debug( - "Registering additional %s local Mamba blocks", len(blocks_data) + # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-read split is + # unnecessary — a single conv desc per block suffices. Consider + # adding a fast path that falls back to the standard 2-region + # registration (register_blocks mamba=True) when no hetero-TP + # remote has been seen. Currently we always register 4 regions + # because local descs are created before knowing the remote TP. + logger.debug("Registering local Mamba descriptors (4 regions/layer)") + blocks_data.extend( + self._build_mamba_local(local_base_addresses, block_size_ratio) ) - register_blocks(blocks_data, mamba=True) descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) # NIXL_INIT_AGENT to be used for preparations of local descs. @@ -1880,6 +2062,9 @@ class NixlConnectorWorker: Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0 so that the whole cache is shared by "tp_ratio" D TP workers. + + For Mamba hetero-TP, both tp_ratio > 0 (D_TP > P_TP) and + tp_ratio < 0 (P_TP > D_TP) are supported by the 3-read transfer. """ # noqa: E501 engine_id = nixl_agent_meta.engine_id # TODO re-evaluate refreshing for scaling/recovery @@ -1915,6 +2100,10 @@ class NixlConnectorWorker: if engine_id not in self.dst_num_blocks: self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks + if self._has_mamba: + self._mamba_phys_ratio[engine_id] = compute_mamba_phys_ratio( + nixl_agent_meta.ssm_sizes, nixl_agent_meta.block_lens[0] + ) # Keep track of remote agent kv caches base addresses. self.kv_caches_base_addr[engine_id][remote_tp_rank] = ( @@ -1931,6 +2120,21 @@ class NixlConnectorWorker: not self.kv_topo.replicates_kv_cache(engine_id) and tp_ratio > 0 ) + # Create transfer config (single source of truth for descriptor sizes). + if self._has_mamba and engine_id not in self._transfer_configs: + self._transfer_configs[engine_id] = HeteroTPTransferConfig( + tp_ratio=tp_ratio, + K=kv_topo.total_num_kv_heads, + d_tp=self.world_size, + p_tp=remote_tp_size, + d_rank=self.tp_rank, + use_mla=self.use_mla, + d_block_len=self.block_len_per_layer[0], + p_block_len=nixl_agent_meta.block_lens[0], + is_blocks_first=kv_topo.is_kv_layout_blocks_first, + ) + logger.info("Created %s", self._transfer_configs[engine_id].describe()) + logger.debug( "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s", engine_id, @@ -1947,21 +2151,48 @@ class NixlConnectorWorker: # Remote tp_size > local tp_size: read from multiple remote ranks. # Logically "split" own regions into |tp_ratio| chunks. Mind that # we only do this once per remote tp_size (replica-friendly). + abs_tp = -tp_ratio self.src_xfer_handles_by_tp_ratio[tp_ratio] = [] - for i in range(-tp_ratio): - blocks_data = [] - for memory_region in self.src_blocks_data: - addr, local_block_len, own_tp_rank = memory_region - # Computing block len layer by layer allows for different - # block sizes to be used. - remote_block_len = local_block_len // (-tp_ratio) - addr = addr + i * remote_block_len - blocks_data.append((addr, remote_block_len, own_tp_rank)) - descs = self.nixl_wrapper.get_xfer_descs( - blocks_data, self.nixl_memory_type - ) - handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) + + if self._has_mamba: + transfer_cfg = self._transfer_configs.get(engine_id) + assert transfer_cfg is not None + if transfer_cfg.needs_split_handles: + # Mamba-HMA: FA and Mamba use different split factors. + for handle_data in transfer_cfg.compute_split_handle_data( + self.src_blocks_data, self.num_descs, abs_tp + ): + descs = self.nixl_wrapper.get_xfer_descs( + handle_data, self.nixl_memory_type + ) + handle = self.nixl_wrapper.prep_xfer_dlist( + "NIXL_INIT_AGENT", descs + ) + self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) + + logger.info( + "Mamba-HMA split handles: targets=%s, fa_reads=%s, " + "fa_entry=%s, mamba_reads=%s, num_descs=%s", + transfer_cfg.transfer_targets, + transfer_cfg.physical_fa_num_reads, + transfer_cfg.fa_entry_size, + transfer_cfg.mamba_num_reads, + self.num_descs, + ) + else: + # Original path: uniform divide by abs_tp (non-Mamba-HMA). + for i in range(abs_tp): + blocks_data = [] + for memory_region in self.src_blocks_data: + addr, local_block_len, own_tp_rank = memory_region + remote_block_len = local_block_len // abs_tp + addr = addr + i * remote_block_len + blocks_data.append((addr, remote_block_len, own_tp_rank)) + descs = self.nixl_wrapper.get_xfer_descs( + blocks_data, self.nixl_memory_type + ) + handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs) + self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle) ### Register remote agent memory regions blocks_data = [] @@ -2044,13 +2275,33 @@ class NixlConnectorWorker: self.tp_rank, ) - register_remote_blocks(blocks_data, mamba=False) if self._has_mamba: - # Create extra descs for the Mamba "view" of the same KV cache tensors. + # Mamba-HMA: separate FA registration with GQA-aware sizing, + # plus mamba 3-read registration for the Mamba "view" of the + # same KV cache tensors. logger.debug( - "Registering additional %s remote Mamba blocks", len(blocks_data) + "Registering remote Mamba blocks for engine %s rank %s", + engine_id, + remote_tp_rank, ) - register_remote_blocks(blocks_data, mamba=True) + transfer_cfg = self._transfer_configs.get(engine_id) + assert transfer_cfg is not None + blocks_data.extend( + self._build_fa_remote_for_mamba( + nixl_agent_meta, + transfer_cfg, + block_size_ratio, + kv_topo, + ) + ) + blocks_data.extend( + self._build_mamba_remote( + nixl_agent_meta, + tp_ratio, + ) + ) + else: + register_remote_blocks(blocks_data, mamba=False) # Register with NIXL. descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) @@ -2083,17 +2334,17 @@ class NixlConnectorWorker: block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id( remote_engine_id ) - # Num kv_heads > tp_size and P TP > D TP case, not supported - assert not (tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id)) + # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba. + # Mamba models can have replicated FA KV with tp_ratio < 0. + if not self._has_mamba: + assert not ( + tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id) + ) if self._is_hma_required: assert block_size_ratio == 1, ( "HMA does not support different remote block size yet" ) - # Mamba additional constraints - if self._has_mamba: - assert tp_ratio == 1, "Mamba does not support heterogeneous TP yet" - kv_cache_layout = ( self.kv_cache_layout if not self.use_host_buffer @@ -2138,11 +2389,14 @@ class NixlConnectorWorker: remote_block_len = nixl_agent_meta.block_lens[0] if self.use_mla or self.kv_topo.is_kv_replicated(remote_engine_id): # With replicated KV cache, only the number of blocks can differ. - for i in range(len(self.block_len_per_layer)): - assert ( - self.block_len_per_layer[i] // block_size_ratio - == nixl_agent_meta.block_lens[i] - ), "KV cache sizes must match between P and D when replicated" + # TODO (ZhanqiuHu): For mamba models, validate FA and mamba + # block_lens separately. + if not self._has_mamba: + for i in range(len(self.block_len_per_layer)): + assert ( + self.block_len_per_layer[i] // block_size_ratio + == nixl_agent_meta.block_lens[i] + ), "KV cache sizes must match between P and D when replicated" else: # When MLA is not used, this is a list of the same block length for block_len in nixl_agent_meta.block_lens: @@ -2150,25 +2404,31 @@ class NixlConnectorWorker: "All remote layers must have the same block size" ) - if tp_ratio > 0: - # Remote tp is smaller: remote block_len size is bigger - assert ( - remote_block_len - == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio - ), ( - "Remote P worker KV layer cache must be of shape [2, N, " - "local_kv_heads*tp_ratio, page_size, head_dim] and same dtype." - ) # noqa: E501 - else: - assert block_size_ratio == 1, ( - "Different local/remote block sizes are not supported when" - " P TP > D TP." - ) - # Remote tp is bigger: remote block_len size is smaller - assert remote_block_len == self.block_len_per_layer[0] // (-tp_ratio), ( - "Remote P worker KV layer cache must be of shape [2, N, " - "local_kv_heads/tp_ratio, page_size, head_dim] and same dtype." - ) # noqa: E501 + # HMA hybrid models (mamba+attention) pad block_len to + # max(attn_page, mamba_page), so the linear tp_ratio scaling + # assumption only holds for pure-attention models. + if not self._has_mamba: + if tp_ratio > 0: + assert ( + remote_block_len + == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio + ), ( + "Remote P worker KV layer cache must be of shape [2, N," + " local_kv_heads*tp_ratio, page_size, head_dim] and " + "same dtype." + ) + else: + assert block_size_ratio == 1, ( + "Different local/remote block sizes are not supported" + " when P TP > D TP." + ) + assert remote_block_len == self.block_len_per_layer[0] // ( + -tp_ratio + ), ( + "Remote P worker KV layer cache must be of shape [2, N," + " local_kv_heads/tp_ratio, page_size, head_dim] and " + "same dtype." + ) # TP workers that handhshake with same remote have same #blocks. assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks @@ -2471,9 +2731,8 @@ class NixlConnectorWorker: meta.local_block_ids ) assert meta.remote is not None - meta.remote.block_ids = self._logical_to_kernel_block_ids( - meta.remote.block_ids - ) + # Remote block IDs are kept logical here; expanded in + # _read_blocks_for_req using the remote engine's phys ratio. remote_engine_id = meta.remote.engine_id logger.debug( "start_load_kv for request %s from remote engine %s. " @@ -2525,6 +2784,13 @@ class NixlConnectorWorker: meta.remote.engine_id ) tp_ratio = self.kv_topo.tp_ratio_from_engine_id(meta.remote.engine_id) + + if self._has_mamba: + # Expand remote logical → kernel block IDs. + meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( + meta.remote.block_ids, + self._mamba_phys_ratio[meta.remote.engine_id], + ) # D may have to perform multiple reads from different remote ranks. for i, remote_rank in enumerate(remote_ranks): if self.use_mla and tp_ratio < 0 and i > 0: @@ -2558,12 +2824,26 @@ class NixlConnectorWorker: remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][ remote_rank ] + + local_ids: BlockIds = meta.local_physical_block_ids + remote_ids: BlockIds = meta.remote.block_ids + if self._has_mamba: + # Mamba-HMA: zero out FA groups for P ranks outside fa_read_targets. + transfer_cfg = self._transfer_configs.get(meta.remote.engine_id) + assert transfer_cfg is not None + local_ids, remote_ids = transfer_cfg.filter_block_ids_for_rank( + remote_rank, + local_ids, + remote_ids, + self._is_mamba_group, + ) + self._read_blocks( request_id=req_id, dst_engine_id=meta.remote.engine_id, remote_request_id=meta.remote.request_id, - local_block_ids=meta.local_physical_block_ids, - remote_block_ids=meta.remote.block_ids, + local_block_ids=local_ids, + remote_block_ids=remote_ids, remote_rank=remote_rank, local_xfer_side_handle=local_xfer_side_handle, remote_xfer_side_handle=remote_xfer_side_handle, @@ -2663,9 +2943,12 @@ class NixlConnectorWorker: for i, remote_group in enumerate(remote_block_ids): num_remote_blocks = len(remote_group) num_local_blocks = len(local_block_ids[i]) - assert num_local_blocks <= num_remote_blocks + if not self._is_mamba_group[i]: + assert num_local_blocks <= num_remote_blocks # Partial prefix cache hit: just read uncomputed blocks. - if num_local_blocks < num_remote_blocks: + # Skip mamba groups — their blocks represent full state (conv+ssm), + # not per-token data, so trimming would corrupt the transfer. + if num_local_blocks < num_remote_blocks and not self._is_mamba_group[i]: remote_block_ids[i] = remote_group[-num_local_blocks:] # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from @@ -2781,16 +3064,22 @@ class NixlConnectorWorker: # This is like having two "low-level views" of the same storage. # `num_fa_descs` offset must be computed per-engine since P and D can # have different num_blocks (and thus different FA descs counts). - ratio = self._physical_blocks_per_logical_kv_block - # SSM may register fewer num_blocks than FA + ratio = self._mamba_phys_ratio[engine_id] logical_blocks = num_blocks // ratio num_fa_descs = self.num_regions * num_blocks + # 3-read mamba: 4 regions per unique cache tensor (x, B, C, ssm). + mamba_region_ids = np.arange(len(self.block_len_per_layer) * 4)[:, None] all_descs = [] for i, group in enumerate(block_ids): - stride = logical_blocks if self._is_mamba_group[i] else num_blocks group_arr = np.asarray(group)[None, :] - offset = num_fa_descs if self._is_mamba_group[i] else 0 - all_descs.append((region_ids * stride + group_arr + offset).flatten()) + if self._is_mamba_group[i]: + all_descs.append( + ( + mamba_region_ids * logical_blocks + group_arr + num_fa_descs + ).flatten() + ) + else: + all_descs.append((region_ids * num_blocks + group_arr).flatten()) return np.concatenate(all_descs) def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds: @@ -2818,6 +3107,36 @@ class NixlConnectorWorker: for i, group in enumerate(block_ids) ] + def _logical_to_remote_kernel_block_ids( + self, block_ids: BlockIds, remote_ratio: int + ) -> BlockIds: + """Map logical block IDs to physical kernel block IDs on the remote. + + Args: + block_ids: per-group lists of logical block IDs. + remote_ratio: remote engine's physical blocks per logical block. + + Returns: + Same structure with FA groups expanded (each logical block L + becomes kernel blocks [L*remote_ratio .. L*remote_ratio + + local_ratio - 1]). Mamba groups are passed through unchanged. + """ + local_ratio = self._physical_blocks_per_logical_kv_block + if remote_ratio == 1: + return block_ids + local_arange = np.arange(local_ratio).reshape(1, -1) + group_specs = self.kv_cache_config.kv_cache_groups + result: list[list[int]] = [] + for i, group in enumerate(block_ids): + if not isinstance(group_specs[i].kv_cache_spec, MambaSpec): + arr = np.array(group).reshape(-1, 1) + expanded = (arr * remote_ratio + local_arange).flatten() + result.append(expanded.tolist()) + else: + # Mamba blocks are 1:1 logical-to-physical (no expansion). + result.append(group) + return result + def get_backend_aware_kv_block_len( self, layer_idx: int, first_split: bool = True, mamba_view: bool = False ) -> int: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py new file mode 100644 index 00000000000..6d65e006e1b --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Mamba conv-state sub-projection decomposition for the 3-read transfer. + +With DS conv state layout (dim, state_len), x/B/C sub-projections are +contiguous in memory. Each D rank reads its x, B, C slices via 3 +separate RDMA transfers — no P-side permutation needed. +""" + +import math +from dataclasses import dataclass + +import torch + +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.v1.kv_cache_interface import MambaSpec + + +@dataclass(frozen=True) +class MambaConvSplitInfo: + """Per-rank byte sizes of x, B, C sub-projections in the Mamba conv state. + + Used by both P and D sides for NIXL descriptor registration. + All fields are LOCAL to this engine's TP (already divided by TP size). + + DS memory layout within one page (contiguous in memory): + |--- x (x_local * conv_rows) ---|- B (b_local * conv_rows) -|- C -| + """ + + conv_rows: int # conv_kernel - 1 (typically 3) + x_local: int # intermediate_size / TP (columns for x) + b_local: int # groups_ss / TP (columns for B; C is same size) + conv_dtype_size: int # bytes per element (e.g. 2 for float16) + + @property + def conv_dim_local(self) -> int: + """Total conv columns per rank: x + B + C.""" + return self.x_local + 2 * self.b_local + + @property + def x_bytes(self) -> int: + """Byte size of the x sub-projection for one rank.""" + return self.x_local * self.conv_rows * self.conv_dtype_size + + @property + def b_bytes(self) -> int: + """Byte size of the B (or C) sub-projection for one rank.""" + return self.b_local * self.conv_rows * self.conv_dtype_size + + @property + def local_conv_offsets(self) -> list[tuple[int, int]]: + """(byte_offset, byte_size) of x, B, C within this engine's page. + + Used by both P and D for local descriptor registration. + """ + xb = self.x_bytes + bb = self.b_bytes + return [(0, xb), (xb, bb), (xb + bb, bb)] + + def remote_conv_offsets( + self, local_rank_offset: int, tp_ratio: int + ) -> list[tuple[int, int]]: + """(byte_offset, byte_size) of this D rank's x, B, C slice within + one P page. + + Used by D side only, during remote descriptor registration. + + Args: + local_rank_offset: which slice this D rank reads. + tp_ratio > 0: tp_rank % tp_ratio (selects slice of P's page). + tp_ratio < 0: always 0 (read P's full page). + tp_ratio: effective ratio (>= 1 when D_TP > P_TP, 1 when + P_TP > D_TP since each P rank is read in full). + """ + xb = self.x_bytes + bb = self.b_bytes + xr = xb * tp_ratio # full remote x section in bytes + br = bb * tp_ratio # full remote B section in bytes + return [ + (local_rank_offset * xb, xb), + (xr + local_rank_offset * bb, bb), + (xr + br + local_rank_offset * bb, bb), + ] + + +def derive_mamba_conv_split( + mamba_spec: MambaSpec, + local_tp: int, +) -> MambaConvSplitInfo: + """Derive per-rank x/B/C byte sizes from a MambaSpec. + + Called once at init on both P and D. Decomposes the conv dimension + (= intermediate_size + 2 * groups_ss) into its x, B, C parts. + + Args: + mamba_spec: MambaSpec whose shapes are: + shapes[0] = conv state: (conv_dim_local, conv_rows) in DS layout. + shapes[1] = SSM temporal: (local_num_heads, head_dim). + local_tp: this engine's tensor-parallel size. + + Returns: + MambaConvSplitInfo with per-rank x_local, b_local, conv_rows, and + conv_dtype_size. + """ + if mamba_spec.mamba_type != "mamba2": + raise NotImplementedError( + f"3-read conv transfer only supports Mamba2 models, " + f"got mamba_type={mamba_spec.mamba_type!r}. " + f"Mamba1 SSM temporal shape is (intermediate_size // tp, state_size) " + f"which cannot be used to reconstruct intermediate_size." + ) + + conv_shape = mamba_spec.shapes[0] + assert len(conv_shape) == 2, f"Expected 2D conv state shape, got {conv_shape}" + + # NOTE (ZhanqiuHu): 3-read requires DS layout, which is already asserted + # in nixl_connector __init__. Use it directly instead of heuristic detection. + assert is_conv_state_dim_first(), "3-read requires DS conv state layout" + local_conv_dim = conv_shape[0] # DS: (conv_dim_local, conv_rows) + conv_rows = conv_shape[1] + + # NOTE (ZhanqiuHu): intermediate_size (= global x dim) is not stored + # in MambaSpec, so we reconstruct it from the SSM temporal state shape: + # shapes[1] = (local_num_heads, head_dim), already divided by TP. + head_dim = mamba_spec.shapes[1][1] + local_num_heads = mamba_spec.shapes[1][0] + intermediate_size = local_num_heads * local_tp * head_dim + + # NOTE (ZhanqiuHu): global conv dim = intermediate_size + 2 * groups_ss, + # where groups_ss is the B (= C) dimension. B and C are always the same + # size, so we recover groups_ss from the remainder after subtracting x. + remainder = local_conv_dim * local_tp - intermediate_size + assert remainder > 0 and remainder % 2 == 0, ( + f"Conv dim ({local_conv_dim}*tp={local_tp}) doesn't decompose into " + f"intermediate_size={intermediate_size} + 2*groups_ss. " + f"remainder={remainder}" + ) + groups_ss = remainder // 2 + + conv_dtype_size = torch.tensor( + [], + dtype=mamba_spec.dtypes[0], # type: ignore[misc] + ).element_size() + + # Divide by TP to get per-rank column counts. + return MambaConvSplitInfo( + conv_rows=conv_rows, + x_local=intermediate_size // local_tp, + b_local=groups_ss // local_tp, + conv_dtype_size=conv_dtype_size, + ) + + +def compute_mamba_phys_ratio(ssm_sizes: tuple[int, ...], block_len: int) -> int: + """Derive _physical_blocks_per_logical_kv_block from remote metadata. + + The remote engine's ratio is not sent directly in the handshake, so we + reconstruct it: total mamba state per logical block / block_len. + + Args: + ssm_sizes: (conv_state_bytes, ssm_state_bytes) from NixlAgentMetadata. + block_len: the engine's block_len in bytes (from block_lens[0]). + """ + return math.ceil((ssm_sizes[0] + ssm_sizes[1]) / block_len) From f01482408c9e1f8a7e1647aab96d339ba3234cca Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Mon, 6 Apr 2026 13:17:23 -0400 Subject: [PATCH 17/39] [MoE Refactor][Test] FusedMoE layer test (#24675) Signed-off-by: Bill Nell Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .buildkite/test_areas/kernels.yaml | 18 + tests/kernels/moe/conftest.py | 14 + .../modular_kernel_tools/parallel_utils.py | 32 +- tests/kernels/moe/test_moe_layer.py | 1727 +++++++++++++++++ tests/kernels/moe/utils.py | 121 +- .../fused_moe/experts/trtllm_bf16_moe.py | 1 - 6 files changed, 1858 insertions(+), 55 deletions(-) create mode 100644 tests/kernels/moe/conftest.py create mode 100644 tests/kernels/moe/test_moe_layer.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 5fd081699d1..a05ee886f5b 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -180,3 +180,21 @@ steps: - pytest -v -s kernels/moe/test_flashinfer_moe.py - pytest -v -s kernels/moe/test_nvfp4_moe.py - pytest -v -s kernels/moe/test_ocp_mx_moe.py + + +- label: Kernels FusedMoE Layer Test (2 H100s) + timeout_in_minutes: 90 + device: h100 + num_devices: 2 + optional: true + commands: + - pytest -v -s kernels/moe/test_moe_layer.py + + +- label: Kernels FusedMoE Layer Test (2 B200s) + timeout_in_minutes: 90 + device: b200 + num_devices: 2 + optional: true + commands: + - pytest -v -s kernels/moe/test_moe_layer.py diff --git a/tests/kernels/moe/conftest.py b/tests/kernels/moe/conftest.py new file mode 100644 index 00000000000..a217fe684eb --- /dev/null +++ b/tests/kernels/moe/conftest.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + + +def pytest_addoption(parser): + parser.addoption( + "--subtests", action="store", type=str, default=None, help="subtest ids" + ) + + +@pytest.fixture +def subtests(request): + return request.config.getoption("--subtests") diff --git a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py index 3ff2ce3b3c0..10a226bcd97 100644 --- a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py +++ b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py @@ -11,7 +11,11 @@ from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUs from typing_extensions import ParamSpec from vllm.config import VllmConfig, set_current_vllm_config -from vllm.distributed import init_distributed_environment, initialize_model_parallel +from vllm.distributed import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, +) from vllm.utils.network_utils import get_open_port ## Parallel Processes Utils @@ -36,10 +40,17 @@ def _set_vllm_config( temp_file = tempfile.mkstemp()[1] + # When DP is enabled, processes are organized as: + # rank = dp_rank * tp_pp_world_size + tp_pp_rank + tp_pp_world_size = vllm_config.parallel_config.world_size + vllm_config.parallel_config.data_parallel_rank = rank // tp_pp_world_size + tp_pp_rank = rank % tp_pp_world_size + vllm_config.parallel_config.rank = tp_pp_rank + with set_current_vllm_config(vllm_config): init_distributed_environment( - world_size=world_size, - rank=rank, + world_size=tp_pp_world_size, + rank=tp_pp_rank, distributed_init_method=f"file://{temp_file}", local_rank=local_rank, backend="nccl", @@ -59,11 +70,11 @@ def _worker_parallel_launch( world_local_size: int, node_rank: int, init_method: str, - worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig | None, Any, P], None], + worker: Callable[..., None], vllm_config: VllmConfig | None, env_dict: dict | None, - *args: P.args, - **kwargs: P.kwargs, + worker_kwargs: dict[str, Any], + *args: Any, ) -> None: rank = node_rank * world_local_size + local_rank torch.accelerator.set_device_index(local_rank) @@ -98,14 +109,17 @@ def _worker_parallel_launch( vllm_config, cpu_group, *args, - **kwargs, + **worker_kwargs, ) except Exception as ex: print(ex) traceback.print_exc() raise finally: - torch.distributed.destroy_process_group() + if vllm_config is not None: + cleanup_dist_env_and_memory() + else: + torch.distributed.destroy_process_group() def parallel_launch_with_config( @@ -116,7 +130,6 @@ def parallel_launch_with_config( *args: P.args, **kwargs: P.kwargs, ) -> None: - assert not kwargs spawn( _worker_parallel_launch, args=( @@ -127,6 +140,7 @@ def parallel_launch_with_config( worker, vllm_config, env_dict, + kwargs, ) + args, nprocs=world_size, diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py new file mode 100644 index 00000000000..7b31edd3360 --- /dev/null +++ b/tests/kernels/moe/test_moe_layer.py @@ -0,0 +1,1727 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the MOE layer. + +Run `pytest tests/kernels/test_moe_layer.py`. +""" + +import functools +import os +import traceback +import types +from collections.abc import Callable +from dataclasses import astuple, dataclass, fields +from itertools import product +from typing import get_args + +import pytest +import torch + +from tests.kernels.moe.modular_kernel_tools.parallel_utils import ( + ProcessGroupInfo, + _set_vllm_config, + parallel_launch_with_config, +) +from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_weights +from vllm.config import ( + CompilationConfig, + ParallelConfig, + VllmConfig, + set_current_vllm_config, +) +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator +from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace +from vllm.distributed.parallel_state import ( + get_ep_group, + get_eplb_group, +) +from vllm.forward_context import set_forward_context +from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE, fused_experts +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.router.router_factory import ( + create_fused_moe_router, +) +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptFp8Config, + ModelOptNvFp4Config, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import ( + has_flashinfer_nvlink_one_sided, + has_flashinfer_nvlink_two_sided, +) +from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep +from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import set_random_seed +from vllm.v1.worker.workspace import ( + init_workspace_manager, + is_workspace_manager_initialized, +) + +fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype + +SHAPE_COMBOS = [ + (1, 128, 256), + (32, 1024, 512), + (222, 2048, 2048), # should be big enough to exercise DP chunking +] + +NUM_EXPERTS = [8, 64] +TOP_KS = [2, 6] + +# dp_size, tp_size, use_ep +# Note: DP+TP is not yet supported in the FusedMoE layer. +PARALLEL_COMBOS = [ + [1, 2, False], + [1, 4, False], + [2, 1, True], + [4, 1, True], +] + +# TODO: should this even be set manually? let oracles handle this +BACKENDS = ["allgather_reducescatter"] + +if has_mori(): + BACKENDS += ["mori"] + +if has_flashinfer_nvlink_two_sided(): + BACKENDS += ["flashinfer_nvlink_two_sided"] + +if has_flashinfer_nvlink_one_sided(): + BACKENDS += ["flashinfer_nvlink_one_sided"] + +if has_deep_ep(): + BACKENDS += ["deepep_low_latency", "deepep_high_throughput"] + +if has_nixl_ep(): + BACKENDS += ["nixl_ep"] + +QUANT_METHODS = [ + None, + "fp8", + "modelopt_fp8", + "modelopt_fp4", +] + +# Which quantization methods each backend supports. +# fmt: off +BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = { + "allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, + "mori": {None, "fp8", "modelopt_fp8"}, + "flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"}, + "flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"}, + "deepep_low_latency": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, + "deepep_high_throughput": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, + "nixl_ep": {None, "fp8", "modelopt_fp8"}, +} +# fmt: on + +# Which quantization methods support EPLB. +# ModelOptFp8MoEMethod inherits supports_eplb=False from FusedMoEMethodBase. +# TODO: double check modelopt fp8 +# modelopt_fp4 excluded: get_expert_weights() can't handle NvFP4 packed format. +EPLB_SUPPORTED_QUANTS: list[str | None] = [None, "fp8"] + +# Which backends support EPLB. +# deepep backends fail in get_expert_weights / rearrange_expert_weights_inplace. +# TODO(bnell): check this +EPLB_SUPPORTED_BACKENDS: list[str] = ["allgather_reducescatter"] + + +def maybe_roundup_layer_hidden_size( + hidden_size: int, + act_dtype: torch.dtype, + backend: str | None, +) -> int: + """ + Given layer hidden size and MoE configurations, round up hidden_size + if necessary. + + Args: + hidden_size: Layer hidden-size + act_dtype: Data type of the layer activations. + moe_parallel_config: Fused MoE parallelization strategy configuration. + + Return: + Rounded up hidden_size if rounding up is required based on the configs + and all2all backend. + Original hidden size otherwise. + """ + if backend == "deepep_high_throughput": + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import ( + DeepEPHTPrepareAndFinalize, + ) + + hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size( + hidden_size, act_dtype + ) + + if backend == "deepep_low_latency": + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( + DeepEPLLPrepareAndFinalize, + ) + + hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size( + hidden_size + ) + + return hidden_size + + +def rank_chunk(num: int, r: int, w: int) -> int: + rem = num % w + return (num // w) + (1 if r < rem else 0) + + +def chunk_by_rank( + t: torch.Tensor, + r: int, + w: int, + dim: int = 0, + device: torch.device | None = None, +) -> torch.Tensor: + chunk = cdiv(t.shape[dim], w) + t = t.narrow(dim, r * chunk, chunk) + if device is not None: + t = t.to(device) + return t + + +def maybe_chunk_by_rank( + t: torch.Tensor | None, + r: int, + w: int, + dim: int = 0, + device: torch.device | None = None, +) -> torch.Tensor | None: + if t is not None: + return chunk_by_rank(t, r, w, dim, device) + else: + return t + + +def tp_chunk_gate_up( + w: torch.Tensor, + tp_rank: int, + tp_size: int, + dim: int, + device: torch.device | int | None = None, +) -> torch.Tensor: + """TP-chunk a combined [gate; up] weight, splitting each half separately + so every rank gets a portion of both gate and up.""" + half = w.shape[dim] // 2 + gate = chunk_by_rank( + w.narrow(dim, 0, half), tp_rank, tp_size, dim=dim, device=device + ) + up = chunk_by_rank( + w.narrow(dim, half, half), tp_rank, tp_size, dim=dim, device=device + ) + return torch.cat([gate, up], dim=dim) + + +@dataclass +class MoETestConfig: + m: int + n: int + k: int + num_experts: int + top_k: int + in_dtype: torch.dtype + quantization: str | None + use_shared_experts: bool + use_gate: bool + use_routed_input_transform: bool + enable_eplb: bool = False + reduce_results: bool = False + backend: str | None = None + ep_size: int = 1 + dp_size: int = 1 + tp_size: int = 1 + + # TODO: add more error messages + def id(self) -> str: + def proc(s: str) -> str: + return s.removeprefix("torch.") + + id_str = "-".join([proc(str(f)) for f in astuple(self)]) + return f"[{id_str}]" + + # TODO: add more error messages + @staticmethod + def from_id(id: str) -> "MoETestConfig": + id = id[1:-1] + str_values = id.split("-") + + def convert(v: str, ty): + if isinstance(ty, types.UnionType): + sub_ty = list(get_args(ty)) + assert len(sub_ty) == 2 and types.NoneType in sub_ty + sub_ty.remove(types.NoneType) + return sub_ty[0](v) if v != "None" else None + elif ty is torch.dtype: + ty_val = getattr(torch, v, None) + assert isinstance(ty_val, torch.dtype) + return ty_val + elif ty is bool: + return v == "True" + else: + return ty(v) + + values = tuple( + [convert(v, f.type) for v, f in zip(str_values, fields(MoETestConfig))] + ) + return MoETestConfig(*values) + + +def generate_valid_test_configs( + backend: str, + ep_size: int, + dp_size: int, + tp_size: int, + enable_eplb: bool, + verbosity: int = 0, +) -> list[MoETestConfig]: + configs: list[MoETestConfig] = [] + + for ( + shape, + num_experts, + top_k, + quantization, + use_shared_experts, + use_gate, + use_routed_input_transform, + reduce_results, + ) in product( + SHAPE_COMBOS, + NUM_EXPERTS, + TOP_KS, + QUANT_METHODS, + [False, True], # shared + [False, True], # gate + [False, True], # routed input exform + [False, True], # reduce results + ): + config = MoETestConfig( + shape[0], # m + shape[1], # n + shape[2], # k + num_experts, + top_k, + torch.bfloat16, + quantization, + use_shared_experts, + use_gate, + use_routed_input_transform, + enable_eplb, + reduce_results, + backend, + ep_size, + dp_size, + tp_size, + ) + + valid, reason = is_valid_config(config) + if valid: + configs.append(config) + elif verbosity > 1: + print(f"Skipping invalid config {config} - {reason}") + + return configs + + +# TODO: break this up into sections +def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: + # routed_input_transform only makes sense with shared_experts (latent MoE) + # TODO: not sure this is true + if config.use_routed_input_transform and not config.use_shared_experts: + return False, "routed_input_transform requires shared_experts" + + # TODO: disable for now + if config.use_routed_input_transform and config.enable_eplb: + return False, "routed_input_transform not supported with EPLB." + + # TODO: disable for now + if config.use_routed_input_transform and config.use_gate: + return ( + False, + "routed_input_transform not supported with gate because of " + "padding problems", + ) + + # TODO: disable for now + if config.use_routed_input_transform and config.backend in [ + "deepep_low_latency", + "deepep_high_throughput", + ]: + return ( + False, + "routed_input_transform not supported with DeepEP backends because " + "of padding problems", + ) + + # routed_input_transform + quantization + high hidden dimensions + # TODO: Disable >= 2048 w/fp8 + deepep LL for now due to insane errors. + if ( + (config.use_routed_input_transform or config.backend == "deepep_low_latency") + and config.quantization is not None + and config.k >= 2048 + ): + return ( + False, + "routed_input_transform + quantization + higher hidden dimensions " + "leads to large differences.", + ) + + # gate requires shared_experts (use_overlapped mode) + # TODO: also not sure this is true + if config.use_gate and not config.use_shared_experts: + return False, "gate requires shared_experts (use_overlapped mode)" + + # Skip modelopt_fp4 if not on B100+ (compute capability 10.0+) + if ( + config.quantization == "modelopt_fp4" + and not current_platform.has_device_capability(100) + ): + return False, "modelopt_fp4 not supported on H100+ GPUs" + + # Skip flashinfer_nvlink if not on H100+ (compute capability 10.0+) + if ( + config.backend is not None + and config.backend.startswith("flashinfer_nvlink") + and not current_platform.has_device_capability(90) + ): + return False, "flashinfer_nvlink needs an H100+ GPUs" + + # reduce_results incompatibilities + if config.reduce_results and config.use_shared_experts: + return False, "reduce_results=True is not compatible with shared_experts=True" + + if config.reduce_results and config.quantization is not None: + return ( + False, + "reduce_results=True only tested with unquantized data types in " + "order to limit number of tests run", + ) + + # Backend-specific checks + if config.backend is not None: + supported_quants = BACKEND_SUPPORTED_QUANTS.get(config.backend) + if supported_quants is not None and config.quantization not in supported_quants: + return ( + False, + f"{config.backend} does not support quantization={config.quantization}", + ) + + if config.backend == "deepep_low_latency": + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( # noqa: E501 + DeepEPLLPrepareAndFinalize, + ) + + if config.k not in DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES: + return ( + False, + f"Skipping unsupported K {config.k} in {config.backend} w/o EP.", + ) + + if config.backend == "nixl_ep": + from vllm.model_executor.layers.fused_moe.nixl_ep_prepare_finalize import ( # noqa: E501 + NixlEPPrepareAndFinalize, + ) + + if config.k not in NixlEPPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES: + return ( + False, + f"Skipping unsupported K {config.k} in {config.backend} w/o EP.", + ) + + if config.enable_eplb and config.ep_size == 1: + return False, "EPLB requires EP." + + if config.enable_eplb and config.quantization not in EPLB_SUPPORTED_QUANTS: + return False, f"EPLB not supported with {config.quantization} quantization." + + if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS: + return False, f"EPLB not supported with {config.backend}." + + world_size = config.tp_size * config.dp_size + if config.reduce_results and world_size == 1: + return False, "reduce_results=True only makes sense for multi-GPU tests" + + if ( + config.backend is not None + and config.backend.startswith("flashinfer_nvlink") + and config.ep_size > 1 + ): + return False, "flashinfer_nvlink EP not yet supported." + + if config.enable_eplb and config.num_experts % config.dp_size != 0: + return False, "EPLB requires num_experts divisible by ep_size" + + if config.enable_eplb and config.ep_size == 1: + return False, "EPLB only works with EP+DP" + + return True, None + + +def chunk_scales_by_rank( + t: torch.Tensor | None, + r: int, + w: int, + device: torch.device | None = None, +) -> torch.Tensor | None: + if t is not None and t.numel() > 1: + # Calculate start index by summing chunk sizes for all previous ranks + # start = sum(rank_chunk(t.shape[0], i, w) for i in range(r)) + # chunk = rank_chunk(t.shape[0], r, w) + # t = t[start:(start + chunk)] + chunk = rank_chunk(t.shape[0], r, w) + t = t[(r * chunk) : max(t.shape[0], (r + 1) * chunk)] + + if t is not None and device is not None: + t = t.to(device) + + return t + + +def chunk_scales( + t: torch.Tensor | None, + start: int, + end: int, + device: torch.device | None = None, +) -> torch.Tensor | None: + if t is not None and t.numel() > 1: + t = t[start:end] + + if t is not None and device is not None: + t = t.to(device) + + return t + + +@dataclass +class QuantizedWeights: + w13_weight: torch.Tensor + w2_weight: torch.Tensor + w13_weight_scale: torch.Tensor | None = None + w2_weight_scale: torch.Tensor | None = None + w13_weight_scale_2: torch.Tensor | None = None + w2_weight_scale_2: torch.Tensor | None = None + w13_input_scale: torch.Tensor | None = None + w2_input_scale: torch.Tensor | None = None + + +def _quantize_fp8_halves( + w1: torch.Tensor, + w2: torch.Tensor, +) -> QuantizedWeights: + """Quantize w13 gate/up halves separately to FP8, producing per-shard scales.""" + half = w1.shape[1] // 2 + w1q_a, w1s_a, _ = moe_quantize_weights( + w1[:, :half, :], None, fp8_dtype, False, None + ) + w1q_b, w1s_b, _ = moe_quantize_weights( + w1[:, half:, :], None, fp8_dtype, False, None + ) + assert w1s_a is not None and w1s_b is not None + + w2q, w2s, _ = moe_quantize_weights(w2, None, fp8_dtype, False, None) + assert w2s is not None + + return QuantizedWeights( + w13_weight=torch.cat([w1q_a, w1q_b], dim=1), + w2_weight=w2q, + # Each w1s_x is (E, 1, 1) -> reshape to (E, 1), cat to (E, 2) + w13_weight_scale=torch.cat([w1s_a.view(-1, 1), w1s_b.view(-1, 1)], dim=1), + # w2s is (E, 1, 1) -> reshape to (E,) + w2_weight_scale=w2s.view(-1), + ) + + +def quantization_to_quant_dtype( + quantization: str | None, +) -> torch.dtype | str | None: + if quantization is None: + return None + elif quantization in ["fp8", "modelopt_fp8"]: + return fp8_dtype + elif quantization in ["modelopt_fp4"]: + return "nvfp4" + else: + raise NotImplementedError(f"Unsupported quantization: {quantization}") + + +def make_quant_config( + quantization: str | None, + w1: torch.Tensor, + w2: torch.Tensor, + num_experts: int, +) -> tuple[QuantizationConfig | None, QuantizedWeights]: + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + if quantization is None: + return None, QuantizedWeights(w13_weight=w1, w2_weight=w2) + + if quantization == "fp8": + return Fp8Config(True), _quantize_fp8_halves(w1, w2) + + if quantization == "modelopt_fp8": + qw = _quantize_fp8_halves(w1, w2) + # why? + qw.w13_input_scale = torch.ones( + num_experts, dtype=torch.float32, device=w1.device + ) + # why? + qw.w2_input_scale = torch.ones( + num_experts, dtype=torch.float32, device=w2.device + ) + quant_config = ModelOptFp8Config( + quant_method="FP8", + is_checkpoint_fp8_serialized=True, + kv_cache_quant_method=None, + exclude_modules=[], + ) + return quant_config, qw + + if quantization == "modelopt_fp4": + # Quantize full w13 at once so both gate/up halves share the same + # global scale per expert. process_weights_after_loading uses + # w13_weight_scale_2[:, 0] for the entire tensor, so the two shard + # scales must match. + w1q, w1s, w1gs = moe_quantize_weights(w1, None, "nvfp4", False, None) + assert w1s is not None and w1gs is not None + + w2q, w2s, w2gs = moe_quantize_weights(w2, None, "nvfp4", False, None) + assert w2s is not None and w2gs is not None + + qw = QuantizedWeights( + w13_weight=w1q, + w2_weight=w2q, + w13_weight_scale=w1s, + w2_weight_scale=w2s, + # weight_scale_2 = 1/w_gs: the kernel computes + # g_alphas = a_scale * w_scale_2, and correct dequant needs 1/w_gs. + # Expand per-expert scalar to (E, 2) for the two shards. + w13_weight_scale_2=(1.0 / w1gs).unsqueeze(1).expand(-1, 2).contiguous(), + w2_weight_scale_2=1.0 / w2gs, + w13_input_scale=torch.ones( + (num_experts, 2), dtype=torch.float32, device=w1.device + ), + w2_input_scale=torch.ones( + num_experts, dtype=torch.float32, device=w2.device + ), + ) + quant_config = ModelOptNvFp4Config( + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + ) + return quant_config, qw + + raise NotImplementedError(f"Unsupported quantization: {quantization}") + + +@dataclass +class SharedExpertsConfig: + w1: torch.Tensor + w2: torch.Tensor + w1_s: torch.Tensor | None = None + w2_s: torch.Tensor | None = None + quant_dtype: torch.dtype | str | None = None + + +@dataclass +class MoETestData: + """Container for MOE test data and transforms.""" + + w1: torch.Tensor + w2: torch.Tensor + hidden_states: torch.Tensor + router_logits: torch.Tensor + shared_experts_config: SharedExpertsConfig | None + gate: torch.nn.Module | None + routed_input_transform: torch.nn.Module | None + routed_output_transform: torch.nn.Module | None + routed_expert_hidden_size: int + + +class SimpleGate(torch.nn.Module): + """Simple gate module for testing: computes router logits from hidden states.""" + + def __init__( + self, + hidden_size: int, + num_experts: int, + dtype: torch.dtype, + device: str = "cuda", + ): + super().__init__() + self.weight = torch.nn.Parameter( + torch.randn(num_experts, hidden_size, device=device, dtype=dtype) / 10 + ) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: + """Returns (router_logits, None) to match expected signature.""" + router_logits = torch.nn.functional.linear(hidden_states, self.weight) + return router_logits, None + + +class SimpleRoutedInputTransform(torch.nn.Module): + """Simple linear transform for testing routed input transform + (e.g., latent projection). + """ + + def __init__( + self, + in_features: int, + out_features: int, + dtype: torch.dtype, + device: str = "cuda", + ): + super().__init__() + self.weight = torch.nn.Parameter( + torch.randn(out_features, in_features, device=device, dtype=dtype) / 10 + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.linear(x, self.weight) + + +def create_shared_experts_from_config( + shared_experts_config: SharedExpertsConfig | None, + in_dtype: torch.dtype, + tp_size: int = 1, + tp_rank: int = 0, + device: torch.device | str | None = None, +) -> TestMLP | None: + """Create TestMLP for shared experts from config. + + Args: + shared_experts_config: Configuration for shared experts + in_dtype: Output data type + tp_size: Tensor parallel size (for weight chunking) + tp_rank: Tensor parallel rank (for weight chunking) + device: Device to move weights to (optional) + + Returns: + TestMLP instance or None if config is None + """ + if shared_experts_config is None: + return None + + s_w1 = shared_experts_config.w1 + s_w2 = shared_experts_config.w2 + + # Apply TP chunking if needed + if tp_size > 1: + s_w1 = tp_chunk_gate_up(s_w1, tp_rank, tp_size, dim=1, device=device) + s_w2 = chunk_by_rank(s_w2, tp_rank, tp_size, dim=0, device=device) + else: + s_w1 = s_w1.to(device) + s_w2 = s_w2.to(device) + + return TestMLP(w1=s_w1, w2=s_w2, out_dtype=in_dtype) + + +# Make version that takes a MoETestConfig? +def setup_moe_test_data( + m: int, + k: int, + n: int, + num_experts: int, + in_dtype: torch.dtype, + use_shared_experts: bool, + use_gate: bool, + use_routed_input_transform: bool, + backend: str | None, + device: str = "cuda", +) -> MoETestData: + """Setup test data and transforms for MOE tests. + + Args: + m: Number of tokens + k: Hidden size + n: Intermediate size + num_experts: Number of experts + in_dtype: Data type for tensors + use_shared_experts: Whether to create shared experts config + use_gate: Whether to create gate module + use_routed_input_transform: Whether to create routed input/output transforms + device: Device to create tensors on ("cuda" or "cpu") + + Returns: + MoETestData containing all test data and transforms + """ + # For latent MoE: latent_size = k // 2 + latent_size = k // 2 + + # k = maybe_roundup_layer_hidden_size(k, in_dtype, backend) + # latent_size = maybe_roundup_layer_hidden_size(latent_size, in_dtype, backend) + + # Determine dimensions for routed experts (may be transformed) + # For latent MoE, routed experts operate entirely in latent space + # (k//2). The routed_output_transform then projects back to k before + # adding with shared experts. + # w1: (E, 2*N, latent_size) - input latent_size + # w2: (E, latent_size, N) - output latent_size (fused_experts returns + # same shape as input) + routed_expert_hidden_size = latent_size if use_routed_input_transform else k + + # Create expert weights + (w1, _, _, _), (w2, _, _, _) = make_test_weights( + num_experts, + n, + routed_expert_hidden_size, # Both w1 input and w2 output use latent_size + in_dtype=in_dtype, + ) + + # Create shared experts config if needed + if use_shared_experts: + shared_experts_config = SharedExpertsConfig( + w1=torch.randn((k, n * 2), device=device, dtype=in_dtype) / 15, + w2=torch.randn((n, k), device=device, dtype=in_dtype) / 15, + ) + else: + shared_experts_config = None + + # Create routed input transform if needed + routed_input_transform = ( + SimpleRoutedInputTransform(k, latent_size, in_dtype, device=device) + if use_routed_input_transform + else None + ) + + # Create gate if needed + # Note: gate is called AFTER routed_input_transform, so it should expect + # the transformed dimension (latent_size) when routed_input_transform is used + gate_input_dim = latent_size if use_routed_input_transform else k + gate = ( + SimpleGate(gate_input_dim, num_experts, in_dtype, device=device) + if use_gate + else None + ) + + # Create routed output transform if needed (projects latent space back to original) + routed_output_transform = ( + SimpleRoutedInputTransform(latent_size, k, in_dtype, device=device) + if use_routed_input_transform + else None + ) + + # Create test inputs + hidden_states = torch.randn((m, k), device=device, dtype=in_dtype) / 10 + router_logits = torch.randn((m, num_experts), device=device, dtype=in_dtype) + + return MoETestData( + w1=w1, + w2=w2, + hidden_states=hidden_states, + router_logits=router_logits, + shared_experts_config=shared_experts_config, + gate=gate, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + routed_expert_hidden_size=routed_expert_hidden_size, + ) + + +def make_fused_moe_layer( + quantization: str | None, + use_ep: bool, + hidden_size: int, + intermediate_size: int, + in_dtype: torch.dtype, + tp_size: int, + ep_size: int, + dp_size: int, + reduce_results: bool, + w1: torch.Tensor, + w2: torch.Tensor, + top_k: int, + global_num_experts: int, + renormalize: bool = False, + shared_experts: torch.nn.Module | None = None, + use_grouped_topk: bool = False, + topk_group: int | None = None, + num_expert_group: int | None = None, + custom_routing_function: Callable | None = None, + scoring_func: str = "softmax", + routed_scaling_factor: float = 1.0, + e_score_correction_bias: torch.Tensor | None = None, + apply_router_weight_on_input: bool = False, + activation: str = "silu", + indices_type: torch.dtype | None = None, + expert_map: torch.Tensor | None = None, + enable_eplb: bool = False, + expert_load_view: torch.Tensor | None = None, + logical_to_physical_map: torch.Tensor | None = None, + logical_replica_count: torch.Tensor | None = None, + num_redundant_experts: int = 0, + has_bias: bool = False, + gate: torch.nn.Module | None = None, + routed_input_transform: torch.nn.Module | None = None, + routed_output_transform: torch.nn.Module | None = None, + pcp_size: int | None = 1, +) -> tuple[Callable, FusedMoE]: + quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts) + + kwargs = dict() + if shared_experts is None: + builder = FusedMoE + else: + builder = SharedFusedMoE + kwargs["shared_experts"] = shared_experts + + # Add gate and routed_input_transform if provided + if gate is not None: + kwargs["gate"] = gate + if routed_input_transform is not None: + kwargs["routed_input_transform"] = routed_input_transform + + layer = builder( + num_experts=global_num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + params_dtype=in_dtype, + reduce_results=reduce_results, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + quant_config=quant_config, + tp_size=tp_size, + ep_size=ep_size, + dp_size=dp_size, + pcp_size=pcp_size, + prefix="from_forward_context", + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + apply_router_weight_on_input=apply_router_weight_on_input, + activation=activation, + enable_eplb=enable_eplb, + num_redundant_experts=num_redundant_experts, + has_bias=has_bias, + **kwargs, + ) + + for name, value in [ + ("w13_weight", qw.w13_weight), + ("w2_weight", qw.w2_weight), + ("w13_weight_scale", qw.w13_weight_scale), + ("w2_weight_scale", qw.w2_weight_scale), + ("w13_weight_scale_2", qw.w13_weight_scale_2), + ("w2_weight_scale_2", qw.w2_weight_scale_2), + ("w13_input_scale", qw.w13_input_scale), + ("w2_input_scale", qw.w2_input_scale), + ]: + if value is not None: + layer.register_parameter( + name, torch.nn.Parameter(value, requires_grad=False) + ) + + layer.quant_method.process_weights_after_loading(layer) + + def _moe( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + if shared_experts is None: + final_shared_states = None + final_hidden_states = layer(hidden_states, router_logits) + else: + final_shared_states, final_hidden_states = layer( + hidden_states, router_logits + ) + + # Apply routed output transform if provided + # (e.g., latent space -> original space) + if routed_output_transform is not None: + final_hidden_states = routed_output_transform(final_hidden_states) + + if shared_experts is not None: + assert not reduce_results + assert final_shared_states is not None + final_hidden_states += final_shared_states + + if not reduce_results and layer.tp_size > 1: + final_hidden_states = layer.maybe_all_reduce_tensor_model_parallel( + final_hidden_states + ) + + return final_hidden_states + + return _moe, layer + + +def make_fake_moe_layer( + w1: torch.Tensor, + w2: torch.Tensor, + top_k: int, + global_num_experts: int, + in_dtype: torch.dtype, + quant_dtype: torch.dtype | None, + renormalize: bool = False, + shared_experts_config: SharedExpertsConfig | None = None, + use_grouped_topk: bool = False, + topk_group: int | None = None, + num_expert_group: int | None = None, + custom_routing_function: Callable | None = None, + scoring_func: str = "softmax", + routed_scaling_factor: float = 1.0, + e_score_correction_bias: torch.Tensor | None = None, + apply_router_weight_on_input: bool = False, + activation: str = "silu", + indices_type: torch.dtype | None = None, + expert_map: torch.Tensor | None = None, + enable_eplb: bool = False, + expert_load_view: torch.Tensor | None = None, + logical_to_physical_map: torch.Tensor | None = None, + logical_replica_count: torch.Tensor | None = None, + gate: torch.nn.Module | None = None, + routed_input_transform: torch.nn.Module | None = None, + routed_output_transform: torch.nn.Module | None = None, + use_ep: bool = False, + tp_size: int = 1, + dp_size: int = 1, + ep_size: int = 1, + reduce_results: bool = False, +) -> Callable: + activation = MoEActivation.from_str(activation) + + router = create_fused_moe_router( + top_k=top_k, + global_num_experts=global_num_experts, + # eplb_state=None, # TODO + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + custom_routing_function=custom_routing_function, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + num_fused_shared_experts=0, # TODO + enable_eplb=enable_eplb, + # TODO(bnell): once we can construct the MK at init time, we + # can make this a value. + indices_type_getter=lambda: indices_type, + ) + + if quant_dtype is not None: + w1, w1_s, _ = moe_quantize_weights(w1, None, quant_dtype, False, None) + w2, w2_s, _ = moe_quantize_weights(w2, None, quant_dtype, False, None) + else: + w1_s = None + w2_s = None + + shared_experts = create_shared_experts_from_config( + shared_experts_config, in_dtype, 1, 0, "cuda" + ) + + quant_config = FusedMoEQuantConfig.make( + quant_dtype, + w1_scale=w1_s, + w2_scale=w2_s, + ) + + def _moe( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + # Save original hidden_states for shared experts (before transform) + original_hidden_states = hidden_states + + # Apply routed input transform if provided + if routed_input_transform is not None: + hidden_states = routed_input_transform(hidden_states) + + # If gate provided, compute router_logits from hidden_states + # Note: gate operates on transformed hidden_states (after + # routed_input_transform) + if gate is not None: + router_logits, _ = gate(hidden_states) + + topk_weights, topk_ids = router.select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + ) + + # Shared experts use original (untransformed) hidden_states + if shared_experts is not None: + shared_output = shared_experts(original_hidden_states) + else: + shared_output = None + + # Routed experts use transformed hidden_states + output = fused_experts( + hidden_states=hidden_states, + w1=w1, + w2=w2, + quant_config=quant_config, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=False, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + global_num_experts=global_num_experts, + expert_map=expert_map, + ) + + # Apply routed output transform if provided + # (e.g., latent space -> original space) + if routed_output_transform is not None: + output = routed_output_transform(output) + + if shared_experts is not None: + assert shared_output is not None + output += shared_output + + # Apply TP/DP reduction if not already reduced + # if (tp_size > 1 or dp_size > 1): + # output = tensor_model_parallel_all_reduce(output) + + return output + + return _moe + + +def _test_body_regular( + moe_fn: Callable, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + vllm_config: VllmConfig, + num_tokens: int, + num_tokens_across_dp: torch.Tensor, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + """Regular MoE test body: compare layer output to baseline.""" + baseline_output = kwargs["baseline_output"] + + with set_forward_context( + None, + vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + ): + output = moe_fn(hidden_states, router_logits) + + return baseline_output, output + + +def _test_body_eplb( + moe_fn: Callable, + moe_layer: FusedMoE, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + vllm_config: VllmConfig, + num_tokens: int, + num_tokens_across_dp: torch.Tensor, + cpu_group, + in_dtype: torch.dtype, + quantization: str | None, + use_ep: bool, + tp_size: int, + ep_size: int, + dp_size: int, + w1: torch.Tensor, + w2: torch.Tensor, + num_experts: int, + k: int, + n: int, + top_k: int, + shared_experts, + reduce_results: bool, + gate: torch.nn.Module | None, + routed_input_transform: torch.nn.Module | None, + routed_output_transform: torch.nn.Module | None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + device = torch.accelerator.current_accelerator() + + """EPLB test body: compare output before and after expert weight rearrangement.""" + # Get "before" output with original weight arrangement + with set_forward_context( + None, + vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + ): + output_before = moe_fn(hidden_states, router_logits) + + # Create a fresh FusedMoE layer with enable_eplb=True + # Delete the original layer's registration so the constructor can + # re-use the same "from_forward_context" prefix + cc = vllm_config.compilation_config + del cc.static_forward_context["from_forward_context"] + cc.static_all_moe_layers.remove("from_forward_context") + + # Determine hidden size for MoE layer + # When using routed_input_transform, experts operate in latent space + hidden_size_for_layer = k // 2 if routed_input_transform is not None else k + + moe_fn, moe_layer = make_fused_moe_layer( + quantization=quantization, + use_ep=use_ep, + hidden_size=hidden_size_for_layer, + intermediate_size=n, + in_dtype=in_dtype, + tp_size=tp_size, + ep_size=ep_size, + dp_size=dp_size, + reduce_results=reduce_results, + w1=w1, + w2=w2, + top_k=top_k, + global_num_experts=num_experts, + shared_experts=shared_experts, + enable_eplb=True, + gate=gate, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + ) + + # Necessary? + if moe_layer._expert_map is not None: + moe_layer._expert_map = moe_layer._expert_map.to(device) + + # All ranks must generate the same permutation + initial_indices = torch.arange(num_experts, dtype=torch.long) + shuffled_indices = initial_indices[torch.randperm(num_experts)] + + expert_weights = [list(moe_layer.get_expert_weights())] + + communicator = create_eplb_communicator( + group_coordinator=get_eplb_group(), + backend=vllm_config.parallel_config.eplb_config.communicator, + expert_weights=expert_weights[0], + ) + + # Rearrange expert weights across EP ranks + rearrange_expert_weights_inplace( + old_global_expert_indices=initial_indices.unsqueeze(0), + new_global_expert_indices=shuffled_indices.unsqueeze(0), + expert_weights=expert_weights, + ep_group=cpu_group, + communicator=communicator, + ) + + # Build logical_to_physical_map from shuffled_indices + # shuffled_indices[physical] = logical, we need the inverse + logical_to_physical = torch.empty(num_experts, dtype=torch.int32, device=device) + logical_to_physical[shuffled_indices.to(device)] = torch.arange( + num_experts, dtype=torch.int32, device=device + ) + + moe_layer.set_eplb_state( + moe_layer_idx=0, + expert_load_view=torch.zeros( + (1, num_experts), + dtype=torch.int32, + device=device, + ), + logical_to_physical_map=logical_to_physical.reshape(num_experts, 1).unsqueeze( + 0 + ), + logical_replica_count=torch.ones( + (1, num_experts), + dtype=torch.int32, + device=device, + ), + ) + + moe_layer.eplb_state.should_record_tensor = torch.ones( + (), dtype=torch.bool, device=device + ) + + # Get "after" output with rearranged weights and EPLB routing + with set_forward_context( + None, + vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + ): + output_after = moe_fn(hidden_states, router_logits) + + return output_before, output_after + + +# TODO: make this take a MoETestConfig +def _run_one_config( + vllm_config: VllmConfig, + ep_size: int, + dp_size: int, + tp_size: int, + dp_rank: int, + tp_rank: int, + m: int, + n: int, + k: int, + num_experts: int, + top_k: int, + quantization: str | None, + reduce_results: bool, + backend: str | None, + test_body_fn: Callable, + use_shared_experts: bool, + use_gate: bool, + use_routed_input_transform: bool, + **kwargs, +) -> None: + set_random_seed(7) + + """Generic test loop that sets up environment and delegates to test_body_fn. + + This function is called directly by test_moe_layer and test_moe_layer_eplb + via parallel_launch_with_config, passing either _test_body_regular or + _test_body_eplb as the test_body_fn parameter. + """ + world_size = tp_size * dp_size + use_ep = ep_size > 1 + + assert vllm_config.parallel_config.enable_expert_parallel == use_ep + + in_dtype = torch.bfloat16 + device = torch.accelerator.current_accelerator() + + if not is_workspace_manager_initialized(): + init_workspace_manager(device) + + # Create test data and transforms + test_data = setup_moe_test_data( + m=m, + k=k, + n=n, + num_experts=num_experts, + in_dtype=in_dtype, + use_shared_experts=use_shared_experts, + use_gate=use_gate, + use_routed_input_transform=use_routed_input_transform, + backend=backend, + device=device, + ) + + # Extract data from test_data + hidden_states = test_data.hidden_states + router_logits = test_data.router_logits + w1 = test_data.w1 + w2 = test_data.w2 + shared_experts_config = test_data.shared_experts_config + gate = test_data.gate + routed_input_transform = test_data.routed_input_transform + routed_output_transform = test_data.routed_output_transform + + baseline_layer = make_fake_moe_layer( + w1=w1, + w2=w2, + top_k=top_k, + global_num_experts=num_experts, + in_dtype=in_dtype, + quant_dtype=None, # quantization_to_quant_dtype(quantization), + renormalize=False, + shared_experts_config=shared_experts_config, + gate=gate, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + use_ep=use_ep, + tp_size=tp_size, + ep_size=ep_size, + dp_size=dp_size, + reduce_results=reduce_results, + ) + + baseline_output = baseline_layer(hidden_states, router_logits) + + del baseline_layer + torch.accelerator.empty_cache() + + with set_current_vllm_config(vllm_config): + # Chunk weights for EP/TP (after baseline is created) + if ep_size > 1: + w1 = chunk_by_rank(w1, dp_rank, dp_size, dim=0, device=device) + w2 = chunk_by_rank(w2, dp_rank, dp_size, dim=0, device=device) + + if tp_size > 1: + w1 = tp_chunk_gate_up(w1, tp_rank, tp_size, dim=1, device=device) + w2 = chunk_by_rank(w2, tp_rank, tp_size, dim=2, device=device) + + # Setup shared experts if needed + shared_experts = create_shared_experts_from_config( + shared_experts_config, in_dtype, tp_size, tp_rank, device + ) + + # Determine hidden size for MoE layer + # When using routed_input_transform, experts operate in latent space + hidden_size_for_layer = k // 2 if routed_input_transform is not None else k + + # Create initial MoE layer + moe_fn, moe_layer = make_fused_moe_layer( + quantization=quantization, + use_ep=use_ep, + hidden_size=hidden_size_for_layer, + intermediate_size=n, + in_dtype=in_dtype, + tp_size=tp_size, + ep_size=ep_size, + dp_size=dp_size, + reduce_results=reduce_results, + w1=w1, + w2=w2, + top_k=top_k, + global_num_experts=num_experts, + shared_experts=shared_experts, + gate=gate, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + ) + + # Necessary? + if moe_layer._expert_map is not None: + moe_layer._expert_map = moe_layer._expert_map.to(device) + + num_tokens = m + num_tokens_across_dp = torch.tensor( + [num_tokens] * world_size, + device=device, + dtype=torch.int, + ) + + # Call the test body function with all necessary context + expected, actual = test_body_fn( + moe_fn=moe_fn, + moe_layer=moe_layer, + hidden_states=hidden_states, + router_logits=router_logits, + vllm_config=vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp, + in_dtype=in_dtype, + quantization=quantization, + use_ep=use_ep, + tp_size=tp_size, + ep_size=ep_size, + dp_size=dp_size, + w1=w1, + w2=w2, + num_experts=num_experts, + k=k, + n=n, + m=m, + top_k=top_k, + shared_experts=shared_experts, + reduce_results=reduce_results, + gate=gate, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + baseline_output=baseline_output, + **kwargs, + ) + + # Common tolerance logic + # TODO: consider associating tolerances with quant methods. + if quantization is None: + if k >= 2048: + atol, rtol = 7.6e-2, 7.6e-2 + else: + atol, rtol = 3.5e-2, 3.5e-2 + elif quantization in ("fp8", "modelopt_fp8"): + if k >= 2048: + atol, rtol = 7.6e-2, 7.6e-2 + else: + atol, rtol = 6e-2, 6e-2 + elif quantization == "modelopt_fp4": + atol = rtol = 1e-1 + k * 5e-4 + else: + atol, rtol = 6e-2, 6e-2 + + torch.accelerator.synchronize() # TODO: Is this needed? + torch.testing.assert_close(expected, actual, atol=atol, rtol=rtol) + + +# Test for non-parallel cases (world_size == 1) - backend doesn't matter +@pytest.mark.parametrize("m, n, k", SHAPE_COMBOS) +@pytest.mark.parametrize("num_experts", NUM_EXPERTS) +@pytest.mark.parametrize("top_k", TOP_KS) +@pytest.mark.parametrize("quantization", QUANT_METHODS) +@pytest.mark.parametrize("use_shared_experts", [False, True]) +@pytest.mark.parametrize("use_gate", [False, True]) +@pytest.mark.parametrize("use_routed_input_transform", [False, True]) +def test_moe_layer_no_parallel( + m: int, + n: int, + k: int, + num_experts: int, + top_k: int, + quantization: str | None, + use_shared_experts: bool, + use_gate: bool, + use_routed_input_transform: bool, + monkeypatch, +): + """Test MoE layer without parallelism (dp_size=1, tp_size=1, use_ep=False).""" + + if os.environ.get("VLLM_LOGGING_LEVEL") is None: + monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR") + + test_config = MoETestConfig( + m, + n, + k, + num_experts, + top_k, + torch.bfloat16, + quantization, + use_shared_experts, + use_gate, + use_routed_input_transform, + ) + + valid, reason = is_valid_config(test_config) + if not valid: + pytest.skip(reason) + + set_random_seed(7) + + parallel_config = ParallelConfig() + compilation_config = CompilationConfig() + compilation_config.pass_config.fuse_allreduce_rms = False + + vllm_config = VllmConfig( + parallel_config=parallel_config, compilation_config=compilation_config + ) + + # Initialize distributed environment for single GPU + _set_vllm_config(vllm_config, 1, rank=0, local_rank=0) + + _run_one_config( + vllm_config, + test_config.ep_size, + test_config.dp_size, + test_config.tp_size, + 0, + 0, + test_config.m, + test_config.n, + test_config.k, + test_config.num_experts, + test_config.top_k, + test_config.quantization, + test_config.reduce_results, + test_config.backend, + _test_body_regular, + use_shared_experts=test_config.use_shared_experts, + use_gate=test_config.use_gate, + use_routed_input_transform=test_config.use_routed_input_transform, + ) + + +def _test_body_config(test_config: MoETestConfig, cpu_group, **kwargs): + if not test_config.enable_eplb: + return _test_body_regular(**kwargs) + else: + return _test_body_eplb(**kwargs, cpu_group=cpu_group) + + +def _parallel_worker( + pgi: ProcessGroupInfo, + vllm_config: VllmConfig, + cpu_group, + test_configs: list[MoETestConfig], + verbosity: int, + **kwargs, +) -> None: + set_random_seed(7) + + total = 0 + passed = 0 + failed = 0 + fail_ids = [] + + dp_rank = vllm_config.parallel_config.data_parallel_rank + + for test_config in test_configs: + cc = vllm_config.compilation_config + if "from_forward_context" in cc.static_forward_context: + del cc.static_forward_context["from_forward_context"] + cc.static_all_moe_layers.remove("from_forward_context") + + tp_rank = pgi.rank % test_config.tp_size + + if verbosity > 0: + print(f"subtest: {test_config.id()}", end="") + + try: + _run_one_config( + vllm_config, + test_config.ep_size, + test_config.dp_size, + test_config.tp_size, + dp_rank, + tp_rank, + test_config.m, + test_config.n, + test_config.k, + test_config.num_experts, + test_config.top_k, + test_config.quantization, + test_config.reduce_results, + test_config.backend, + functools.partial( + _test_body_config, test_config=test_config, cpu_group=cpu_group + ), + use_shared_experts=test_config.use_shared_experts, + use_gate=test_config.use_gate, + use_routed_input_transform=test_config.use_routed_input_transform, + ) + if verbosity > 0: + print(" PASSED") + else: + print(".", end="") + passed = passed + 1 + except Exception as ex: + fail_ids.append(test_config.id()) + failed = failed + 1 + if verbosity > 0: + traceback.print_exc() + print(f"\n{str(ex)}\nFAILED {ex.__class__}") + else: + print("F", end="") + finally: + # Note: for some reason DeepEP buffers don't seem to be + # entirely reusable on B200. In order to work around this + # we clear the all2all manager's cache after each testpoint. + cap = current_platform.get_device_capability() + if ( + cap is not None + and cap.major == 10 + and ( + test_config.backend == "deepep_low_latency" + or test_config.backend == "deepep_high_throughput" + ) + ): + torch.accelerator.synchronize() + all2all_manager = get_ep_group().device_communicator.all2all_manager + if all2all_manager is not None: + all2all_manager.destroy() + total = total + 1 + + skipped = total - (passed + failed) + + fails = f"{failed} failed" if failed > 0 else "" + sep = ", " if fails != "" else "" + skips = f"{sep}{skipped} skipped" if skipped > 0 else "" + sep = ", " if skips != "" or fails != "" else "" + passes = f"{sep}{passed} passed" if passed > 0 else "" + + report = ( + f"============= {fails}{skips}{passes} of {total} total tests =============" + ) + + sep = "\n" if verbosity == 0 else "" + print(f"{sep}{report}") + + if failed > 0: + fail_ids_str = "\n".join(fail_ids) + raise RuntimeError( + f"\n============= Failed subtests =============\n{fail_ids_str}\n{report}" + ) + + +# TODO: add cudagraphs/torch.compile tests +@pytest.mark.parametrize("dp_size, tp_size, use_ep", PARALLEL_COMBOS) +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("enable_eplb", [False, True]) +def test_moe_layer( + dp_size: int, + tp_size: int, + use_ep: bool, + backend: str, + enable_eplb: bool, + monkeypatch, + pytestconfig, + subtests, +): + """Test MoE layer with parallelism (multi-GPU or TP/EP enabled). + + For non-parallel cases (world_size == 1), use test_moe_layer_no_parallel instead. + """ + num_gpus = current_platform.device_count() + world_size = tp_size * dp_size + ep_size = 1 if not use_ep else world_size # or dp_size? + assert world_size > 1 + + # Check if enough GPUs available + if world_size is not None and num_gpus is not None and world_size > num_gpus: + pytest.skip(f"Not enough GPUs got {num_gpus}, expected {world_size}.") + + if enable_eplb and not use_ep: + pytest.skip("EPLB requires EP.") + + verbosity = pytestconfig.getoption("verbose") + + test_env = dict() + test_env["VLLM_MOE_DP_CHUNK_SIZE"] = "128" + monkeypatch.setenv("VLLM_MOE_DP_CHUNK_SIZE", "128") + if os.environ.get("VLLM_LOGGING_LEVEL") is None: + monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR") + + # TODO + # VLLM_FLASHINFER_MOE_BACKEND=latency + # VLLM_USE_FLASHINFER_MOE_FP16=1 + # VLLM_USE_FLASHINFER_MOE_FP8 + # VLLM_USE_FLASHINFER_MOE_FP4 + # VLLM_USE_FLASHINFER_MOE_INT4 + + parallel_config = ParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=dp_size, + tensor_parallel_size=tp_size, + enable_expert_parallel=use_ep, + all2all_backend=backend, + enable_eplb=enable_eplb, + ) + + compilation_config = CompilationConfig() + # compilation_config.mode = CompilationMode.NONE # for now + compilation_config.pass_config.fuse_allreduce_rms = False # for now + + vllm_config = VllmConfig( + parallel_config=parallel_config, compilation_config=compilation_config + ) + + test_configs = generate_valid_test_configs( + backend, ep_size, dp_size, tp_size, enable_eplb, verbosity + ) + + if subtests is not None: + new_test_configs = [] + for subtest in subtests.split(","): + sub_test_config = MoETestConfig.from_id(subtest) + if sub_test_config in test_configs: + new_test_configs.append(sub_test_config) + else: + pytest.skip( + f"subtest config {subtest} does not match any valid test " + "configuration" + ) + test_configs = new_test_configs + + if len(test_configs) == 0: + pytest.skip("No supported configs found for this testpoint.") + + try: + parallel_launch_with_config( + world_size, + _parallel_worker, + vllm_config, + test_env, + test_configs, + verbosity, + ) + finally: + torch.accelerator.synchronize() # TODO: Is this needed? + torch.accelerator.empty_cache() diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 2ef4424c2ba..8763ad68351 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -248,7 +248,7 @@ def make_quantized_test_activations( return a, a_q, a_scale -def moe_quantize_weights( +def moe_quantize_weights_2d( w: torch.Tensor, w_s: torch.Tensor | None, quant_dtype: torch.dtype | str | None, @@ -293,6 +293,40 @@ def moe_quantize_weights( return w, w_s, w_gs +def moe_quantize_weights( + w: torch.Tensor, + w_s: torch.Tensor | None, + quant_dtype: torch.dtype | str | None, + per_token_quant: bool, + block_shape: list[int] | None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + assert w.dim() == 3 + e, rows, cols = w.shape + w_l = [None] * e + w_s_l = [None] * e + w_gs_l = [None] * e + for idx in range(e): + w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights_2d( + w[idx], None, quant_dtype, per_token_quant, block_shape + ) + + w = torch.stack(w_l) + w_s = torch.stack(w_s_l) + w_gs = torch.stack(w_gs_l) if e > 0 and w_gs_l[0] is not None else None + + if w_s.ndim == 2: + assert w_s.shape[-1] == 1 + w_s = w_s.view(-1, 1, 1) + + if block_shape is not None: + block_n, block_k = block_shape + n_tiles = (rows + block_n - 1) // block_n + k_tiles = (cols + block_k - 1) // block_k + assert w_s.shape == (e, n_tiles, k_tiles) + + return w, w_s, w_gs + + def make_test_weight( e: int, rows: int, @@ -303,30 +337,11 @@ def make_test_weight( per_out_ch_quant: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15 - w_gs = None if quant_dtype is not None: - w_l = [None] * e - w_s_l = [None] * e - w_gs_l = [None] * e - for idx in range(e): - w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights( - w_16[idx], None, quant_dtype, per_out_ch_quant, block_shape - ) - - w = torch.stack(w_l) - w_s = torch.stack(w_s_l) - if e > 0 and w_gs_l[0] is not None: - w_gs = torch.stack(w_gs_l) - if w_s.ndim == 2: - assert w_s.shape[-1] == 1 - w_s = w_s.view(-1, 1, 1) - - if block_shape is not None: - block_n, block_k = block_shape - n_tiles = (rows + block_n - 1) // block_n - k_tiles = (cols + block_k - 1) // block_k - assert w_s.shape == (e, n_tiles, k_tiles) + w, w_s, w_gs = moe_quantize_weights( + w_16, None, quant_dtype, per_out_ch_quant, block_shape + ) else: w = w_16 w_s = None @@ -454,7 +469,6 @@ def fused_moe( ) -# CustomOp? class BaselineMM(torch.nn.Module): def __init__( self, @@ -462,13 +476,22 @@ class BaselineMM(torch.nn.Module): out_dtype: torch.dtype, ): super().__init__() - self.b = b.to(dtype=torch.float32) + self.b = torch.nn.Parameter(b.to(dtype=torch.float32)) self.out_dtype = out_dtype def forward(self, a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: return torch.mm(a.to(dtype=torch.float32), self.b).to(self.out_dtype), None +class BaselineSiluAndMul(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + return torch.nn.functional.silu(x[..., :d]) * x[..., d:] + + class TestMLP(torch.nn.Module): def __init__( self, @@ -479,7 +502,7 @@ class TestMLP(torch.nn.Module): super().__init__() self.gate_up_proj = BaselineMM(w1, out_dtype) self.down_proj = BaselineMM(w2, out_dtype) - self.act_fn = SiluAndMul() + self.act_fn = BaselineSiluAndMul() def forward(self, x): x, _ = self.gate_up_proj(x) @@ -564,35 +587,24 @@ class RealMLP(torch.nn.Module): return x -def make_shared_experts( +def make_shared_experts_with_weights( N: int, K: int, - in_dtype: torch.dtype = torch.bfloat16, + in_dtype: torch.dtype, + w1: torch.Tensor, + w2: torch.Tensor, + w1_s: torch.Tensor | None = None, + w2_s: torch.Tensor | None = None, quant_dtype: torch.dtype | str | None = None, ) -> torch.nn.Module: - from vllm.model_executor.layers.quantization.fp8 import Fp8Config - - (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( - 1, - N, - K, - in_dtype=in_dtype, - quant_dtype=quant_dtype, - ) old_dtype = torch.get_default_dtype() try: torch.set_default_dtype(in_dtype) if quant_dtype == torch.float8_e4m3fn: - w1 = w1[0].transpose(0, 1) - w2 = w2[0].transpose(0, 1) - w1_s = w1_s[0].transpose(0, 1) if w1_s is not None else None - w2_s = w2_s[0].transpose(0, 1) if w2_s is not None else None + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + quant_config = Fp8Config(True) else: - w1 = w1[0] - w2 = w2[0] - w1_s = None - w2_s = None quant_config = None return RealMLP(K, N, w1, w2, "silu", quant_config, w1_s=w1_s, w2_s=w2_s) @@ -614,3 +626,22 @@ def modular_triton_fused_moe( TritonExperts(moe_config, quant_config), inplace=False, ) + + +def make_shared_experts( + N: int, + K: int, + in_dtype: torch.dtype = torch.bfloat16, + quant_dtype: torch.dtype | str | None = None, +) -> torch.nn.Module: + (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( + 1, + N, + K, + in_dtype=in_dtype, + quant_dtype=quant_dtype, + ) + + return make_shared_experts_with_weights( + N, K, in_dtype, w1, w2, w1_s=w1_s, w2_s=w2_s, quant_dtype=quant_dtype + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index ee6df1af110..0b679b78c92 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -76,7 +76,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic): activation_key: QuantKey | None, ) -> bool: return routing_method in [ - RoutingMethodType.Default, RoutingMethodType.DeepSeekV3, RoutingMethodType.Llama4, RoutingMethodType.Renormalize, From 419e73cdfab3b4a2d2ea6753382b345575e02983 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 6 Apr 2026 13:31:19 -0400 Subject: [PATCH 18/39] [Bug] Fix mistral version dependency (#39086) Signed-off-by: yewentao256 --- requirements/nightly_torch_test.txt | 2 +- requirements/rocm-test.in | 2 +- requirements/test.in | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/requirements/nightly_torch_test.txt b/requirements/nightly_torch_test.txt index ca9c5bd1cac..e0eb7e11411 100644 --- a/requirements/nightly_torch_test.txt +++ b/requirements/nightly_torch_test.txt @@ -23,7 +23,7 @@ jiwer # required for audio tests timm # required for internvl test transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio] >= 1.9.1 # required for voxtral test +mistral_common[image,audio] >= 1.11.0 # required for voxtral test num2words # required for smolvlm test opencv-python-headless >= 4.13.0 # required for video test datamodel_code_generator # required for minicpm3 test diff --git a/requirements/rocm-test.in b/requirements/rocm-test.in index b8978f1f226..23c3a0f91e0 100644 --- a/requirements/rocm-test.in +++ b/requirements/rocm-test.in @@ -31,7 +31,7 @@ tblib # for pickling test exceptions timm>=1.0.17 # required for internvl and gemma3n-mm test transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio]>=1.10.0 # required for voxtral test +mistral_common[image,audio]>=1.11.0 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py opencv-python-headless>=4.13.0 # required for video test diff --git a/requirements/test.in b/requirements/test.in index dc47976b1a6..e21f89d2d80 100644 --- a/requirements/test.in +++ b/requirements/test.in @@ -32,7 +32,7 @@ torchaudio==2.10.0 torchvision==0.25.0 transformers_stream_generator # required for qwen-vl test matplotlib # required for qwen-vl test -mistral_common[image,audio] >= 1.9.1 # required for voxtral test +mistral_common[image,audio] >= 1.11.0 # required for voxtral test num2words # required for smolvlm test open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py opencv-python-headless >= 4.13.0 # required for video test From 94fbb09894a00533a41ce2d976d9aa2f06e7e000 Mon Sep 17 00:00:00 2001 From: namgyu-youn Date: Tue, 7 Apr 2026 03:05:39 +0900 Subject: [PATCH 19/39] [EASY] Drop duplicate KV-cache initialization (#38799) Signed-off-by: namgyu-youn --- vllm/model_executor/layers/attention/attention.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 3ff4ec62a6b..a49415a3df8 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -131,9 +131,6 @@ def _init_kv_cache_quant( quant_config: Optional quantization configuration. prefix: Layer name prefix for quantization method lookup. """ - quant_method = ( - quant_config.get_quant_method(layer, prefix=prefix) if quant_config else None - ) # Note [Register q/k/v/prob scales in state dict] # When calling model.to(device), only parameters/buffers in state dict are From e8ebbdde8304a8cf89bbd4e101ebdfc25118b125 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 6 Apr 2026 14:57:53 -0400 Subject: [PATCH 20/39] [Quantization] Add FlashInfer CuteDSL batched experts backend for NVFP4 MoE (#38251) Signed-off-by: Yongye Zhu Co-authored-by: Michael Goin Co-authored-by: Roger Wang --- tests/kernels/moe/test_cutedsl_moe.py | 2 +- .../experts/flashinfer_cutedsl_batched_moe.py | 353 ++++++++++++++++++ .../experts/flashinfer_cutedsl_moe.py | 302 +++------------ .../layers/fused_moe/oracle/nvfp4.py | 48 ++- .../quantization/utils/flashinfer_fp4_moe.py | 96 ++++- vllm/utils/flashinfer.py | 18 + 6 files changed, 574 insertions(+), 245 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py diff --git a/tests/kernels/moe/test_cutedsl_moe.py b/tests/kernels/moe/test_cutedsl_moe.py index bca3eba0f91..2a6f83695c4 100644 --- a/tests/kernels/moe/test_cutedsl_moe.py +++ b/tests/kernels/moe/test_cutedsl_moe.py @@ -17,7 +17,7 @@ from flashinfer import fp4_quantize from torch.nn import functional as F from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import ( # noqa: E501 flashinfer_cutedsl_moe_masked, ) from vllm.utils.flashinfer import ( diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py new file mode 100644 index 00000000000..5eaaf46739f --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import envs +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceDelegate, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import ( + flashinfer_cutedsl_grouped_gemm_nt_masked, + has_flashinfer_cutedsl_grouped_gemm_nt_masked, + scaled_fp4_grouped_quantize, + silu_and_mul_scaled_nvfp4_experts_quantize, +) + +logger = init_logger(__name__) + + +class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + max_num_tokens: int, + num_dispatchers: int, + ): + super().__init__( + moe_config=moe_config, + quant_config=quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=num_dispatchers, + ) + assert quant_config.quant_dtype == "nvfp4", ( + "Only nvfp4 quantization are currently supported." + ) + self.out_dtype = moe_config.in_dtype + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.BatchedExperts + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + return ( + p.is_cuda() + and p.is_device_capability_family(100) + and has_flashinfer_cutedsl_grouped_gemm_nt_masked() + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kNvfp4Static, kNvfp4Dynamic), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SILU + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + # Let PrepareAndFinalize::finalize() decide the impl. + return TopKWeightAndReduceDelegate() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """ + Compute the shapes for the temporary and final outputs of the two gemms + and activation in the fused expert function. Since the gemms are + independent, the workspace for the first gemm can be shared with the + workspace for the last gemm. + + Returns a tuple of: + - workspace13 shape tuple: must be large enough to hold the + result of either expert gemm. + - workspace2 shape tuple: must be large enough to hold the + result of the activation function. + - output shape tuple: must be exact size of the final gemm output. + - Workspace type: The dtype to use for the workspace tensors. + - Note: in order for activation chunking to work, the first dimension + of each tuple must be the number of tokens. + """ + + # We use global_num_experts due to how moe_align_block_size handles + # expert_maps. + K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K + output_shape = (local_num_experts, M, K_dim) + workspace2 = (local_num_experts, M, N) + workspace1 = output_shape + return (workspace1, workspace2, output_shape) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, # Not used + workspace13: torch.Tensor | None, + workspace2: torch.Tensor | None, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool | None, + ): + assert self.quant_dtype == "nvfp4", ( + "Only nvfp4 quantization are currently supported." + ) + # Ensure w1_scale and w2_scale are not None before calling view + assert self.w1_scale is not None and self.w2_scale is not None, ( + "w1_scale and w2_scale must not be None for FlashInferExperts" + ) + assert expert_tokens_meta is not None + expert_num_tokens = expert_tokens_meta.expert_num_tokens + assert hidden_states.ndim == 3 + assert self.w1_scale.ndim == 3 + assert self.w2_scale.ndim == 3 + + input_global_scale = ( + None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale + ) + flashinfer_hidden_states = ( + (hidden_states, a1q_scale) + if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH + else hidden_states + ) + flashinfer_cutedsl_moe_masked( + hidden_states=flashinfer_hidden_states, + input_global_scale=input_global_scale, + w1=w1, + w1_blockscale=self.w1_scale, + w1_alpha=self.g1_alphas, + w2=w2, + a2_global_scale=self.a2_gscale, + w2_blockscale=self.w2_scale, + w2_alpha=self.g2_alphas, + masked_m=expert_num_tokens, + workspace=workspace2, + out=output, + ) + + +def get_cute_dtype(input: torch.Tensor) -> str: + if input.dtype == torch.bfloat16: + return "bfloat16" + elif input.dtype == torch.float16: + return "float16" + elif input.dtype == torch.float32: + return "float32" + else: + raise ValueError(f"Unsupported cute dtype {input.dtype}") + + +def flashinfer_cutedsl_moe_masked( + hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + input_global_scale: torch.Tensor, + w1: torch.Tensor, + w1_blockscale: torch.Tensor, + w1_alpha, + w2: torch.Tensor, + a2_global_scale: torch.Tensor, + w2_blockscale: torch.Tensor, + w2_alpha, + masked_m: torch.Tensor, + workspace: torch.Tensor, + out: torch.Tensor, +): + """ + Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL + kernels. + + Args: + hidden_states: Either of the following case + * torch.Tensor: [num_experts, m, k], bf16 + * tuple[torch.Tensor, torch.Tensor]: [num_experts, m, k // 2], + uint8, [num_experts, m, k // 16], float8_e4m3fn + input_global_scale (torch.Tensor): (l,) + w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8 + w1_blockscale (torch.Tensor): blockscale factors, e4m3, + w1_alpha (torch.Tensor): (l,) + w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8 + a2_global_scale (torch.Tensor): (l,) + w2_blockscale (torch.Tensor): blockscale factors, e4m3, + w2_alpha (torch.Tensor): (l,) + masked_m (torch.Tensor): Masked dimension indices + workspace (torch.Tensor): For gateup_output + + Notes: + - Assumes max(masked_m) <= m. + """ + + # === Assertions on dtypes === + assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}" + assert w1_blockscale.dtype == torch.float8_e4m3fn, ( + f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}" + ) + assert w1_alpha.dtype == torch.float32, ( + f"w1_alpha must be float32, got {w1_alpha.dtype}" + ) + assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}" + assert a2_global_scale.dtype == torch.float32, ( + f"a2_global_scale must be float32, got {a2_global_scale.dtype}" + ) + assert w2_blockscale.dtype == torch.float8_e4m3fn, ( + f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}" + ) + assert w2_alpha.dtype == torch.float32, ( + f"w2_alpha must be float32, got {w2_alpha.dtype}" + ) + + # === Assertions on shapes === + n = w2.shape[-1] * 2 # intermediate dimension + if isinstance(hidden_states, tuple): + assert input_global_scale is None, ( + "input_global_scale is needed when input needs quant" + ) + + aq = hidden_states[0].view(torch.uint8) + aq_sf = hidden_states[1].view(torch.float8_e4m3fn) + # m, k_by_2, num_experts = aq.shape + num_experts, m, k_by_2 = aq.shape + k = k_by_2 * 2 + aq = aq.permute(1, 2, 0) + else: + num_experts, m, k = hidden_states.shape + + assert input_global_scale.dtype == torch.float32, ( + f"input_global_scale must be float32, got {input_global_scale.dtype}" + ) + assert input_global_scale.shape == (num_experts,), ( + f"input_global_scale must be (l,), got {input_global_scale.shape}" + ) + + aq, aq_sf = scaled_fp4_grouped_quantize( + hidden_states, + masked_m, + input_global_scale, + ) + + assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}" + assert w1.shape[-1] * 2 == k, ( + f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}" + ) + assert w2.shape[-2:] == ( + k, + n // 2, + ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}" + + assert w1_alpha.shape == (num_experts,), ( + f"w1_alpha must be (l,), got {w1_alpha.shape}" + ) + assert a2_global_scale.shape == (num_experts,), ( + f"a2_global_scale must be (l,), got {a2_global_scale.shape}" + ) + assert w2_alpha.shape == (num_experts,), ( + f"w2_alpha must be (l,), got {w2_alpha.shape}" + ) + + workspace = workspace.permute(1, 2, 0) # requirement of kernel + sf_vec_size = 16 + assert aq_sf.dtype == torch.float8_e4m3fn + assert aq.dtype == torch.uint8 + ab_dtype = "float4_e2m1fn" + sf_dtype = "float8_e4m3fn" + + if isinstance(hidden_states, tuple): + c_dtype = "bfloat16" + else: + c_dtype = get_cute_dtype(hidden_states) + + # Gemm1 + flashinfer_cutedsl_grouped_gemm_nt_masked( + (aq, aq_sf), + (w1.permute(1, 2, 0), w1_blockscale), + workspace, + masked_m, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + alpha=w1_alpha.view(1, 1, num_experts), + alpha_dtype=get_cute_dtype(w1_alpha), + ) # in logical [m, n, l] + + # SILU and quantization + diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize( + workspace.permute(2, 0, 1), + masked_m, + a2_global_scale, + ) + + # Gemm2 + out = out.permute(1, 2, 0) # requirement of kernel + flashinfer_cutedsl_grouped_gemm_nt_masked( + (diq, diq_sf), + (w2.permute(1, 2, 0), w2_blockscale), + out, + masked_m, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + alpha=w2_alpha.view(1, 1, num_experts), + alpha_dtype=get_cute_dtype(w2_alpha), + ) # in logical [m, k, l] + out = out.permute(2, 0, 1) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index a1db2661938..5ce58220b07 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -4,8 +4,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm import envs -from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -13,7 +11,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceDelegate, + TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -22,33 +20,42 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( - flashinfer_cutedsl_grouped_gemm_nt_masked, - has_flashinfer_cutedsl_grouped_gemm_nt_masked, - scaled_fp4_grouped_quantize, - silu_and_mul_scaled_nvfp4_experts_quantize, + flashinfer_cute_dsl_fused_moe_nvfp4, + has_flashinfer_cutedsl_moe_nvfp4, ) -logger = init_logger(__name__) - class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): + """ + CuteDSL NvFP4 MoE experts using the FlashInfer functional API. + + Uses Standard activation format (non-batched). The kernel handles + routing, expert computation, and reduction internally. + Supports expert parallelism natively. + """ + def __init__( self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, - max_num_tokens: int, - num_dispatchers: int, ): super().__init__( moe_config=moe_config, quant_config=quant_config, - max_num_tokens=max_num_tokens, - num_dispatchers=num_dispatchers, ) assert quant_config.quant_dtype == "nvfp4", ( - "Only nvfp4 quantization are currently supported." + "Only nvfp4 quantization is currently supported." ) self.out_dtype = moe_config.in_dtype + self.hidden_dim = moe_config.hidden_dim + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.topk = moe_config.experts_per_token + self.local_num_experts = moe_config.num_local_experts + self.global_num_experts = moe_config.num_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + self.local_expert_offset = self.ep_rank * self.local_num_experts def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) @@ -56,7 +63,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.BatchedExperts + return mk.FusedMoEActivationFormat.Standard @staticmethod def _supports_current_device() -> bool: @@ -64,7 +71,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): return ( p.is_cuda() and p.is_device_capability_family(100) - and has_flashinfer_cutedsl_grouped_gemm_nt_masked() + and has_flashinfer_cutedsl_moe_nvfp4() ) @staticmethod @@ -86,15 +93,16 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): return activation == MoEActivation.SILU @staticmethod - def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: return True def supports_expert_map(self) -> bool: return False def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - # Let PrepareAndFinalize::finalize() decide the impl. - return TopKWeightAndReduceDelegate() + return TopKWeightAndReduceNoOP() def workspace_shapes( self, @@ -107,29 +115,12 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): expert_tokens_meta: mk.ExpertTokensMetadata | None, activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - # We use global_num_experts due to how moe_align_block_size handles - # expert_maps. - """ - Compute the shapes for the temporary and final outputs of the two gemms - and activation in the fused expert function. Since the gemms are - independent, the workspace for the first gemm can be shared with the - workspace for the last gemm. - - Returns a tuple of: - - workspace13 shape tuple: must be large enough to hold the - result of either expert gemm. - - workspace2 shape tuple: must be large enough to hold the - result of the activation function. - - output shape tuple: must be exact size of the final gemm output. - - Workspace type: The dtype to use for the workspace tensors. - - Note: in order for activation chunking to work, the first dimension - of each tuple must be the number of tokens. - """ - K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K - output_shape = (local_num_experts, M, K_dim) - workspace2 = (local_num_experts, M, N) - workspace1 = output_shape - return (workspace1, workspace2, output_shape) + workspace1 = (0,) + workspace2 = (0,) + # K is packed (K//2 for uint8), so output uses hidden_dim. + assert self.hidden_dim == K * 2 + output = (M, self.hidden_dim) + return (workspace1, workspace2, output) def apply( self, @@ -143,210 +134,39 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, # Not used + a2_scale: torch.Tensor | None, workspace13: torch.Tensor | None, workspace2: torch.Tensor | None, expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool | None, ): - assert self.quant_dtype == "nvfp4", ( - "Only nvfp4 quantization are currently supported." - ) - # Ensure w1_scale and w2_scale are not None before calling view - assert self.w1_scale is not None and self.w2_scale is not None, ( - "w1_scale and w2_scale must not be None for FlashInferExperts" - ) - assert expert_tokens_meta is not None - expert_num_tokens = expert_tokens_meta.expert_num_tokens - assert hidden_states.ndim == 3 - assert self.w1_scale.ndim == 3 - assert self.w2_scale.ndim == 3 + assert self.quant_dtype == "nvfp4" + assert a1q_scale is not None + assert self.w1_scale is not None + assert self.w2_scale is not None - input_global_scale = ( - None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale - ) - flashinfer_hidden_states = ( - (hidden_states, a1q_scale) - if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH - else hidden_states - ) - flashinfer_cutedsl_moe_masked( - hidden_states=flashinfer_hidden_states, - input_global_scale=input_global_scale, - w1=w1, - w1_blockscale=self.w1_scale, - w1_alpha=self.g1_alphas, - w2=w2, - a2_global_scale=self.a2_gscale, - w2_blockscale=self.w2_scale, - w2_alpha=self.g2_alphas, - masked_m=expert_num_tokens, - workspace=workspace2, - out=output, - ) + # a1q_scale is (M, K//16) float8_e4m3fn from fp4_quantize. + # The functional API expects x_sf with trailing dim: (M, K//16, 1). + x_sf = a1q_scale.unsqueeze(-1) + from vllm.utils.flashinfer import _is_fi_autotuning, autotune -def get_cute_dtype(input: torch.Tensor) -> str: - if input.dtype == torch.bfloat16: - return "bfloat16" - elif input.dtype == torch.float16: - return "float16" - elif input.dtype == torch.float32: - return "float32" - else: - raise ValueError(f"Unsupported cute dtype {input.dtype}") - - -def flashinfer_cutedsl_moe_masked( - hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - input_global_scale: torch.Tensor, - w1: torch.Tensor, - w1_blockscale: torch.Tensor, - w1_alpha, - w2: torch.Tensor, - a2_global_scale: torch.Tensor, - w2_blockscale: torch.Tensor, - w2_alpha, - masked_m: torch.Tensor, - workspace: torch.Tensor, - out: torch.Tensor, -): - """ - Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL - kernels. - - Args: - hidden_states: Either of the following case - * torch.Tensor: [num_experts, m, k], bf16 - * tuple[torch.Tensor, torch.Tensor]: [num_experts, m, k // 2], - uint8, [num_experts, m, k // 16], float8_e4m3fn - input_global_scale (torch.Tensor): (l,) - w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8 - w1_blockscale (torch.Tensor): blockscale factors, e4m3, - w1_alpha (torch.Tensor): (l,) - w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8 - a2_global_scale (torch.Tensor): (l,) - w2_blockscale (torch.Tensor): blockscale factors, e4m3, - w2_alpha (torch.Tensor): (l,) - masked_m (torch.Tensor): Masked dimension indices - workspace (torch.Tensor): For gateup_output - - Notes: - - Assumes max(masked_m) <= m. - """ - - # === Assertions on dtypes === - assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}" - assert w1_blockscale.dtype == torch.float8_e4m3fn, ( - f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}" - ) - assert w1_alpha.dtype == torch.float32, ( - f"w1_alpha must be float32, got {w1_alpha.dtype}" - ) - assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}" - assert a2_global_scale.dtype == torch.float32, ( - f"a2_global_scale must be float32, got {a2_global_scale.dtype}" - ) - assert w2_blockscale.dtype == torch.float8_e4m3fn, ( - f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}" - ) - assert w2_alpha.dtype == torch.float32, ( - f"w2_alpha must be float32, got {w2_alpha.dtype}" - ) - - # === Assertions on shapes === - n = w2.shape[-1] * 2 # intermediate dimension - if isinstance(hidden_states, tuple): - assert input_global_scale is None, ( - "input_global_scale is needed when input needs quant" - ) - - aq = hidden_states[0].view(torch.uint8) - aq_sf = hidden_states[1].view(torch.float8_e4m3fn) - # m, k_by_2, num_experts = aq.shape - num_experts, m, k_by_2 = aq.shape - k = k_by_2 * 2 - aq = aq.permute(1, 2, 0) - else: - num_experts, m, k = hidden_states.shape - - assert input_global_scale.dtype == torch.float32, ( - f"input_global_scale must be float32, got {input_global_scale.dtype}" - ) - assert input_global_scale.shape == (num_experts,), ( - f"input_global_scale must be (l,), got {input_global_scale.shape}" - ) - - aq, aq_sf = scaled_fp4_grouped_quantize( - hidden_states, - masked_m, - input_global_scale, - ) - - assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}" - assert w1.shape[-1] * 2 == k, ( - f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}" - ) - assert w2.shape[-2:] == ( - k, - n // 2, - ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}" - - assert w1_alpha.shape == (num_experts,), ( - f"w1_alpha must be (l,), got {w1_alpha.shape}" - ) - assert a2_global_scale.shape == (num_experts,), ( - f"a2_global_scale must be (l,), got {a2_global_scale.shape}" - ) - assert w2_alpha.shape == (num_experts,), ( - f"w2_alpha must be (l,), got {w2_alpha.shape}" - ) - - workspace = workspace.permute(1, 2, 0) # requirement of kernel - sf_vec_size = 16 - assert aq_sf.dtype == torch.float8_e4m3fn - assert aq.dtype == torch.uint8 - ab_dtype = "float4_e2m1fn" - sf_dtype = "float8_e4m3fn" - - if isinstance(hidden_states, tuple): - c_dtype = "bfloat16" - else: - c_dtype = get_cute_dtype(hidden_states) - - # Gemm1 - flashinfer_cutedsl_grouped_gemm_nt_masked( - (aq, aq_sf), - (w1.permute(1, 2, 0), w1_blockscale), - workspace, - masked_m, - ab_dtype=ab_dtype, - sf_dtype=sf_dtype, - c_dtype=c_dtype, - sf_vec_size=sf_vec_size, - alpha=w1_alpha.view(1, 1, num_experts), - alpha_dtype=get_cute_dtype(w1_alpha), - ) # in logical [m, n, l] - - # SILU and quantization - diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize( - workspace.permute(2, 0, 1), - masked_m, - a2_global_scale, - ) - - # Gemm2 - out = out.permute(1, 2, 0) # requirement of kernel - flashinfer_cutedsl_grouped_gemm_nt_masked( - (diq, diq_sf), - (w2.permute(1, 2, 0), w2_blockscale), - out, - masked_m, - ab_dtype=ab_dtype, - sf_dtype=sf_dtype, - c_dtype=c_dtype, - sf_vec_size=sf_vec_size, - alpha=w2_alpha.view(1, 1, num_experts), - alpha_dtype=get_cute_dtype(w2_alpha), - ) # in logical [m, k, l] - out = out.permute(2, 0, 1) + with autotune(_is_fi_autotuning): + flashinfer_cute_dsl_fused_moe_nvfp4( + x=hidden_states, + x_sf=x_sf, + token_selected_experts=topk_ids.to(torch.int32), + token_final_scales=topk_weights.float(), + w1_weight=w1, + w1_weight_sf=self.w1_scale, + w1_alpha=self.g1_alphas, + fc2_input_scale=self.a2_gscale, + w2_weight=w2, + w2_weight_sf=self.w2_scale, + w2_alpha=self.g2_alphas, + num_experts=self.global_num_experts, + top_k=self.topk, + num_local_experts=self.local_num_experts, + local_expert_offset=self.local_expert_offset, + moe_output=output, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index d946c5eb53c..597d784d3b6 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( ) from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( prepare_nvfp4_moe_layer_for_fi_or_cutlass, + prepare_nvfp4_moe_layer_for_flashinfer_cutedsl, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( FlashinferMoeBackend, @@ -41,6 +42,7 @@ class NvFp4MoeBackend(Enum): FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS" FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL" + FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED" VLLM_CUTLASS = "VLLM_CUTLASS" MARLIN = "MARLIN" @@ -49,6 +51,7 @@ FLASHINFER_NVFP4_MOE_BACKENDS = [ NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTEDSL, + NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, ] fi_2_vllm_backend_map: dict[FlashinferMoeBackend, NvFp4MoeBackend] = { @@ -95,6 +98,13 @@ def backend_to_kernel_cls( return [FlashInferCuteDSLExperts] + elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED: + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import ( # noqa: E501 + FlashInferCuteDSLBatchedExperts, + ) + + return [FlashInferCuteDSLBatchedExperts] + elif backend == NvFp4MoeBackend.VLLM_CUTLASS: from vllm.model_executor.layers.fused_moe.cutlass_moe import ( CutlassExpertsFp4, @@ -143,6 +153,7 @@ def select_nvfp4_moe_backend( AVAILABLE_BACKENDS = [ NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTEDSL, + NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, @@ -198,6 +209,12 @@ def select_nvfp4_moe_backend( runner_backend = config.moe_backend if runner_backend != "auto": requested_backend = map_nvfp4_backend(runner_backend) + # For batched activation format, use batched variant if available. + if ( + activation_format == mk.FusedMoEActivationFormat.BatchedExperts + and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL + ): + requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED return _return_or_raise( requested_backend, config, weight_key, activation_key, activation_format ) @@ -288,7 +305,28 @@ def convert_to_nvfp4_moe_kernel_format( torch.Tensor, torch.Tensor, ]: - if ( + if nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL: + ( + w13, + w13_scale, + w13_scale_2, + a13_scale, + w2, + w2_scale, + w2_scale_2, + a2_scale, + ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( + layer=layer, + w13=w13, + w13_scale=w13_scale, + w13_scale_2=w13_scale_2, + a13_scale=a13_scale, + w2=w2, + w2_scale=w2_scale, + w2_scale_2=w2_scale_2, + a2_scale=a2_scale, + ) + elif ( nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS or nvfp4_backend == NvFp4MoeBackend.VLLM_CUTLASS ): @@ -380,7 +418,13 @@ def make_nvfp4_moe_quant_config( # NOTE(rob): this is a hack until the MoE kernels # create their own quant configs. TRTLLM kernel # does not accept swizzled input quant scales. - is_nvfp4_scale_swizzled=(backend != NvFp4MoeBackend.FLASHINFER_TRTLLM), + is_nvfp4_scale_swizzled=( + backend + not in ( + NvFp4MoeBackend.FLASHINFER_TRTLLM, + NvFp4MoeBackend.FLASHINFER_CUTEDSL, + ) + ), ) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index d16d4a3d261..397442aeced 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -60,6 +60,100 @@ def reorder_w1w3_to_w3w1( ) +def interleave_linear_and_gate( + x: torch.Tensor, + group_size: int = 64, + dim: int = -1, +) -> torch.Tensor: + """Interleave gate and linear weight rows for CuteDSL wrapper.""" + sizes = x.size() + dim = dim % x.dim() + assert sizes[dim] % (group_size * 2) == 0, ( + f"dim {dim} size {sizes[dim]} must be divisible by {group_size * 2}" + ) + prev_sizes = sizes[:dim] + post_sizes = sizes[dim + 1 :] + x = x.view(*prev_sizes, 2, sizes[dim] // (group_size * 2), group_size, *post_sizes) + x = x.transpose(dim, dim + 1).contiguous().view(*sizes) + return x + + +def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( + layer: "FusedMoE", + w13: torch.Tensor, + w13_scale: torch.Tensor, + w13_scale_2: torch.Tensor, + a13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + w2_scale_2: torch.Tensor, + a2_scale: torch.Tensor, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend. + + Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper, + and interleaves w13 gate/linear rows. + """ + from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout + + # Global scaling factors (same as other FlashInfer backends). + num_experts = w13.shape[0] + a13_scale = a13_scale.max().to(torch.float32).expand(num_experts) + a2_scale = a2_scale.max().to(torch.float32).expand(num_experts) + + half = w13.shape[1] // 2 + w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) + w13_scale = torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1) + + # Interleave up/gate rows for w13 weights and scales. + w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) + w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) + + # Convert w13 scale factors: linear → swizzled → MMA layout. + w13_scale = swizzle_blockscale(w13_scale) + E, M_padded, K_sf_padded = w13_scale.shape + w13_scale_flat = w13_scale.reshape(E * M_padded, K_sf_padded) + w13_scale = convert_sf_to_mma_layout( + w13_scale_flat, + m=M_padded, + k=K_sf_padded * 16, + num_groups=E, + sf_vec_size=16, + ) + + # Convert w2 scale factors: linear → swizzled → MMA layout. + w2_scale = swizzle_blockscale(w2_scale) + E, M_padded, K_sf_padded = w2_scale.shape + w2_scale_flat = w2_scale.reshape(E * M_padded, K_sf_padded) + w2_scale = convert_sf_to_mma_layout( + w2_scale_flat, + m=M_padded, + k=K_sf_padded * 16, + num_groups=E, + sf_vec_size=16, + ) + + return ( + w13, + w13_scale, + w13_scale_2, + a13_scale, + w2, + w2_scale, + w2_scale_2, + a2_scale, + ) + + def prepare_static_weights_for_trtllm_fp4_moe( # args_dequant, # args, @@ -221,7 +315,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_TRTLLM, - NvFp4MoeBackend.FLASHINFER_CUTEDSL, + NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, ] # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels. diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 8ffac48cc87..373134e655d 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -128,6 +128,12 @@ scaled_fp4_grouped_quantize = _lazy_import_wrapper( nvfp4_block_scale_interleave = _lazy_import_wrapper( "flashinfer.fp4_quantization", "block_scale_interleave" ) +flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper( + "flashinfer", "cute_dsl_fused_moe_nvfp4" +) +flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper( + "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout" +) trtllm_fp4_block_scale_moe = _lazy_import_wrapper( "flashinfer", "trtllm_fp4_block_scale_moe" ) @@ -252,6 +258,15 @@ def has_flashinfer_cutedsl_grouped_gemm_nt_masked() -> bool: return True +@functools.cache +def has_flashinfer_cutedsl_moe_nvfp4() -> bool: + """Return ``True`` if FlashInfer cute_dsl_fused_moe_nvfp4 is available.""" + if not has_flashinfer_cutedsl(): + return False + mod = _get_submodule("flashinfer") + return mod is not None and hasattr(mod, "cute_dsl_fused_moe_nvfp4") + + @functools.cache def has_nvidia_artifactory() -> bool: """Return `True` if NVIDIA's artifactory is accessible. @@ -768,6 +783,8 @@ __all__ = [ "silu_and_mul_scaled_nvfp4_experts_quantize", "scaled_fp4_grouped_quantize", "nvfp4_block_scale_interleave", + "flashinfer_cute_dsl_fused_moe_nvfp4", + "flashinfer_convert_sf_to_mma_layout", "trtllm_fp4_block_scale_moe", "autotune", "has_flashinfer_moe", @@ -776,6 +793,7 @@ __all__ = [ "has_flashinfer_nvlink_one_sided", "has_flashinfer_cutlass_fused_moe", "has_flashinfer_cutedsl_grouped_gemm_nt_masked", + "has_flashinfer_cutedsl_moe_nvfp4", "has_flashinfer_fp8_blockscale_gemm", "has_nvidia_artifactory", "supports_trtllm_attention", From dfa5062a8f372ea78e48197939289c15247c1840 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Mon, 6 Apr 2026 22:47:46 +0300 Subject: [PATCH 21/39] NemotronH default mamba_ssm_cache_dtype=float32; enable auto-hook for NemotronHNanoVLV2Config (#39032) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- vllm/model_executor/models/config.py | 33 ++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 7b4fa9252b5..22d300a7ebf 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -7,7 +7,9 @@ from vllm.logger import init_logger from vllm.utils.math_utils import round_up if TYPE_CHECKING: - from vllm.config import ModelConfig, VllmConfig + from transformers import PretrainedConfig + + from vllm.config import CacheConfig, ModelConfig, VllmConfig logger = init_logger(__name__) @@ -346,17 +348,20 @@ class MambaModelConfig(VerifyAndUpdateConfig): class NemotronHForCausalLMConfig(VerifyAndUpdateConfig): - @staticmethod - def verify_and_update_config(vllm_config: "VllmConfig") -> None: + DEFAULT_MAMBA_SSM_CACHE_DTYPE = "float32" + """Only `float32` is known to have no accuracy issues by default.""" + + @classmethod + def update_mamba_ssm_cache_dtype( + cls, *, cache_config: "CacheConfig", hf_config: "PretrainedConfig" + ) -> None: """Update mamba_ssm_cache_dtype for NemotronH models when set to 'auto' (or not explicitly set), to the value specified in the HF config, or to - float16 if not specified. + `float32` if not specified. """ - cache_config = vllm_config.cache_config if cache_config.mamba_ssm_cache_dtype == "auto": - hf_config = vllm_config.model_config.hf_config mamba_ssm_cache_dtype = getattr( - hf_config, "mamba_ssm_cache_dtype", "float16" + hf_config, "mamba_ssm_cache_dtype", cls.DEFAULT_MAMBA_SSM_CACHE_DTYPE ) logger.info( "Updating mamba_ssm_cache_dtype to '%s' for NemotronH model", @@ -364,8 +369,22 @@ class NemotronHForCausalLMConfig(VerifyAndUpdateConfig): ) cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype + @classmethod + def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: + cls.update_mamba_ssm_cache_dtype( + cache_config=vllm_config.cache_config, + hf_config=vllm_config.model_config.hf_config, + ) + class NemotronHNanoVLV2Config(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, + ) + @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: mm_config = model_config.multimodal_config From f186cfe75e452aeb76f5233da7392d51ee34d3ef Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 6 Apr 2026 12:55:13 -0700 Subject: [PATCH 22/39] [MRV2] Fix hanging issue with DeepSeek V3.2 by setting `skip_attn=False` (#39098) Signed-off-by: WoosukKwon Signed-off-by: Woosuk Kwon --- vllm/v1/worker/gpu/model_runner.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a2f83c52e95..56df70fc0c9 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -391,12 +391,17 @@ class GPUModelRunner(LoRAModelRunnerMixin): self, num_tokens: int, *args, - skip_attn: bool = True, + skip_attn: bool = False, uniform_decode: bool = False, skip_eplb: bool = False, is_profile: bool = False, **kwargs, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if skip_attn and not is_profile: + raise ValueError( + "skip_attn must only be True for initial memory profiling." + ) + # Create a dummy scheduler output. num_reqs = min(num_tokens, self.max_num_reqs) if uniform_decode: @@ -988,6 +993,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): if not skip_attn_for_dummy_run: block_tables, slot_mappings = self.prepare_dummy_attn(input_batch) else: + assert batch_desc.cg_mode != CUDAGraphMode.FULL, ( + "Attention metadata must be prepared for dummy runs when using " + "FULL cudagraph mode." + ) block_tables = None slot_mappings = None # FIXME(woosuk): Fix warmup for LoRA. From 9c81f35b1ae6c70681661e11c461ddbb7e417aff Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Mon, 6 Apr 2026 17:51:46 -0400 Subject: [PATCH 23/39] [Attention][MLA] Re-enable FA4 as default MLA prefill backend (#38819) --- vllm/config/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 014bb9b2260..1da647a6d6f 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -30,7 +30,7 @@ class AttentionConfig: use_cudnn_prefill: bool = False """Whether to use cudnn prefill.""" - use_trtllm_ragged_deepseek_prefill: bool = True + use_trtllm_ragged_deepseek_prefill: bool = False """Whether to use TRTLLM ragged deepseek prefill.""" use_trtllm_attention: bool | None = None From 00d7b497b33679a3ce641db9119eb09dd4ed5e30 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Tue, 7 Apr 2026 00:18:27 +0200 Subject: [PATCH 24/39] [NVFP4] Support NVFP4 dense models from `modelopt` and `compressed-tensors` on AMD Instinct MI300, MI355X and Hopper through emulation (#35733) Signed-off-by: Felix Marty Signed-off-by: fxmarty-amd Co-authored-by: Kyle Sayers --- tests/models/quantization/test_nvfp4.py | 19 ++- tests/quantization/test_compressed_tensors.py | 5 +- vllm/envs.py | 5 + .../schemes/compressed_tensors_w4a4_nvfp4.py | 23 +++ .../layers/quantization/modelopt.py | 19 +++ .../quantization/utils/marlin_utils_fp4.py | 2 +- .../utils/nvfp4_emulation_utils.py | 31 ++-- .../layers/quantization/utils/nvfp4_utils.py | 139 +++++++++++++----- vllm/platforms/rocm.py | 1 + vllm/utils/import_utils.py | 5 + 10 files changed, 191 insertions(+), 58 deletions(-) diff --git a/tests/models/quantization/test_nvfp4.py b/tests/models/quantization/test_nvfp4.py index b73462bfd19..30f69f62130 100644 --- a/tests/models/quantization/test_nvfp4.py +++ b/tests/models/quantization/test_nvfp4.py @@ -89,22 +89,33 @@ def test_models(example_prompts, model_name) -> None: EAGER = [True, False] +SM_100_NVFP4_BACKENDS = [ + "flashinfer-cudnn", + "flashinfer-trtllm", + "flashinfer-cutlass", +] + -@pytest.mark.skipif( - not current_platform.has_device_capability(100), - reason="modelopt_fp4 is not supported on this GPU type.", -) @pytest.mark.parametrize("model", ["nvidia/Llama-3.1-8B-Instruct-NVFP4"]) @pytest.mark.parametrize("eager", EAGER) @pytest.mark.parametrize( "backend", [ + "emulation", "flashinfer-cudnn", "flashinfer-trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used "flashinfer-cutlass", ], ) def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch): + if ( + not current_platform.has_device_capability(100) + and backend in SM_100_NVFP4_BACKENDS + ): + pytest.skip( + f"The backend {backend} is not supported with current_platform.has_device_capability(100) == False" + ) + monkeypatch.setenv("VLLM_NVFP4_GEMM_BACKEND", backend) with vllm_runner(model, enforce_eager=eager) as llm: output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2) diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index f23506b00a7..badcac00546 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -366,9 +366,6 @@ def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner): assert output -@pytest.mark.skipif( - not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." -) @pytest.mark.parametrize( "args", [ @@ -398,7 +395,7 @@ def test_compressed_tensors_nvfp4(vllm_runner, args): assert qkv_proj.scheme.group_size == 16 llm.apply_model(check_model) - output = llm.generate_greedy("Hello my name is", max_tokens=4) + output = llm.generate_greedy(["Hello my name is"], max_tokens=4) print(output) assert output diff --git a/vllm/envs.py b/vllm/envs.py index c2f8ca8c580..d2af9e64d66 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1464,6 +1464,10 @@ environment_variables: dict[str, Callable[[], Any]] = { # - "flashinfer-trtllm": use flashinfer trtllm GEMM backend # - "flashinfer-cutlass": use flashinfer cutlass GEMM backend # - "marlin": use marlin GEMM backend (for GPUs without native FP4 support) + # - "emulation": + # use BF16/FP16 GEMM, dequantizing weights and running QDQ on activations. + # This is only meant for research purposes to run on devices where NVFP4 + # GEMM kernels are not available. # - : automatically pick an available backend "VLLM_NVFP4_GEMM_BACKEND": env_with_choices( "VLLM_NVFP4_GEMM_BACKEND", @@ -1474,6 +1478,7 @@ environment_variables: dict[str, Callable[[], Any]] = { "flashinfer-cutlass", "cutlass", "marlin", + "emulation", ], ), # Controls garbage collection during CUDA graph capture. 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 a3b53626bf6..fff7387260e 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 @@ -5,10 +5,12 @@ from collections.abc import Callable import torch from torch.nn.parameter import Parameter +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsScheme, ) from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + NvFp4LinearBackend, apply_nvfp4_linear, convert_to_nvfp4_linear_kernel_format, select_nvfp4_linear_backend, @@ -19,6 +21,9 @@ from vllm.model_executor.parameter import ( PerTensorScaleParameter, ) +logger = init_logger(__name__) + + __all__ = ["CompressedTensorsW4A4Fp4"] @@ -27,6 +32,10 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme): self.backend = select_nvfp4_linear_backend() self.group_size = 16 + self.swizzle = None + if self.backend == NvFp4LinearBackend.EMULATION: + self.swizzle = False + @classmethod def get_min_capability(cls) -> int: return 75 @@ -89,6 +98,19 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme): # Rename CT checkpoint names to standardized names layer.weight = layer.weight_packed del layer.weight_packed + + if ( + torch.unique(layer.input_global_scale).numel() != 1 + or torch.unique(layer.weight_global_scale).numel() != 1 + ): + logger.warning_once( + "In NVFP4 linear, the global scale for input or weight are different" + " for parallel layers (e.g. q_proj, k_proj, v_proj). This " + " will likely result in reduced accuracy. Please verify the model" + " accuracy. Consider using a checkpoint with a shared global NVFP4" + " scale for fused layers." + ) + # Process global scales (CT stores as divisors, i.e. 1/scale) input_global_scale_inv = layer.input_global_scale.max().to(torch.float32) layer.input_global_scale = Parameter( @@ -121,4 +143,5 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme): layer=layer, x=x, bias=bias, + swizzle=self.swizzle, ) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 7871b774e0c..53b15950d92 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -71,6 +71,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( mxfp8_e4m3_quantize, ) from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + NvFp4LinearBackend, apply_nvfp4_linear, convert_to_nvfp4_linear_kernel_format, select_nvfp4_linear_backend, @@ -1074,6 +1075,10 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase): self.marlin_input_dtype = None self.backend = select_nvfp4_linear_backend() + self.swizzle = None + if self.backend == NvFp4LinearBackend.EMULATION: + self.swizzle = False + def create_weights( self, layer: torch.nn.Module, @@ -1149,10 +1154,23 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase): layer.register_parameter("weight_scale", weight_scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if ( + torch.unique(layer.input_scale).numel() != 1 + or torch.unique(layer.weight_scale_2).numel() != 1 + ): + logger.warning_once( + "In NVFP4 linear, the global scale for input or weight are different" + " for parallel layers (e.g. q_proj, k_proj, v_proj). This " + " will likely results in reduce accuracy. Please verify the model" + " accuracy. Consider using a checkpoint with a shared global NVFP4" + " scale for parallel layers." + ) + # Rename ModelOpt checkpoint names to standardized names input_global_scale = layer.input_scale.max().to(torch.float32) layer.input_global_scale = Parameter(input_global_scale, requires_grad=False) del layer.input_scale + weight_global_scale = layer.weight_scale_2.max().to(torch.float32) layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False) del layer.weight_scale_2 @@ -1179,6 +1197,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase): layer=layer, x=x, bias=bias, + swizzle=self.swizzle, ) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index 4fd484edeb3..19473fa3273 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -24,7 +24,7 @@ logger = init_logger(__name__) def is_fp4_marlin_supported(): - return current_platform.has_device_capability(75) + return current_platform.is_cuda() and current_platform.has_device_capability(75) def _nvfp4_compute_scale_factor( diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index 62b480210fc..9a0c52b62c1 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import torch from vllm.scalar_type import scalar_types @@ -11,9 +13,10 @@ __all__ = [ ] FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max() +FLOAT4_E2M1_MAX_RECIPROCAL = 1 / FLOAT4_E2M1_MAX -kE2M1ToFloat = torch.tensor( - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 +kE2M1ToFloat_handle = SimpleNamespace( + val=torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32) ) @@ -29,8 +32,9 @@ def break_fp4_bytes(a, dtype): # Vectorized sign and magnitude extraction signs = (combined & 0x08).to(torch.bool) # Sign bits abs_vals = (combined & 0x07).to(torch.long) + + kE2M1 = kE2M1ToFloat_handle.val # Device-aware lookup and sign application - kE2M1 = kE2M1ToFloat.to(device=a.device) values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0) # Reshape to final form return values.reshape(m, n * 2).to(dtype=dtype) @@ -47,7 +51,12 @@ def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size): def dequantize_to_dtype( - tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16 + tensor_fp4: torch.Tensor, + tensor_sf: torch.Tensor, + global_scale: torch.Tensor | float, + dtype: torch.dtype, + block_size: int = 16, + swizzle: bool | None = True, ): """Dequantize the fp4 tensor back to high precision.""" # Two fp4 values are packed into one uint8. @@ -57,8 +66,10 @@ def dequantize_to_dtype( tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32) tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size) tensor_sf = tensor_sf.view(torch.float8_e4m3fn) - tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size) - tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale + + if swizzle: + tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size) + tensor_sf_dtype = tensor_sf.to(torch.float32) * global_scale # scale the tensor out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k) @@ -67,7 +78,8 @@ def dequantize_to_dtype( def get_reciprocal(x): if isinstance(x, torch.Tensor): - return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x) + # torch.where yields operation not permitted when stream is capturing. + return 1.0 / (x + (x == 0) * 1e8) elif isinstance(x, (float, int)): return 0.0 if x == 0 else 1.0 / x else: @@ -94,7 +106,7 @@ def ref_nvfp4_quant(x, global_scale, block_size): m, n = x.shape x = torch.reshape(x, (m, n // block_size, block_size)) vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32) - scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX)) + scale = global_scale * (vec_max * FLOAT4_E2M1_MAX_RECIPROCAL) scale = torch.clamp(scale, max=448, min=-448) scale = scale.to(torch.float8_e4m3fn).to(torch.float32) output_scale = get_reciprocal(scale * get_reciprocal(global_scale)) @@ -111,6 +123,7 @@ def run_nvfp4_emulations( weight: torch.Tensor, weight_scale_swizzled: torch.Tensor, weight_global_scale: torch.Tensor, + swizzle: bool | None = True, ): group_size = 16 x_m, x_k = x.shape @@ -132,8 +145,8 @@ def run_nvfp4_emulations( weight_scale_swizzled.data, weight_global_scale, output_dtype, - x.device, group_size, + swizzle=swizzle, ) # matmul diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py index f21f2ef23f4..4796cc6c95c 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py @@ -17,31 +17,99 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( prepare_fp4_layer_for_marlin, ) from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + kE2M1ToFloat_handle, run_nvfp4_emulations, ) from vllm.platforms import current_platform from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer +from vllm.utils.import_utils import has_fbgemm_gpu from vllm.utils.math_utils import round_up logger = init_logger(__name__) +# NOTE: This is ordered by preferred backend. +# Example: if both are available, FLASHINFER_CUTLASS is preferred to VLLM_CUTLASS. class NvFp4LinearBackend(Enum): - VLLM_CUTLASS = "cutlass" FLASHINFER_CUTLASS = "flashinfer-cutlass" + VLLM_CUTLASS = "cutlass" + MARLIN = "marlin" FLASHINFER_TRTLLM = "flashinfer-trtllm" FLASHINFER_CUDNN = "flashinfer-cudnn" FBGEMM = "fbgemm" - MARLIN = "marlin" EMULATION = "emulation" +NVFP4_LINEAR_BACKENDS = list(NvFp4LinearBackend) + + +def is_backend_supported(backend: NvFp4LinearBackend) -> tuple[bool, str | None]: + reason = None + supported = True + + if backend == NvFp4LinearBackend.FLASHINFER_CUTLASS: + # cutlass_fp4_supported() checks that the vLLM NVFP4 kernels (both + # quantization and GEMM) were compiled for the current SM version. + # FlashInfer backends still rely on the vLLM quantization kernels, + # so we gate them on the same check. + supported = ( + cutlass_fp4_supported() + and current_platform.has_device_capability(100) + and has_flashinfer() + ) + + if not supported: + reason = "FlashInfer is required, >=sm_100 is required" + elif backend == NvFp4LinearBackend.VLLM_CUTLASS: + supported = cutlass_fp4_supported() + if not supported: + reason = "Cutlass is required" + elif backend == NvFp4LinearBackend.MARLIN: + supported = is_fp4_marlin_supported() + if not supported: + reason = "Marlin is required" + elif backend in [ + NvFp4LinearBackend.FLASHINFER_TRTLLM, + NvFp4LinearBackend.FLASHINFER_CUDNN, + ]: + supported = has_flashinfer() + if not supported: + reason = "FlashInfer is required" + elif backend == NvFp4LinearBackend.FBGEMM: + supported = has_fbgemm_gpu() + if not supported: + reason = "fbgemm_gpu is required" + elif backend == NvFp4LinearBackend.EMULATION: + # e.g. AMD Instinct does not support native NVFP4. + unsupported_reasons = {} + for other_backend in NVFP4_LINEAR_BACKENDS: + if other_backend == NvFp4LinearBackend.EMULATION: + continue + other_supported, other_reason = is_backend_supported(other_backend) + if not other_supported: + unsupported_reasons[other_backend] = other_reason + + if unsupported_reasons: + unsupported_reasons_str = "\n - ".join( + [f"{b.value}: {r}" for b, r in unsupported_reasons.items()] + ) + logger.warning_once( + f"NVFP4 linear falling back to the slow and unoptimized " + f"backend=NvFp4LinearBackend.EMULATION as no optimized backend is " + f"available (unavailable reasons:\n - {unsupported_reasons_str}\n). " + "In case you expect one of these backend to be used, " + "please verify your environment." + ) + + return supported, reason + + def select_nvfp4_linear_backend() -> NvFp4LinearBackend: """ Select the best available NVFP4 GEMM backend based on environment configuration and platform capabilities. """ - backend: NvFp4LinearBackend | None = None + selected_backend: NvFp4LinearBackend | None = None if envs.VLLM_USE_FBGEMM: try: @@ -51,51 +119,36 @@ def select_nvfp4_linear_backend() -> NvFp4LinearBackend: "Backend fbgemm requires fbgemm.f4f4bf16 operator, " "Please install with: pip install fbgemm-gpu-genai" ) from exc - backend = NvFp4LinearBackend.FBGEMM + selected_backend = NvFp4LinearBackend.FBGEMM elif envs.VLLM_USE_NVFP4_CT_EMULATIONS: - backend = NvFp4LinearBackend.EMULATION + selected_backend = NvFp4LinearBackend.EMULATION elif envs.VLLM_NVFP4_GEMM_BACKEND is None: - # Auto-select best available backend. - # cutlass_fp4_supported() checks that the vLLM NVFP4 kernels (both - # quantization and GEMM) were compiled for the current SM version. - # FlashInfer backends still rely on the vLLM quantization kernels, - # so we gate them on the same check. - if ( - cutlass_fp4_supported() - and current_platform.has_device_capability(100) - and has_flashinfer() - ): - backend = NvFp4LinearBackend.FLASHINFER_CUTLASS - elif cutlass_fp4_supported(): - backend = NvFp4LinearBackend.VLLM_CUTLASS - elif is_fp4_marlin_supported(): - backend = NvFp4LinearBackend.MARLIN + for backend in NVFP4_LINEAR_BACKENDS: + supported, reason = is_backend_supported(backend) + if supported: + selected_backend = backend + break else: - backend = NvFp4LinearBackend(envs.VLLM_NVFP4_GEMM_BACKEND) + selected_backend = NvFp4LinearBackend(envs.VLLM_NVFP4_GEMM_BACKEND) - # Validate that the backend is supported - if backend in ( - NvFp4LinearBackend.FLASHINFER_CUTLASS, - NvFp4LinearBackend.FLASHINFER_TRTLLM, - NvFp4LinearBackend.FLASHINFER_CUDNN, - ): - assert has_flashinfer(), f"FlashInfer is required for {backend}" - assert cutlass_fp4_supported(), ( - f"{backend} requires vLLM NVFP4 quantization kernels compiled " - f"for the current GPU (SM {current_platform.get_device_capability()})" - ) - elif backend == NvFp4LinearBackend.VLLM_CUTLASS: - assert cutlass_fp4_supported(), f"Cutlass is required for {backend}" - elif backend == NvFp4LinearBackend.MARLIN: - assert is_fp4_marlin_supported(), f"Marlin is required for {backend}" - elif backend is None: + if selected_backend is None: raise ValueError( f"No NVFP4 GEMM backend selected, " - f"available backends: {list(NvFp4LinearBackend)}" + f"available backends: {NVFP4_LINEAR_BACKENDS}" ) - logger.info_once(f"Using {backend} for NVFP4 GEMM") - return backend + supported, reason = is_backend_supported(selected_backend) + + if not supported: + raise ValueError( + f"The selected backend={selected_backend} is not supported in current " + f"environment. Reason: {reason}. Current environment: " + f"{envs.VLLM_USE_FBGEMM=}, {envs.VLLM_USE_NVFP4_CT_EMULATIONS=}, " + f"{envs.VLLM_NVFP4_GEMM_BACKEND}." + ) + + logger.info_once(f"Using {selected_backend} for NVFP4 GEMM") + return selected_backend def prepare_weights_for_nvfp4_flashinfer_trtllm( @@ -183,6 +236,10 @@ def convert_to_nvfp4_linear_kernel_format( layer.weight = torch.nn.Parameter(weight, requires_grad=False) layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) layer.weights_padding_cols = weights_padding_cols + elif backend == NvFp4LinearBackend.EMULATION: + # We can not call `.to(device)` during cuda graph capture - do it here instead. + # (operation not permitted when stream is capturing) + kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device) def apply_nvfp4_linear( @@ -190,6 +247,7 @@ def apply_nvfp4_linear( layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, + swizzle: bool | None = None, ) -> torch.Tensor: """ Apply NVFP4 linear transformation using the specified backend. @@ -220,6 +278,7 @@ def apply_nvfp4_linear( weight=weight, weight_scale_swizzled=weight_scale, weight_global_scale=weight_global_scale, + swizzle=swizzle, ) if bias is not None: out = out + bias diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 7b713536602..2ba4ef3fe8a 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -409,6 +409,7 @@ class RocmPlatform(Platform): "mxfp4", "torchao", "bitsandbytes", + "modelopt_fp4", ] @classmethod diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index e7f966b275e..b72108b4be6 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -461,3 +461,8 @@ def has_aiter() -> bool: def has_mori() -> bool: """Whether the optional `mori` package is available.""" return _has_module("mori") + + +def has_fbgemm_gpu() -> bool: + """Whether the optional `fbgemm_gpu` package is available.""" + return _has_module("fbgemm_gpu") From b2b2c5239ec65fc6ba1109b0d06cb7462d99b70e Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Mon, 6 Apr 2026 20:07:54 -0400 Subject: [PATCH 25/39] [MoE Refactor] Split up compressed_tensors_moe.py (#38960) Signed-off-by: Bill Nell --- docs/design/moe_kernel_features.md | 4 +- .../compressed_tensors_moe.py | 2541 ----------------- .../compressed_tensors_moe/__init__.py | 10 + .../compressed_tensors_moe.py | 175 ++ .../compressed_tensors_moe_w4a4_mxfp4.py | 168 ++ .../compressed_tensors_moe_w4a4_nvfp4.py | 306 ++ .../compressed_tensors_moe_w4a8_fp8.py | 343 +++ .../compressed_tensors_moe_w4a8_int8.py | 349 +++ .../compressed_tensors_moe_w8a8_fp8.py | 414 +++ .../compressed_tensors_moe_w8a8_int8.py | 161 ++ .../compressed_tensors_moe_wna16.py | 267 ++ .../compressed_tensors_moe_wna16_marlin.py | 575 ++++ 12 files changed, 2770 insertions(+), 2543 deletions(-) delete mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/__init__.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py create mode 100644 vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 03d25a9b1cb..7cf7b76d6ed 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -57,8 +57,8 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes. - [`ModelOptFp8MoEMethod`][vllm.model_executor.layers.quantization.modelopt.ModelOptFp8MoEMethod] - [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod] -- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW4A4Nvfp4MoEMethod] -- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW8A8Fp8MoEMethod] +- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod] +- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod] - [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod] - [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod] diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py deleted file mode 100644 index bce63bcbeb7..00000000000 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ /dev/null @@ -1,2541 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import enum -from enum import Enum - -import torch -from compressed_tensors import CompressionFormat -from compressed_tensors.quantization import ( - ActivationOrdering, - QuantizationArgs, - QuantizationStrategy, -) - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm import _custom_ops as ops -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, - FusedMoEActivationFormat, - FusedMoEExpertsModular, - FusedMoEMethodBase, - FusedMoeWeightScaleSupported, - UnquantizedFusedMoEMethod, -) -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.config import ( - FusedMoEConfig, - FusedMoEQuantConfig, - int4_w4a16_moe_quant_config, - int4_w4afp8_moe_quant_config, - int8_w8a8_moe_quant_config, - int8_w8a16_moe_quant_config, -) -from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts -from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( - BatchedMarlinExperts, - MarlinExperts, - fused_marlin_moe, -) -from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( - convert_to_fp8_moe_kernel_format, - make_fp8_moe_kernel, - make_fp8_moe_quant_config, - select_fp8_moe_backend, -) -from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( - Mxfp4MoeBackend, - make_mxfp4_moe_kernel, - make_mxfp4_moe_quant_config, -) -from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( - convert_to_nvfp4_moe_kernel_format, - is_global_sf_supported_for_nvfp4_backend, - make_nvfp4_moe_kernel, - make_nvfp4_moe_quant_config, - select_nvfp4_moe_backend, -) -from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa - WNA16_SUPPORTED_BITS, - WNA16_SUPPORTED_TYPES_MAP, -) -from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( - flashinfer_trtllm_mxint4_moe, - is_flashinfer_mxint4_moe_available, - prepare_static_weights_for_trtllm_mxint4_moe, -) -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - process_fp8_input_tensor_strategy_moe, - process_fp8_weight_tensor_strategy_moe, -) -from vllm.model_executor.layers.quantization.utils.marlin_utils import ( - check_moe_marlin_supports_layer, - get_marlin_input_dtype, - marlin_act_int8_process_scales, - marlin_make_workspace_new, - marlin_moe_permute_scales, -) -from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( - prepare_moe_fp4_layer_for_marlin, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - convert_bf16_scales_to_fp8, - convert_packed_uint4b8_to_signed_int4_inplace, - kFp8Dynamic128Sym, - kFp8DynamicTokenSym, - kFp8Static128BlockSym, - kFp8StaticChannelSym, - kFp8StaticTensorSym, - kNvfp4Dynamic, - kNvfp4Static, -) -from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( - normalize_e4m3fn_to_e4m3fnuz, -) -from vllm.model_executor.utils import replace_parameter, set_weight_attrs -from vllm.platforms import CpuArchEnum, current_platform - -logger = init_logger(__name__) - - -class GPTQMarlinState(Enum): - REPACK = enum.auto() - READY = enum.auto() - - -__all__ = [ - "CompressedTensorsMoEMethod", - "CompressedTensorsW8A8Fp8MoEMethod", - "CompressedTensorsW8A8Int8MoEMethod", - "CompressedTensorsWNA16MarlinMoEMethod", - "CompressedTensorsWNA16MoEMethod", - "CompressedTensorsW4A4Nvfp4MoEMethod", - "CompressedTensorsW4A8Int8MoEMethod", -] - - -class CompressedTensorsMoEMethod(FusedMoEMethodBase): - @staticmethod - def get_moe_method( - quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501 - layer: torch.nn.Module, - layer_name: str, - ) -> FusedMoEMethodBase: - # FusedMoE was made by combining multiple Linears so need to - # make sure quantization config for Linear can target it - quant_config._add_fused_moe_to_target_scheme_map() - unfused_names = [ - layer_name + proj_name - for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"] - ] - # TODO: refactor this to use expert_mapping and check all layer numbers - all_scheme_dicts = [ - quant_config.get_scheme_dict(layer, name) for name in unfused_names - ] - scheme_dict = all_scheme_dicts.pop() - - # multiple schemes found - if not all([cur_dict == scheme_dict for cur_dict in all_scheme_dicts]): - raise ValueError( - "All MoE projections need to have same " - "quantization scheme but found multiple" - ) - - if scheme_dict is None: # ignored layer - return UnquantizedFusedMoEMethod(layer.moe_config) - - # TODO: @dsikka: refactor this to use schemes as other kernels - # are supported + check if the layer is being ignored. - weight_quant = scheme_dict.get("weights") - input_quant = scheme_dict.get("input_activations") - format = scheme_dict.get("format") - - if quant_config._is_mxfp4(weight_quant): - return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config) - - if quant_config._is_wNa16_group_channel(weight_quant, input_quant): - # group_size=None means channelwise - group_size = weight_quant.group_size or -1 - - valid_format_and_bits = ( - weight_quant.num_bits in WNA16_SUPPORTED_BITS - and format == CompressionFormat.pack_quantized.value - ) - - if not valid_format_and_bits: - raise ValueError( - "For Fused MoE layers, only format: ", - f"{CompressionFormat.pack_quantized.value} ", - f" and bits: {WNA16_SUPPORTED_BITS} is supported ", - f"but got format: {CompressionFormat.pack_quantized.value} " - f" and bits: {weight_quant.num_bits}", - ) - - # Prefer to use the MarlinMoE kernel when it is supported. - if ( - not check_moe_marlin_supports_layer(layer, group_size) - or current_platform.is_rocm() - ): - if ( - weight_quant.strategy == QuantizationStrategy.GROUP - and weight_quant.actorder - in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC) - ): - raise ValueError( - "WNA16MoE is not supported with actorder=group/dynamic." - ) - logger.info_once("Using CompressedTensorsWNA16MoEMethod") - return CompressedTensorsWNA16MoEMethod( - weight_quant, input_quant, layer.moe_config - ) - else: - logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod") - return CompressedTensorsWNA16MarlinMoEMethod( - weight_quant, input_quant, layer.moe_config - ) - elif quant_config._is_nvfp4_format(weight_quant): - _is_valid_nvfp4_activations = ( - quant_config._is_nvfp4_format(input_quant) or input_quant is None - ) - if not _is_valid_nvfp4_activations: - raise ValueError( - "For NVFP4 weights, input quantization must also be NVFP4 format ", - f"or None for NVFP4A16, found {input_quant}", - ) - return CompressedTensorsW4A4Nvfp4MoEMethod( - layer.moe_config, layer_name, use_a16=(input_quant is None) - ) - elif ( - quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant) - or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant) - or quant_config._is_fp8_w8a8(weight_quant, input_quant) - ): - return CompressedTensorsW8A8Fp8MoEMethod( - weight_quant, input_quant, layer.moe_config - ) - elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant): - return CompressedTensorsW8A8Int8MoEMethod( - weight_quant, input_quant, layer.moe_config - ) - elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant): - logger.info_once("Using CompressedTensorsW4A8Fp8MoEMethod") - return CompressedTensorsW4A8Fp8MoEMethod( - weight_quant, input_quant, layer.moe_config - ) - elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant): - return CompressedTensorsW4A8Int8MoEMethod( - weight_quant, input_quant, layer.moe_config - ) - else: - raise RuntimeError( - f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}" - ) - - -class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): - def __init__(self, moe): - super().__init__(moe) - self.group_size = 32 - self.mxfp4_backend = Mxfp4MoeBackend.MARLIN - self.experts_cls = MarlinExperts - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - layer.num_experts = num_experts - layer.params_dtype = params_dtype - - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - # 2 fp4 items are packed in the input dimension - hidden_size // 2, - requires_grad=False, - dtype=torch.uint8, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_packed", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - # 2 fp4 items are packed in the input dimension - intermediate_size_per_partition // 2, - dtype=torch.uint8, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_packed", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - w13_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - # 2 fp4 items are packed in the input dimension - hidden_size // self.group_size, - dtype=torch.uint8, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} - ) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - - w2_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - # 2 fp4 items are packed in the input dimension - intermediate_size_per_partition // self.group_size, - dtype=torch.uint8, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - return make_mxfp4_moe_quant_config( - mxfp4_backend=self.mxfp4_backend, - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - ) - - def process_weights_after_loading(self, layer: FusedMoE) -> 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") - - logger.warning_once( - "Your GPU does not have native support for FP4 computation but " - "FP4 quantization is being used. Weight-only FP4 compression " - "will be used leveraging the Marlin kernel. This may degrade " - "performance for compute-heavy workloads." - ) - prepare_moe_fp4_layer_for_marlin(layer) - - self.moe_quant_config = self.get_fused_moe_quant_config(layer) - if self.moe_quant_config is not None: - self.moe_kernel = make_mxfp4_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - experts_cls=self.experts_cls, - mxfp4_backend=self.mxfp4_backend, - shared_experts=layer.shared_experts, - routing_tables=layer._maybe_init_expert_routing_tables(), - ) - - def apply( - self, - layer: FusedMoE, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - assert self.moe_kernel is not None - return self.moe_kernel.apply( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - shared_experts_input=shared_experts_input, - ) - - -class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): - def __init__( - self, - moe: FusedMoEConfig, - layer_name: str | None = None, - use_a16: bool = False, - ): - super().__init__(moe) - self.group_size = 16 - - # 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, - ) - - self.use_global_sf = is_global_sf_supported_for_nvfp4_backend( - self.nvfp4_backend - ) - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - layer.num_experts = num_experts - layer.params_dtype = params_dtype - w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - w13_num_shards * intermediate_size_per_partition, - # 2 fp4 items are packed in the input dimension - hidden_size // 2, - requires_grad=False, - dtype=torch.uint8, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_packed", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - # 2 fp4 items are packed in the input dimension - intermediate_size_per_partition // 2, - dtype=torch.uint8, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_packed", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - # Weight Scales - w13_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, - w13_num_shards * intermediate_size_per_partition, - # 2 fp4 items are packed in the input dimension - hidden_size // self.group_size, - dtype=torch.float8_e4m3fn, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} - ) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - - w2_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - # 2 fp4 items are packed in the input dimension - intermediate_size_per_partition // self.group_size, - dtype=torch.float8_e4m3fn, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} - ) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - - # Weight Global Scales - w13_weight_scale_2 = torch.nn.Parameter( - torch.empty(num_experts, w13_num_shards, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w13_weight_global_scale", w13_weight_scale_2) - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} - ) - set_weight_attrs(w13_weight_scale_2, extra_weight_attrs) - - w2_weight_scale_2 = torch.nn.Parameter( - torch.empty(num_experts, dtype=torch.float32), requires_grad=False - ) - layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2) - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} - ) - set_weight_attrs(w2_weight_scale_2, extra_weight_attrs) - - # Input Global Scales - w13_input_scale = torch.nn.Parameter( - torch.empty(num_experts, w13_num_shards, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w13_input_global_scale", w13_input_scale) - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} - ) - set_weight_attrs(w13_input_scale, extra_weight_attrs) - - w2_input_scale = torch.nn.Parameter( - torch.empty(num_experts, dtype=torch.float32), requires_grad=False - ) - layer.register_parameter("w2_input_global_scale", w2_input_scale) - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} - ) - set_weight_attrs(w2_input_scale, extra_weight_attrs) - - def process_weights_after_loading(self, layer: FusedMoE) -> 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") - - # Use a single gscale for w13. - if self.moe.is_act_and_mul and not torch.allclose( - layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1] - ): - logger.warning_once( - "w1_weight_global_scale must match w3_weight_global_scale. " - "Accuracy may be affected.", - ) - w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous() - - # Shuffle weights into the NvFp4 kernel format. - ( - w13, - w13_scale, - w13_scale_2, - a13_scale, - w2, - w2_scale, - w2_scale_2, - a2_scale, - ) = convert_to_nvfp4_moe_kernel_format( - nvfp4_backend=self.nvfp4_backend, - layer=layer, - 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), - 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), - is_act_and_mul=self.moe.is_act_and_mul, - ) - - replace_parameter(layer, "w13_weight", w13) - replace_parameter(layer, "w13_weight_scale", w13_scale) - replace_parameter(layer, "w2_weight", w2) - replace_parameter(layer, "w2_weight_scale", w2_scale) - layer.w13_weight_scale_2 = w13_scale_2 - layer.w2_weight_scale_2 = w2_scale_2 - layer.w13_input_scale = a13_scale - layer.w2_input_scale = a2_scale - - # Setup modular kernel. - self.moe_quant_config = self.get_fused_moe_quant_config(layer) - assert self.experts_cls is not None - self.moe_kernel = make_nvfp4_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - experts_cls=self.experts_cls, - shared_experts=layer.shared_experts, - routing_tables=layer._maybe_init_expert_routing_tables(), - ) - self.moe_kernel.fused_experts.process_weights_after_loading(layer) - - def maybe_make_prepare_finalize( - self, - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalizeModular | None: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic. This function should not be called." - ) - - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: - return make_nvfp4_moe_quant_config( - backend=self.nvfp4_backend, - w13_scale=layer.w13_weight_scale, - 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, - ) - - def apply_monolithic( - self, - layer: FusedMoE, - x: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - assert self.is_monolithic - assert self.moe_kernel is not None - return self.moe_kernel.apply_monolithic( - x, - layer.w13_weight, - layer.w2_weight, - router_logits, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - e_score_correction_bias=layer.e_score_correction_bias, - routed_scaling_factor=layer.routed_scaling_factor, - ) - - def apply( - self, - layer: FusedMoE, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - assert self.moe_kernel is not None - return self.moe_kernel.apply( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - shared_experts_input=shared_experts_input, - ) - - -class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): - """W8A8 FP8 MoE quantization using compressed tensors.""" - - def __init__( - self, - weight_quant: QuantizationArgs, - input_quant: QuantizationArgs, - moe: FusedMoEConfig, - layer_name: str | None = None, - ): - super().__init__(moe) - self.weight_quant = weight_quant - self.input_quant = input_quant - - per_tensor = ( - self.weight_quant.strategy == QuantizationStrategy.TENSOR - and self.input_quant.strategy == QuantizationStrategy.TENSOR - ) - per_channel = ( - self.weight_quant.strategy == QuantizationStrategy.CHANNEL - and self.input_quant.strategy == QuantizationStrategy.TOKEN - ) - if not (per_tensor or per_channel): - assert self.weight_quant.strategy == QuantizationStrategy.BLOCK - self.weight_block_size = self.weight_quant.block_structure - assert self.weight_quant.dynamic is not None - else: - self.weight_block_size = None - self.block_quant = self.weight_block_size is not None - - self.static_input_scales = not self.input_quant.dynamic - if self.static_input_scales and per_channel: - raise ValueError( - "For FP8 Fused MoE layer, we require either per tensor or " - "channelwise, dynamic per token quantization." - ) - - ct2vllm_weight = { - QuantizationStrategy.CHANNEL: kFp8StaticChannelSym, - QuantizationStrategy.TENSOR: kFp8StaticTensorSym, - QuantizationStrategy.BLOCK: kFp8Static128BlockSym, - } - ct2vllm_act = { - QuantizationStrategy.TOKEN: kFp8DynamicTokenSym, - QuantizationStrategy.TENSOR: ( - kFp8StaticTensorSym if self.static_input_scales else kFp8Dynamic128Sym - ), - } - weight_key = ct2vllm_weight[self.weight_quant.strategy] - if weight_key == kFp8Static128BlockSym: - activation_key = kFp8Dynamic128Sym - else: - activation_key = ct2vllm_act[self.input_quant.strategy] - - # Select Fp8 MoE backend - self.fp8_backend, self.experts_cls = select_fp8_moe_backend( - config=self.moe, - weight_key=weight_key, - activation_key=activation_key, - allow_vllm_cutlass=True, - ) - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - layer.num_experts = num_experts - layer.orig_dtype = params_dtype - layer.weight_block_size = None - - params_dtype = torch.float8_e4m3fn - w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - - if self.block_quant: - assert self.weight_block_size is not None - layer.weight_block_size = self.weight_block_size - tp_size = get_tensor_model_parallel_world_size() - block_n, block_k = ( - self.weight_block_size[0], - self.weight_block_size[1], - ) - # NOTE: To ensure proper alignment of the block-wise quantization - # scales, the output_size of the weights for both the gate and up - # layers must be divisible by block_n. - # Required by column parallel or enabling merged weights - if intermediate_size_per_partition % block_n != 0: - raise ValueError( - f"The output_size of gate's and up's weight = " - f"{intermediate_size_per_partition} is not divisible by " - f"weight quantization block_n = {block_n}." - ) - if tp_size > 1 and intermediate_size_per_partition % block_k != 0: - # Required by row parallel - raise ValueError( - f"The input_size of down's weight = " - f"{intermediate_size_per_partition} is not divisible by " - f"weight quantization block_k = {block_k}." - ) - - # WEIGHTS - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - w13_num_shards * intermediate_size_per_partition, - hidden_size, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - # WEIGHT_SCALES - if self.weight_quant.strategy == QuantizationStrategy.TENSOR: - # For gated MoE, allocate 2 scales for w1 and w3 respectively. - # They will be combined to a single scale after weight loading. - # For non-gated MoE, allocate 1 scale for w13. - w13_weight_scale = torch.nn.Parameter( - torch.ones(num_experts, w13_num_shards, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - w2_weight_scale = torch.nn.Parameter( - torch.ones(num_experts, dtype=torch.float32), requires_grad=False - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - # Add PER-TENSOR quantization for FusedMoE.weight_loader. - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} - ) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - - elif self.weight_quant.strategy == QuantizationStrategy.CHANNEL: - w13_weight_scale = torch.nn.Parameter( - torch.ones( - num_experts, - w13_num_shards * intermediate_size_per_partition, - 1, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - w2_weight_scale = torch.nn.Parameter( - torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - # Add PER-CHANNEL quantization for FusedMoE.weight_loader. - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} - ) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - - elif self.weight_quant.strategy == QuantizationStrategy.BLOCK: - w13_weight_scale = torch.nn.Parameter( - torch.ones( - num_experts, - w13_num_shards - * ((intermediate_size_per_partition + block_n - 1) // block_n), - (hidden_size + block_k - 1) // block_k, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - w2_weight_scale = torch.nn.Parameter( - torch.ones( - num_experts, - (hidden_size + block_n - 1) // block_n, - (intermediate_size_per_partition + block_k - 1) // block_k, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - # Add PER-CHANNEL quantization for FusedMoE.weight_loader. - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} - ) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - - # INPUT_SCALES - if self.static_input_scales: - w13_input_scale = torch.nn.Parameter( - torch.ones(num_experts, dtype=torch.float32), requires_grad=False - ) - layer.register_parameter("w13_input_scale", w13_input_scale) - set_weight_attrs(w13_input_scale, extra_weight_attrs) - - w2_input_scale = torch.nn.Parameter( - torch.ones(num_experts, dtype=torch.float32), requires_grad=False - ) - layer.register_parameter("w2_input_scale", w2_input_scale) - set_weight_attrs(w2_input_scale, extra_weight_attrs) - else: - layer.w13_input_scale = None - layer.w2_input_scale = None - - def process_weights_after_loading(self, layer: FusedMoE) -> None: - # Allow for accessing weights and scales in standard way. - w13 = layer.w13_weight - w2 = layer.w2_weight - w13_scale = layer.w13_weight_scale - w2_scale = layer.w2_weight_scale - w13_input_scale = layer.w13_input_scale - w2_input_scale = layer.w2_input_scale - - # MI300x and MI325x use FNUZ format for FP8. Convert if needed. - if current_platform.is_fp8_fnuz(): - w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz( - w13, w13_scale, w13_input_scale - ) - w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz( - w2, w2_scale, w2_input_scale - ) - - # Per tensor kernels require single activation scale. Use the max. - if self.static_input_scales: - assert self.input_quant.strategy == QuantizationStrategy.TENSOR - assert w13_input_scale is not None and w2_input_scale is not None - w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe( - w13_input_scale, w2_input_scale - ) - replace_parameter(layer, "w13_input_scale", w13_input_scale) - replace_parameter(layer, "w2_input_scale", w2_input_scale) - - # Per-tensor kernels use a single scale, for W13, but on disk there - # is a separate scale for W1 and W3. Requantize with the max scale. - if self.weight_quant.strategy == QuantizationStrategy.TENSOR: - w13, w13_scale = process_fp8_weight_tensor_strategy_moe( - w13, - w13_scale, - shard_size=layer.intermediate_size_per_partition, - num_experts=layer.local_num_experts, - is_act_and_mul=self.moe.is_act_and_mul, - ) - - w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( - fp8_backend=self.fp8_backend, - layer=layer, - w13=w13, - w2=w2, - w13_scale=w13_scale, - w2_scale=w2_scale, - w13_input_scale=w13_input_scale, - w2_input_scale=w2_input_scale, - ) - - # Replace parameters with updated versions. Note that this helper - # function ensures the replacement is compatible with RL weight reloads. - replace_parameter(layer, "w13_weight", w13) - replace_parameter(layer, "w2_weight", w2) - replace_parameter(layer, "w13_weight_scale", w13_scale) - replace_parameter(layer, "w2_weight_scale", w2_scale) - - # Setup modular kernel for TP case and naive DP/EP case. - # In non-naive DP/EP case, we will create a ModularKernelMethod. - # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel - # in both cases. - self.moe_quant_config = self.get_fused_moe_quant_config(layer) - if self.moe_quant_config: - assert self.experts_cls is not None - self.moe_kernel = make_fp8_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - fp8_backend=self.fp8_backend, - experts_cls=self.experts_cls, - routing_tables=layer._maybe_init_expert_routing_tables(), - shared_experts=layer.shared_experts, - ) - - def maybe_make_prepare_finalize( - self, - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalizeModular | None: - raise ValueError( - f"{self.__class__.__name__} uses the new modular kernel initialization " - "logic. This function should not be called." - ) - - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: - is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN - return make_fp8_moe_quant_config( - fp8_backend=self.fp8_backend, - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - a1_scale=layer.w13_input_scale, - a2_scale=layer.w2_input_scale, - per_act_token_quant=is_per_token, - per_out_ch_quant=is_per_token, - block_shape=self.weight_block_size, - ) - - def apply_monolithic( - self, - layer: FusedMoE, - x: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - assert self.moe_kernel is not None - return self.moe_kernel.apply_monolithic( - x, - layer.w13_weight, - layer.w2_weight, - router_logits, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - e_score_correction_bias=layer.e_score_correction_bias, - routed_scaling_factor=layer.routed_scaling_factor, - ) - - def apply( - self, - layer: FusedMoE, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - assert not self.is_monolithic - assert self.moe_kernel is not None - return self.moe_kernel.apply( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - # TODO(rob): investigate the disable_expert_map introduced by: - # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501 - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - shared_experts_input=shared_experts_input, - ) - - @property - def supports_eplb(self) -> bool: - return True - - -class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): - def __init__( - self, - weight_quant: QuantizationArgs, - input_quant: QuantizationArgs, - moe: FusedMoEConfig, - layer_name: str | None = None, - ): - super().__init__(moe) - self.weight_quant = weight_quant - self.input_quant = input_quant - - per_channel = ( - self.weight_quant.strategy == QuantizationStrategy.CHANNEL - and self.input_quant.strategy == QuantizationStrategy.TOKEN - ) - if not per_channel: - raise ValueError( - "For INT8 Fused MoE layers, we require channelwise, " - "dynamic per token quantization. Found " - f"{self.weight_quant}, {self.input_quant}" - ) - - self.static_input_scales = not self.input_quant.dynamic - if self.static_input_scales: - raise ValueError( - "For INT8 Fused MoE layers, we require channelwise, " - "dynamic per token quantization. Found static input scales." - ) - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - params_dtype = torch.int8 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - - # WEIGHTS - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - w13_num_shards * intermediate_size_per_partition, - hidden_size, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - # WEIGHT_SCALES - assert self.weight_quant.strategy == QuantizationStrategy.CHANNEL - w13_weight_scale = torch.nn.Parameter( - torch.ones( - num_experts, - w13_num_shards * intermediate_size_per_partition, - 1, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - w2_weight_scale = torch.nn.Parameter( - torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - # Add PER-CHANNEL quantization for FusedMoE.weight_loader. - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} - ) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - - # INPUT_SCALES - assert not self.static_input_scales - layer.w13_input_scale = None - layer.w2_input_scale = None - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - pass - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - return int8_w8a8_moe_quant_config( - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - a1_scale=layer.w13_input_scale, - a2_scale=layer.w2_input_scale, - per_act_token_quant=True, - ) - - def apply( - self, - layer: FusedMoE, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe import fused_experts - - return fused_experts( - hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, - topk_weights=topk_weights, - topk_ids=topk_ids, - inplace=not self.moe.disable_inplace, - activation=layer.activation, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - quant_config=self.moe_quant_config, - ) - - -class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): - def __init__( - self, - weight_quant: QuantizationArgs, - input_quant: QuantizationArgs | None, - moe: FusedMoEConfig, - layer_name: str | None = None, - ): - super().__init__(moe) - self.weight_quant = weight_quant - self.input_quant = input_quant - assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE" - ) - # Extract properties from weight_quant - self.num_bits = weight_quant.num_bits - self.packed_factor = 32 // weight_quant.num_bits - self.strategy = weight_quant.strategy - self.group_size = weight_quant.group_size - self.actorder = weight_quant.actorder - - self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits] - - self.marlin_input_dtype = get_marlin_input_dtype(layer_name) - self.use_flashinfer_mxint4_moe = ( - is_flashinfer_mxint4_moe_available() - and self.group_size == 32 - and weight_quant.num_bits == 4 - ) - self.kernel_backend = ( - "Flashinfer" if self.use_flashinfer_mxint4_moe else "Marlin" - ) - logger.info_once( - f"Using {self.kernel_backend} backend for WNA16 MoE " - f"(group_size={self.group_size}, num_bits={self.num_bits})", - scope="local", - ) - - def get_weight_shape( - self, - weight_name: str, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - num_groups_w2: int | None = None, - num_groups_w13: int | None = None, - ) -> tuple[int, int, int]: - """ - Get the shape of the weight based on the weight name, number of experts - hidden size, intermediate size per partition, number of groups for w2, - and number of groups for w13. Pass in num_groups_w2 and num_groups_w13 - for weight scales. - """ - if weight_name == "w13_scale": - assert num_groups_w13 is not None, ( - "num_groups_w13 must be provided for weight scales" - ) - if weight_name == "w2_scale": - assert num_groups_w2 is not None, ( - "num_groups_w2 must be provided for weight scales" - ) - w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - shape_map = { - "w13_weight": { - "Flashinfer": ( - num_experts, - w13_num_shards * intermediate_size_per_partition, - hidden_size // self.packed_factor, - ), - "Marlin": ( - num_experts, - hidden_size // self.packed_factor, - w13_num_shards * intermediate_size_per_partition, - ), - }, - "w13_scale": { - "Flashinfer": ( - num_experts, - w13_num_shards * intermediate_size_per_partition, - num_groups_w13, - ), - "Marlin": ( - num_experts, - num_groups_w13, - w13_num_shards * intermediate_size_per_partition, - ), - }, - "w2_weight": { - "Flashinfer": ( - num_experts, - hidden_size, - intermediate_size_per_partition // self.packed_factor, - ), - "Marlin": ( - num_experts, - intermediate_size_per_partition // self.packed_factor, - hidden_size, - ), - }, - "w2_scale": { - "Flashinfer": (num_experts, hidden_size, num_groups_w2), - "Marlin": (num_experts, num_groups_w2, hidden_size), - }, - } - return shape_map[weight_name][self.kernel_backend] - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full") - - # Will transpose the loaded weight along the - # intermediate and hidden dim sizes. Will - # shard for TP along the transposed dims - is_transposed = self.kernel_backend != "Flashinfer" - extra_weight_attrs.update( - {"is_transposed": is_transposed, "quant_method": self.strategy} - ) - - w13_weight = torch.nn.Parameter( - torch.empty( - *self.get_weight_shape( - "w13_weight", - num_experts, - hidden_size, - intermediate_size_per_partition, - ), - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_packed", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - *self.get_weight_shape( - "w2_weight", - num_experts, - hidden_size, - intermediate_size_per_partition, - ), - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_packed", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - # In the case where we have actorder/g_idx, - # we do not partition the w2 scales - load_full_w2 = self.actorder and self.group_size != -1 - w2_scales_size = ( - intermediate_size_full if load_full_w2 else intermediate_size_per_partition - ) - - self.is_k_full = (not self.actorder) or ( - intermediate_size_per_partition == intermediate_size_full - ) - - if self.strategy == "channel": - num_groups_w2 = num_groups_w13 = 1 - self.group_size = -1 - else: - num_groups_w2 = w2_scales_size // self.group_size - num_groups_w13 = hidden_size // self.group_size - - layer.num_groups_w13 = num_groups_w13 - layer.num_groups_w2 = num_groups_w2 - - w13_scale = torch.nn.Parameter( - torch.ones( - *self.get_weight_shape( - "w13_scale", - num_experts, - hidden_size, - intermediate_size_per_partition, - num_groups_w13=num_groups_w13, - ), - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_scale) - set_weight_attrs(w13_scale, extra_weight_attrs) - - w2_scale = torch.nn.Parameter( - torch.ones( - *self.get_weight_shape( - "w2_scale", - num_experts, - hidden_size, - intermediate_size_per_partition, - num_groups_w2=num_groups_w2, - ), - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_scale) - set_weight_attrs(w2_scale, extra_weight_attrs) - set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) - - w2_weight_shape = torch.nn.Parameter( - torch.empty(num_experts, 2), requires_grad=False - ) - layer.register_parameter("w2_weight_shape", w2_weight_shape) - set_weight_attrs(w2_weight_shape, extra_weight_attrs) - w13_weight_shape = torch.nn.Parameter( - torch.empty(num_experts, 2), requires_grad=False - ) - - layer.register_parameter("w13_weight_shape", w13_weight_shape) - set_weight_attrs(w13_weight_shape, extra_weight_attrs) - - w13_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_g_idx", w13_g_idx) - set_weight_attrs(w13_g_idx, extra_weight_attrs) - - w2_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_g_idx", w2_g_idx) - set_weight_attrs(w2_g_idx, extra_weight_attrs) - - w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) - - w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) - - layer.a13_scale = None - layer.a2_scale = None - layer.marlin_state = GPTQMarlinState.REPACK - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - num_experts = layer.w13_weight_g_idx.shape[0] - device = layer.w13_weight_g_idx.device - if self.kernel_backend == "Flashinfer": - dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe( - layer.w13_weight_packed, - layer.w13_weight_scale, - layer.w2_weight_packed, - layer.w2_weight_scale, - ) - replace_parameter( - layer, "w13_weight_packed", dict_weights_mxint4["gemm1_weights"] - ) - replace_parameter( - layer, "w13_weight_scale", dict_weights_mxint4["gemm1_scales"] - ) - replace_parameter( - layer, "w2_weight_packed", dict_weights_mxint4["gemm2_weights"] - ) - replace_parameter( - layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"] - ) - return None - - is_a_8bit = ( - self.marlin_input_dtype is not None - and self.marlin_input_dtype.itemsize == 1 - ) - - if self.marlin_input_dtype == torch.float8_e4m3fn: - # NOTE: for non-zp quantization format only - ops.marlin_int4_fp8_preprocess(layer.w13_weight_packed, inplace=True) - ops.marlin_int4_fp8_preprocess(layer.w2_weight_packed, inplace=True) - layer.w13_weight_scale.data = layer.w13_weight_scale.data * 512 - layer.w2_weight_scale.data = layer.w2_weight_scale.data * 512 - - # when running models with grouped act order, - # resort to g_idx values provided in checkpoint - if self.actorder == "group": - w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx) - w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx) - w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx) - w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx) - - for e in range(num_experts): - w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to( - torch.int32 - ) - w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to( - torch.int32 - ) - w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][ - w13_g_idx_sort_indices[e] - ] - w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]] - - replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx) - replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx) - replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) - - else: - layer.w13_weight_g_idx = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - layer.w2_weight_g_idx = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - layer.w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - layer.w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - - marlin_w13_qweight = ops.gptq_marlin_moe_repack( - layer.w13_weight_packed, - layer.w13_g_idx_sort_indices, - layer.w13_weight_packed.shape[1] * self.packed_factor, - layer.w13_weight_packed.shape[2], - self.num_bits, - is_a_8bit=is_a_8bit, - ) - replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight) - - marlin_w2_qweight = ops.gptq_marlin_moe_repack( - layer.w2_weight_packed, - layer.w2_g_idx_sort_indices, - layer.w2_weight_packed.shape[1] * self.packed_factor, - layer.w2_weight_packed.shape[2], - self.num_bits, - is_a_8bit=is_a_8bit, - ) - replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight) - - # Repack scales - marlin_w13_scales = marlin_moe_permute_scales( - s=layer.w13_weight_scale, - size_k=layer.w13_weight_packed.shape[2], - size_n=layer.w13_weight_scale.shape[2], - group_size=self.group_size, - is_a_8bit=is_a_8bit, - ) - if self.marlin_input_dtype == torch.int8 and layer.num_groups_w13 > 1: - marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales( - marlin_w13_scales - ) - layer.register_parameter( - "w13_input_global_scale", - torch.nn.Parameter(w13_input_global_scale, requires_grad=False), - ) - replace_parameter(layer, "w13_weight_scale", marlin_w13_scales) - - marlin_w2_scales = marlin_moe_permute_scales( - s=layer.w2_weight_scale, - size_k=layer.w2_weight_scale.shape[1] - * (self.group_size if self.group_size != -1 else self.packed_factor), - size_n=layer.w2_weight_scale.shape[2], - group_size=self.group_size, - is_a_8bit=is_a_8bit, - ) - if self.marlin_input_dtype == torch.int8 and layer.num_groups_w2 > 1: - marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales( - marlin_w2_scales - ) - layer.register_parameter( - "w2_input_global_scale", - torch.nn.Parameter(w2_input_global_scale, requires_grad=False), - ) - replace_parameter(layer, "w2_weight_scale", marlin_w2_scales) - - layer.workspace = marlin_make_workspace_new(device, 4) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - if self.num_bits != 4: - return None - return int4_w4a16_moe_quant_config( - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - w1_zp=None, - w2_zp=None, - block_shape=[0, self.group_size], - ) - - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - assert self.num_bits == 4, "only supporting w4" - layer.w13_weight = layer.w13_weight_packed - layer.w2_weight = layer.w2_weight_packed - assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]]) - assert self.moe_quant_config is not None - if ( - prepare_finalize.activation_format - == mk.FusedMoEActivationFormat.BatchedExperts - ): - max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank() - assert max_num_tokens_per_rank is not None - return BatchedMarlinExperts( - max_num_tokens=max_num_tokens_per_rank, - num_dispatchers=prepare_finalize.num_dispatchers(), - moe_config=self.moe, - quant_config=self.moe_quant_config, - w13_g_idx=layer.w13_weight_g_idx, - w2_g_idx=layer.w2_weight_g_idx, - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, - is_k_full=self.is_k_full, - ) - else: - return MarlinExperts( - moe_config=self.moe, - quant_config=self.moe_quant_config, - w13_g_idx=layer.w13_weight_g_idx, - w2_g_idx=layer.w2_weight_g_idx, - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, - is_k_full=self.is_k_full, - ) - - @property - def is_monolithic(self) -> bool: - return self.kernel_backend == "Flashinfer" - - def apply_monolithic( - self, - layer: FusedMoE, - x: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - assert self.kernel_backend == "Flashinfer" - return flashinfer_trtllm_mxint4_moe( - x=x, - router_logits=router_logits, - w13_weight_packed=layer.w13_weight_packed, - w13_weight_scale=layer.w13_weight_scale, - w2_weight_packed=layer.w2_weight_packed, - w2_weight_scale=layer.w2_weight_scale, - global_num_experts=layer.global_num_experts, - top_k=layer.top_k, - intermediate_size_per_partition=layer.intermediate_size_per_partition, - local_num_experts=layer.local_num_experts, - ep_rank=layer.ep_rank, - num_expert_group=layer.num_expert_group, - topk_group=layer.topk_group, - e_score_correction_bias=layer.e_score_correction_bias, - routing_method_type=layer.routing_method_type, - ) - - def apply( - self, - layer: FusedMoE, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - assert self.kernel_backend == "Marlin" - return fused_marlin_moe( - x, - layer.w13_weight_packed, - layer.w2_weight_packed, - None, - None, - layer.w13_weight_scale, - layer.w2_weight_scale, - topk_weights, - topk_ids, - input_global_scale1=getattr(layer, "w13_input_global_scale", None), - input_global_scale2=getattr(layer, "w2_input_global_scale", None), - quant_type_id=self.quant_type.id, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=layer.global_num_experts, - activation=layer.activation, - expert_map=layer.expert_map, - g_idx1=layer.w13_weight_g_idx, - g_idx2=layer.w2_weight_g_idx, - sort_indices1=layer.w13_g_idx_sort_indices, - sort_indices2=layer.w2_g_idx_sort_indices, - workspace=layer.workspace, - input_dtype=self.marlin_input_dtype, - is_k_full=self.is_k_full, - inplace=not self.moe.disable_inplace, - ) - - -class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): - def __init__( - self, - weight_quant: QuantizationArgs, - input_quant: QuantizationArgs | None, - moe: FusedMoEConfig, - layer_name: str | None = None, - ): - super().__init__(moe) - self.weight_quant = weight_quant - self.input_quant = input_quant - # Extract properties from weight_quant - self.num_bits = weight_quant.num_bits - self.packed_factor = 32 // weight_quant.num_bits - self.strategy = weight_quant.strategy - # channelwise is not supported by this kernel - assert weight_quant.strategy == "group" - self.group_size = weight_quant.group_size - # grouped actorder isn't supported by this kernel - assert weight_quant.actorder != "group" - assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE" - ) - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - # Will transpose the loaded weight along the - # intermediate and hidden dim sizes. Will - # shard for TP along the transposed dims - extra_weight_attrs.update( - {"is_transposed": True, "quant_method": self.strategy} - ) - w13_num_shards = 2 if self.moe.is_act_and_mul else 1 - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size // self.packed_factor, - w13_num_shards * intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_packed", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition // self.packed_factor, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_packed", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - w2_scales_size = intermediate_size_per_partition - - if self.strategy == "channel": - num_groups_w2 = num_groups_w13 = 1 - self.group_size = -1 - else: - num_groups_w2 = w2_scales_size // self.group_size - num_groups_w13 = hidden_size // self.group_size - - w13_scale = torch.nn.Parameter( - torch.ones( - num_experts, - num_groups_w13, - w13_num_shards * intermediate_size_per_partition, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_scale) - set_weight_attrs(w13_scale, extra_weight_attrs) - - w2_scale = torch.nn.Parameter( - torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_scale) - set_weight_attrs(w2_scale, extra_weight_attrs) - set_weight_attrs(w2_scale, {"load_full_w2": False}) - - w2_weight_shape = torch.nn.Parameter( - torch.empty(num_experts, 2), requires_grad=False - ) - layer.register_parameter("w2_weight_shape", w2_weight_shape) - set_weight_attrs(w2_weight_shape, extra_weight_attrs) - w13_weight_shape = torch.nn.Parameter( - torch.empty(num_experts, 2), requires_grad=False - ) - - layer.register_parameter("w13_weight_shape", w13_weight_shape) - set_weight_attrs(w13_weight_shape, extra_weight_attrs) - - w13_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_g_idx", w13_g_idx) - set_weight_attrs(w13_g_idx, extra_weight_attrs) - - w2_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_g_idx", w2_g_idx) - set_weight_attrs(w2_g_idx, extra_weight_attrs) - - w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) - set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) - - w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty( - num_experts, - intermediate_size_per_partition, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) - set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) - - layer.a13_scale = None - layer.a2_scale = None - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Reconfigure packed weights and scales to match moe_wna16 format - layer.w13_weight_packed = torch.nn.Parameter( - layer.w13_weight_packed.transpose(1, 2).contiguous().view(torch.uint8), - requires_grad=False, - ) - layer.w2_weight_packed = torch.nn.Parameter( - layer.w2_weight_packed.transpose(1, 2).contiguous().view(torch.uint8), - requires_grad=False, - ) - layer.w13_weight_scale = torch.nn.Parameter( - layer.w13_weight_scale.transpose(1, 2).contiguous(), requires_grad=False - ) - layer.w2_weight_scale = torch.nn.Parameter( - layer.w2_weight_scale.transpose(1, 2).contiguous(), requires_grad=False - ) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - assert self.num_bits == 4 or self.num_bits == 8 - config_builder = ( - int4_w4a16_moe_quant_config - if self.num_bits == 4 - else int8_w8a16_moe_quant_config - ) - - return config_builder( - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - w1_zp=None, - w2_zp=None, - block_shape=[0, self.group_size], - ) - - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - if self.moe.is_lora_enabled: - assert self.moe_quant_config is not None - from vllm.triton_utils import HAS_TRITON - - if HAS_TRITON: - from vllm.model_executor.layers.fused_moe import TritonWNA16Experts - - layer.w13_weight = layer.w13_weight_packed - layer.w2_weight = layer.w2_weight_packed - return TritonWNA16Experts( - moe_config=self.moe, quant_config=self.moe_quant_config - ) - else: - raise NotImplementedError( - "TritonExperts requires Triton. " - "Install triton or disable LoRA for MoE." - ) - - raise NotImplementedError - - def apply( - self, - layer: FusedMoE, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe import fused_experts - - return fused_experts( - x, - layer.w13_weight_packed, - layer.w2_weight_packed, - topk_weights=topk_weights, - topk_ids=topk_ids, - inplace=not self.moe.disable_inplace, - activation=layer.activation, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - quant_config=self.moe_quant_config, - ) - - @property - def supports_eplb(self) -> bool: - return True - - -class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod): - """ - CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform - - Weights: int4 (stored as int8 values in [-8,7], packed to uint8 nibbles) - - Scales: Fp32 for Channelwise , bf16 for groupwise quantization - - Bias: Same data type as original weights - - Activations: FP32/Bf16 dynamic per-token (A8 Int), - quantized inside the kernel - """ - - def __init__( - self, - weight_quant: QuantizationArgs, - input_quant: QuantizationArgs, - moe: FusedMoEConfig, - layer_name: str | None = None, - ): - super().__init__(moe) - self.has_bias = self.moe.has_bias - self.weight_quant = weight_quant - self.input_quant = input_quant - - # Validate scheme: weights=W4 (channel or group), - # activations=dynamic TOKEN (A8) - - # Must be dynamic per-token activations - if ( - input_quant.strategy != QuantizationStrategy.TOKEN - or not input_quant.dynamic - ): - raise ValueError( - "W4A8-int MoE needs dynamic per-token activation quantization." - ) - - # Weight can be channel-wise (group_size=None) or group-wise - self.group_size = ( - weight_quant.group_size if (weight_quant.group_size is not None) else -1 - ) - if weight_quant.num_bits != 4: - raise ValueError("This method only supports 4-bit weights (num_bits=4).") - - # CPU only - if not current_platform.is_cpu(): - raise ValueError("CompressedTensorsW4A8Int8MoEMethod is CPU-only.") - - # Arm: check _dyn ops availability - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: - try: - _ = torch.ops.aten._dyn_quant_matmul_4bit - _ = torch.ops.aten._dyn_quant_pack_4bit_weight - except AttributeError as err: - raise RuntimeError( - f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops; - install a newer build.""" - ) from err - self.static_input_scales = False # always dynamic per token - - # ---- parameter creation ---- - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - # Shapes per local rank (TP/EP): - # w13: [E, 2*I_local, H] int8 (int4 values in [-8,7]) - # w2 : [E, H, I_local] int8 - # Scales: - # channel-wise: group_size=-1 -> per-output-row, single scale per row - # group-wise : group_size=g -> - # per-output-row, (in_features/g) scales - - E = num_experts - H = hidden_size - IN = intermediate_size_per_partition - g = self.group_size - - # Per-row scale columns - def _n_scale_cols(in_features: int) -> int: - return 1 if g == -1 else (in_features // g) - - # Register unpacked int4-as-int8 weights the loader will fill. - w13 = torch.nn.Parameter( - torch.empty(E, 2 * IN, H, dtype=torch.int8), requires_grad=False - ) - set_weight_attrs(w13, extra_weight_attrs) - layer.register_parameter("w13_weight", w13) - - w2 = torch.nn.Parameter( - torch.empty(E, H, IN, dtype=torch.int8), requires_grad=False - ) - set_weight_attrs(w2, extra_weight_attrs) - layer.register_parameter("w2_weight", w2) - - # Register scales - # KleidiAI groupwise kernels accepts float32 scales - # KleidiAI groupwise kernels accepts bfloat16 scales - scale_dtype = torch.float32 if g == -1 else torch.bfloat16 - - w13_s = torch.nn.Parameter( - torch.ones(E, 2 * IN, _n_scale_cols(H), dtype=scale_dtype), - requires_grad=False, - ) - set_weight_attrs( - w13_s, - {"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs}, - ) - layer.register_parameter("w13_weight_scale", w13_s) - - w2_s = torch.nn.Parameter( - torch.ones(E, H, _n_scale_cols(IN), dtype=scale_dtype), requires_grad=False - ) - set_weight_attrs( - w2_s, - {"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs}, - ) - layer.register_parameter("w2_weight_scale", w2_s) - - if self.has_bias: - w13_bias = torch.nn.Parameter( - torch.zeros(E, 2 * IN, dtype=params_dtype), requires_grad=False - ) - layer.register_parameter("w13_bias", w13_bias) - set_weight_attrs(w13_bias, extra_weight_attrs) - - w2_bias = torch.nn.Parameter( - torch.zeros(num_experts, hidden_size, dtype=params_dtype), - requires_grad=False, - ) - layer.register_parameter("w2_bias", w2_bias) - set_weight_attrs(w2_bias, extra_weight_attrs) - - # Placeholders for packed weights (will be replaced after packing) - layer.register_parameter( - "w13_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False) - ) - set_weight_attrs(layer.w13_weight_packed, extra_weight_attrs) - - layer.register_parameter( - "w2_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False) - ) - set_weight_attrs(layer.w2_weight_packed, extra_weight_attrs) - - # dims for 4 bit fused matmuls - layer.w13_in_features = H - layer.w13_out_features = 2 * IN - layer.w2_in_features = IN - layer.w2_out_features = H - layer.group_size = g - - # post-load packing to dyn-4bit KleidiAI kernel's format - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - E = layer.w13_weight.shape[0] - H = layer.w13_in_features - I2 = layer.w13_out_features - IN = layer.w2_in_features - g = layer.group_size - - def _pack_matrix( - int4_as_int8_2d: torch.Tensor, - scales_2d: torch.Tensor, - bias_1d: torch.Tensor | None, - in_features: int, - out_features: int, - ) -> torch.Tensor: - # int4 values are stored as int8 in [-8,7]. - # Shift to unsigned nibble and pack pairs along input-dim. - tmp = int4_as_int8_2d.add(8) # [out, in] - uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to( - torch.uint8 - ) # [out, in//2] - - # KleidiAI groupwise kernels accepts float32 scales - # KleidiAI groupwise kernels accepts bfloat16 scales - scale_dtype = torch.float32 if g == -1 else torch.bfloat16 - scales = scales_2d.to(scale_dtype) - bias = None if bias_1d is None else bias_1d.to(torch.float32) - return torch.ops.aten._dyn_quant_pack_4bit_weight( - uint8_nibbles, - scales, - bias, - g if g != -1 else in_features, - in_features, - out_features, - ) - - # Pack per expert - w13_packed_list = [] - w2_packed_list = [] - - has_w13_bias = hasattr(layer, "w13_bias") and layer.w13_bias is not None - has_w2_bias = hasattr(layer, "w2_bias") and layer.w2_bias is not None - - for e in range(E): - w13_packed_list.append( - _pack_matrix( - layer.w13_weight[e], # [2I, H] - layer.w13_weight_scale[e], # [2I, H/g or 1] - layer.w13_bias[e] if has_w13_bias else None, # [2I] - H, - I2, - ) - ) - w2_packed_list.append( - _pack_matrix( - # w2 shape is [H, IN]; we need [out, in] == [H, IN]. - layer.w2_weight[e], # [H, IN] - layer.w2_weight_scale[e], # [H, IN/g or 1] - layer.w2_bias[e] if has_w2_bias else None, # [H] - IN, - layer.w2_out_features, # in_features=IN, out_features=H - ) - ) - - # each packed tensor has identical shape per expert; stack on dim 0 - w13_packed = torch.stack(w13_packed_list, dim=0) - w2_packed = torch.stack(w2_packed_list, dim=0) - - 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), - ) - - # free raw tensors/scales/bias now that they're packed into the payload. - 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(torch.empty(0), requires_grad=False), - ) - replace_parameter( - layer, - "w2_weight_scale", - torch.nn.Parameter(torch.empty(0), requires_grad=False), - ) - if has_w13_bias: - replace_parameter( - layer, - "w13_bias", - torch.nn.Parameter(torch.empty(0), requires_grad=False), - ) - if has_w2_bias: - replace_parameter( - layer, - "w2_bias", - torch.nn.Parameter(torch.empty(0), requires_grad=False), - ) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - # CPU dynamic 4-bit MoE path does not use modular kernels or - # fused_experts; quant config is not needed. - return None - - @property - def is_monolithic(self) -> bool: - return True - - def apply_monolithic( - self, - layer: FusedMoE, - x: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor: - assert not layer.enable_eplb, "EPLB not supported for W4A8-int MoE yet." - assert layer.activation in ( - MoEActivation.SILU, - MoEActivation.SWIGLUOAI, - MoEActivation.SWIGLUSTEP, - ), "Only SiLU/SwiGLUGU/SwiGLUUG are supported." - assert layer.expert_map is None, """expert_map/EP not implemented - for CPU dyn-4bit MoE.""" - - def _act_kind(s: MoEActivation) -> int: - # 0 = SwiGLU_Gu (SiLU(g)*u), 1 = SwiGLU_Ug (SiLU(u)*g), 2 = SiLU - if s == MoEActivation.SWIGLUSTEP: - return 0 - if s == MoEActivation.SWIGLUOAI: - return 1 - if s == MoEActivation.SILU: - return 2 - raise ValueError(f"Unknown activation '{s}'") - - # Apply topk softmax on router output - topk_weights, topk_ids = select_experts( - hidden_states=x, - router_logits=router_logits, - top_k=layer.top_k, - use_grouped_topk=layer.use_grouped_topk, - renormalize=layer.renormalize, - ) - - return torch.ops._C.dynamic_4bit_int_moe( - x, - topk_ids.to(torch.long), - topk_weights, - layer.w13_weight_packed, - layer.w2_weight_packed, - layer.w2_out_features, - layer.w2_in_features, - layer.w13_out_features, - layer.group_size, - layer.apply_router_weight_on_input, - int(_act_kind(layer.activation)), - ) - - -class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod): - def __init__( - self, - weight_quant: QuantizationArgs, - input_quant: QuantizationArgs, - moe: FusedMoEConfig, - layer_name: str | None = None, - ): - super().__init__(moe) - self.weight_quant = weight_quant - self.input_quant = input_quant - - self.group_size = self.weight_quant.group_size - self.num_bits = self.weight_quant.num_bits - self.packed_factor = 32 // self.num_bits - - assert self.weight_quant.symmetric, ( - "Only symmetric quantization is supported for W4A8 MoE" - ) - assert self.weight_quant.actorder != "group" - assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE" - - self.disable_expert_map = False - self.layer_name = layer_name - - from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 - from vllm.model_executor.layers.quantization.utils.quant_utils import ( - GroupShape, - ) - - self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN) - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - layer.num_experts = num_experts - layer.orig_dtype = params_dtype - layer.weight_block_size = None - - # requirement for CUTLASS reorder_tensor - assert hidden_size % 256 == 0, f"{hidden_size=} must be divisible by 256" - assert intermediate_size_per_partition % 256 == 0, ( - f"{intermediate_size_per_partition=} must be divisible by 256" - ) - # storage type, pack 8xint4 into int32 - params_dtype = torch.int32 - - # WEIGHTS - w13_weight_packed = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.packed_factor, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_packed", w13_weight_packed) - set_weight_attrs(w13_weight_packed, extra_weight_attrs) - - w2_weight_packed = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition // self.packed_factor, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_packed", w2_weight_packed) - set_weight_attrs(w2_weight_packed, extra_weight_attrs) - - # SCALES - # weight_scale refers to the group-wise scales - # they are initially loaded as bf16, we will convert to fp8 - # after loading - w13_weight_scale = torch.nn.Parameter( - torch.ones( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.group_size, - dtype=layer.orig_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - - w2_weight_scale = torch.nn.Parameter( - torch.ones( - num_experts, - hidden_size, - intermediate_size_per_partition // self.group_size, - dtype=layer.orig_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - # Add PER-GROUP quantization for FusedMoE.weight_loader. - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} - ) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - - # weight shapes - w2_weight_shape = torch.nn.Parameter( - torch.empty(num_experts, 2), requires_grad=False - ) - layer.register_parameter("w2_weight_shape", w2_weight_shape) - set_weight_attrs(w2_weight_shape, extra_weight_attrs) - w13_weight_shape = torch.nn.Parameter( - torch.empty(num_experts, 2), requires_grad=False - ) - layer.register_parameter("w13_weight_shape", w13_weight_shape) - set_weight_attrs(w13_weight_shape, extra_weight_attrs) - - # don't use input scales - layer.w13_input_scale = None - layer.w2_input_scale = None - - def process_weights_after_loading(self, layer): - device = layer.w13_weight_packed.device - - # STRIDES - # A, C - self.a_strides1_c_strides2 = torch.full( - (layer.local_num_experts,), - layer.hidden_size, - device=device, - dtype=torch.int64, - ) - self.a_strides2 = torch.full( - (layer.local_num_experts,), - layer.intermediate_size_per_partition, - device=device, - dtype=torch.int64, - ) - self.c_strides1 = torch.full( - (layer.local_num_experts,), - 2 * layer.intermediate_size_per_partition, - device=device, - dtype=torch.int64, - ) - - # S (group-wise scales) - # sizeof(StrideS) = 16 bytes, so we need to use 2xint64 to encode it - self.s_strides1 = torch.zeros( - (layer.local_num_experts, 2), device=device, dtype=torch.int64 - ) - self.s_strides1[:, 0] = 2 * layer.intermediate_size_per_partition - - self.s_strides2 = torch.zeros( - (layer.local_num_experts, 2), device=device, dtype=torch.int64 - ) - self.s_strides2[:, 0] = layer.hidden_size - - # encode and reorder weight tensors, and get the layout to pass to - # the grouped gemm kernel. `b_strides1/2` specifies the entire layout - convert_packed_uint4b8_to_signed_int4_inplace(layer.w13_weight_packed) - w13_weight_shuffled, self.b_strides1 = ( - ops.cutlass_encode_and_reorder_int4b_grouped(layer.w13_weight_packed) - ) - replace_parameter(layer, "w13_weight_packed", w13_weight_shuffled) - convert_packed_uint4b8_to_signed_int4_inplace(layer.w2_weight_packed) - w2_weight_shuffled, self.b_strides2 = ( - ops.cutlass_encode_and_reorder_int4b_grouped(layer.w2_weight_packed) - ) - replace_parameter(layer, "w2_weight_packed", w2_weight_shuffled) - - # convert bf16 scales to (fp8_scales, channel_scales) - w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8( - self.quant_fp8, layer.w13_weight_scale - ) - w2_weight_scale, w2_weight_chan_scale = convert_bf16_scales_to_fp8( - self.quant_fp8, layer.w2_weight_scale - ) - - # register channel scales - layer.register_parameter( - "w13_weight_chan_scale", - torch.nn.Parameter(w13_weight_chan_scale, requires_grad=False), - ) - layer.register_parameter( - "w2_weight_chan_scale", - torch.nn.Parameter(w2_weight_chan_scale, requires_grad=False), - ) - - # The scales are stored as (E, N, K // 128) but the kernel expects - # (E, K // 128, N) in row-major format, so we need to permute the last 2 dims - # and make it contiguous - w13_weight_scale_packed = ops.cutlass_pack_scale_fp8( - w13_weight_scale.permute(0, 2, 1).contiguous() - ) - replace_parameter(layer, "w13_weight_scale", w13_weight_scale_packed) - w2_weight_scale_packed = ops.cutlass_pack_scale_fp8( - w2_weight_scale.permute(0, 2, 1).contiguous() - ) - replace_parameter(layer, "w2_weight_scale", w2_weight_scale_packed) - - def maybe_make_prepare_finalize( - self, - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - ) -> mk.FusedMoEPrepareAndFinalizeModular | None: - return super().maybe_make_prepare_finalize(routing_tables) - - def get_fused_moe_quant_config( - self, layer: torch.nn.Module - ) -> FusedMoEQuantConfig | None: - # Store quantization scales; both per-group and per-channel - # Note we haven't specified the group size here because - # the quant config logic assumes group-wise scaling - # and channel-wise scaling are exclusive. - return int4_w4afp8_moe_quant_config( - w1_scale=layer.w13_weight_scale, # group scale - w2_scale=layer.w2_weight_scale, # group scale - g1_alphas=layer.w13_weight_chan_scale, - g2_alphas=layer.w2_weight_chan_scale, - per_act_token_quant=True, # always use dynamic per-token - per_out_ch_quant=True, # always use per-channel - ) - - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - assert self.moe_quant_config is not None - assert ( - prepare_finalize.activation_format == FusedMoEActivationFormat.Standard - ), "BatchedExperts not supported" - - from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8 - - experts: FusedMoEExpertsModular - - logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__) - experts = CutlassExpertsW4A8Fp8( - out_dtype=self.moe.in_dtype, - a_strides1=self.a_strides1_c_strides2, - a_strides2=self.a_strides2, - b_strides1=self.b_strides1, - b_strides2=self.b_strides2, - c_strides1=self.c_strides1, - c_strides2=self.a_strides1_c_strides2, - s_strides1=self.s_strides1, - s_strides2=self.s_strides2, - moe_config=self.moe, - quant_config=self.moe_quant_config, - group_size=self.group_size, - ) - - num_dispatchers = prepare_finalize.num_dispatchers() - self.disable_expert_map = ( - num_dispatchers > 1 or not experts.supports_expert_map() - ) - - return experts - - def apply( - self, - layer: FusedMoE, - x: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor: - if layer.enable_eplb: - raise NotImplementedError( - "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet." - ) - assert self.moe_quant_config is not None - - from vllm.model_executor.layers.fused_moe.cutlass_moe import ( - cutlass_moe_w4a8_fp8, - ) - - return cutlass_moe_w4a8_fp8( - x, - layer.w13_weight_packed, - layer.w2_weight_packed, - topk_weights, - topk_ids, - moe_config=self.moe, - quant_config=self.moe_quant_config, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=None if self.disable_expert_map else layer.expert_map, - a_strides1=self.a_strides1_c_strides2, - a_strides2=self.a_strides2, - b_strides1=self.b_strides1, - b_strides2=self.b_strides2, - c_strides1=self.c_strides1, - c_strides2=self.a_strides1_c_strides2, - s_strides1=self.s_strides1, - s_strides2=self.s_strides2, - group_size=self.group_size, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - ) - - @property - def supports_eplb(self) -> bool: - return False diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/__init__.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/__init__.py new file mode 100644 index 00000000000..39c9113f65f --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe import ( # noqa: E501 + CompressedTensorsMoEMethod, +) + +__all__ = [ + "CompressedTensorsMoEMethod", +] diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py new file mode 100644 index 00000000000..9ee8df9daba --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from compressed_tensors import CompressionFormat +from compressed_tensors.quantization import ( + ActivationOrdering, + QuantizationStrategy, +) + +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoEMethodBase, + UnquantizedFusedMoEMethod, +) +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa + WNA16_SUPPORTED_BITS, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + check_moe_marlin_supports_layer, +) +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +class CompressedTensorsMoEMethod(FusedMoEMethodBase): + @staticmethod + def get_moe_method( + quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501 + layer: torch.nn.Module, + layer_name: str, + ) -> FusedMoEMethodBase: + # FusedMoE was made by combining multiple Linears so need to + # make sure quantization config for Linear can target it + quant_config._add_fused_moe_to_target_scheme_map() + unfused_names = [ + layer_name + proj_name + for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"] + ] + # TODO: refactor this to use expert_mapping and check all layer numbers + all_scheme_dicts = [ + quant_config.get_scheme_dict(layer, name) for name in unfused_names + ] + scheme_dict = all_scheme_dicts.pop() + + # multiple schemes found + if not all([cur_dict == scheme_dict for cur_dict in all_scheme_dicts]): + raise ValueError( + "All MoE projections need to have same " + "quantization scheme but found multiple" + ) + + if scheme_dict is None: # ignored layer + return UnquantizedFusedMoEMethod(layer.moe_config) + + # TODO: @dsikka: refactor this to use schemes as other kernels + # are supported + check if the layer is being ignored. + weight_quant = scheme_dict.get("weights") + input_quant = scheme_dict.get("input_activations") + format = scheme_dict.get("format") + + if quant_config._is_mxfp4(weight_quant): + from .compressed_tensors_moe_w4a4_mxfp4 import ( + CompressedTensorsW4A4Mxfp4MoEMethod, + ) + + return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config) + + if quant_config._is_wNa16_group_channel(weight_quant, input_quant): + # group_size=None means channelwise + group_size = weight_quant.group_size or -1 + + valid_format_and_bits = ( + weight_quant.num_bits in WNA16_SUPPORTED_BITS + and format == CompressionFormat.pack_quantized.value + ) + + if not valid_format_and_bits: + raise ValueError( + "For Fused MoE layers, only format: ", + f"{CompressionFormat.pack_quantized.value} ", + f" and bits: {WNA16_SUPPORTED_BITS} is supported ", + f"but got format: {CompressionFormat.pack_quantized.value} " + f" and bits: {weight_quant.num_bits}", + ) + + # Prefer to use the MarlinMoE kernel when it is supported. + if ( + not check_moe_marlin_supports_layer(layer, group_size) + or current_platform.is_rocm() + ): + from .compressed_tensors_moe_wna16 import ( + CompressedTensorsWNA16MoEMethod, + ) + + if ( + weight_quant.strategy == QuantizationStrategy.GROUP + and weight_quant.actorder + in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC) + ): + raise ValueError( + "WNA16MoE is not supported with actorder=group/dynamic." + ) + logger.info_once("Using CompressedTensorsWNA16MoEMethod") + return CompressedTensorsWNA16MoEMethod( + weight_quant, input_quant, layer.moe_config + ) + else: + from .compressed_tensors_moe_wna16_marlin import ( + CompressedTensorsWNA16MarlinMoEMethod, + ) + + logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod") + return CompressedTensorsWNA16MarlinMoEMethod( + weight_quant, input_quant, layer.moe_config + ) + elif quant_config._is_nvfp4_format(weight_quant): + from .compressed_tensors_moe_w4a4_nvfp4 import ( + CompressedTensorsW4A4Nvfp4MoEMethod, + ) + + _is_valid_nvfp4_activations = ( + quant_config._is_nvfp4_format(input_quant) or input_quant is None + ) + if not _is_valid_nvfp4_activations: + raise ValueError( + "For NVFP4 weights, input quantization must also be NVFP4 format ", + f"or None for NVFP4A16, found {input_quant}", + ) + return CompressedTensorsW4A4Nvfp4MoEMethod( + layer.moe_config, layer_name, use_a16=(input_quant is None) + ) + elif ( + quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant) + or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant) + or quant_config._is_fp8_w8a8(weight_quant, input_quant) + ): + from .compressed_tensors_moe_w8a8_fp8 import ( + CompressedTensorsW8A8Fp8MoEMethod, + ) + + return CompressedTensorsW8A8Fp8MoEMethod( + weight_quant, input_quant, layer.moe_config + ) + elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant): + from .compressed_tensors_moe_w8a8_int8 import ( + CompressedTensorsW8A8Int8MoEMethod, + ) + + return CompressedTensorsW8A8Int8MoEMethod( + weight_quant, input_quant, layer.moe_config + ) + elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant): + from .compressed_tensors_moe_w4a8_fp8 import ( + CompressedTensorsW4A8Fp8MoEMethod, + ) + + logger.info_once("Using CompressedTensorsW4A8Fp8MoEMethod") + return CompressedTensorsW4A8Fp8MoEMethod( + weight_quant, input_quant, layer.moe_config + ) + elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant): + from .compressed_tensors_moe_w4a8_int8 import ( + CompressedTensorsW4A8Int8MoEMethod, + ) + + return CompressedTensorsW4A8Int8MoEMethod( + weight_quant, input_quant, layer.moe_config + ) + else: + raise RuntimeError( + f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}" + ) 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 new file mode 100644 index 00000000000..8cc6d17bcc1 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + FusedMoeWeightScaleSupported, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( + MarlinExperts, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + make_mxfp4_moe_kernel, + make_mxfp4_moe_quant_config, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( + prepare_moe_fp4_layer_for_marlin, +) +from vllm.model_executor.utils import set_weight_attrs + +logger = init_logger(__name__) + + +class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): + def __init__(self, moe): + super().__init__(moe) + self.group_size = 32 + self.mxfp4_backend = Mxfp4MoeBackend.MARLIN + self.experts_cls = MarlinExperts + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.params_dtype = params_dtype + + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + # 2 fp4 items are packed in the input dimension + hidden_size // 2, + requires_grad=False, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + # 2 fp4 items are packed in the input dimension + intermediate_size_per_partition // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + # 2 fp4 items are packed in the input dimension + hidden_size // self.group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + w2_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + # 2 fp4 items are packed in the input dimension + intermediate_size_per_partition // self.group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + return make_mxfp4_moe_quant_config( + mxfp4_backend=self.mxfp4_backend, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + ) + + def process_weights_after_loading(self, layer: FusedMoE) -> 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") + + logger.warning_once( + "Your GPU does not have native support for FP4 computation but " + "FP4 quantization is being used. Weight-only FP4 compression " + "will be used leveraging the Marlin kernel. This may degrade " + "performance for compute-heavy workloads." + ) + prepare_moe_fp4_layer_for_marlin(layer) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config is not None: + self.moe_kernel = make_mxfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + mxfp4_backend=self.mxfp4_backend, + shared_experts=layer.shared_experts, + routing_tables=layer._maybe_init_expert_routing_tables(), + ) + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, + ) 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 new file mode 100644 index 00000000000..09a216fd2cb --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + FusedMoeWeightScaleSupported, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + convert_to_nvfp4_moe_kernel_format, + is_global_sf_supported_for_nvfp4_backend, + make_nvfp4_moe_kernel, + make_nvfp4_moe_quant_config, + select_nvfp4_moe_backend, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs + +logger = init_logger(__name__) + + +class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): + def __init__( + self, + moe: FusedMoEConfig, + layer_name: str | None = None, + use_a16: bool = False, + ): + super().__init__(moe) + self.group_size = 16 + + # 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, + ) + + self.use_global_sf = is_global_sf_supported_for_nvfp4_backend( + self.nvfp4_backend + ) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.params_dtype = params_dtype + w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + w13_num_shards * intermediate_size_per_partition, + # 2 fp4 items are packed in the input dimension + hidden_size // 2, + requires_grad=False, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + # 2 fp4 items are packed in the input dimension + intermediate_size_per_partition // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # Weight Scales + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + w13_num_shards * intermediate_size_per_partition, + # 2 fp4 items are packed in the input dimension + hidden_size // self.group_size, + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + w2_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + # 2 fp4 items are packed in the input dimension + intermediate_size_per_partition // self.group_size, + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} + ) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # Weight Global Scales + w13_weight_scale_2 = torch.nn.Parameter( + torch.empty(num_experts, w13_num_shards, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_global_scale", w13_weight_scale_2) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + set_weight_attrs(w13_weight_scale_2, extra_weight_attrs) + + w2_weight_scale_2 = torch.nn.Parameter( + torch.empty(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + set_weight_attrs(w2_weight_scale_2, extra_weight_attrs) + + # Input Global Scales + w13_input_scale = torch.nn.Parameter( + torch.empty(num_experts, w13_num_shards, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_input_global_scale", w13_input_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.empty(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_input_global_scale", w2_input_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + + def process_weights_after_loading(self, layer: FusedMoE) -> 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") + + # Use a single gscale for w13. + if self.moe.is_act_and_mul and not torch.allclose( + layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1] + ): + logger.warning_once( + "w1_weight_global_scale must match w3_weight_global_scale. " + "Accuracy may be affected.", + ) + w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous() + + # Shuffle weights into the NvFp4 kernel format. + ( + w13, + w13_scale, + w13_scale_2, + a13_scale, + w2, + w2_scale, + w2_scale_2, + a2_scale, + ) = convert_to_nvfp4_moe_kernel_format( + nvfp4_backend=self.nvfp4_backend, + layer=layer, + 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), + 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), + is_act_and_mul=self.moe.is_act_and_mul, + ) + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w2_weight_scale", w2_scale) + layer.w13_weight_scale_2 = w13_scale_2 + layer.w2_weight_scale_2 = w2_scale_2 + layer.w13_input_scale = a13_scale + layer.w2_input_scale = a2_scale + + # Setup modular kernel. + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.experts_cls is not None + self.moe_kernel = make_nvfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + shared_experts=layer.shared_experts, + routing_tables=layer._maybe_init_expert_routing_tables(), + ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel initialization " + "logic. This function should not be called." + ) + + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: + return make_nvfp4_moe_quant_config( + backend=self.nvfp4_backend, + w13_scale=layer.w13_weight_scale, + 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, + ) + + def apply_monolithic( + self, + layer: FusedMoE, + x: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, + ) 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 new file mode 100644 index 00000000000..74cb0b4f6e1 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from compressed_tensors.quantization import ( + QuantizationArgs, +) + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import _custom_ops as ops +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + FusedMoEActivationFormat, + FusedMoEExpertsModular, + FusedMoeWeightScaleSupported, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + int4_w4afp8_moe_quant_config, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + convert_bf16_scales_to_fp8, + convert_packed_uint4b8_to_signed_int4_inplace, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs + +logger = init_logger(__name__) + + +class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod): + def __init__( + self, + weight_quant: QuantizationArgs, + input_quant: QuantizationArgs, + moe: FusedMoEConfig, + layer_name: str | None = None, + ): + super().__init__(moe) + self.weight_quant = weight_quant + self.input_quant = input_quant + + self.group_size = self.weight_quant.group_size + self.num_bits = self.weight_quant.num_bits + self.packed_factor = 32 // self.num_bits + + assert self.weight_quant.symmetric, ( + "Only symmetric quantization is supported for W4A8 MoE" + ) + assert self.weight_quant.actorder != "group" + assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE" + + self.disable_expert_map = False + self.layer_name = layer_name + + from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + ) + + self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + # requirement for CUTLASS reorder_tensor + assert hidden_size % 256 == 0, f"{hidden_size=} must be divisible by 256" + assert intermediate_size_per_partition % 256 == 0, ( + f"{intermediate_size_per_partition=} must be divisible by 256" + ) + # storage type, pack 8xint4 into int32 + params_dtype = torch.int32 + + # WEIGHTS + w13_weight_packed = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.packed_factor, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight_packed) + set_weight_attrs(w13_weight_packed, extra_weight_attrs) + + w2_weight_packed = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.packed_factor, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight_packed) + set_weight_attrs(w2_weight_packed, extra_weight_attrs) + + # SCALES + # weight_scale refers to the group-wise scales + # they are initially loaded as bf16, we will convert to fp8 + # after loading + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.group_size, + dtype=layer.orig_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + + w2_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + hidden_size, + intermediate_size_per_partition // self.group_size, + dtype=layer.orig_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + # Add PER-GROUP quantization for FusedMoE.weight_loader. + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # weight shapes + w2_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + layer.register_parameter("w2_weight_shape", w2_weight_shape) + set_weight_attrs(w2_weight_shape, extra_weight_attrs) + w13_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + layer.register_parameter("w13_weight_shape", w13_weight_shape) + set_weight_attrs(w13_weight_shape, extra_weight_attrs) + + # don't use input scales + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer): + device = layer.w13_weight_packed.device + + # STRIDES + # A, C + self.a_strides1_c_strides2 = torch.full( + (layer.local_num_experts,), + layer.hidden_size, + device=device, + dtype=torch.int64, + ) + self.a_strides2 = torch.full( + (layer.local_num_experts,), + layer.intermediate_size_per_partition, + device=device, + dtype=torch.int64, + ) + self.c_strides1 = torch.full( + (layer.local_num_experts,), + 2 * layer.intermediate_size_per_partition, + device=device, + dtype=torch.int64, + ) + + # S (group-wise scales) + # sizeof(StrideS) = 16 bytes, so we need to use 2xint64 to encode it + self.s_strides1 = torch.zeros( + (layer.local_num_experts, 2), device=device, dtype=torch.int64 + ) + self.s_strides1[:, 0] = 2 * layer.intermediate_size_per_partition + + self.s_strides2 = torch.zeros( + (layer.local_num_experts, 2), device=device, dtype=torch.int64 + ) + self.s_strides2[:, 0] = layer.hidden_size + + # encode and reorder weight tensors, and get the layout to pass to + # the grouped gemm kernel. `b_strides1/2` specifies the entire layout + convert_packed_uint4b8_to_signed_int4_inplace(layer.w13_weight_packed) + w13_weight_shuffled, self.b_strides1 = ( + ops.cutlass_encode_and_reorder_int4b_grouped(layer.w13_weight_packed) + ) + replace_parameter(layer, "w13_weight_packed", w13_weight_shuffled) + convert_packed_uint4b8_to_signed_int4_inplace(layer.w2_weight_packed) + w2_weight_shuffled, self.b_strides2 = ( + ops.cutlass_encode_and_reorder_int4b_grouped(layer.w2_weight_packed) + ) + replace_parameter(layer, "w2_weight_packed", w2_weight_shuffled) + + # convert bf16 scales to (fp8_scales, channel_scales) + w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8( + self.quant_fp8, layer.w13_weight_scale + ) + w2_weight_scale, w2_weight_chan_scale = convert_bf16_scales_to_fp8( + self.quant_fp8, layer.w2_weight_scale + ) + + # register channel scales + layer.register_parameter( + "w13_weight_chan_scale", + torch.nn.Parameter(w13_weight_chan_scale, requires_grad=False), + ) + layer.register_parameter( + "w2_weight_chan_scale", + torch.nn.Parameter(w2_weight_chan_scale, requires_grad=False), + ) + + # The scales are stored as (E, N, K // 128) but the kernel expects + # (E, K // 128, N) in row-major format, so we need to permute the last 2 dims + # and make it contiguous + w13_weight_scale_packed = ops.cutlass_pack_scale_fp8( + w13_weight_scale.permute(0, 2, 1).contiguous() + ) + replace_parameter(layer, "w13_weight_scale", w13_weight_scale_packed) + w2_weight_scale_packed = ops.cutlass_pack_scale_fp8( + w2_weight_scale.permute(0, 2, 1).contiguous() + ) + replace_parameter(layer, "w2_weight_scale", w2_weight_scale_packed) + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: + return super().maybe_make_prepare_finalize(routing_tables) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + # Store quantization scales; both per-group and per-channel + # Note we haven't specified the group size here because + # the quant config logic assumes group-wise scaling + # and channel-wise scaling are exclusive. + return int4_w4afp8_moe_quant_config( + w1_scale=layer.w13_weight_scale, # group scale + w2_scale=layer.w2_weight_scale, # group scale + g1_alphas=layer.w13_weight_chan_scale, + g2_alphas=layer.w2_weight_chan_scale, + per_act_token_quant=True, # always use dynamic per-token + per_out_ch_quant=True, # always use per-channel + ) + + def select_gemm_impl( + self, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, + layer: torch.nn.Module, + ) -> mk.FusedMoEExpertsModular: + assert self.moe_quant_config is not None + assert ( + prepare_finalize.activation_format == FusedMoEActivationFormat.Standard + ), "BatchedExperts not supported" + + from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8 + + experts: FusedMoEExpertsModular + + logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__) + experts = CutlassExpertsW4A8Fp8( + out_dtype=self.moe.in_dtype, + a_strides1=self.a_strides1_c_strides2, + a_strides2=self.a_strides2, + b_strides1=self.b_strides1, + b_strides2=self.b_strides2, + c_strides1=self.c_strides1, + c_strides2=self.a_strides1_c_strides2, + s_strides1=self.s_strides1, + s_strides2=self.s_strides2, + moe_config=self.moe, + quant_config=self.moe_quant_config, + group_size=self.group_size, + ) + + num_dispatchers = prepare_finalize.num_dispatchers() + self.disable_expert_map = ( + num_dispatchers > 1 or not experts.supports_expert_map() + ) + + return experts + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + if layer.enable_eplb: + raise NotImplementedError( + "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet." + ) + assert self.moe_quant_config is not None + + from vllm.model_executor.layers.fused_moe.cutlass_moe import ( + cutlass_moe_w4a8_fp8, + ) + + return cutlass_moe_w4a8_fp8( + x, + layer.w13_weight_packed, + layer.w2_weight_packed, + topk_weights, + topk_ids, + moe_config=self.moe, + quant_config=self.moe_quant_config, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=None if self.disable_expert_map else layer.expert_map, + a_strides1=self.a_strides1_c_strides2, + a_strides2=self.a_strides2, + b_strides1=self.b_strides1, + b_strides2=self.b_strides2, + c_strides1=self.c_strides1, + c_strides2=self.a_strides1_c_strides2, + s_strides1=self.s_strides1, + s_strides2=self.s_strides2, + group_size=self.group_size, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ) + + @property + def supports_eplb(self) -> bool: + return False 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 new file mode 100644 index 00000000000..2e8a935cad6 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py @@ -0,0 +1,349 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, +) + +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, +) +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts +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.platforms import CpuArchEnum, current_platform + +logger = init_logger(__name__) + + +class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod): + """ + CPU-only MoE method using dynamic 4-bit matmul kernels on Arm Platform + - Weights: int4 (stored as int8 values in [-8,7], packed to uint8 nibbles) + - Scales: Fp32 for Channelwise , bf16 for groupwise quantization + - Bias: Same data type as original weights + - Activations: FP32/Bf16 dynamic per-token (A8 Int), + quantized inside the kernel + """ + + def __init__( + self, + weight_quant: QuantizationArgs, + input_quant: QuantizationArgs, + moe: FusedMoEConfig, + layer_name: str | None = None, + ): + super().__init__(moe) + self.has_bias = self.moe.has_bias + self.weight_quant = weight_quant + self.input_quant = input_quant + + # Validate scheme: weights=W4 (channel or group), + # activations=dynamic TOKEN (A8) + + # Must be dynamic per-token activations + if ( + input_quant.strategy != QuantizationStrategy.TOKEN + or not input_quant.dynamic + ): + raise ValueError( + "W4A8-int MoE needs dynamic per-token activation quantization." + ) + + # Weight can be channel-wise (group_size=None) or group-wise + self.group_size = ( + weight_quant.group_size if (weight_quant.group_size is not None) else -1 + ) + if weight_quant.num_bits != 4: + raise ValueError("This method only supports 4-bit weights (num_bits=4).") + + # CPU only + if not current_platform.is_cpu(): + raise ValueError("CompressedTensorsW4A8Int8MoEMethod is CPU-only.") + + # Arm: check _dyn ops availability + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM: + try: + _ = torch.ops.aten._dyn_quant_matmul_4bit + _ = torch.ops.aten._dyn_quant_pack_4bit_weight + except AttributeError as err: + raise RuntimeError( + f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops; + install a newer build.""" + ) from err + self.static_input_scales = False # always dynamic per token + + # ---- parameter creation ---- + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + # Shapes per local rank (TP/EP): + # w13: [E, 2*I_local, H] int8 (int4 values in [-8,7]) + # w2 : [E, H, I_local] int8 + # Scales: + # channel-wise: group_size=-1 -> per-output-row, single scale per row + # group-wise : group_size=g -> + # per-output-row, (in_features/g) scales + + E = num_experts + H = hidden_size + IN = intermediate_size_per_partition + g = self.group_size + + # Per-row scale columns + def _n_scale_cols(in_features: int) -> int: + return 1 if g == -1 else (in_features // g) + + # Register unpacked int4-as-int8 weights the loader will fill. + w13 = torch.nn.Parameter( + torch.empty(E, 2 * IN, H, dtype=torch.int8), requires_grad=False + ) + set_weight_attrs(w13, extra_weight_attrs) + layer.register_parameter("w13_weight", w13) + + w2 = torch.nn.Parameter( + torch.empty(E, H, IN, dtype=torch.int8), requires_grad=False + ) + set_weight_attrs(w2, extra_weight_attrs) + layer.register_parameter("w2_weight", w2) + + # Register scales + # KleidiAI groupwise kernels accepts float32 scales + # KleidiAI groupwise kernels accepts bfloat16 scales + scale_dtype = torch.float32 if g == -1 else torch.bfloat16 + + w13_s = torch.nn.Parameter( + torch.ones(E, 2 * IN, _n_scale_cols(H), dtype=scale_dtype), + requires_grad=False, + ) + set_weight_attrs( + w13_s, + {"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs}, + ) + layer.register_parameter("w13_weight_scale", w13_s) + + w2_s = torch.nn.Parameter( + torch.ones(E, H, _n_scale_cols(IN), dtype=scale_dtype), requires_grad=False + ) + set_weight_attrs( + w2_s, + {"quant_method": "channel" if g == -1 else "group", **extra_weight_attrs}, + ) + layer.register_parameter("w2_weight_scale", w2_s) + + if self.has_bias: + w13_bias = torch.nn.Parameter( + torch.zeros(E, 2 * IN, dtype=params_dtype), requires_grad=False + ) + layer.register_parameter("w13_bias", w13_bias) + set_weight_attrs(w13_bias, extra_weight_attrs) + + w2_bias = torch.nn.Parameter( + torch.zeros(num_experts, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) + + # Placeholders for packed weights (will be replaced after packing) + layer.register_parameter( + "w13_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False) + ) + set_weight_attrs(layer.w13_weight_packed, extra_weight_attrs) + + layer.register_parameter( + "w2_weight_packed", torch.nn.Parameter(torch.empty(0), requires_grad=False) + ) + set_weight_attrs(layer.w2_weight_packed, extra_weight_attrs) + + # dims for 4 bit fused matmuls + layer.w13_in_features = H + layer.w13_out_features = 2 * IN + layer.w2_in_features = IN + layer.w2_out_features = H + layer.group_size = g + + # post-load packing to dyn-4bit KleidiAI kernel's format + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + E = layer.w13_weight.shape[0] + H = layer.w13_in_features + I2 = layer.w13_out_features + IN = layer.w2_in_features + g = layer.group_size + + def _pack_matrix( + int4_as_int8_2d: torch.Tensor, + scales_2d: torch.Tensor, + bias_1d: torch.Tensor | None, + in_features: int, + out_features: int, + ) -> torch.Tensor: + # int4 values are stored as int8 in [-8,7]. + # Shift to unsigned nibble and pack pairs along input-dim. + tmp = int4_as_int8_2d.add(8) # [out, in] + uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to( + torch.uint8 + ) # [out, in//2] + + # KleidiAI groupwise kernels accepts float32 scales + # KleidiAI groupwise kernels accepts bfloat16 scales + scale_dtype = torch.float32 if g == -1 else torch.bfloat16 + scales = scales_2d.to(scale_dtype) + bias = None if bias_1d is None else bias_1d.to(torch.float32) + return torch.ops.aten._dyn_quant_pack_4bit_weight( + uint8_nibbles, + scales, + bias, + g if g != -1 else in_features, + in_features, + out_features, + ) + + # Pack per expert + w13_packed_list = [] + w2_packed_list = [] + + has_w13_bias = hasattr(layer, "w13_bias") and layer.w13_bias is not None + has_w2_bias = hasattr(layer, "w2_bias") and layer.w2_bias is not None + + for e in range(E): + w13_packed_list.append( + _pack_matrix( + layer.w13_weight[e], # [2I, H] + layer.w13_weight_scale[e], # [2I, H/g or 1] + layer.w13_bias[e] if has_w13_bias else None, # [2I] + H, + I2, + ) + ) + w2_packed_list.append( + _pack_matrix( + # w2 shape is [H, IN]; we need [out, in] == [H, IN]. + layer.w2_weight[e], # [H, IN] + layer.w2_weight_scale[e], # [H, IN/g or 1] + layer.w2_bias[e] if has_w2_bias else None, # [H] + IN, + layer.w2_out_features, # in_features=IN, out_features=H + ) + ) + + # each packed tensor has identical shape per expert; stack on dim 0 + w13_packed = torch.stack(w13_packed_list, dim=0) + w2_packed = torch.stack(w2_packed_list, dim=0) + + 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), + ) + + # free raw tensors/scales/bias now that they're packed into the payload. + 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(torch.empty(0), requires_grad=False), + ) + replace_parameter( + layer, + "w2_weight_scale", + torch.nn.Parameter(torch.empty(0), requires_grad=False), + ) + if has_w13_bias: + replace_parameter( + layer, + "w13_bias", + torch.nn.Parameter(torch.empty(0), requires_grad=False), + ) + if has_w2_bias: + replace_parameter( + layer, + "w2_bias", + torch.nn.Parameter(torch.empty(0), requires_grad=False), + ) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + # CPU dynamic 4-bit MoE path does not use modular kernels or + # fused_experts; quant config is not needed. + return None + + @property + def is_monolithic(self) -> bool: + return True + + def apply_monolithic( + self, + layer: FusedMoE, + x: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + assert not layer.enable_eplb, "EPLB not supported for W4A8-int MoE yet." + assert layer.activation in ( + MoEActivation.SILU, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUSTEP, + ), "Only SiLU/SwiGLUGU/SwiGLUUG are supported." + assert layer.expert_map is None, """expert_map/EP not implemented + for CPU dyn-4bit MoE.""" + + def _act_kind(s: MoEActivation) -> int: + # 0 = SwiGLU_Gu (SiLU(g)*u), 1 = SwiGLU_Ug (SiLU(u)*g), 2 = SiLU + if s == MoEActivation.SWIGLUSTEP: + return 0 + if s == MoEActivation.SWIGLUOAI: + return 1 + if s == MoEActivation.SILU: + return 2 + raise ValueError(f"Unknown activation '{s}'") + + # Apply topk softmax on router output + topk_weights, topk_ids = select_experts( + hidden_states=x, + router_logits=router_logits, + top_k=layer.top_k, + use_grouped_topk=layer.use_grouped_topk, + renormalize=layer.renormalize, + ) + + return torch.ops._C.dynamic_4bit_int_moe( + x, + topk_ids.to(torch.long), + topk_weights, + layer.w13_weight_packed, + layer.w2_weight_packed, + layer.w2_out_features, + layer.w2_in_features, + layer.w13_out_features, + layer.group_size, + layer.apply_router_weight_on_input, + int(_act_kind(layer.activation)), + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py new file mode 100644 index 00000000000..ed8ed79c50c --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, +) + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + FusedMoeWeightScaleSupported, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + convert_to_fp8_moe_kernel_format, + make_fp8_moe_kernel, + make_fp8_moe_quant_config, + select_fp8_moe_backend, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + process_fp8_input_tensor_strategy_moe, + process_fp8_weight_tensor_strategy_moe, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8DynamicTokenSym, + kFp8Static128BlockSym, + kFp8StaticChannelSym, + kFp8StaticTensorSym, +) +from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( + normalize_e4m3fn_to_e4m3fnuz, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): + """W8A8 FP8 MoE quantization using compressed tensors.""" + + def __init__( + self, + weight_quant: QuantizationArgs, + input_quant: QuantizationArgs, + moe: FusedMoEConfig, + layer_name: str | None = None, + ): + super().__init__(moe) + self.weight_quant = weight_quant + self.input_quant = input_quant + + per_tensor = ( + self.weight_quant.strategy == QuantizationStrategy.TENSOR + and self.input_quant.strategy == QuantizationStrategy.TENSOR + ) + per_channel = ( + self.weight_quant.strategy == QuantizationStrategy.CHANNEL + and self.input_quant.strategy == QuantizationStrategy.TOKEN + ) + if not (per_tensor or per_channel): + assert self.weight_quant.strategy == QuantizationStrategy.BLOCK + self.weight_block_size = self.weight_quant.block_structure + assert self.weight_quant.dynamic is not None + else: + self.weight_block_size = None + self.block_quant = self.weight_block_size is not None + + self.static_input_scales = not self.input_quant.dynamic + if self.static_input_scales and per_channel: + raise ValueError( + "For FP8 Fused MoE layer, we require either per tensor or " + "channelwise, dynamic per token quantization." + ) + + ct2vllm_weight = { + QuantizationStrategy.CHANNEL: kFp8StaticChannelSym, + QuantizationStrategy.TENSOR: kFp8StaticTensorSym, + QuantizationStrategy.BLOCK: kFp8Static128BlockSym, + } + ct2vllm_act = { + QuantizationStrategy.TOKEN: kFp8DynamicTokenSym, + QuantizationStrategy.TENSOR: ( + kFp8StaticTensorSym if self.static_input_scales else kFp8Dynamic128Sym + ), + } + weight_key = ct2vllm_weight[self.weight_quant.strategy] + if weight_key == kFp8Static128BlockSym: + activation_key = kFp8Dynamic128Sym + else: + activation_key = ct2vllm_act[self.input_quant.strategy] + + # Select Fp8 MoE backend + self.fp8_backend, self.experts_cls = select_fp8_moe_backend( + config=self.moe, + weight_key=weight_key, + activation_key=activation_key, + allow_vllm_cutlass=True, + ) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + params_dtype = torch.float8_e4m3fn + w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + + if self.block_quant: + assert self.weight_block_size is not None + layer.weight_block_size = self.weight_block_size + tp_size = get_tensor_model_parallel_world_size() + block_n, block_k = ( + self.weight_block_size[0], + self.weight_block_size[1], + ) + # NOTE: To ensure proper alignment of the block-wise quantization + # scales, the output_size of the weights for both the gate and up + # layers must be divisible by block_n. + # Required by column parallel or enabling merged weights + if intermediate_size_per_partition % block_n != 0: + raise ValueError( + f"The output_size of gate's and up's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + if tp_size > 1 and intermediate_size_per_partition % block_k != 0: + # Required by row parallel + raise ValueError( + f"The input_size of down's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + w13_num_shards * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + if self.weight_quant.strategy == QuantizationStrategy.TENSOR: + # For gated MoE, allocate 2 scales for w1 and w3 respectively. + # They will be combined to a single scale after weight loading. + # For non-gated MoE, allocate 1 scale for w13. + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, w13_num_shards, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + # Add PER-TENSOR quantization for FusedMoE.weight_loader. + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + elif self.weight_quant.strategy == QuantizationStrategy.CHANNEL: + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + w13_num_shards * intermediate_size_per_partition, + 1, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + # Add PER-CHANNEL quantization for FusedMoE.weight_loader. + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + elif self.weight_quant.strategy == QuantizationStrategy.BLOCK: + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + w13_num_shards + * ((intermediate_size_per_partition + block_n - 1) // block_n), + (hidden_size + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + (hidden_size + block_n - 1) // block_n, + (intermediate_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + # Add PER-CHANNEL quantization for FusedMoE.weight_loader. + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + if self.static_input_scales: + w13_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + else: + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: FusedMoE) -> None: + # Allow for accessing weights and scales in standard way. + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = layer.w13_weight_scale + w2_scale = layer.w2_weight_scale + w13_input_scale = layer.w13_input_scale + w2_input_scale = layer.w2_input_scale + + # MI300x and MI325x use FNUZ format for FP8. Convert if needed. + if current_platform.is_fp8_fnuz(): + w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz( + w13, w13_scale, w13_input_scale + ) + w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz( + w2, w2_scale, w2_input_scale + ) + + # Per tensor kernels require single activation scale. Use the max. + if self.static_input_scales: + assert self.input_quant.strategy == QuantizationStrategy.TENSOR + assert w13_input_scale is not None and w2_input_scale is not None + w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe( + w13_input_scale, w2_input_scale + ) + replace_parameter(layer, "w13_input_scale", w13_input_scale) + replace_parameter(layer, "w2_input_scale", w2_input_scale) + + # Per-tensor kernels use a single scale, for W13, but on disk there + # is a separate scale for W1 and W3. Requantize with the max scale. + if self.weight_quant.strategy == QuantizationStrategy.TENSOR: + w13, w13_scale = process_fp8_weight_tensor_strategy_moe( + w13, + w13_scale, + shard_size=layer.intermediate_size_per_partition, + num_experts=layer.local_num_experts, + is_act_and_mul=self.moe.is_act_and_mul, + ) + + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=self.fp8_backend, + layer=layer, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + w13_input_scale=w13_input_scale, + w2_input_scale=w2_input_scale, + ) + + # Replace parameters with updated versions. Note that this helper + # function ensures the replacement is compatible with RL weight reloads. + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) + + # Setup modular kernel for TP case and naive DP/EP case. + # In non-naive DP/EP case, we will create a ModularKernelMethod. + # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel + # in both cases. + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config: + assert self.experts_cls is not None + self.moe_kernel = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + experts_cls=self.experts_cls, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, + ) + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel initialization " + "logic. This function should not be called." + ) + + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: + is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN + return make_fp8_moe_quant_config( + fp8_backend=self.fp8_backend, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + per_act_token_quant=is_per_token, + per_out_ch_quant=is_per_token, + block_shape=self.weight_block_size, + ) + + def apply_monolithic( + self, + layer: FusedMoE, + x: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + # TODO(rob): investigate the disable_expert_map introduced by: + # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501 + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, + ) + + @property + def supports_eplb(self) -> bool: + return True diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py new file mode 100644 index 00000000000..de155f9e179 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, +) + +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + FusedMoeWeightScaleSupported, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + int8_w8a8_moe_quant_config, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.utils import set_weight_attrs + +logger = init_logger(__name__) + + +class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): + def __init__( + self, + weight_quant: QuantizationArgs, + input_quant: QuantizationArgs, + moe: FusedMoEConfig, + layer_name: str | None = None, + ): + super().__init__(moe) + self.weight_quant = weight_quant + self.input_quant = input_quant + + per_channel = ( + self.weight_quant.strategy == QuantizationStrategy.CHANNEL + and self.input_quant.strategy == QuantizationStrategy.TOKEN + ) + if not per_channel: + raise ValueError( + "For INT8 Fused MoE layers, we require channelwise, " + "dynamic per token quantization. Found " + f"{self.weight_quant}, {self.input_quant}" + ) + + self.static_input_scales = not self.input_quant.dynamic + if self.static_input_scales: + raise ValueError( + "For INT8 Fused MoE layers, we require channelwise, " + "dynamic per token quantization. Found static input scales." + ) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + params_dtype = torch.int8 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + w13_num_shards * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + assert self.weight_quant.strategy == QuantizationStrategy.CHANNEL + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + w13_num_shards * intermediate_size_per_partition, + 1, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + # Add PER-CHANNEL quantization for FusedMoE.weight_loader. + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + assert not self.static_input_scales + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + return int8_w8a8_moe_quant_config( + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + per_act_token_quant=True, + ) + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + from vllm.model_executor.layers.fused_moe import fused_experts + + return fused_experts( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=not self.moe.disable_inplace, + activation=layer.activation, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + quant_config=self.moe_quant_config, + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py new file mode 100644 index 00000000000..f530a1a1df2 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from compressed_tensors.quantization import ( + QuantizationArgs, +) + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + int4_w4a16_moe_quant_config, + int8_w8a16_moe_quant_config, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.utils import set_weight_attrs + +logger = init_logger(__name__) + + +class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): + def __init__( + self, + weight_quant: QuantizationArgs, + input_quant: QuantizationArgs | None, + moe: FusedMoEConfig, + layer_name: str | None = None, + ): + super().__init__(moe) + self.weight_quant = weight_quant + self.input_quant = input_quant + # Extract properties from weight_quant + self.num_bits = weight_quant.num_bits + self.packed_factor = 32 // weight_quant.num_bits + self.strategy = weight_quant.strategy + # channelwise is not supported by this kernel + assert weight_quant.strategy == "group" + self.group_size = weight_quant.group_size + # grouped actorder isn't supported by this kernel + assert weight_quant.actorder != "group" + assert weight_quant.symmetric, ( + "Only symmetric quantization is supported for MoE" + ) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + # Will transpose the loaded weight along the + # intermediate and hidden dim sizes. Will + # shard for TP along the transposed dims + extra_weight_attrs.update( + {"is_transposed": True, "quant_method": self.strategy} + ) + w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size // self.packed_factor, + w13_num_shards * intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition // self.packed_factor, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w2_scales_size = intermediate_size_per_partition + + if self.strategy == "channel": + num_groups_w2 = num_groups_w13 = 1 + self.group_size = -1 + else: + num_groups_w2 = w2_scales_size // self.group_size + num_groups_w13 = hidden_size // self.group_size + + w13_scale = torch.nn.Parameter( + torch.ones( + num_experts, + num_groups_w13, + w13_num_shards * intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_scale) + set_weight_attrs(w13_scale, extra_weight_attrs) + + w2_scale = torch.nn.Parameter( + torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_scale) + set_weight_attrs(w2_scale, extra_weight_attrs) + set_weight_attrs(w2_scale, {"load_full_w2": False}) + + w2_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + layer.register_parameter("w2_weight_shape", w2_weight_shape) + set_weight_attrs(w2_weight_shape, extra_weight_attrs) + w13_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + + layer.register_parameter("w13_weight_shape", w13_weight_shape) + set_weight_attrs(w13_weight_shape, extra_weight_attrs) + + w13_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_g_idx", w13_g_idx) + set_weight_attrs(w13_g_idx, extra_weight_attrs) + + w2_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_g_idx", w2_g_idx) + set_weight_attrs(w2_g_idx, extra_weight_attrs) + + w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) + + w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + + layer.a13_scale = None + layer.a2_scale = None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Reconfigure packed weights and scales to match moe_wna16 format + layer.w13_weight_packed = torch.nn.Parameter( + layer.w13_weight_packed.transpose(1, 2).contiguous().view(torch.uint8), + requires_grad=False, + ) + layer.w2_weight_packed = torch.nn.Parameter( + layer.w2_weight_packed.transpose(1, 2).contiguous().view(torch.uint8), + requires_grad=False, + ) + layer.w13_weight_scale = torch.nn.Parameter( + layer.w13_weight_scale.transpose(1, 2).contiguous(), requires_grad=False + ) + layer.w2_weight_scale = torch.nn.Parameter( + layer.w2_weight_scale.transpose(1, 2).contiguous(), requires_grad=False + ) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + assert self.num_bits == 4 or self.num_bits == 8 + config_builder = ( + int4_w4a16_moe_quant_config + if self.num_bits == 4 + else int8_w8a16_moe_quant_config + ) + + return config_builder( + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, self.group_size], + ) + + def select_gemm_impl( + self, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, + layer: torch.nn.Module, + ) -> mk.FusedMoEExpertsModular: + if self.moe.is_lora_enabled: + assert self.moe_quant_config is not None + from vllm.triton_utils import HAS_TRITON + + if HAS_TRITON: + from vllm.model_executor.layers.fused_moe import TritonWNA16Experts + + layer.w13_weight = layer.w13_weight_packed + layer.w2_weight = layer.w2_weight_packed + return TritonWNA16Experts( + moe_config=self.moe, quant_config=self.moe_quant_config + ) + else: + raise NotImplementedError( + "TritonExperts requires Triton. " + "Install triton or disable LoRA for MoE." + ) + + raise NotImplementedError + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + from vllm.model_executor.layers.fused_moe import fused_experts + + return fused_experts( + x, + layer.w13_weight_packed, + layer.w2_weight_packed, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=not self.moe.disable_inplace, + activation=layer.activation, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + quant_config=self.moe_quant_config, + ) + + @property + def supports_eplb(self) -> bool: + return True diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py new file mode 100644 index 00000000000..216eed6372a --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -0,0 +1,575 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import enum +from enum import Enum + +import torch +from compressed_tensors.quantization import ( + QuantizationArgs, +) + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import _custom_ops as ops +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + int4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( + BatchedMarlinExperts, + MarlinExperts, + fused_marlin_moe, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa + WNA16_SUPPORTED_TYPES_MAP, +) +from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( + flashinfer_trtllm_mxint4_moe, + is_flashinfer_mxint4_moe_available, + prepare_static_weights_for_trtllm_mxint4_moe, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + get_marlin_input_dtype, + marlin_act_int8_process_scales, + marlin_make_workspace_new, + marlin_moe_permute_scales, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs + +logger = init_logger(__name__) + + +class GPTQMarlinState(Enum): + REPACK = enum.auto() + READY = enum.auto() + + +class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): + def __init__( + self, + weight_quant: QuantizationArgs, + input_quant: QuantizationArgs | None, + moe: FusedMoEConfig, + layer_name: str | None = None, + ): + super().__init__(moe) + self.weight_quant = weight_quant + self.input_quant = input_quant + assert weight_quant.symmetric, ( + "Only symmetric quantization is supported for MoE" + ) + # Extract properties from weight_quant + self.num_bits = weight_quant.num_bits + self.packed_factor = 32 // weight_quant.num_bits + self.strategy = weight_quant.strategy + self.group_size = weight_quant.group_size + self.actorder = weight_quant.actorder + + self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits] + + self.marlin_input_dtype = get_marlin_input_dtype(layer_name) + self.use_flashinfer_mxint4_moe = ( + is_flashinfer_mxint4_moe_available() + and self.group_size == 32 + and weight_quant.num_bits == 4 + ) + self.kernel_backend = ( + "Flashinfer" if self.use_flashinfer_mxint4_moe else "Marlin" + ) + logger.info_once( + f"Using {self.kernel_backend} backend for WNA16 MoE " + f"(group_size={self.group_size}, num_bits={self.num_bits})", + scope="local", + ) + + def get_weight_shape( + self, + weight_name: str, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + num_groups_w2: int | None = None, + num_groups_w13: int | None = None, + ) -> tuple[int, int, int]: + """ + Get the shape of the weight based on the weight name, number of experts + hidden size, intermediate size per partition, number of groups for w2, + and number of groups for w13. Pass in num_groups_w2 and num_groups_w13 + for weight scales. + """ + if weight_name == "w13_scale": + assert num_groups_w13 is not None, ( + "num_groups_w13 must be provided for weight scales" + ) + if weight_name == "w2_scale": + assert num_groups_w2 is not None, ( + "num_groups_w2 must be provided for weight scales" + ) + w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + shape_map = { + "w13_weight": { + "Flashinfer": ( + num_experts, + w13_num_shards * intermediate_size_per_partition, + hidden_size // self.packed_factor, + ), + "Marlin": ( + num_experts, + hidden_size // self.packed_factor, + w13_num_shards * intermediate_size_per_partition, + ), + }, + "w13_scale": { + "Flashinfer": ( + num_experts, + w13_num_shards * intermediate_size_per_partition, + num_groups_w13, + ), + "Marlin": ( + num_experts, + num_groups_w13, + w13_num_shards * intermediate_size_per_partition, + ), + }, + "w2_weight": { + "Flashinfer": ( + num_experts, + hidden_size, + intermediate_size_per_partition // self.packed_factor, + ), + "Marlin": ( + num_experts, + intermediate_size_per_partition // self.packed_factor, + hidden_size, + ), + }, + "w2_scale": { + "Flashinfer": (num_experts, hidden_size, num_groups_w2), + "Marlin": (num_experts, num_groups_w2, hidden_size), + }, + } + return shape_map[weight_name][self.kernel_backend] + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full") + + # Will transpose the loaded weight along the + # intermediate and hidden dim sizes. Will + # shard for TP along the transposed dims + is_transposed = self.kernel_backend != "Flashinfer" + extra_weight_attrs.update( + {"is_transposed": is_transposed, "quant_method": self.strategy} + ) + + w13_weight = torch.nn.Parameter( + torch.empty( + *self.get_weight_shape( + "w13_weight", + num_experts, + hidden_size, + intermediate_size_per_partition, + ), + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + *self.get_weight_shape( + "w2_weight", + num_experts, + hidden_size, + intermediate_size_per_partition, + ), + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # In the case where we have actorder/g_idx, + # we do not partition the w2 scales + load_full_w2 = self.actorder and self.group_size != -1 + w2_scales_size = ( + intermediate_size_full if load_full_w2 else intermediate_size_per_partition + ) + + self.is_k_full = (not self.actorder) or ( + intermediate_size_per_partition == intermediate_size_full + ) + + if self.strategy == "channel": + num_groups_w2 = num_groups_w13 = 1 + self.group_size = -1 + else: + num_groups_w2 = w2_scales_size // self.group_size + num_groups_w13 = hidden_size // self.group_size + + layer.num_groups_w13 = num_groups_w13 + layer.num_groups_w2 = num_groups_w2 + + w13_scale = torch.nn.Parameter( + torch.ones( + *self.get_weight_shape( + "w13_scale", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w13=num_groups_w13, + ), + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_scale) + set_weight_attrs(w13_scale, extra_weight_attrs) + + w2_scale = torch.nn.Parameter( + torch.ones( + *self.get_weight_shape( + "w2_scale", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w2=num_groups_w2, + ), + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_scale) + set_weight_attrs(w2_scale, extra_weight_attrs) + set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) + + w2_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + layer.register_parameter("w2_weight_shape", w2_weight_shape) + set_weight_attrs(w2_weight_shape, extra_weight_attrs) + w13_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + + layer.register_parameter("w13_weight_shape", w13_weight_shape) + set_weight_attrs(w13_weight_shape, extra_weight_attrs) + + w13_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_g_idx", w13_g_idx) + set_weight_attrs(w13_g_idx, extra_weight_attrs) + + w2_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_g_idx", w2_g_idx) + set_weight_attrs(w2_g_idx, extra_weight_attrs) + + w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) + + w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + + layer.a13_scale = None + layer.a2_scale = None + layer.marlin_state = GPTQMarlinState.REPACK + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + num_experts = layer.w13_weight_g_idx.shape[0] + device = layer.w13_weight_g_idx.device + if self.kernel_backend == "Flashinfer": + dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe( + layer.w13_weight_packed, + layer.w13_weight_scale, + layer.w2_weight_packed, + layer.w2_weight_scale, + ) + replace_parameter( + layer, "w13_weight_packed", dict_weights_mxint4["gemm1_weights"] + ) + replace_parameter( + layer, "w13_weight_scale", dict_weights_mxint4["gemm1_scales"] + ) + replace_parameter( + layer, "w2_weight_packed", dict_weights_mxint4["gemm2_weights"] + ) + replace_parameter( + layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"] + ) + return None + + is_a_8bit = ( + self.marlin_input_dtype is not None + and self.marlin_input_dtype.itemsize == 1 + ) + + if self.marlin_input_dtype == torch.float8_e4m3fn: + # NOTE: for non-zp quantization format only + ops.marlin_int4_fp8_preprocess(layer.w13_weight_packed, inplace=True) + ops.marlin_int4_fp8_preprocess(layer.w2_weight_packed, inplace=True) + layer.w13_weight_scale.data = layer.w13_weight_scale.data * 512 + layer.w2_weight_scale.data = layer.w2_weight_scale.data * 512 + + # when running models with grouped act order, + # resort to g_idx values provided in checkpoint + if self.actorder == "group": + w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx) + w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx) + w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx) + w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx) + + for e in range(num_experts): + w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to( + torch.int32 + ) + w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to( + torch.int32 + ) + w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][ + w13_g_idx_sort_indices[e] + ] + w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]] + + replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx) + replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx) + replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + + else: + layer.w13_weight_g_idx = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_weight_g_idx = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + + marlin_w13_qweight = ops.gptq_marlin_moe_repack( + layer.w13_weight_packed, + layer.w13_g_idx_sort_indices, + layer.w13_weight_packed.shape[1] * self.packed_factor, + layer.w13_weight_packed.shape[2], + self.num_bits, + is_a_8bit=is_a_8bit, + ) + replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight) + + marlin_w2_qweight = ops.gptq_marlin_moe_repack( + layer.w2_weight_packed, + layer.w2_g_idx_sort_indices, + layer.w2_weight_packed.shape[1] * self.packed_factor, + layer.w2_weight_packed.shape[2], + self.num_bits, + is_a_8bit=is_a_8bit, + ) + replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight) + + # Repack scales + marlin_w13_scales = marlin_moe_permute_scales( + s=layer.w13_weight_scale, + size_k=layer.w13_weight_packed.shape[2], + size_n=layer.w13_weight_scale.shape[2], + group_size=self.group_size, + is_a_8bit=is_a_8bit, + ) + if self.marlin_input_dtype == torch.int8 and layer.num_groups_w13 > 1: + marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales( + marlin_w13_scales + ) + layer.register_parameter( + "w13_input_global_scale", + torch.nn.Parameter(w13_input_global_scale, requires_grad=False), + ) + replace_parameter(layer, "w13_weight_scale", marlin_w13_scales) + + marlin_w2_scales = marlin_moe_permute_scales( + s=layer.w2_weight_scale, + size_k=layer.w2_weight_scale.shape[1] + * (self.group_size if self.group_size != -1 else self.packed_factor), + size_n=layer.w2_weight_scale.shape[2], + group_size=self.group_size, + is_a_8bit=is_a_8bit, + ) + if self.marlin_input_dtype == torch.int8 and layer.num_groups_w2 > 1: + marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales( + marlin_w2_scales + ) + layer.register_parameter( + "w2_input_global_scale", + torch.nn.Parameter(w2_input_global_scale, requires_grad=False), + ) + replace_parameter(layer, "w2_weight_scale", marlin_w2_scales) + + layer.workspace = marlin_make_workspace_new(device, 4) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + if self.num_bits != 4: + return None + return int4_w4a16_moe_quant_config( + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, self.group_size], + ) + + def select_gemm_impl( + self, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, + layer: torch.nn.Module, + ) -> mk.FusedMoEExpertsModular: + assert self.num_bits == 4, "only supporting w4" + layer.w13_weight = layer.w13_weight_packed + layer.w2_weight = layer.w2_weight_packed + assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]]) + assert self.moe_quant_config is not None + if ( + prepare_finalize.activation_format + == mk.FusedMoEActivationFormat.BatchedExperts + ): + max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank() + assert max_num_tokens_per_rank is not None + return BatchedMarlinExperts( + max_num_tokens=max_num_tokens_per_rank, + num_dispatchers=prepare_finalize.num_dispatchers(), + moe_config=self.moe, + quant_config=self.moe_quant_config, + w13_g_idx=layer.w13_weight_g_idx, + w2_g_idx=layer.w2_weight_g_idx, + w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, + w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, + is_k_full=self.is_k_full, + ) + else: + return MarlinExperts( + moe_config=self.moe, + quant_config=self.moe_quant_config, + w13_g_idx=layer.w13_weight_g_idx, + w2_g_idx=layer.w2_weight_g_idx, + w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, + w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, + is_k_full=self.is_k_full, + ) + + @property + def is_monolithic(self) -> bool: + return self.kernel_backend == "Flashinfer" + + def apply_monolithic( + self, + layer: FusedMoE, + x: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + assert self.kernel_backend == "Flashinfer" + return flashinfer_trtllm_mxint4_moe( + x=x, + router_logits=router_logits, + w13_weight_packed=layer.w13_weight_packed, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_packed=layer.w2_weight_packed, + w2_weight_scale=layer.w2_weight_scale, + global_num_experts=layer.global_num_experts, + top_k=layer.top_k, + intermediate_size_per_partition=layer.intermediate_size_per_partition, + local_num_experts=layer.local_num_experts, + ep_rank=layer.ep_rank, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routing_method_type=layer.routing_method_type, + ) + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert self.kernel_backend == "Marlin" + return fused_marlin_moe( + x, + layer.w13_weight_packed, + layer.w2_weight_packed, + None, + None, + layer.w13_weight_scale, + layer.w2_weight_scale, + topk_weights, + topk_ids, + input_global_scale1=getattr(layer, "w13_input_global_scale", None), + input_global_scale2=getattr(layer, "w2_input_global_scale", None), + quant_type_id=self.quant_type.id, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + global_num_experts=layer.global_num_experts, + activation=layer.activation, + expert_map=layer.expert_map, + g_idx1=layer.w13_weight_g_idx, + g_idx2=layer.w2_weight_g_idx, + sort_indices1=layer.w13_g_idx_sort_indices, + sort_indices2=layer.w2_g_idx_sort_indices, + workspace=layer.workspace, + input_dtype=self.marlin_input_dtype, + is_k_full=self.is_k_full, + inplace=not self.moe.disable_inplace, + ) From 62095e82c1ef7ed91be5223d36c6109958e9aec6 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Mon, 6 Apr 2026 17:21:09 -0700 Subject: [PATCH 26/39] [BugFix][MRV2] Fix cuda event reuse race (#39115) Signed-off-by: Nick Hill --- vllm/v1/worker/gpu/async_utils.py | 6 ++---- vllm/v1/worker/gpu/model_runner.py | 3 --- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index 7f270c2b8c9..b3d6f5e4d90 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -17,7 +17,6 @@ class AsyncOutput(AsyncModelRunnerOutput): num_sampled_tokens: torch.Tensor, main_stream: torch.cuda.Stream, copy_stream: torch.cuda.Stream, - copy_event: torch.cuda.Event, ): # NOTE(woosuk): We must retain references to the GPU tensors, # as the copy operations are performed on a different CUDA stream than @@ -25,7 +24,7 @@ class AsyncOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.sampler_output = sampler_output self.num_sampled_tokens = num_sampled_tokens - self.copy_event = copy_event + self.copy_event = torch.cuda.Event() with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) @@ -78,12 +77,11 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput): is_valid: torch.Tensor | None, main_stream: torch.cuda.Stream, copy_stream: torch.cuda.Stream, - copy_event: torch.cuda.Event, ): self.model_runner_output = model_runner_output self.pooler_output = pooler_output self.is_valid = is_valid - self.copy_event = copy_event + self.copy_event = torch.cuda.Event() with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 56df70fc0c9..f188b061a6c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -130,7 +130,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.use_async_scheduling = self.scheduler_config.async_scheduling self.output_copy_stream = torch.cuda.Stream(self.device) - self.output_copy_event = torch.cuda.Event() # Pipeline parallelism. self.use_pp = self.parallel_config.pipeline_parallel_size > 1 @@ -1180,7 +1179,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): num_sampled_tokens=num_sampled, main_stream=self.main_stream, copy_stream=self.output_copy_stream, - copy_event=self.output_copy_event, ) mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None @@ -1270,7 +1268,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): is_valid=is_valid, main_stream=self.main_stream, copy_stream=self.output_copy_stream, - copy_event=self.output_copy_event, ) self.postprocess_pool(input_batch) From 2df2c85be494d1b98f8a66d0acb7c65f0b119a37 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 6 Apr 2026 21:57:09 -0500 Subject: [PATCH 27/39] [Kernels][MoE] Fix legacy_routing to use bitmatrix-based routing path (#38504) Signed-off-by: Andreas Karatzas --- .../configs/gpt-oss-20b-rocm-baseline.yaml | 2 +- ...t-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml | 4 +- ...-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml | 2 +- ...t-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml | 2 +- .../moe/test_gpt_oss_triton_kernels.py | 69 ++------ .../quantization/test_mxfp4_triton_ep.py | 59 ++----- .../fused_moe/gpt_oss_triton_kernels_moe.py | 162 ++++++------------ 7 files changed, 84 insertions(+), 216 deletions(-) diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-baseline.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-baseline.yaml index 76b1d796230..ec1c2b3922d 100644 --- a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-baseline.yaml +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-baseline.yaml @@ -3,4 +3,4 @@ model_name: openai/gpt-oss-20b metric_threshold: 0.568 reasoning_effort: low -server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN" \ No newline at end of file +server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --tensor-parallel-size 2" \ No newline at end of file diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml index 850a6d28be0..4ff2648ca82 100644 --- a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml @@ -3,6 +3,6 @@ model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16 metric_threshold: 0.568 reasoning_effort: low -server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter" +server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter --tokenizer openai/gpt-oss-20b --tensor-parallel-size 2" env: - VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER: "1" \ No newline at end of file diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml index 903f30e59e7..5ae665a044a 100644 --- a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml @@ -3,4 +3,4 @@ model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16 metric_threshold: 0.568 reasoning_effort: low -server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton" \ No newline at end of file +server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton --tokenizer openai/gpt-oss-20b --tensor-parallel-size 2" \ No newline at end of file diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml index f7dd14784a1..81270e0105f 100644 --- a/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml @@ -3,6 +3,6 @@ model_name: amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8 metric_threshold: 0.568 reasoning_effort: low -server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN" +server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --tensor-parallel-size 2" env: VLLM_ROCM_USE_AITER: "1" \ No newline at end of file diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 172938f18e4..032b4fc047c 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -23,16 +23,12 @@ from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_m from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor from triton_kernels.tensor_details import layout from triton_kernels.testing import assert_close -from triton_kernels.topk import topk as topk_fn from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( - legacy_routing, - make_routing_data, triton_kernel_moe_forward, ) from vllm.utils.math_utils import round_up -from vllm.utils.torch_utils import set_random_seed from .utils import shuffle_weight @@ -97,10 +93,18 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int): if w_dtype != "mx4": pytest.skip("NYI") else: # quantize to mx4 - # careful on the padding here, the activation padding need to be - # multiple of 64, the actual engine is not implemented - w1_bottom_pad = round_up(w1_tri.shape[1], 64) - w1_tri.shape[1] - w1_right_pad = round_up(w1_tri.shape[2], 128) - w1_tri.shape[2] + # Padding alignment depends on the platform. On CDNA4 the scale + # swizzle requires SCALE_K % 8 == 0 (K % 256) and + # SCALE_N % 32 == 0 (2*N % 512), matching the production + # alignment in mxfp4_round_up_hidden_size_and_intermediate_size. + # On CUDA (Hopper) the scale layout pads internally, so the + # original 64/128 alignment is sufficient. + if current_platform.is_rocm(): + k_align, n2_align = 256, 512 + else: + k_align, n2_align = 64, 128 + w1_bottom_pad = round_up(w1_tri.shape[1], k_align) - w1_tri.shape[1] + w1_right_pad = round_up(w1_tri.shape[2], n2_align) - w1_tri.shape[2] w2_bottom_pad = w1_right_pad // 2 w2_right_pad = w1_bottom_pad @@ -367,52 +371,3 @@ def test_unit_shuffle(): ) assert_close(ref=out_ref, tri=out) - - -@pytest.mark.parametrize("num_tokens", [2, 8, 64]) -@pytest.mark.parametrize("num_experts", [32, 128]) -@pytest.mark.parametrize("topk", [1, 4]) -@pytest.mark.parametrize("renormalize", [True, False]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -def test_legacy_routing( - num_tokens: int, num_experts: int, topk: int, renormalize: bool, dtype: torch.dtype -): - set_random_seed(0) - gating_output = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype) - - sm_first = not renormalize - logits = gating_output - if sm_first: - logits = torch.softmax(logits, dim=-1) - topk_result = topk_fn(logits, topk, apply_softmax=not sm_first) - # topk_fn returns SparseMatrix on NVIDIA, plain tuple on ROCm. - if isinstance(topk_result, tuple): - topk_weights, topk_ids_raw, bitmatrix = topk_result - from triton_kernels.routing import routing_from_bitmatrix - - routing_data_ref, gather_indx_ref, scatter_indx_ref = routing_from_bitmatrix( - bitmatrix, topk_weights, topk_ids_raw, num_experts, topk - ) - else: - topk_ids = topk_result.indx.to(torch.long) - topk_weights = topk_result.vals - routing_data_ref, gather_indx_ref, scatter_indx_ref = make_routing_data( - topk_ids, topk_weights, num_experts - ) - - routing_data, gather_indx, scatter_indx = legacy_routing( - gating_output, topk, sm_first=sm_first - ) - - assert_close( - ref=gather_indx_ref.src_indx, tri=gather_indx.src_indx, maxtol=0, rmstol=0 - ) - assert_close( - ref=gather_indx_ref.dst_indx, tri=gather_indx.dst_indx, maxtol=0, rmstol=0 - ) - assert_close( - ref=scatter_indx_ref.src_indx, tri=scatter_indx.src_indx, maxtol=0, rmstol=0 - ) - assert_close( - ref=scatter_indx_ref.dst_indx, tri=scatter_indx.dst_indx, maxtol=0, rmstol=0 - ) diff --git a/tests/kernels/quantization/test_mxfp4_triton_ep.py b/tests/kernels/quantization/test_mxfp4_triton_ep.py index 6c8aebe42c0..045bc63de90 100644 --- a/tests/kernels/quantization/test_mxfp4_triton_ep.py +++ b/tests/kernels/quantization/test_mxfp4_triton_ep.py @@ -4,12 +4,9 @@ Tests that triton_kernel_moe_forward correctly applies expert_map remapping when expert parallelism (EP) is enabled. -Previously, legacy_routing was always used and it produced routing data -with global expert IDs that didn't correspond to local weight indices, -causing illegal memory access with EP. The fix splits routing: when -expert_map is provided, topk selection is performed first, expert_map is -applied to remap global→local IDs, and make_routing_data builds routing -structures from the local IDs. +Both EP and non-EP paths use topk + make_routing_data. When expert_map +is provided, global expert IDs are remapped to local IDs before building +routing structures. """ from unittest.mock import MagicMock, patch @@ -24,21 +21,15 @@ class TestTritonMoeForwardExpertMap: @pytest.mark.parametrize("expert_map_present", [False, True]) def test_routing_path_selection(self, expert_map_present): - """Verify that the EP-aware routing path is taken when expert_map - is present, and the legacy_routing path is taken otherwise.""" + """Verify that both EP and non-EP paths use topk + make_routing_data, + and that expert_map remapping is applied when present.""" device = "cuda" if torch.cuda.is_available() else "cpu" - # This is a structural test: we mock the routing functions to - # verify the correct path is exercised. mock_expert_map = ( torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None ) with ( - patch( - "vllm.model_executor.layers.fused_moe." - "gpt_oss_triton_kernels_moe.legacy_routing" - ) as mock_legacy, patch("triton_kernels.topk.topk") as mock_topk, patch( "vllm.model_executor.layers.fused_moe." @@ -53,27 +44,19 @@ class TestTritonMoeForwardExpertMap: triton_kernel_moe_forward, ) - # Set up return values mock_routing_data = MagicMock() mock_gather = MagicMock() mock_scatter = MagicMock() - if expert_map_present: - sparse_result = MagicMock() - sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32) - sparse_result.vals = torch.tensor([[0.6, 0.4]]) - mock_topk.return_value = sparse_result - mock_make_routing.return_value = ( - mock_routing_data, - mock_gather, - mock_scatter, - ) - else: - mock_legacy.return_value = ( - mock_routing_data, - mock_gather, - mock_scatter, - ) + sparse_result = MagicMock() + sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32) + sparse_result.vals = torch.tensor([[0.6, 0.4]]) + mock_topk.return_value = sparse_result + mock_make_routing.return_value = ( + mock_routing_data, + mock_gather, + mock_scatter, + ) mock_fused_experts.return_value = torch.zeros((1, 8), device=device) @@ -92,20 +75,14 @@ class TestTritonMoeForwardExpertMap: expert_map=mock_expert_map, ) + # Both paths use topk + make_routing_data + mock_topk.assert_called_once() + mock_make_routing.assert_called_once() + if expert_map_present: - # EP path: should use topk + make_routing_data, NOT - # legacy_routing - mock_topk.assert_called_once() - mock_make_routing.assert_called_once() - mock_legacy.assert_not_called() # expert_map should be None in the fused_experts call # (already applied) call_kwargs = mock_fused_experts.call_args assert call_kwargs[1].get("expert_map") is None or ( len(call_kwargs[0]) > 0 ) - else: - # Non-EP path: should use legacy_routing - mock_legacy.assert_called_once() - mock_topk.assert_not_called() - mock_make_routing.assert_not_called() diff --git a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py index e03ecd01ae7..a21ddaba075 100644 --- a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py @@ -47,7 +47,6 @@ if has_triton_kernels(): BIT, Bitmatrix, ) - from triton_kernels.topk import topk try: from triton_kernels.tensor import ( @@ -89,6 +88,7 @@ def pack_bitmatrix( offsets = offsets_m[:, None] * n_expts_act + offsets_k[None, :] mask = (offsets_m < n_rows)[:, None] & (offsets_k < n_expts_act)[None, :] indices = tl.load(topk_ids + offsets, mask=mask, other=-1) + valid = indices >= 0 div = indices // 32 rem = indices % 32 one = tl.cast(1, tl.uint32) @@ -99,8 +99,13 @@ def pack_bitmatrix( offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32) # All topks that need to go into this column has the correct bit set. # Other bits are 0. x is a 2D tensor. + # Guard with `valid` to prevent negative indices from producing + # spurious bits (on HIP, -1 // 32 == 0 and 1 << (-1 % 32) sets + # bit 31). x = tl.where( - div[:, :, None] == offs[None, None, :], (one << rem)[:, :, None], 0 + valid[:, :, None] & (div[:, :, None] == offs[None, None, :]), + (one << rem)[:, :, None], + 0, ) # Reduce x to get a single int32_t bitpack. y = tl.reduce_or(x, axis=1) @@ -108,93 +113,6 @@ def pack_bitmatrix( tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows) -def legacy_routing_from_bitmatrix( - bitmatrix: "Bitmatrix", - expt_scal: torch.Tensor, - expt_indx: torch.Tensor, - n_expts_tot: int, - n_expts_act: int, -) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]: - """ - Replacement for the removed triton_kernels.routing.routing_from_bitmatrix. - Creates routing data from a bitmatrix representation. - """ - if use_legacy_triton_kernels: - from triton_kernels.routing import routing_from_bitmatrix - - return routing_from_bitmatrix( - bitmatrix, expt_scal, expt_indx, n_expts_tot, n_expts_act - ) - sparse_logits = SparseMatrix(indx=expt_indx, vals=expt_scal, mask=bitmatrix) - dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx - combine_indx = sparse_logits.mask_metadata.col_sorted_indx - ragged_batch_metadata = make_ragged_tensor_metadata( - sparse_logits.mask_metadata.col_sum, - dispatch_indx.shape[0], - ) - gate_scal = sparse_logits.vals.flatten()[combine_indx] - routing_data = RoutingData( - gate_scal, - ragged_batch_metadata.block_sizes, - n_expts_tot, - n_expts_act, - ragged_batch_metadata, - ) - gather_idx = GatherIndx(combine_indx, dispatch_indx) - scatter_idx = ScatterIndx(dispatch_indx, combine_indx) - return routing_data, gather_idx, scatter_idx - - -def legacy_routing_from_sparsematrix( - sparse_logits: "SparseMatrix", - n_expts_tot: int, - n_expts_act: int, -) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]: - """ - Creates routing data from a SparseMatrix representation. - """ - dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx - combine_indx = sparse_logits.mask_metadata.col_sorted_indx - ragged_batch_metadata = make_ragged_tensor_metadata( - sparse_logits.mask_metadata.col_sum, - dispatch_indx.shape[0], - ) - gate_scal = sparse_logits.vals.flatten()[combine_indx] - routing_data = RoutingData( - gate_scal, - ragged_batch_metadata.block_sizes, - n_expts_tot, - n_expts_act, - ragged_batch_metadata, - ) - gather_idx = GatherIndx(combine_indx, dispatch_indx) - scatter_idx = ScatterIndx(dispatch_indx, combine_indx) - return routing_data, gather_idx, scatter_idx - - -def legacy_routing( - logits: torch.Tensor, - n_expts_act: int, - sm_first: bool = False, -) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]: - """ - Replacement for the removed triton_kernels.routing.routing function. - Computes routing data from gating logits. - """ - if use_legacy_triton_kernels: - from triton_kernels.routing import routing - - return routing(logits, n_expts_act, sm_first=sm_first) - if sm_first: - logits = torch.softmax(logits, dim=-1) - sparse_logits = topk(logits, n_expts_act, apply_softmax=not sm_first) - return legacy_routing_from_sparsematrix( - sparse_logits, - logits.shape[-1], - n_expts_act, - ) - - def triton_kernel_moe_forward( hidden_states: torch.Tensor, w1, # Tensor or triton_kernels.Tensor @@ -241,26 +159,22 @@ def triton_kernel_moe_forward( unpadded_K_w2=unpadded_K_w2, ) - if expert_map is not None: - # With expert parallelism, legacy_routing produces routing data - # using global expert IDs which don't correspond to local weight - # indices. Split the routing into topk selection + expert_map - # remapping + local routing data construction (matching the - # approach used by OAITritonExperts.apply). - from triton_kernels.topk import topk as topk_fn + from triton_kernels.topk import topk as topk_fn - sm_first = not renormalize - logits = gating_output - if sm_first: - logits = torch.softmax(logits, dim=-1) - topk_result = topk_fn(logits, topk, apply_softmax=not sm_first) - # topk may return a tuple (vals, indx, bitmatrix) or a - # SparseMatrix depending on the triton_kernels version. - if isinstance(topk_result, tuple): - topk_weights, topk_ids_raw, _ = topk_result - else: - topk_weights = topk_result.vals - topk_ids_raw = topk_result.indx + sm_first = not renormalize + logits = gating_output + if sm_first: + logits = torch.softmax(logits, dim=-1) + topk_result = topk_fn(logits, topk, apply_softmax=not sm_first) + # topk may return a tuple (vals, indx, bitmatrix) or a + # SparseMatrix depending on the triton_kernels version. + if isinstance(topk_result, tuple): + topk_weights, topk_ids_raw, _ = topk_result + else: + topk_weights = topk_result.vals + topk_ids_raw = topk_result.indx + + if expert_map is not None: # topk_ids_raw contains global expert IDs - remap to local. topk_ids = expert_map[topk_ids_raw.to(torch.long)] local_num_experts = w1.shape[0] @@ -271,8 +185,9 @@ def triton_kernel_moe_forward( effective_expert_map = None effective_global_num_experts = local_num_experts else: - routing_data, gather_idx, scatter_idx = legacy_routing( - gating_output, topk, sm_first=not renormalize + topk_ids = topk_ids_raw.to(torch.long) + routing_data, gather_idx, scatter_idx = make_routing_data( + topk_ids, topk_weights, gating_output.shape[-1] ) effective_expert_map = expert_map effective_global_num_experts = global_num_experts @@ -539,10 +454,31 @@ def make_routing_data( # matmul_ogs expects invalid topk_weights to be -1s topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights) - routing_data, gather_indx, scatter_indx = legacy_routing_from_bitmatrix( - bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk - ) + if use_legacy_triton_kernels: + from triton_kernels.routing import routing_from_bitmatrix + + return routing_from_bitmatrix( + bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk + ) + + sparse_logits = SparseMatrix(indx=topk_ids, vals=topk_weights, mask=bitmatrix) + dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx + combine_indx = sparse_logits.mask_metadata.col_sorted_indx + ragged_batch_metadata = make_ragged_tensor_metadata( + sparse_logits.mask_metadata.col_sum, + dispatch_indx.shape[0], + ) + gate_scal = sparse_logits.vals.flatten()[combine_indx] + routing_data = RoutingData( + gate_scal, + ragged_batch_metadata.block_sizes, + num_local_experts, + num_topk, + ragged_batch_metadata, + ) + gather_indx = GatherIndx(combine_indx, dispatch_indx) + scatter_indx = ScatterIndx(dispatch_indx, combine_indx) return routing_data, gather_indx, scatter_indx From a435e3108d82eb96d9b3954c1935afbbf4c5f69b Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Tue, 7 Apr 2026 00:36:21 -0500 Subject: [PATCH 28/39] [ROCm][CI] Fix test repo-root assumptions (#39053) Signed-off-by: Andreas Karatzas --- .../ec_connector/integration/run_epd_correctness_test.sh | 5 +++-- .../v1/kv_connector/nixl_integration/run_accuracy_test.sh | 7 +++++-- .../v1/kv_connector/nixl_integration/run_edge_case_test.sh | 5 +++-- .../nixl_integration/run_tpu_disagg_accuracy_test.sh | 5 +++-- .../nixl_integration/run_tpu_edge_case_test.sh | 3 ++- .../nixl_integration/run_xpu_disagg_accuracy_test.sh | 3 ++- .../nixl_integration/spec_decode_acceptance_test.sh | 4 +++- 7 files changed, 21 insertions(+), 11 deletions(-) diff --git a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh index ffe9cac3803..e199a3ecea4 100644 --- a/tests/v1/ec_connector/integration/run_epd_correctness_test.sh +++ b/tests/v1/ec_connector/integration/run_epd_correctness_test.sh @@ -15,8 +15,9 @@ # set -xe -# Find the git repository root directory -GIT_ROOT=$(git rev-parse --show-toplevel) +# Resolve the repository root from the script location instead of `.git`. +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" # Model to test MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}" diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index fe95249602a..fc446a0e765 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -85,8 +85,11 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128} # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048) VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} -# Find the git repository root directory -GIT_ROOT=$(git rev-parse --show-toplevel) +# Resolve the repository root from the script location instead of `.git`. +# The ROCm CI image copies `/vllm-workspace` without the Git metadata, so +# `git rev-parse --show-toplevel` is not reliable at runtime. +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") diff --git a/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh index 703a27fd3f7..9d8e4df8c53 100755 --- a/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh @@ -33,8 +33,9 @@ MODELS=( "Qwen/Qwen3-0.6B" ) -# Find the git repository root directory -GIT_ROOT=$(git rev-parse --show-toplevel) +# Resolve the repository root from the script location instead of `.git`. +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" # Trap the SIGINT signal (triggered by Ctrl+C) trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT diff --git a/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh index 407542eb82b..9274e3c573c 100644 --- a/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh @@ -20,7 +20,8 @@ BLOCK_SIZE=${BLOCK_SIZE:-32} # execution env -GIT_ROOT=$(git rev-parse --show-toplevel) +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration" CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"} CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"} @@ -153,4 +154,4 @@ echo "-----P/D success----" rm "${OUTPUT_FILE}" cleanup -exit 0 \ No newline at end of file +exit 0 diff --git a/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh index f32ef5e764c..5969455025e 100644 --- a/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh @@ -20,7 +20,8 @@ BLOCK_SIZE=${BLOCK_SIZE:-32} # execution env -GIT_ROOT=$(git rev-parse --show-toplevel) +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration" CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"} CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"} diff --git a/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh index 79863123b72..8340720f927 100644 --- a/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh @@ -44,7 +44,8 @@ DECODER_ZE_AFFINITY_MASK=${DECODER_ZE_AFFINITY_MASK:-$(generate_affinity_mask "$ # execution env -GIT_ROOT=$(git rev-parse --show-toplevel) +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration" OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.xpu_accuracy_test_outputs.txt"} diff --git a/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh b/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh index c2c938ebffe..a82dae2d510 100755 --- a/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh +++ b/tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh @@ -52,7 +52,9 @@ DECODER_TP_SIZE=${DECODER_TP_SIZE:-1} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.7} BLOCK_SIZE=${BLOCK_SIZE:-16} -GIT_ROOT=$(git rev-parse --show-toplevel) +# Resolve the repository root from the script location instead of `.git`. +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)" +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}" SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") From 5c35517a3e66740ec728a124b81029b57998d937 Mon Sep 17 00:00:00 2001 From: Andrew Barnes Date: Tue, 7 Apr 2026 03:17:59 -0400 Subject: [PATCH 29/39] [ROCm] Remove unused IS_FNUZ parameter from reshape_and_cache_shuffle_kernel (#39123) Signed-off-by: Bortlesboat --- vllm/v1/attention/backends/rocm_aiter_fa.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 29351fcbf51..a897b33fcab 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -228,7 +228,6 @@ if current_platform.is_rocm(): num_kv_heads, BLOCK_SIZE: tl.constexpr, QUANT: tl.constexpr, - IS_FNUZ: tl.constexpr, ): tid = tl.program_id(0) head_id = tl.program_id(1) @@ -314,7 +313,6 @@ if current_platform.is_rocm(): num_kv_heads, BLOCK_SIZE=head_size, QUANT=QUANT, - IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz, ) From a9a0e0551f038d7306cf8887069264d917c643a5 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Tue, 7 Apr 2026 10:23:29 +0300 Subject: [PATCH 30/39] nano-nemotron-vl: get_mm_max_tokens_per_item for audio, video, image == seq_len (#38727) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- .../model_executor/models/nano_nemotron_vl.py | 62 ++++++++++++++++--- 1 file changed, 52 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 249b2896910..9983015b0ee 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -288,6 +288,35 @@ class NanoNemotronVLProcessingInfo(BaseProcessingInfo): max_num_tiles=max_num_tiles, ) + def get_dummy_image_size_and_max_tokens( + self, mm_counts: Mapping[str, int] + ) -> tuple[tuple[int, int], int]: + processor = self.get_hf_processor() + num_images = mm_counts.get("image", 0) + + if tiler := processor.dynamic_tiler: + budget = tiler.max_num_tokens_available(text_prompt_length=num_images) + target_width, target_height = ( + tiler.width_and_height_for_max_num_tokens_available(budget) + ) + return ( + (target_width, target_height), + tiler._get_num_embeddings(target_width, target_height), + ) + + max_num_tiles = processor.max_num_tiles + target_width, target_height = self.get_image_size_with_most_features( + max_num_tiles + ) + return ( + (target_width, target_height), + processor.get_num_image_tokens( + image_width=target_width, + image_height=target_height, + max_num_tiles=max_num_tiles, + ), + ) + def get_num_frames_with_most_features( self, seq_len: int, @@ -306,6 +335,26 @@ class NanoNemotronVLProcessingInfo(BaseProcessingInfo): max_frames_per_video = max_tubelets_per_video * T return max(max_frames_per_video, 1) + def get_mm_max_tokens_per_item( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> Mapping[str, int]: + mm_max_tokens: dict[str, int] = {} + + if mm_counts.get("image", 0) > 0: + _, mm_max_tokens["image"] = self.get_dummy_image_size_and_max_tokens( + mm_counts + ) + + if mm_counts.get("video", 0) > 0: + assert self.supports_video + mm_max_tokens["video"] = seq_len + + if mm_counts.get("audio", 0) > 0: + assert self.supports_audio + mm_max_tokens["audio"] = seq_len + + return mm_max_tokens + class NanoNemotronVLMultiModalProcessor( BaseMultiModalProcessor[NanoNemotronVLProcessingInfo] @@ -708,17 +757,10 @@ class NanoNemotronVLDummyInputsBuilder( mm_options: Mapping[str, BaseDummyOptions], ) -> MultiModalDataDict: num_images = mm_counts.get("image", 0) + (target_width, target_height), _ = ( + self.info.get_dummy_image_size_and_max_tokens(mm_counts) + ) processor = self.info.get_hf_processor() - if tiler := processor.dynamic_tiler: - budget = tiler.max_num_tokens_available(text_prompt_length=num_images) - target_width, target_height = ( - tiler.width_and_height_for_max_num_tokens_available(budget) - ) - else: - max_num_tiles = 12 - target_width, target_height = self.info.get_image_size_with_most_features( - max_num_tiles - ) image_overrides = mm_options.get("image") From da4c0e4db93c3d3ad18003cca56c645c91010532 Mon Sep 17 00:00:00 2001 From: Rishapveer Singh Date: Tue, 7 Apr 2026 10:25:17 +0200 Subject: [PATCH 31/39] [Model] Use AutoWeightsLoader for FalconH1 (#39092) Signed-off-by: Rishapveer Singh <215205492+rishaps@users.noreply.github.com> --- vllm/model_executor/models/falcon_h1.py | 122 ++++++++++++------------ 1 file changed, 63 insertions(+), 59 deletions(-) diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index fba2e216e3f..b837dc010da 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -50,6 +50,7 @@ from .interfaces import ( SupportsPP, ) from .utils import ( + AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, @@ -495,6 +496,63 @@ class FalconH1Model(nn.Module): hidden_states = self.final_layernorm(hidden_states) return hidden_states + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + if "A_log" in name: + name = name.replace("A_log", "A") + + if "mamba" in name: + name = name.replace("mamba", "mamba.mamba") + + if "scale" in name: + # Remapping the name of kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip layers on other devices. + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params + class FalconH1ForCausalLM( nn.Module, @@ -632,62 +690,8 @@ class FalconH1ForCausalLM( return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("qkv_proj", "q_proj", "q"), - ("qkv_proj", "k_proj", "k"), - ("qkv_proj", "v_proj", "v"), - ("gate_up_proj", "gate_proj", 0), - ("gate_up_proj", "up_proj", 1), - ] - - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - - if "A_log" in name: - name = name.replace("A_log", "A") - - if "mamba" in name: - name = name.replace("mamba", "mamba.mamba") - - if "scale" in name: - # Remapping the name of kv-scale. - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Skip layers on other devices. - if is_pp_missing_parameter(name, self): - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if is_pp_missing_parameter(name, self): - continue - if self.tie_word_embeddings and "lm_head" in name: - continue - - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - if self.tie_word_embeddings: - loaded_params.add("lm_head.weight") - return loaded_params + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None), + ) + return loader.load_weights(weights) From 8060bb0333855739d5bdee0e67f7eb227fa04ebb Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 7 Apr 2026 16:37:00 +0800 Subject: [PATCH 32/39] [vLLM IR] rework gemma_rms_norm (#39014) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: zjy0516 Signed-off-by: Jiangyun Zhu Co-authored-by: Luka Govedič --- tests/kernels/core/test_layernorm.py | 30 +++++++- .../passes/fusion/allreduce_rms_fusion.py | 24 ++++++- .../passes/fusion/rms_quant_fusion.py | 30 +++++++- vllm/ir/ops/layernorm.py | 5 +- vllm/kernels/aiter_ops.py | 10 ++- vllm/kernels/vllm_c.py | 7 +- vllm/kernels/xpu_ops.py | 4 +- vllm/model_executor/layers/layernorm.py | 71 ++++--------------- 8 files changed, 106 insertions(+), 75 deletions(-) diff --git a/tests/kernels/core/test_layernorm.py b/tests/kernels/core/test_layernorm.py index 42da24ccb96..c39d42c7593 100644 --- a/tests/kernels/core/test_layernorm.py +++ b/tests/kernels/core/test_layernorm.py @@ -6,7 +6,7 @@ import torch from tests.kernels.quant_utils import FP8_DTYPE from tests.kernels.utils import opcheck -from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -162,3 +162,31 @@ def test_fused_rms_norm_quant( atol=1e-3, rtol=1e-3, ) + + +@torch.inference_mode() +def test_gemma_rms_norm_mixed_input_weight_dtype(default_vllm_config) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + device = CUDA_DEVICES[0] + torch.set_default_device(device) + + num_tokens, hidden_size = 32, 1024 + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) + layer = GemmaRMSNorm(hidden_size, eps=1e-6).to(device=device) + layer.weight.data.normal_(mean=0.0, std=0.1) + + # Gemma uses fp32 weight parameter while activations can be bf16. + assert layer.weight.dtype == torch.float32 + out = layer(x) + + x_fp32 = x.float() + weight_fp32 = layer.weight.data.float() + 1.0 + variance = x_fp32.pow(2).mean(dim=-1, keepdim=True) + ref = (x_fp32 * torch.rsqrt(variance + layer.variance_epsilon) * weight_fp32).to( + x.dtype + ) + + assert out.dtype == x.dtype + torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2) diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 86cdd7c5e89..09b9a557fe4 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -12,6 +12,9 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized from torch._inductor.pattern_matcher import PatternMatcherPass import vllm.ir.ops +from vllm.compilation.passes.fusion.rms_quant_fusion import ( + _rms_input_weight_dtype_match, +) from vllm.config import VllmConfig from vllm.config.utils import Range from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce @@ -320,7 +323,12 @@ class AllReduceRMSNormPattern(BasePattern): return allreduce[3], allreduce[1] pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_rms_input_weight_dtype_match, ) @@ -459,7 +467,12 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern): return allreduce[4], allreduce[1] pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_rms_input_weight_dtype_match, ) @@ -621,7 +634,12 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern): return allreduce[4], allreduce[1], allreduce[5] pm.register_replacement( - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass + pattern, + replacement, + self.get_inputs(), + pm.fwd_only, + pm_pass, + extra_check=_rms_input_weight_dtype_match, ) diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index 0e5121c7890..850e434a3e7 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -38,6 +38,22 @@ FP8_DTYPE = current_platform.fp8_dtype() FP4_DTYPE = torch.uint8 +_RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default + + +# TODO: extend rmsnorm quant kernels to support mixed input/weight dtypes, +# and remove this check. +def _rms_input_weight_dtype_match(match: pm.Match) -> bool: + """Prevent fusion when rms_norm input and weight dtypes differ.""" + for node in match.nodes: + if node.target == _RMS_NORM_OP: + # rms_norm(x, weight, epsilon, variance_size) + x, weight = node.args[0], node.args[1] + if isinstance(x, fx.Node) and isinstance(weight, fx.Node): + return x.meta["val"].dtype == weight.meta["val"].dtype + return True + + def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor: return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda") @@ -186,7 +202,14 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern): ] pattern(*inputs) - pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass) + pm.register_replacement( + pattern, + replacement, + inputs, + pm.fwd_only, + pm_pass, + extra_check=_rms_input_weight_dtype_match, + ) class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern): @@ -249,6 +272,7 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern): inputs, pm.fwd_only, pm_pass, + extra_check=_rms_input_weight_dtype_match, ) @@ -350,6 +374,7 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern): self.rmsnorm_matcher.inputs() + [scale], pm.fwd_only, pm_pass, + extra_check=_rms_input_weight_dtype_match, ) @@ -445,6 +470,7 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern): ], pm.fwd_only, pm_pass, + extra_check=_rms_input_weight_dtype_match, ) @@ -503,6 +529,7 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern): ], pm.fwd_only, pm_pass, + extra_check=_rms_input_weight_dtype_match, ) @@ -559,6 +586,7 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern): self.rmsnorm_matcher.inputs(), pm.fwd_only, pm_pass, + extra_check=_rms_input_weight_dtype_match, ) diff --git a/vllm/ir/ops/layernorm.py b/vllm/ir/ops/layernorm.py index 8471aa043c8..ac0c38a9e4d 100644 --- a/vllm/ir/ops/layernorm.py +++ b/vllm/ir/ops/layernorm.py @@ -16,7 +16,6 @@ def rms_norm( x_var = x if variance_size is None else x[..., :variance_size] variance = x_var.pow(2).mean(dim=-1, keepdim=True) x = x * torch.rsqrt(variance + epsilon) - x = x.to(orig_dtype) if weight is not None: - x = x * weight - return x + x = x.to(weight.dtype) * weight + return x.to(orig_dtype) diff --git a/vllm/kernels/aiter_ops.py b/vllm/kernels/aiter_ops.py index 1980051dd92..14c2b87fbbd 100644 --- a/vllm/kernels/aiter_ops.py +++ b/vllm/kernels/aiter_ops.py @@ -36,13 +36,11 @@ AITER_SUPPORTED = is_aiter_found() rms_no_var_16bit_only = ( lambda x, weight, epsilon, variance_size=None: variance_size is None - and x.dtype - in ( - torch.float16, - torch.bfloat16, - ) + and x.dtype in (torch.float16, torch.bfloat16) + and (weight is None or weight.dtype == x.dtype) ) -"""AITER rms_norm only supports float16 and bfloat16 acts and no var_size override.""" +"""AITER rms_norm only supports float16 and bfloat16 acts, no var_size override, +and requires weight dtype to match x dtype.""" @ir.ops.rms_norm.register_impl( diff --git a/vllm/kernels/vllm_c.py b/vllm/kernels/vllm_c.py index fabb36d7b43..124b02e4e27 100644 --- a/vllm/kernels/vllm_c.py +++ b/vllm/kernels/vllm_c.py @@ -11,8 +11,11 @@ current_platform.import_kernels() CUDA_ALIKE = current_platform.is_cuda_alike() """Most kernels in this file are supported on all CUDA-alike platforms.""" -rms_no_var_size = lambda x, weight, epsilon, variance_size=None: variance_size is None -"""vLLM kernel does not support variance_size parameter.""" +rms_no_var_size = ( + lambda x, weight, epsilon, variance_size=None: variance_size is None + and (weight is None or weight.dtype == x.dtype) +) +"""vLLM kernel requires no variance_size override and matching input/weight dtype.""" @ir.ops.rms_norm.register_impl( diff --git a/vllm/kernels/xpu_ops.py b/vllm/kernels/xpu_ops.py index 3548fb8682f..c680c542c1d 100644 --- a/vllm/kernels/xpu_ops.py +++ b/vllm/kernels/xpu_ops.py @@ -18,7 +18,9 @@ def is_xpu_kernels_found() -> bool: XPU_KERNELS_SUPPORTED = is_xpu_kernels_found() """Kernels in this file are supported if vLLM XPU kernels are installed.""" -rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None +rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None and ( + weight is None or weight.dtype == x.dtype +) @ir.ops.rms_norm.register_impl( diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 7b222f9c431..9afc4c9c08d 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -376,77 +376,32 @@ class GemmaRMSNorm(CustomOp): self.weight = nn.Parameter(torch.zeros(hidden_size)) self.variance_epsilon = eps - @staticmethod - def _forward_static_no_residual( - weight: torch.Tensor, - variance_epsilon: float, - x: torch.Tensor, - ) -> torch.Tensor: - """PyTorch-native implementation equivalent to forward() without residual.""" - orig_dtype = x.dtype - x = x.float() - variance = x.pow(2).mean(dim=-1, keepdim=True) - x = x * torch.rsqrt(variance + variance_epsilon) - x = x * (1.0 + weight.float()) - x = x.to(orig_dtype) - return x - - @staticmethod - def _forward_static_with_residual( - weight: torch.Tensor, - variance_epsilon: float, - x: torch.Tensor, - residual: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """PyTorch-native implementation equivalent to forward() with residual.""" - orig_dtype = x.dtype - x = ( - x.float() + residual.float() - if orig_dtype == torch.float16 - else x + residual - ) - residual = x - - x = x.float() - variance = x.pow(2).mean(dim=-1, keepdim=True) - x = x * torch.rsqrt(variance + variance_epsilon) - # Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16) - # See https://github.com/huggingface/transformers/pull/29402 - x = x * (1.0 + weight.float()) - x = x.to(orig_dtype) - return x, residual - def forward_native( self, x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """PyTorch-native implementation equivalent to forward().""" - if residual is None: - return self._forward_static_no_residual( - self.weight.data, self.variance_epsilon, x - ) - else: - return self._forward_static_with_residual( - self.weight.data, self.variance_epsilon, x, residual + orig_dtype = x.dtype + weight = self.weight.data.float() + 1.0 + if residual is not None: + x = ( + x.float() + residual.float() + if orig_dtype == torch.float16 + else x + residual ) + residual = x + # ir.ops.rms_norm handles fp32 upcast internally + out = ir.ops.rms_norm(x, weight, self.variance_epsilon) + return ( + out.to(orig_dtype) if residual is None else (out.to(orig_dtype), residual) + ) def forward_cuda( self, x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if torch.compiler.is_compiling(): - return self.forward_native(x, residual) - - if not getattr(self, "_is_compiled", False): - self._forward_static_no_residual = torch.compile( # type: ignore - self._forward_static_no_residual - ) - self._forward_static_with_residual = torch.compile( # type: ignore - self._forward_static_with_residual - ) - self._is_compiled = True return self.forward_native(x, residual) From dd9342e6bc92a52a4674a3e472318c241cb18fe1 Mon Sep 17 00:00:00 2001 From: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Date: Tue, 7 Apr 2026 04:29:23 -0500 Subject: [PATCH 33/39] only patch runtime_env for torch >= 2.10 (#38763) Signed-off-by: Rohan138 --- vllm/env_override.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/env_override.py b/vllm/env_override.py index 432300f4cf5..aa09c4a9d29 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -500,7 +500,7 @@ if is_torch_equal("2.9.0"): # This mirrors the fix in https://github.com/pytorch/pytorch/pull/177558 # and can be removed once torch >=2.12 is the minimum supported version. -if not is_torch_equal_or_newer("2.12.0"): +if is_torch_equal_or_newer("2.10.0") and not is_torch_equal_or_newer("2.12.0"): import builtins as _builtins import pickle From 7b9de7c892b3b5d50942d245400b740767a40bd8 Mon Sep 17 00:00:00 2001 From: Kyle Mylonakis <122286752+KyleMylonakisProtopia@users.noreply.github.com> Date: Tue, 7 Apr 2026 11:24:39 +0100 Subject: [PATCH 34/39] [Bugfix] Correct mistake in chained comparison in static assert logic (#38699) Signed-off-by: Kyle Mylonakis --- csrc/cpu/cpu_attn_vec.hpp | 2 +- csrc/cpu/cpu_attn_vec16.hpp | 2 +- csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/csrc/cpu/cpu_attn_vec.hpp b/csrc/cpu/cpu_attn_vec.hpp index 479313f0e19..f51a232ba95 100644 --- a/csrc/cpu/cpu_attn_vec.hpp +++ b/csrc/cpu/cpu_attn_vec.hpp @@ -53,7 +53,7 @@ class TileGemm82 { const int64_t ldb, const int64_t ldc, const int32_t block_size, const int32_t dynamic_k_size, const bool accum_c) { - static_assert(0 < M <= 8); + static_assert(0 < M && M <= 8); using load_vec_t = typename VecTypeTrait::vec_t; kv_cache_t* __restrict__ curr_b_0 = b_tile; diff --git a/csrc/cpu/cpu_attn_vec16.hpp b/csrc/cpu/cpu_attn_vec16.hpp index 7402312c092..06e4ad7624e 100644 --- a/csrc/cpu/cpu_attn_vec16.hpp +++ b/csrc/cpu/cpu_attn_vec16.hpp @@ -68,7 +68,7 @@ class TileGemm161 { const int64_t ldb, const int64_t ldc, const int32_t block_size, const int32_t dynamic_k_size, const bool accum_c) { - static_assert(0 < M <= 16); + static_assert(0 < M && M <= 16); using load_vec_t = typename VecTypeTrait::vec_t; kv_cache_t* __restrict__ curr_b_0 = b_tile; diff --git a/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp b/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp index bdd3e85a1c5..1c605a2851d 100644 --- a/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp +++ b/csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp @@ -39,7 +39,7 @@ class TileGemm82 { template static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) { - static_assert(0 < M <= 8); + static_assert(0 < M && M <= 8); using load_vec_t = typename cpu_utils::VecTypeTrait::vec_t; scalar_t* __restrict__ curr_b_0 = b_ptr; From 0be9516ea43df0fcb24bf50021e22768a49d61cf Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Tue, 7 Apr 2026 08:04:08 -0400 Subject: [PATCH 35/39] [Bug] Fix Trtllm Fp8 MoE Weight Shuffle Memory Fragamentation (#39054) Signed-off-by: wzhao18 --- .../quantization/utils/flashinfer_utils.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 13c82893dde..0e39dc881f2 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -322,20 +322,23 @@ def _shuffle_deepseek_fp8_moe_weights( block_k = 128 num_experts = w13.shape[0] - w13_shuffled: list[torch.Tensor] = [] - w2_shuffled: list[torch.Tensor] = [] + M13, K13 = w13.shape[1], w13.shape[2] + M2, K2 = w2.shape[1], w2.shape[2] + w13_out = torch.empty( + num_experts, K13 // block_k, M13, block_k, dtype=torch.uint8, device=w13.device + ) + w2_out = torch.empty( + num_experts, K2 // block_k, M2, block_k, dtype=torch.uint8, device=w2.device + ) + for i in range(num_experts): t13 = shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m) - t13 = convert_to_block_layout(t13, block_k) - w13_shuffled.append(t13) + w13_out[i] = convert_to_block_layout(t13, block_k) t2 = shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m) - t2 = convert_to_block_layout(t2, block_k) - w2_shuffled.append(t2) + w2_out[i] = convert_to_block_layout(t2, block_k) - w13_out = torch.stack(w13_shuffled).view(torch.float8_e4m3fn) - w2_out = torch.stack(w2_shuffled).view(torch.float8_e4m3fn) - return w13_out, w2_out + return w13_out.view(torch.float8_e4m3fn), w2_out.view(torch.float8_e4m3fn) def _shuffle_mxfp8_moe_weights( From 7c139ab23f6d2e9b4603b40814956100a1ccf569 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Tue, 7 Apr 2026 15:14:45 +0300 Subject: [PATCH 36/39] [KV Offload] Clean up ARC/LRU refactoring leftovers: group ARC tests and fix stale comment (#38217) Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/test_cpu_manager.py | 510 +++++++++++------------- vllm/v1/kv_offload/reuse_manager.py | 5 +- 2 files changed, 239 insertions(+), 276 deletions(-) diff --git a/tests/v1/kv_offload/test_cpu_manager.py b/tests/v1/kv_offload/test_cpu_manager.py index eea0367bf50..a9a8e21d617 100644 --- a/tests/v1/kv_offload/test_cpu_manager.py +++ b/tests/v1/kv_offload/test_cpu_manager.py @@ -15,6 +15,7 @@ from vllm.v1.kv_offload.abstract import ( from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy from vllm.v1.kv_offload.mediums import CPULoadStoreSpec +from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager @dataclass @@ -243,335 +244,300 @@ def test_cpu_manager(): ) -def test_arc_manager_basic(): - """ - Tests CPUOffloadingManager with arc policy. - Verifies that ARC handles store, load, and lookup operations correctly. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True - ) - arc_policy = arc_manager._policy - assert isinstance(arc_policy, ARCCachePolicy) +class TestARCPolicy: + """Unit tests for CPUOffloadingManager with ARC eviction policy.""" - # prepare store [1, 2] - prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2])) - verify_store_output( - prepare_store_output, - ExpectedPrepareStoreOutput( - block_hashes_to_store=[1, 2], - store_block_ids=[0, 1], - block_hashes_evicted=[], - ), - ) + def _make_manager( + self, num_blocks: int = 4, enable_events: bool = True + ) -> tuple[CPUOffloadingManager, ARCCachePolicy]: + manager = CPUOffloadingManager( + block_size=256, + num_blocks=num_blocks, + cache_policy="arc", + enable_events=enable_events, + ) + policy = manager._policy + assert isinstance(policy, ARCCachePolicy) + return manager, policy - # lookup [1, 2] -> not ready - assert arc_manager.lookup(to_hashes([1, 2])) == 0 + def test_basic(self): + """ + Tests CPUOffloadingManager with arc policy. + Verifies that ARC handles store, load, and lookup operations correctly. + """ + cpu_manager, arc_policy = self._make_manager() - # no events so far - assert list(arc_manager.take_events()) == [] + # prepare store [1, 2] + prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2])) + verify_store_output( + prepare_store_output, + ExpectedPrepareStoreOutput( + block_hashes_to_store=[1, 2], + store_block_ids=[0, 1], + block_hashes_evicted=[], + ), + ) - # complete store [1, 2] - arc_manager.complete_store(to_hashes([1, 2])) - verify_events( - arc_manager.take_events(), block_size=block_size, expected_stores=({1, 2},) - ) + # lookup [1, 2] -> not ready + assert cpu_manager.lookup(to_hashes([1, 2])) == 0 - # lookup [1, 2] - assert arc_manager.lookup(to_hashes([1])) == 1 - assert arc_manager.lookup(to_hashes([1, 2])) == 2 - assert arc_manager.lookup(to_hashes([1, 2, 3])) == 2 + # no events so far + assert list(cpu_manager.take_events()) == [] - # blocks should be in T1 (recent) - assert len(arc_policy.t1) == 2 - assert len(arc_policy.t2) == 0 + # complete store [1, 2] + cpu_manager.complete_store(to_hashes([1, 2])) + verify_events( + cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},) + ) + # lookup [1, 2] + assert cpu_manager.lookup(to_hashes([1])) == 1 + assert cpu_manager.lookup(to_hashes([1, 2])) == 2 + assert cpu_manager.lookup(to_hashes([1, 2, 3])) == 2 -def test_arc_manager_t1_to_t2_promotion(): - """ - Tests that accessing a block in T1 promotes it to T2 (frequent). - This is a key feature of ARC's adaptive behavior. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=False - ) - arc_policy = arc_manager._policy - assert isinstance(arc_policy, ARCCachePolicy) + # blocks should be in T1 (recent) + assert len(arc_policy.t1) == 2 + assert len(arc_policy.t2) == 0 - # store and complete block 1 - arc_manager.prepare_store(to_hashes([1])) - arc_manager.complete_store(to_hashes([1])) + def test_t1_to_t2_promotion(self): + """ + Tests that accessing a block in T1 promotes it to T2 (frequent). + This is a key feature of ARC's adaptive behavior. + """ + cpu_manager, arc_policy = self._make_manager(enable_events=False) - # block 1 starts in T1 (recent) - assert to_hashes([1])[0] in arc_policy.t1 - assert to_hashes([1])[0] not in arc_policy.t2 + # store and complete block 1 + cpu_manager.prepare_store(to_hashes([1])) + cpu_manager.complete_store(to_hashes([1])) - # touch block 1 (simulate second access) - arc_manager.touch(to_hashes([1])) + # block 1 starts in T1 (recent) + assert to_hashes([1])[0] in arc_policy.t1 + assert to_hashes([1])[0] not in arc_policy.t2 - # block 1 should now be in T2 (frequent) - assert to_hashes([1])[0] not in arc_policy.t1 - assert to_hashes([1])[0] in arc_policy.t2 + # touch block 1 (simulate second access) + cpu_manager.touch(to_hashes([1])) + # block 1 should now be in T2 (frequent) + assert to_hashes([1])[0] not in arc_policy.t1 + assert to_hashes([1])[0] in arc_policy.t2 -def test_arc_manager_eviction_with_load(): - """ - Tests ARC eviction behavior similar to LRU test. - Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True - ) + def test_eviction_with_load(self): + """ + Tests ARC eviction behavior similar to LRU test. + Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted. + """ + cpu_manager, _ = self._make_manager() - # prepare and complete store [1, 2, 3, 4] - prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2, 3, 4])) - verify_store_output( - prepare_store_output, - ExpectedPrepareStoreOutput( - block_hashes_to_store=[1, 2, 3, 4], - store_block_ids=[0, 1, 2, 3], - block_hashes_evicted=[], - ), - ) - arc_manager.complete_store(to_hashes([1, 2, 3, 4])) + # prepare and complete store [1, 2, 3, 4] + prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2, 3, 4])) + verify_store_output( + prepare_store_output, + ExpectedPrepareStoreOutput( + block_hashes_to_store=[1, 2, 3, 4], + store_block_ids=[0, 1, 2, 3], + block_hashes_evicted=[], + ), + ) + cpu_manager.complete_store(to_hashes([1, 2, 3, 4])) - # prepare load [2, 3] (increases ref_cnt) - prepare_load_output = arc_manager.prepare_load(to_hashes([2, 3])) - verify_load_output(prepare_load_output, [1, 2]) + # prepare load [2, 3] (increases ref_cnt) + prepare_load_output = cpu_manager.prepare_load(to_hashes([2, 3])) + verify_load_output(prepare_load_output, [1, 2]) - # prepare store [5, 6, 7] with [2, 3] being loaded - # should fail because [2, 3] have ref_cnt > 0 - assert arc_manager.prepare_store(to_hashes([5, 6, 7])) is None + # prepare store [5, 6, 7] with [2, 3] being loaded + # should fail because [2, 3] have ref_cnt > 0 + assert cpu_manager.prepare_store(to_hashes([5, 6, 7])) is None - # complete load [2, 3] - arc_manager.complete_load(to_hashes([2, 3])) + # complete load [2, 3] + cpu_manager.complete_load(to_hashes([2, 3])) - # now prepare store [5, 6, 7] should succeed - # ARC will evict blocks one at a time from T1 as needed - prepare_store_output = arc_manager.prepare_store(to_hashes([5, 6, 7])) - assert prepare_store_output is not None - # Should successfully evict enough blocks to make room (at least 1) - assert len(prepare_store_output.block_hashes_evicted) >= 1 + # now prepare store [5, 6, 7] should succeed + # ARC will evict blocks one at a time from T1 as needed + prepare_store_output = cpu_manager.prepare_store(to_hashes([5, 6, 7])) + assert prepare_store_output is not None + # Should successfully evict enough blocks to make room (at least 1) + assert len(prepare_store_output.block_hashes_evicted) >= 1 + def test_adaptive_target(self): + """ + Tests ARC's adaptive target adjustment via ghost lists. + When a block in B1 (ghost list) is accessed, target_t1_size increases. + When a block in B2 is accessed, target_t1_size decreases. + """ + cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False) -def test_arc_manager_adaptive_target(): - """ - Tests ARC's adaptive target adjustment via ghost lists. - When a block in B1 (ghost list) is accessed, target_t1_size increases. - When a block in B2 is accessed, target_t1_size decreases. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=2, cache_policy="arc", enable_events=False - ) - arc_policy = arc_manager._policy - assert isinstance(arc_policy, ARCCachePolicy) + # store blocks 1, 2 (fills cache) + cpu_manager.prepare_store(to_hashes([1, 2])) + cpu_manager.complete_store(to_hashes([1, 2])) - # store blocks 1, 2 (fills cache) - arc_manager.prepare_store(to_hashes([1, 2])) - arc_manager.complete_store(to_hashes([1, 2])) + initial_target = arc_policy.target_t1_size - initial_target = arc_policy.target_t1_size + # store block 3, evicting block 1 (moves to B1 ghost list) + cpu_manager.prepare_store(to_hashes([3])) + cpu_manager.complete_store(to_hashes([3])) - # store block 3, evicting block 1 (moves to B1 ghost list) - arc_manager.prepare_store(to_hashes([3])) - arc_manager.complete_store(to_hashes([3])) + # block 1 should be in B1 (ghost list) + assert to_hashes([1])[0] in arc_policy.b1 - # block 1 should be in B1 (ghost list) - assert to_hashes([1])[0] in arc_policy.b1 + # touch block 1 (cache miss, but in B1) + # this should increase target_t1_size (favor recency) + cpu_manager.touch(to_hashes([1])) - # touch block 1 (cache miss, but in B1) - # this should increase target_t1_size (favor recency) - arc_manager.touch(to_hashes([1])) + # target should have increased + assert arc_policy.target_t1_size > initial_target - # target should have increased - assert arc_policy.target_t1_size > initial_target + def test_t1_t2_eviction_policy(self): + """ + Tests that ARC evicts from T1 or T2 based on target_t1_size. + If |T1| >= target_t1_size, evict from T1, otherwise from T2. + """ + cpu_manager, arc_policy = self._make_manager(enable_events=False) + # store blocks 1, 2, 3, 4 + cpu_manager.prepare_store(to_hashes([1, 2, 3, 4])) + cpu_manager.complete_store(to_hashes([1, 2, 3, 4])) -def test_arc_manager_t1_t2_eviction_policy(): - """ - Tests that ARC evicts from T1 or T2 based on target_t1_size. - If |T1| >= target_t1_size, evict from T1, otherwise from T2. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=False - ) - arc_policy = arc_manager._policy - assert isinstance(arc_policy, ARCCachePolicy) + # promote blocks 3, 4 to T2 by touching them + cpu_manager.touch(to_hashes([3, 4])) - # store blocks 1, 2, 3, 4 - arc_manager.prepare_store(to_hashes([1, 2, 3, 4])) - arc_manager.complete_store(to_hashes([1, 2, 3, 4])) + # now: T1 = {1, 2}, T2 = {3, 4} + assert len(arc_policy.t1) == 2 + assert len(arc_policy.t2) == 2 - # promote blocks 3, 4 to T2 by touching them - arc_manager.touch(to_hashes([3, 4])) + # set target_t1_size to prefer evicting from T1 + # (when |T1| >= target, evict from T1) + arc_policy.target_t1_size = 1 - # now: T1 = {1, 2}, T2 = {3, 4} - assert len(arc_policy.t1) == 2 - assert len(arc_policy.t2) == 2 + # store block 5, should evict from T1 (block 1, LRU in T1) + output = cpu_manager.prepare_store(to_hashes([5])) + assert output is not None + assert to_hashes([1]) == output.block_hashes_evicted - # set target_t1_size to prefer evicting from T1 - # (when |T1| >= target, evict from T1) - arc_policy.target_t1_size = 1 + cpu_manager.complete_store(to_hashes([5])) - # store block 5, should evict from T1 (block 1, LRU in T1) - output = arc_manager.prepare_store(to_hashes([5])) - assert output is not None - assert to_hashes([1]) == output.block_hashes_evicted + # block 1 should be in B1 (ghost list) + assert to_hashes([1])[0] in arc_policy.b1 + # block 5 should be in T1 + assert to_hashes([5])[0] in arc_policy.t1 - arc_manager.complete_store(to_hashes([5])) + def test_ghost_list_bounds(self): + """ + Tests that ghost lists (B1, B2) don't grow unbounded. + They should be capped at cache_capacity. + """ + cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False) - # block 1 should be in B1 (ghost list) - assert to_hashes([1])[0] in arc_policy.b1 - # block 5 should be in T1 - assert to_hashes([5])[0] in arc_policy.t1 + # fill cache with blocks 1, 2 + cpu_manager.prepare_store(to_hashes([1, 2])) + cpu_manager.complete_store(to_hashes([1, 2])) + # store many blocks to fill ghost lists + for i in range(3, 20): + cpu_manager.prepare_store(to_hashes([i])) + cpu_manager.complete_store(to_hashes([i])) -def test_arc_manager_ghost_list_bounds(): - """ - Tests that ghost lists (B1, B2) don't grow unbounded. - They should be capped at cache_capacity. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=2, cache_policy="arc", enable_events=False - ) - arc_policy = arc_manager._policy - assert isinstance(arc_policy, ARCCachePolicy) + # ghost lists should not exceed cache_capacity + assert len(arc_policy.b1) <= arc_policy.cache_capacity + assert len(arc_policy.b2) <= arc_policy.cache_capacity - # fill cache with blocks 1, 2 - arc_manager.prepare_store(to_hashes([1, 2])) - arc_manager.complete_store(to_hashes([1, 2])) + def test_touch_ordering(self): + """ + Tests that touch() correctly updates access patterns. + Similar to LRU test but verifies T1/T2 ordering. + """ + cpu_manager, arc_policy = self._make_manager() - # store many blocks to fill ghost lists - for i in range(3, 20): - arc_manager.prepare_store(to_hashes([i])) - arc_manager.complete_store(to_hashes([i])) + # store blocks 1, 2, 3, 4 + cpu_manager.prepare_store(to_hashes([1, 2, 3, 4])) + cpu_manager.complete_store(to_hashes([1, 2, 3, 4])) - # ghost lists should not exceed cache_capacity - assert len(arc_policy.b1) <= arc_policy.cache_capacity - assert len(arc_policy.b2) <= arc_policy.cache_capacity + # promote 3, 4 to T2 + cpu_manager.touch(to_hashes([3, 4])) + # T1 = {1, 2}, T2 = {3, 4} + # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2 + cpu_manager.touch(to_hashes([1, 3, 4])) -def test_arc_manager_touch_ordering(): - """ - Tests that touch() correctly updates access patterns. - Similar to LRU test but verifies T1/T2 ordering. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True - ) - arc_policy = arc_manager._policy - assert isinstance(arc_policy, ARCCachePolicy) + # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent) + assert len(arc_policy.t1) == 1 + assert len(arc_policy.t2) == 3 - # store blocks 1, 2, 3, 4 - arc_manager.prepare_store(to_hashes([1, 2, 3, 4])) - arc_manager.complete_store(to_hashes([1, 2, 3, 4])) + # store block 5, should evict from T1 (block 2, only one in T1) + prepare_store_output = cpu_manager.prepare_store(to_hashes([5])) + verify_store_output( + prepare_store_output, + ExpectedPrepareStoreOutput( + block_hashes_to_store=[5], + store_block_ids=[1], # reuses block 2's storage + block_hashes_evicted=[2], + ), + ) - # promote 3, 4 to T2 - arc_manager.touch(to_hashes([3, 4])) + def test_failed_store(self): + """ + Tests that failed store operations clean up correctly. + Similar to LRU test but for ARC. + """ + cpu_manager, arc_policy = self._make_manager() - # T1 = {1, 2}, T2 = {3, 4} - # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2 - arc_manager.touch(to_hashes([1, 3, 4])) + # store blocks 1, 2, 3, 4 + cpu_manager.prepare_store(to_hashes([1, 2, 3, 4])) + cpu_manager.complete_store(to_hashes([1, 2, 3, 4])) - # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent) - assert len(arc_policy.t1) == 1 - assert len(arc_policy.t2) == 3 + # prepare store block 5 (will evict block 1) + prepare_store_output = cpu_manager.prepare_store(to_hashes([5])) + assert prepare_store_output is not None + assert len(prepare_store_output.block_hashes_evicted) == 1 - # store block 5, should evict from T1 (block 2, only one in T1) - prepare_store_output = arc_manager.prepare_store(to_hashes([5])) - verify_store_output( - prepare_store_output, - ExpectedPrepareStoreOutput( - block_hashes_to_store=[5], - store_block_ids=[1], # reuses block 2's storage - block_hashes_evicted=[2], - ), - ) + # complete store with failure + cpu_manager.complete_store(to_hashes([5]), success=False) + # block 5 should not be in cache + assert cpu_manager.lookup(to_hashes([5])) == 0 + # block 5 should not be in T1 or T2 + assert to_hashes([5])[0] not in arc_policy.t1 + assert to_hashes([5])[0] not in arc_policy.t2 -def test_arc_manager_failed_store(): - """ - Tests that failed store operations clean up correctly. - Similar to LRU test but for ARC. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True - ) - arc_policy = arc_manager._policy - assert isinstance(arc_policy, ARCCachePolicy) + # evicted block should still be gone (in B1 ghost list) + evicted_hash = prepare_store_output.block_hashes_evicted[0] + assert evicted_hash in arc_policy.b1 - # store blocks 1, 2, 3, 4 - arc_manager.prepare_store(to_hashes([1, 2, 3, 4])) - arc_manager.complete_store(to_hashes([1, 2, 3, 4])) + def test_full_scenario(self): + """ + Comprehensive test covering multiple ARC operations in sequence. + Similar to the full LRU test but adapted for ARC behavior. + """ + cpu_manager, arc_policy = self._make_manager() - # prepare store block 5 (will evict block 1) - prepare_store_output = arc_manager.prepare_store(to_hashes([5])) - assert prepare_store_output is not None - assert len(prepare_store_output.block_hashes_evicted) == 1 + # store [1, 2] + cpu_manager.prepare_store(to_hashes([1, 2])) + cpu_manager.complete_store(to_hashes([1, 2])) - # complete store with failure - arc_manager.complete_store(to_hashes([5]), success=False) + # store [3, 4, 5] -> evicts [1] + prepare_store_output = cpu_manager.prepare_store(to_hashes([3, 4, 5])) + assert prepare_store_output is not None + assert len(prepare_store_output.block_hashes_evicted) == 1 + cpu_manager.complete_store(to_hashes([3, 4, 5])) - # block 5 should not be in cache - assert arc_manager.lookup(to_hashes([5])) == 0 - # block 5 should not be in T1 or T2 - assert to_hashes([5])[0] not in arc_policy.t1 - assert to_hashes([5])[0] not in arc_policy.t2 + # promote some blocks to T2 + cpu_manager.touch(to_hashes([2, 3])) - # evicted block should still be gone (in B1 ghost list) - evicted_hash = prepare_store_output.block_hashes_evicted[0] - assert evicted_hash in arc_policy.b1 + # T1 has {4, 5}, T2 has {2, 3} + assert len(arc_policy.t1) == 2 + assert len(arc_policy.t2) == 2 + # store [6] -> should evict from T1 (4 is oldest in T1) + prepare_store_output = cpu_manager.prepare_store(to_hashes([6])) + assert prepare_store_output is not None + cpu_manager.complete_store(to_hashes([6])) -def test_arc_manager_full_scenario(): - """ - Comprehensive test covering multiple ARC operations in sequence. - Similar to the full LRU test but adapted for ARC behavior. - """ - block_size = 256 - arc_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True - ) - arc_policy = arc_manager._policy - assert isinstance(arc_policy, ARCCachePolicy) + # verify blocks 2, 3 (in T2) are still present + assert cpu_manager.lookup(to_hashes([2])) == 1 + assert cpu_manager.lookup(to_hashes([3])) == 1 - # store [1, 2] - arc_manager.prepare_store(to_hashes([1, 2])) - arc_manager.complete_store(to_hashes([1, 2])) - - # store [3, 4, 5] -> evicts [1] - prepare_store_output = arc_manager.prepare_store(to_hashes([3, 4, 5])) - assert prepare_store_output is not None - assert len(prepare_store_output.block_hashes_evicted) == 1 - arc_manager.complete_store(to_hashes([3, 4, 5])) - - # promote some blocks to T2 - arc_manager.touch(to_hashes([2, 3])) - - # T1 has {4, 5}, T2 has {2, 3} - assert len(arc_policy.t1) == 2 - assert len(arc_policy.t2) == 2 - - # store [6] -> should evict from T1 (4 is oldest in T1) - prepare_store_output = arc_manager.prepare_store(to_hashes([6])) - assert prepare_store_output is not None - arc_manager.complete_store(to_hashes([6])) - - # verify blocks 2, 3 (in T2) are still present - assert arc_manager.lookup(to_hashes([2])) == 1 - assert arc_manager.lookup(to_hashes([3])) == 1 - - # verify events - events = list(arc_manager.take_events()) - assert len(events) > 0 # should have store and eviction events + # verify events + events = list(cpu_manager.take_events()) + assert len(events) > 0 # should have store and eviction events def test_filter_reused_manager(): @@ -583,8 +549,6 @@ def test_filter_reused_manager(): block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True ) - from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager - manager = FilterReusedOffloadingManager( backing=lru_manager, store_threshold=2, max_tracker_size=3 ) diff --git a/vllm/v1/kv_offload/reuse_manager.py b/vllm/v1/kv_offload/reuse_manager.py index daf6c65cd2d..3c372c5d9eb 100644 --- a/vllm/v1/kv_offload/reuse_manager.py +++ b/vllm/v1/kv_offload/reuse_manager.py @@ -93,9 +93,8 @@ class FilterReusedOffloadingManager(OffloadingManager): ] # Delegate to the backing manager with only the eligible hashes. - # Passing an empty list is intentional and safe — both - # LRUOffloadingManager and ARCOffloadingManager handle it correctly, - # returning a PrepareStoreOutput with empty lists. + # Passing an empty list is intentional and safe — CPUOffloadingManager + # handles it correctly, returning a PrepareStoreOutput with empty lists. return self._backing.prepare_store(eligible) # ------------------------------------------------------------------ From 79df4a794d6f2c933cf824b95afa0e9ae42723a9 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 7 Apr 2026 15:21:18 +0200 Subject: [PATCH 37/39] Automatically add links to API docs for matching strings in docs (#37434) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/mkdocs/hooks/autoref_code.py | 167 ++++++++++++++++++++++++++++++ mkdocs.yaml | 1 + 2 files changed, 168 insertions(+) create mode 100644 docs/mkdocs/hooks/autoref_code.py diff --git a/docs/mkdocs/hooks/autoref_code.py b/docs/mkdocs/hooks/autoref_code.py new file mode 100644 index 00000000000..647f74f202d --- /dev/null +++ b/docs/mkdocs/hooks/autoref_code.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +MkDocs hook to automatically convert inline code references to API doc links. + +For example, `WeightTransferConfig` becomes +[`WeightTransferConfig`][vllm.config.WeightTransferConfig] + +This works with the `autorefs` plugin to create clickable cross-references +to API documentation pages generated by `mkdocstrings`. + +The hook builds an index of all documented public Python names (classes and +functions with docstrings) from the vllm package at startup using AST parsing, +then substitutes matching inline code spans on each page. Names without +docstrings are excluded because mkdocstrings will not generate a page for them. +""" + +import ast +import logging +from pathlib import Path + +import regex as re +from mkdocs.config.defaults import MkDocsConfig +from mkdocs.structure.files import Files +from mkdocs.structure.pages import Page + +logger = logging.getLogger("mkdocs") + +ROOT_DIR = Path(__file__).parent.parent.parent.parent.resolve() +VLLM_DIR = ROOT_DIR / "vllm" + +# Maps short name -> qualified name (e.g. "ModelConfig" -> "vllm.config.ModelConfig") +_name_index: dict[str, str] = {} + +# Fenced code block pattern (``` or ~~~, with optional language specifier). +_FENCED_BLOCK = re.compile( + r"(?:^|\n)(?P`{3,}|~{3,})[^\n]*\n.*?(?:\n(?P=fence))", re.DOTALL +) + +# Inline code that is NOT already part of a markdown link. +# Matches `Name` but not [`Name`] and not [`Name`][...] or [`Name`](...). +_INLINE_CODE = re.compile( + r"(?[A-Za-z0-9_]*)`" # `UpperCamelCase` or `UPPER_SNAKE` + r"(?!\])" # not followed by ] +) + + +def _has_docstring(node: ast.AST) -> bool: + """Check if a class or function node has a docstring.""" + if not isinstance(node, ast.ClassDef | ast.FunctionDef | ast.AsyncFunctionDef): + return False + return ast.get_docstring(node, clean=False) is not None + + +def _module_path(filepath: Path) -> str: + """Convert a filesystem path to a dotted module path.""" + rel = filepath.relative_to(ROOT_DIR) + parts = list(rel.with_suffix("").parts) + if parts[-1] == "__init__": + parts = parts[:-1] + return ".".join(parts) + + +def _index_file(filepath: Path) -> dict[str, str]: + """Extract documented public names from a Python file using AST parsing. + + Only classes and functions with docstrings are included, since + mkdocstrings won't generate a page for undocumented symbols. + """ + names: dict[str, str] = {} + try: + source = filepath.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(filepath)) + except (SyntaxError, UnicodeDecodeError): + return names + + module = _module_path(filepath) + + for node in ast.iter_child_nodes(tree): + if ( + # Class definitions (with docstring) + isinstance(node, ast.ClassDef) + and not node.name.startswith("_") + and _has_docstring(node) + ) or ( + # Function definitions (with docstring, only uppercase/CamelCase) + isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef) + and not node.name.startswith("_") + and node.name[0].isupper() + and _has_docstring(node) + ): + names[node.name] = f"{module}.{node.name}" + + return names + + +def _build_index() -> dict[str, str]: + """Walk the vllm package and build a name -> qualified path index.""" + index: dict[str, str] = {} + # Track conflicts: if multiple modules define the same name, + # prefer shallower modules (more likely to be the public API). + depth: dict[str, int] = {} + + for filepath in sorted(VLLM_DIR.rglob("*.py")): + # Skip internal/private modules + if any(part.startswith("_") and part != "__init__" for part in filepath.parts): + continue + # Skip third-party vendored code + rel = filepath.relative_to(VLLM_DIR) + if rel.parts and rel.parts[0] in ("third_party", "vllm_flash_attn"): + continue + + module_depth = len(filepath.relative_to(ROOT_DIR).parts) + file_names = _index_file(filepath) + + for name, qualified in file_names.items(): + if name not in index or module_depth < depth[name]: + index[name] = qualified + depth[name] = module_depth + + return index + + +def on_startup(*, command: str, dirty: bool) -> None: + """Build the name index once at startup.""" + global _name_index + _name_index = _build_index() + logger.info("autoref_code: indexed %d names from vllm/", len(_name_index)) + + +def on_page_markdown( + markdown: str, *, page: Page, config: MkDocsConfig, files: Files +) -> str: + """Replace inline code references with autoref links.""" + if not _name_index: + return markdown + + # Skip API reference pages to avoid circular/redundant links. + if page.file.src_path.startswith("api/"): + return markdown + + # Step 1: Mask fenced code blocks so we don't touch code inside them. + masks: list[str] = [] + + def _mask_block(match: re.Match) -> str: + masks.append(match.group(0)) + return f"\ue000CODEBLOCK{len(masks) - 1}\ue000" + + masked = _FENCED_BLOCK.sub(_mask_block, markdown) + + # Step 2: Replace inline code references. + def _replace(match: re.Match) -> str: + name = match.group("name") + qualified = _name_index.get(name) + if qualified is None: + return match.group(0) + logger.debug("autoref_code: linking `%s` to [%s]", name, qualified) + return f"[`{name}`][{qualified}]" + + result = _INLINE_CODE.sub(_replace, masked) + + # Step 3: Restore masked code blocks. + result = re.sub( + r"\ue000CODEBLOCK(\d+)\ue000", lambda m: masks[int(m.group(1))], result + ) + return result diff --git a/mkdocs.yaml b/mkdocs.yaml index e37ae9b879a..4b06b31ebe3 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -54,6 +54,7 @@ hooks: - docs/mkdocs/hooks/generate_argparse.py - docs/mkdocs/hooks/generate_metrics.py - docs/mkdocs/hooks/url_schemes.py + - docs/mkdocs/hooks/autoref_code.py plugins: - meta From edcc37a8cee26813fe868b9fc267c3cba5818ff7 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 7 Apr 2026 15:23:33 +0200 Subject: [PATCH 38/39] Fix Mistral yarn warning in Transformers v5 (#37292) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Julien Denize <40604584+juliendenize@users.noreply.github.com> --- vllm/transformers_utils/configs/mistral.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vllm/transformers_utils/configs/mistral.py b/vllm/transformers_utils/configs/mistral.py index bdeadec1bf0..2b079669147 100644 --- a/vllm/transformers_utils/configs/mistral.py +++ b/vllm/transformers_utils/configs/mistral.py @@ -2,7 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any +from packaging.version import Version from transformers import PretrainedConfig, WhisperConfig +from transformers import __version__ as TRANSFORMERS_VERSION from vllm.logger import init_logger @@ -134,6 +136,10 @@ def _remap_mistral_yarn_args(config: dict) -> dict: # Cast to remove Transformers > v5 type warnings config["rope_parameters"][new_name] = cast(yarn_config.pop(old_name)) + # Ignore apply_yarn_scaling in Transformers > v5 RoPE validation to remove warnings + if Version(TRANSFORMERS_VERSION) >= Version("5.3.0.dev0"): + config["ignore_keys_at_rope_validation"] = {"apply_yarn_scaling"} + assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}" return config From 6e1100889e6a675d17ad82815acf8f02f1cc419e Mon Sep 17 00:00:00 2001 From: Ilya Boytsov Date: Tue, 7 Apr 2026 16:40:55 +0200 Subject: [PATCH 39/39] fix(test): recompute Jina ColBERT rotary inv_freq cleared by transformers v5 weight loader (#39176) Signed-off-by: Ilya Boytsov --- tests/models/language/pooling/test_colbert.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index e19e1f75da9..10c229fe063 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -109,6 +109,14 @@ def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device): **extra, ).to(device) model.eval() + + # Transformers 5.0 weight materialization can clear non-persistent + # buffers (e.g. rotary inv_freq) that were registered with + # persistent=False. Re-compute them so the model produces valid output. + for mod in model.modules(): + if hasattr(mod, "_compute_inv_freq") and hasattr(mod, "inv_freq"): + mod.inv_freq = mod._compute_inv_freq(device=device) + return model