fuse more

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

.buildkite/test_areas/distributed.yaml

@@ -196,6 +196,8 @@ steps:
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
     - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
     - pytest -v -s tests/v1/distributed/test_dbo.py
+    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
+    - pytest -v -s tests/distributed/test_packed_tensor.py
 
 - label: Distributed Tests (2 GPUs)(B200)
   device: b200

.github/CODEOWNERS

@@ -44,8 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/pooling_params.py @noooop @DarkLight1337
 /vllm/tokenizers @DarkLight1337 @njhill
 /vllm/renderers @DarkLight1337 @njhill
-/vllm/reasoning @aarnphm @chaunceyjiang
-/vllm/tool_parsers @aarnphm @chaunceyjiang
+/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33
+/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33
+/vllm/parser @aarnphm @chaunceyjiang @sfeng33
 
 # vLLM V1
 /vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -91,7 +92,10 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/v1/kv_connector/nixl_integration @NickLucche
 /tests/v1/kv_connector @ApostaC @orozery
 /tests/v1/kv_offload @ApostaC @orozery
-/tests/v1/determinism @yewentao256 
+/tests/v1/determinism @yewentao256
+/tests/reasoning @aarnphm @chaunceyjiang @sfeng33
+/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33
+/tests/tool_use @aarnphm @chaunceyjiang @sfeng33
 
 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor

docs/governance/committers.md

@@ -31,6 +31,7 @@ Sorted alphabetically by GitHub handle:
 - [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance
 - [@luccafong](https://github.com/luccafong): Llama models, speculative decoding, distributed
 - [@markmc](https://github.com/markmc): Observability
+- [@MatthewBonanni](https://github.com/MatthewBonanni): Kernels and performance
 - [@mgoin](https://github.com/mgoin): Quantization and performance
 - [@NickLucche](https://github.com/NickLucche): KV connector
 - [@njhill](https://github.com/njhill): Distributed, API server, engine core
@@ -41,6 +42,7 @@ Sorted alphabetically by GitHub handle:
 - [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg
 - [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support
 - [@russellb](https://github.com/russellb): Structured output, engine core, security
+- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser
 - [@sighingnow](https://github.com/sighingnow): Qwen models, new model support
 - [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community
 - [@tdoublep](https://github.com/tdoublep): State space models
@@ -86,7 +88,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints
     - @robertgshaw2-redhat, @njhill, @russellb
 - ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation
-    - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg
+    - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni
 - KV Connector: Connector interface and implementation for KV cache offload and transfer
     - @robertgshaw2-redhat, @njhill, @KuntaiDu, @NickLucche, @ApostaC
 - Distributed, Parallelism, Process Management: Process launchers managing each worker, and assign them to the right DP/TP/PP/EP ranks
@@ -105,7 +107,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - Custom Layers: Utility layers in vLLM such as rotary embedding and rms norms
     - @ProExpertProg
 - Attention: Attention interface for paged attention
-    - @WoosukKwon, @LucasWilkinson, @heheda12345
+    - @WoosukKwon, @LucasWilkinson, @heheda12345, @MatthewBonanni
 - FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB
     - @tlrmchlsmth
 - Quantization: Various quantization config, weight loading, and kernel.
@@ -119,7 +121,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - State space models: The state space models implementation in vLLM
     - @tdoublep, @tlrmchlsmth
 - Reasoning and tool calling parsers
-    - @chaunceyjiang, @aarnphm
+    - @chaunceyjiang, @aarnphm, @sfeng33
 
 ### Entrypoints
 
@@ -133,7 +135,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 ### Features
 
 - Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP.
-    - @WoosukKwon, @benchislett, @luccafong
+    - @WoosukKwon, @benchislett, @luccafong, @MatthewBonanni
 - Structured Output: The structured output implementation
     - @russellb, @aarnphm
 - RL: The RL related features such as collective rpc, sleep mode, etc.
@@ -153,8 +155,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 
 ### External Kernels Integration
 
-- FlashAttention: @LucasWilkinson
-- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon
+- FlashAttention: @LucasWilkinson, @MatthewBonanni
+- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni
 - Blackwell Kernels: @mgoin, @yewentao256
 - DeepEP/DeepGEMM: @mgoin, @yewentao256
 

examples/basic/offline_inference/basic.py

@@ -16,7 +16,7 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
 def main():
     # Create an LLM.
-    llm = LLM(model="facebook/opt-125m")
+    llm = LLM(model="nvidia/DeepSeek-V3.2-NVFP4", enforce_eager=True, tensor_parallel_size=4, kernel_config={"enable_flashinfer_autotune": False})
     # Generate texts from the prompts.
     # The output is a list of RequestOutput objects
     # that contain the prompt, generated text, and other information.

examples/offline_inference/simple_profiling.py

@@ -2,31 +2,46 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import time
+import os
+
+os.environ["VLLM_USE_SPECIALIZED_MODELS"] = "1"
+os.environ["VLLM_USE_V2_MODEL_RUNNER"] = "1"
 
 from vllm import LLM, SamplingParams
 
 # Sample prompts.
 prompts = [
-    "Hello, my name is",
-    "The president of the United States is",
-    "The capital of France is",
-    "The future of AI is",
+    [0] * 10_000,
+    [1] * 10_000,
+    [2] * 10_000,
+    [3] * 10_000,
+    [4] * 10_000,
+    [5] * 10_000,
+    [6] * 10_000,
+    [7] * 10_000,
 ]
 # Create a sampling params object.
-sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+sampling_params = SamplingParams(temperature=0.0)
 
 
 def main():
     # Create an LLM.
     llm = LLM(
-        model="facebook/opt-125m",
-        tensor_parallel_size=1,
+        model="nvidia/DeepSeek-V3.2-NVFP4",
+        tensor_parallel_size=4,
+        kernel_config={"enable_flashinfer_autotune": False},
         profiler_config={
             "profiler": "torch",
-            "torch_profiler_dir": "./vllm_profile",
+            "torch_profiler_dir": f"./vllm_profile/bsz{len(prompts)}/",
         },
+        enable_prefix_caching=False,
+        load_format="dummy",
+        compilation_config={"max_cudagraph_capture_size": 64},
+        speculative_config={"method": "mtp", "num_speculative_tokens": 3},
+        max_num_batched_tokens=32768,
     )
 
+    outputs = llm.generate(prompts, sampling_params)
     llm.start_profile()
 
     # Generate texts from the prompts. The output is a list of RequestOutput

tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+import vllm.config
+from tests.compile.backend import TestBackend
+from vllm.compilation.passes.vllm_inductor_pass import (
+    VllmFusionPatternMatcherPass,
+    VllmPatternMatcherPass,
+    VllmPatternReplacement,
+)
+from vllm.config import CompilationConfig, CompilationMode, VllmConfig
+from vllm.platforms import current_platform
+
+
+class ReluToAbsPattern(VllmPatternReplacement):
+    """Replaces relu(x) with abs(x) — a minimal test fixture."""
+
+    @property
+    def pattern(self):
+        def _pattern(x: torch.Tensor) -> torch.Tensor:
+            return torch.ops.aten.relu.default(x)
+
+        return _pattern
+
+    @property
+    def replacement(self):
+        def _replacement(x: torch.Tensor) -> torch.Tensor:
+            return torch.ops.aten.abs.default(x)
+
+        return _replacement
+
+    def get_inputs(self) -> list[torch.Tensor]:
+        return [self.empty_fp32(4)]
+
+
+class ExpToSqrtPattern(VllmPatternReplacement):
+    """A second distinct pattern type — used to test uuid differentiation."""
+
+    @property
+    def pattern(self):
+        def _pattern(x: torch.Tensor) -> torch.Tensor:
+            return torch.ops.aten.exp.default(x)
+
+        return _pattern
+
+    @property
+    def replacement(self):
+        def _replacement(x: torch.Tensor) -> torch.Tensor:
+            return torch.ops.aten.sqrt.default(x)
+
+        return _replacement
+
+    def get_inputs(self) -> list[torch.Tensor]:
+        return [self.empty_fp32(4)]
+
+
+class ReluFusionPass(VllmFusionPatternMatcherPass):
+    def __init__(self, config: VllmConfig) -> None:
+        super().__init__(config, "test_relu_fusion")
+        self.register(ReluToAbsPattern())
+
+
+class TwoPatternFusionPass(VllmFusionPatternMatcherPass):
+    def __init__(self, config: VllmConfig) -> None:
+        super().__init__(config, "test_two_pattern_fusion")
+        self.register(ReluToAbsPattern())
+        self.register(ExpToSqrtPattern())
+
+
+@pytest.fixture
+def vllm_config():
+    return VllmConfig(
+        compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE),
+    )
+
+
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
+def test_register_tracks_patterns(vllm_config):
+    """register() appends each VllmPatternReplacement to _pattern_replacements."""
+    with vllm.config.set_current_vllm_config(vllm_config):
+        single = ReluFusionPass(vllm_config)
+        two = TwoPatternFusionPass(vllm_config)
+
+    assert len(single._pattern_replacements) == 1
+    assert len(two._pattern_replacements) == 2
+
+
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
+def test_uuid_stable(vllm_config):
+    """Two instances of the same pass class produce identical uuids."""
+    with vllm.config.set_current_vllm_config(vllm_config):
+        p1 = ReluFusionPass(vllm_config)
+        p2 = ReluFusionPass(vllm_config)
+        p3 = TwoPatternFusionPass(vllm_config)
+
+    assert p1.uuid() == p2.uuid()
+    assert p1.uuid() != p3.uuid()
+    assert p2.uuid() != p3.uuid()
+
+
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
+@pytest.mark.parametrize("N", [1, 2, 4])
+def test_matched_count_and_match_table(vllm_config, N):
+    """matched_count and match_table reflect the number of matched patterns."""
+
+    class Model(torch.nn.Module):
+        def forward(self, *inputs):
+            # N independent relus
+            return sum(torch.relu(x) for x in inputs)
+
+    with vllm.config.set_current_vllm_config(vllm_config):
+        torch.set_default_device("cuda")
+        torch.set_default_dtype(torch.float32)
+
+        fusion_pass = ReluFusionPass(vllm_config)
+        backend = TestBackend(fusion_pass)
+        model = torch.compile(Model(), backend=backend)
+
+        inputs = [torch.rand(8) for _ in range(N)]
+        model(*inputs)
+
+    assert fusion_pass.matched_count == N
+    assert VllmPatternMatcherPass.match_table["test_relu_fusion"] >= N

tests/distributed/test_weight_transfer.py

@@ -41,6 +41,7 @@ def create_mock_parallel_config(
     config.rank = rank
     config.world_size = world_size
     config.data_parallel_rank = dp_rank
+    config.data_parallel_index = dp_rank
     return config
 
 
@@ -283,6 +284,7 @@ def inference_receive_tensor(
     parallel_config.rank = 0
     parallel_config.world_size = 1
     parallel_config.data_parallel_rank = 0
+    parallel_config.data_parallel_index = 0
 
     engine = NCCLWeightTransferEngine(config, parallel_config)
 
@@ -666,6 +668,7 @@ def inference_receive_ipc_tensor(
     parallel_config.rank = 0
     parallel_config.world_size = 1
     parallel_config.data_parallel_rank = 0
+    parallel_config.data_parallel_index = 0
 
     engine = IPCWeightTransferEngine(config, parallel_config)
 

tests/kernels/mamba/test_ssu_dispatch.py

@@ -13,6 +13,11 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     selective_state_update,
 )
 from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+)
 
 try:
     import flashinfer.mamba  # noqa: F401
@@ -22,22 +27,40 @@ except ImportError:
     HAS_FLASHINFER = False
 
 
+def _kv_cache_config_with_ssu(mamba_type: str = "mamba2") -> KVCacheConfig:
+    spec = MambaSpec(
+        block_size=16,
+        shapes=((16, 64),),
+        dtypes=(torch.float16,),
+        mamba_type=mamba_type,
+    )
+    return KVCacheConfig(
+        num_blocks=1,
+        kv_cache_tensors=[],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["l0"], kv_cache_spec=spec)],
+    )
+
+
 def test_default_backend_is_triton():
-    initialize_mamba_ssu_backend(MambaConfig())
+    initialize_mamba_ssu_backend(MambaConfig(), _kv_cache_config_with_ssu())
     backend = get_mamba_ssu_backend()
     assert isinstance(backend, TritonSSUBackend)
     assert backend.name == "triton"
 
 
 def test_explicit_triton_backend():
-    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
+    initialize_mamba_ssu_backend(
+        MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
+    )
     backend = get_mamba_ssu_backend()
     assert isinstance(backend, TritonSSUBackend)
 
 
 @pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed")
 def test_flashinfer_backend_init():
-    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
+    initialize_mamba_ssu_backend(
+        MambaConfig(backend=MambaBackendEnum.FLASHINFER), _kv_cache_config_with_ssu()
+    )
     backend = get_mamba_ssu_backend()
     assert isinstance(backend, FlashInferSSUBackend)
     assert backend.name == "flashinfer"
@@ -53,6 +76,25 @@ def test_uninitialized_backend_raises():
     mod._mamba_ssu_backend = old
 
 
+@pytest.mark.parametrize(
+    "mamba_type", ["linear_attention", "gdn_attention", "short_conv"]
+)
+def test_init_is_noop_for_non_ssu_mamba_type(mamba_type):
+    import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
+
+    old = mod._mamba_ssu_backend
+    mod._mamba_ssu_backend = None
+    try:
+        initialize_mamba_ssu_backend(
+            MambaConfig(), _kv_cache_config_with_ssu(mamba_type)
+        )
+        assert mod._mamba_ssu_backend is None
+        with pytest.raises(RuntimeError, match="not been initialized"):
+            get_mamba_ssu_backend()
+    finally:
+        mod._mamba_ssu_backend = old
+
+
 @pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed")
 def test_flashinfer_import_error():
     with pytest.raises(ImportError, match="FlashInfer is required"):
@@ -61,7 +103,9 @@ def test_flashinfer_import_error():
 
 def test_triton_basic_call():
     set_random_seed(0)
-    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
+    initialize_mamba_ssu_backend(
+        MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
+    )
     device = "cuda"
     batch_size = 2
     dim = 64

tests/kernels/moe/test_moe_layer.py

@@ -465,6 +465,14 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
     if config.enable_eplb and config.ep_size == 1:
         return False, "EPLB only works with EP+DP"
 
+    # Disable fp4 tests until flashinfer is updated or the Dockerfile is
+    # modified to install cublasLt.h. See #39525.
+    if (
+        config.quantization == "modelopt_fp4"
+        and current_platform.is_device_capability_family(100)
+    ):
+        return False, "Temporarily skip until #39525 is resolved"
+
     return True, None
 
 

tests/models/multimodal/processing/test_audioflamingo3.py

@@ -140,88 +140,3 @@ def test_audio_token_count_matches_hf_processor_math():
         _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 0) == 1499
     )
     assert _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 1) == 375
-
-
-def test_audio_feature_pipeline_matches_hf_small_config():
-    from transformers.models.audioflamingo3 import (
-        modeling_audioflamingo3 as hf_audioflamingo3_modeling,
-    )
-    from transformers.models.audioflamingo3.configuration_audioflamingo3 import (
-        AudioFlamingo3Config,
-    )
-
-    from vllm.model_executor.models.audioflamingo3 import (
-        AudioFlamingo3Encoder,
-        AudioFlamingo3MultiModalProjector,
-        _build_audio_encoder_attention_mask,
-        _flatten_valid_audio_embeddings,
-    )
-
-    text_config = {
-        "model_type": "qwen2",
-        "intermediate_size": 64,
-        "initializer_range": 0.02,
-        "hidden_size": 32,
-        "max_position_embeddings": 1024,
-        "num_hidden_layers": 2,
-        "num_attention_heads": 4,
-        "num_key_value_heads": 2,
-        "vocab_size": 128,
-        "pad_token_id": 1,
-        "use_mrope": False,
-    }
-    audio_config = {
-        "hidden_size": 16,
-        "num_attention_heads": 4,
-        "intermediate_size": 32,
-        "num_hidden_layers": 2,
-        "num_mel_bins": 80,
-        "max_source_positions": 1500,
-        "dropout": 0.0,
-        "attention_dropout": 0.0,
-        "activation_dropout": 0.0,
-        "encoder_layerdrop": 0.0,
-    }
-
-    torch.manual_seed(0)
-    config = AudioFlamingo3Config(
-        text_config=text_config,
-        audio_config=audio_config,
-        audio_token_id=0,
-    )
-    hf_model = hf_audioflamingo3_modeling.AudioFlamingo3ForConditionalGeneration(
-        config
-    ).eval()
-
-    vllm_encoder = AudioFlamingo3Encoder(config.audio_config).eval()
-    vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
-
-    vllm_projector = AudioFlamingo3MultiModalProjector(config).eval()
-    vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
-
-    input_features = torch.randn(3, 80, 3000)
-    feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
-    feature_attention_mask[0, :3000] = True
-    feature_attention_mask[1, :2500] = True
-    feature_attention_mask[2, :1500] = True
-
-    hf_output = hf_model.get_audio_features(
-        input_features,
-        feature_attention_mask,
-        return_dict=True,
-    ).pooler_output
-    vllm_attention_mask = _build_audio_encoder_attention_mask(
-        feature_attention_mask,
-        dtype=vllm_encoder.conv1.weight.dtype,
-        device=vllm_encoder.conv1.weight.device,
-    )
-    vllm_hidden_states = vllm_encoder(
-        input_features,
-        attention_mask=vllm_attention_mask,
-    )
-    vllm_output, _ = _flatten_valid_audio_embeddings(
-        vllm_projector(vllm_hidden_states),
-        feature_attention_mask,
-    )
-
-    torch.testing.assert_close(vllm_output, hf_output)

tests/quantization/test_experts_int8.py

@@ -38,6 +38,5 @@ def test_model_experts_int8_startup(
         dtype=dtype,
         enforce_eager=True,
         quantization="experts_int8",
-        allow_deprecated_quantization=True,
     ) as vllm_model:
         vllm_model.generate_greedy(example_prompts, max_tokens)

vllm/benchmarks/lib/endpoint_request_func.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """The request function for API endpoints."""
 
+import codecs
 import io
 import json
 import os
@@ -25,11 +26,12 @@ class StreamedResponseHandler:
 
     def __init__(self):
         self.buffer = ""
+        self._decoder = codecs.getincrementaldecoder("utf-8")()
 
     def add_chunk(self, chunk_bytes: bytes) -> list[str]:
         """Add a chunk of bytes to the buffer and return any complete
         messages."""
-        chunk_str = chunk_bytes.decode("utf-8")
+        chunk_str = self._decoder.decode(chunk_bytes)
         self.buffer += chunk_str
 
         messages = []

vllm/benchmarks/sweep/plot.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
 from types import TracebackType
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 from typing_extensions import Self, override
 
@@ -17,20 +17,8 @@ from vllm.utils.import_utils import PlaceholderModule
 
 from .utils import sanitize_filename
 
-try:
-    import matplotlib.pyplot as plt
-except ImportError:
-    plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
-
-try:
+if TYPE_CHECKING:
     import pandas as pd
-except ImportError:
-    pd = PlaceholderModule("pandas")
-
-try:
-    import seaborn as sns
-except ImportError:
-    seaborn = PlaceholderModule("seaborn")
 
 
 @dataclass
@@ -265,6 +253,20 @@ def _plot_fig(
     fig_height: float,
     fig_dpi: int,
 ):
+    # Lazy-import matplotlib/pandas/seaborn
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
+    try:
+        import pandas as pd
+    except ImportError:
+        pd = PlaceholderModule("pandas")
+    try:
+        import seaborn as sns
+    except ImportError:
+        sns = PlaceholderModule("seaborn")
+
     fig_group, fig_data = fig_group_data
 
     row_groups = full_groupby(

vllm/benchmarks/sweep/plot_pareto.py

@@ -6,7 +6,7 @@ from concurrent.futures import ProcessPoolExecutor
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
 from vllm.utils.collection_utils import full_groupby
 from vllm.utils.import_utils import PlaceholderModule
@@ -14,20 +14,8 @@ from vllm.utils.import_utils import PlaceholderModule
 from .plot import DummyExecutor, _json_load_bytes
 from .utils import sanitize_filename
 
-try:
-    import matplotlib.pyplot as plt
-except ImportError:
-    plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
-
-try:
+if TYPE_CHECKING:
     import pandas as pd
-except ImportError:
-    pd = PlaceholderModule("pandas")
-
-try:
-    import seaborn as sns
-except ImportError:
-    seaborn = PlaceholderModule("seaborn")
 
 
 def _first_present(run_data: dict[str, object], keys: list[str]):
@@ -195,6 +183,20 @@ def _plot_fig(
         print("[END FIGURE]")
         return
 
+    # Lazy-import matplotlib/pandas/seaborn
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        plt = PlaceholderModule("matplotlib").placeholder_attr("pyplot")
+    try:
+        import pandas as pd
+    except ImportError:
+        pd = PlaceholderModule("pandas")
+    try:
+        import seaborn as sns
+    except ImportError:
+        sns = PlaceholderModule("seaborn")
+
     df = pd.DataFrame.from_records(fig_data)
     df = df.dropna(subset=["tokens_per_user", "tokens_per_gpu"])
 

vllm/config/compilation.py

@@ -737,6 +737,8 @@ class CompilationConfig:
         "vllm::kda_attention",
         "vllm::sparse_attn_indexer",
         "vllm::rocm_aiter_sparse_attn_indexer",
+        # For specialized models
+        "vllm::monolithic_attn",
     ]
 
     def compute_hash(self) -> str:

vllm/config/quantization.py

@@ -19,6 +19,10 @@ class OnlineQuantScheme(Enum):
     # blocks of 128x128 elements (popularized by DeepSeek)
     FP8_PER_BLOCK = "fp8_per_block"
 
+    # int8, weight-only per-channel quantization for MoE expert weights.
+    # Linear layers remain unquantized.
+    INT8_PER_CHANNEL_WEIGHT_ONLY = "int8_per_channel_weight_only"
+
     # TODO(future PRs): add more online quant schemes here: mxfp8, etc
 
 

vllm/entrypoints/chat_utils.py

@@ -1550,6 +1550,18 @@ def _parse_chat_message_content(
             parsed_msg = _ToolParser(message)
             if "tool_call_id" in parsed_msg:
                 result_msg["tool_call_id"] = parsed_msg["tool_call_id"]
+            # Normalize tool message content from OpenAI array format to plain
+            # string. Clients like Claude Code / Cursor send tool results as
+            # [{"type": "text", "text": "..."}], but most chat templates only
+            # handle string content for tool messages.
+            msg_content = result_msg.get("content")
+            if isinstance(msg_content, list):
+                texts = [
+                    item.get("text", "")
+                    for item in msg_content
+                    if isinstance(item, dict) and item.get("type") == "text"
+                ]
+                result_msg["content"] = "\n".join(texts) if texts else ""
 
         if "name" in message and isinstance(message["name"], str):
             result_msg["name"] = message["name"]

vllm/entrypoints/cli/__init__.py

@@ -1,19 +1,2 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from vllm.entrypoints.cli.benchmark.latency import BenchmarkLatencySubcommand
-from vllm.entrypoints.cli.benchmark.mm_processor import (
-    BenchmarkMMProcessorSubcommand,
-)
-from vllm.entrypoints.cli.benchmark.serve import BenchmarkServingSubcommand
-from vllm.entrypoints.cli.benchmark.startup import BenchmarkStartupSubcommand
-from vllm.entrypoints.cli.benchmark.sweep import BenchmarkSweepSubcommand
-from vllm.entrypoints.cli.benchmark.throughput import BenchmarkThroughputSubcommand
-
-__all__: list[str] = [
-    "BenchmarkLatencySubcommand",
-    "BenchmarkMMProcessorSubcommand",
-    "BenchmarkServingSubcommand",
-    "BenchmarkStartupSubcommand",
-    "BenchmarkSweepSubcommand",
-    "BenchmarkThroughputSubcommand",
-]

vllm/entrypoints/cli/benchmark/main.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
+import sys
 import typing
 
 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
@@ -14,6 +15,17 @@ else:
     FlexibleArgumentParser = argparse.ArgumentParser
 
 
+def _import_bench_subcommand_modules() -> None:
+    # Imported lazily so `BenchmarkSubcommandBase` subclasses register only
+    # when `vllm bench` is actually invoked.
+    import vllm.entrypoints.cli.benchmark.latency  # noqa: F401
+    import vllm.entrypoints.cli.benchmark.mm_processor  # noqa: F401
+    import vllm.entrypoints.cli.benchmark.serve  # noqa: F401
+    import vllm.entrypoints.cli.benchmark.startup  # noqa: F401
+    import vllm.entrypoints.cli.benchmark.sweep  # noqa: F401
+    import vllm.entrypoints.cli.benchmark.throughput  # noqa: F401
+
+
 class BenchmarkSubcommand(CLISubcommand):
     """The `bench` subcommand for the vLLM CLI."""
 
@@ -38,18 +50,28 @@ class BenchmarkSubcommand(CLISubcommand):
         )
         bench_subparsers = bench_parser.add_subparsers(required=True, dest="bench_type")
 
-        for cmd_cls in BenchmarkSubcommandBase.__subclasses__():
-            cmd_subparser = bench_subparsers.add_parser(
-                cmd_cls.name,
-                help=cmd_cls.help,
-                description=cmd_cls.help,
-                usage=f"vllm {self.name} {cmd_cls.name} [options]",
-            )
-            cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd)
-            cmd_cls.add_cli_args(cmd_subparser)
-            cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(
-                subcmd=f"{self.name} {cmd_cls.name}"
-            )
+        # Only build the nested bench subparsers when the user is actually
+        # invoking `bench`; otherwise we'd drag in imports
+        # unnecessarily on every `vllm --help` and `vllm serve`.
+        # Scan for the first positional arg so global flags (e.g. `-v`)
+        # before the subcommand don't break detection.
+        first_positional = next(
+            (arg for arg in sys.argv[1:] if not arg.startswith("-")), None
+        )
+        if first_positional == self.name:
+            _import_bench_subcommand_modules()
+            for cmd_cls in BenchmarkSubcommandBase.__subclasses__():
+                cmd_subparser = bench_subparsers.add_parser(
+                    cmd_cls.name,
+                    help=cmd_cls.help,
+                    description=cmd_cls.help,
+                    usage=f"vllm {self.name} {cmd_cls.name} [options]",
+                )
+                cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd)
+                cmd_cls.add_cli_args(cmd_subparser)
+                cmd_subparser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(
+                    subcmd=f"{self.name} {cmd_cls.name}"
+                )
         return bench_parser
 
 

vllm/env_override.py

@@ -100,10 +100,9 @@ logger = init_logger(__name__)
 # it avoids unintentional cuda initialization from torch.cuda.is_available()
 os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
 
-# see https://github.com/vllm-project/vllm/issues/10480
+# see https://github.com/vllm-project/vllm/issues/10480 and
+# https://github.com/vllm-project/vllm/issues/10619.
 os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
-# see https://github.com/vllm-project/vllm/issues/10619
-torch._inductor.config.compile_threads = 1
 
 # Enable Triton autotuning result caching to disk by default.
 # Without this, Triton re-runs autotuning on every process restart,

vllm/envs.py

@@ -216,6 +216,7 @@ if TYPE_CHECKING:
     VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False
     VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
     VLLM_ALLREDUCE_USE_FLASHINFER: bool = False
+    VLLM_USE_SPECIALIZED_MODELS: bool = False
     VLLM_TUNED_CONFIG_FOLDER: str | None = None
     VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set()
     VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False
@@ -1520,6 +1521,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool(
         int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "0"))
     ),
+    # Whether to enable specialized model implementations when available.
+    "VLLM_USE_SPECIALIZED_MODELS": lambda: bool(
+        int(os.getenv("VLLM_USE_SPECIALIZED_MODELS", "0"))
+    ),
     # Experimental: use this to enable MCP tool calling for non harmony models
     "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool(
         int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0"))

vllm/model_executor/layers/fused_moe/oracle/int8.py

@@ -0,0 +1,84 @@
+# 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.all2all_utils import (
+    maybe_make_prepare_finalize,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+    int8_w8a16_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
+    SharedExperts,
+)
+
+logger = init_logger(__name__)
+
+
+def select_int8_moe_backend(
+    config: FusedMoEConfig,
+) -> type[mk.FusedMoEExperts]:
+    from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
+
+    supported, reason = TritonExperts.is_supported_config(
+        TritonExperts,
+        config,
+        None,
+        None,
+        mk.FusedMoEActivationFormat.Standard,
+    )
+    if not supported:
+        raise ValueError(
+            f"INT8 Triton MoE backend does not support the "
+            f"deployment configuration: {reason}"
+        )
+
+    logger.info_once("Using Triton INT8 MoE backend", scope="local")
+    return TritonExperts
+
+
+def make_int8_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+) -> FusedMoEQuantConfig:
+    return int8_w8a16_moe_quant_config(
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=None,
+        w2_zp=None,
+    )
+
+
+def make_int8_moe_kernel(
+    moe_quant_config: FusedMoEQuantConfig,
+    moe_config: FusedMoEConfig,
+    experts_cls: type[mk.FusedMoEExperts],
+    routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+    shared_experts: SharedExperts | None = None,
+) -> mk.FusedMoEKernel:
+    prepare_finalize = maybe_make_prepare_finalize(
+        moe=moe_config,
+        quant_config=moe_quant_config,
+        routing_tables=routing_tables,
+        allow_new_interface=True,
+    )
+    assert prepare_finalize is not None
+
+    logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
+
+    experts = experts_cls(
+        moe_config=moe_config,
+        quant_config=moe_quant_config,
+    )
+
+    return mk.FusedMoEKernel(
+        prepare_finalize,
+        experts,
+        shared_experts=shared_experts,
+        inplace=not moe_config.disable_inplace,
+    )

vllm/model_executor/layers/mamba/ops/ssu_dispatch.py

@@ -15,6 +15,7 @@ import torch
 from vllm.config.mamba import MambaBackendEnum, MambaConfig
 from vllm.logger import init_logger
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
 
 logger = init_logger(__name__)
 
@@ -188,12 +189,22 @@ _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
 _mamba_ssu_backend: MambaSSUBackend | None = None
 
 
-def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None:
+def initialize_mamba_ssu_backend(
+    mamba_config: MambaConfig,
+    kv_cache_config: KVCacheConfig,
+) -> None:
     """Initialize the global Mamba SSU backend.
 
-    Args:
-        mamba_config: Mamba configuration.
+    No-op if `kv_cache_config` contains no specs that call
+    selective_state_update.
     """
+    if not any(
+        isinstance(g.kv_cache_spec, MambaSpec)
+        and g.kv_cache_spec.mamba_type in ("mamba1", "mamba2")
+        for g in kv_cache_config.kv_cache_groups
+    ):
+        return
+
     global _mamba_ssu_backend
 
     backend = mamba_config.backend
@@ -203,7 +214,11 @@ def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None:
             f"Valid options: {list(_BACKEND_REGISTRY.keys())}"
         )
 
-    _mamba_ssu_backend = _BACKEND_REGISTRY[backend](mamba_config)
+    backend_cls = _BACKEND_REGISTRY[backend]
+    if isinstance(_mamba_ssu_backend, backend_cls):
+        return
+
+    _mamba_ssu_backend = backend_cls(mamba_config)
     logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
 
 

vllm/model_executor/layers/quantization/__init__.py

@@ -40,6 +40,7 @@ QuantizationMethods = Literal[
     # shorthand for creating a more complicated online quant config object
     "fp8_per_tensor",
     "fp8_per_block",
+    "int8_per_channel_weight_only",
 ]
 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
 
@@ -47,7 +48,6 @@ DEPRECATED_QUANTIZATION_METHODS = [
     "tpu_int8",
     "fbgemm_fp8",
     "fp_quant",
-    "experts_int8",
 ]
 
 # The customized quantization methods which will be added to this dict.

vllm/model_executor/layers/quantization/experts_int8.py

@@ -5,27 +5,25 @@ from typing import Any
 
 import torch
 
-from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
-from vllm.model_executor.layers.fused_moe import (
-    FusedMoE,
-    FusedMoEConfig,
-    FusedMoEMethodBase,
-)
-from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEQuantConfig,
-    int8_w8a16_moe_quant_config,
-)
+from vllm.model_executor.layers.fused_moe import FusedMoE
 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization import QuantizationMethods
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
-from vllm.model_executor.utils import set_weight_attrs
+from vllm.model_executor.layers.quantization.online.int8 import (
+    Int8OnlineMoEMethod,
+)
 
 
 class ExpertsInt8Config(QuantizationConfig):
-    """Config class for Int8 experts quantization."""
+    """Online int8 quantization for MoE expert weights.
+    Linear layers are left unquantized.
+
+    Backward-compatible config for ``--quantization experts_int8``.
+    Prefer ``--quantization int8_per_channel``
+    """
 
     def __init__(self) -> None:
         super().__init__()
@@ -56,149 +54,5 @@ class ExpertsInt8Config(QuantizationConfig):
         if isinstance(layer, LinearBase):
             return UnquantizedLinearMethod()
         elif isinstance(layer, FusedMoE):
-            return ExpertsInt8MoEMethod(self, layer.moe_config)
+            return Int8OnlineMoEMethod(layer=layer)
         return None
-
-
-class ExpertsInt8MoEMethod(FusedMoEMethodBase):
-    def __init__(
-        self,
-        quant_config: ExpertsInt8Config,
-        moe: FusedMoEConfig,
-    ):
-        super().__init__(moe)
-        self.quant_config = quant_config
-
-    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,
-    ):
-        int8_dtype = torch.int8
-
-        assert "weight_loader" in extra_weight_attrs
-        weight_loader = extra_weight_attrs["weight_loader"]
-        wrapped_weight_loader = ExpertsInt8MoEMethod.quantizing_weight_loader(
-            layer, weight_loader
-        )
-        extra_weight_attrs["weight_loader"] = wrapped_weight_loader
-
-        # Fused gate_up_proj (column parallel)
-        w13_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                2 * intermediate_size_per_partition,
-                hidden_size,
-                dtype=int8_dtype,
-            ),
-            requires_grad=False,
-        )
-        layer.register_parameter("w13_weight", w13_weight)
-        set_weight_attrs(w13_weight, extra_weight_attrs)
-
-        # down_proj (row parallel)
-        w2_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                hidden_size,
-                intermediate_size_per_partition,
-                dtype=int8_dtype,
-            ),
-            requires_grad=False,
-        )
-        layer.register_parameter("w2_weight", w2_weight)
-        set_weight_attrs(w2_weight, extra_weight_attrs)
-
-        w13_scale = torch.nn.Parameter(
-            torch.zeros(
-                num_experts, 2 * intermediate_size_per_partition, dtype=torch.float32
-            ),
-            requires_grad=False,
-        )
-        layer.register_parameter("w13_scale", w13_scale)
-
-        w2_scale = torch.nn.Parameter(
-            torch.zeros(num_experts, hidden_size, dtype=torch.float32),
-            requires_grad=False,
-        )
-        layer.register_parameter("w2_scale", w2_scale)
-
-    def get_fused_moe_quant_config(
-        self, layer: torch.nn.Module
-    ) -> FusedMoEQuantConfig | None:
-        return int8_w8a16_moe_quant_config(
-            w1_scale=layer.w13_scale, w2_scale=layer.w2_scale, w1_zp=None, w2_zp=None
-        )
-
-    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,
-            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,
-        )
-
-    @staticmethod
-    def quantizing_weight_loader(layer, weight_loader):
-        def quantize_and_call_weight_loader(
-            param: torch.nn.Parameter,
-            loaded_weight: torch.Tensor,
-            weight_name: str,
-            shard_id: int,
-            expert_id: int,
-        ):
-            tp_rank = get_tensor_model_parallel_rank()
-            shard_size = layer.intermediate_size_per_partition
-            shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
-            device = get_tp_group().device
-            loaded_weight = loaded_weight.to(device)
-            # w1, gate_proj case: Load into first shard of w13.
-            if shard_id == "w1":
-                scales = quantize_in_place_and_get_scales(loaded_weight[shard, :])
-                layer.w13_scale.data[expert_id, 0:shard_size].copy_(scales[:, 0])
-            # w3, up_proj case: Load into second shard of w13.
-            elif shard_id == "w3":
-                scales = quantize_in_place_and_get_scales(loaded_weight[shard, :])
-                layer.w13_scale.data[expert_id, shard_size : 2 * shard_size].copy_(
-                    scales[:, 0]
-                )
-            # w2, down_proj case: Load into only shard of w2.
-            elif shard_id == "w2":
-                scales = quantize_in_place_and_get_scales(loaded_weight[:, shard])
-                layer.w2_scale.data[expert_id, :].copy_(scales[:, 0])
-            else:
-                raise ValueError(f"Shard id must be in [0,1,2] but got {shard_id}")
-            weight_loader(param, loaded_weight, weight_name, shard_id, expert_id)
-
-        return quantize_and_call_weight_loader
-
-
-def quantize_in_place_and_get_scales(weight: torch.Tensor) -> torch.Tensor:
-    vmax = torch.iinfo(torch.int8).max
-    scales = torch.max(torch.abs(weight), dim=1, keepdim=True)[0] / vmax
-
-    weight.div_(scales)
-    weight.round_()
-    weight.clamp_(-vmax, vmax)
-
-    return scales

vllm/model_executor/layers/quantization/online/base.py

@@ -9,6 +9,7 @@ from vllm.config.quantization import (
     OnlineQuantizationConfigArgs,
     OnlineQuantScheme,
 )
+from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
 )
@@ -33,6 +34,11 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
     Fp8PerTensorOnlineLinearMethod,
     Fp8PerTensorOnlineMoEMethod,
 )
+from vllm.model_executor.layers.quantization.online.int8 import (
+    Int8OnlineMoEMethod,
+)
+
+logger = init_logger(__name__)
 
 
 class OnlineQuantizationConfig(QuantizationConfig):
@@ -96,7 +102,13 @@ class OnlineQuantizationConfig(QuantizationConfig):
                 return UnquantizedLinearMethod()
 
             linear_scheme = self.args.linear_scheme_override or self.args.global_scheme
-            if linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
+            if linear_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
+                logger.warning_once(
+                    "INT8 online quantization only quantizes MoE expert "
+                    "weights. linear layers remain in full precision."
+                )
+                return UnquantizedLinearMethod()
+            elif linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
                 return Fp8PerBlockOnlineLinearMethod()
             else:
                 return Fp8PerTensorOnlineLinearMethod()
@@ -109,7 +121,9 @@ class OnlineQuantizationConfig(QuantizationConfig):
                 return UnquantizedFusedMoEMethod(layer.moe_config)
 
             moe_scheme = self.args.moe_scheme_override or self.args.global_scheme
-            if moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
+            if moe_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
+                return Int8OnlineMoEMethod(layer=layer)
+            elif moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
                 return Fp8PerBlockOnlineMoEMethod(layer=layer)
             else:
                 return Fp8PerTensorOnlineMoEMethod(layer=layer)

vllm/model_executor/layers/quantization/online/fp8.py

@@ -10,7 +10,6 @@ if TYPE_CHECKING:
     import vllm.model_executor.layers.fused_moe.modular_kernel as mk
     from vllm.model_executor.layers.fused_moe import FusedMoE
     from vllm.model_executor.layers.fused_moe.config import (
-        FusedMoEConfig,
         FusedMoEQuantConfig,
     )
     from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
@@ -19,15 +18,15 @@ import vllm.envs as envs
 from vllm import _custom_ops as ops
 from vllm.config import get_current_vllm_config
 from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
-from vllm.model_executor.layers.fused_moe import (
-    FusedMoEMethodBase,
-)
 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
     select_fp8_moe_backend,
 )
 from vllm.model_executor.layers.linear import (
     LinearMethodBase,
 )
+from vllm.model_executor.layers.quantization.online.moe_base import (
+    OnlineMoEMethodBase,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
     create_fp8_quant_key,
@@ -44,7 +43,7 @@ from vllm.model_executor.model_loader.reload.layerwise import (
     initialize_online_processing,
 )
 from vllm.model_executor.parameter import ModelWeightParameter
-from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import per_block_cast_to_fp8
 
@@ -268,21 +267,15 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
 # ---------------------------------------------------------------------------
 
 
-class _Fp8OnlineMoEBase(FusedMoEMethodBase):
+class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
     """Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint
     weights onto meta device and materializes them just-in-time."""
 
-    uses_meta_device: bool = True
-
     # Declared here for mypy; actual values are set in __init__.
     fp8_backend: "Fp8MoeBackend"
     experts_cls: "type[mk.FusedMoEExperts] | None"
     weight_scale_name: str
     weight_block_size: list[int] | None
-    moe: "FusedMoEConfig"
-    is_monolithic: bool
-    moe_quant_config: "FusedMoEQuantConfig | None"
-    moe_kernel: "mk.FusedMoEKernel | None"
 
     def __init__(
         self,
@@ -313,77 +306,6 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
             allow_vllm_cutlass=False,
         )
 
-    def create_weights(
-        self,
-        layer: 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
-
-        # WEIGHTS
-        w13_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                2 * intermediate_size_per_partition,
-                hidden_size,
-                device="meta",
-                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,
-                device="meta",  # materialized and processed during loading
-                dtype=params_dtype,
-            ),
-            requires_grad=False,
-        )
-        layer.register_parameter("w2_weight", w2_weight)
-        set_weight_attrs(w2_weight, extra_weight_attrs)
-
-        # BIASES (for models like GPT-OSS that have biased MoE)
-        if self.moe.has_bias:
-            w13_bias = torch.nn.Parameter(
-                torch.zeros(
-                    num_experts,
-                    2 * intermediate_size_per_partition,
-                    device="meta",  # materialized and processed during loading
-                    dtype=layer.orig_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,
-                    device="meta",  # materialized and processed during loading
-                    dtype=layer.orig_dtype,
-                ),
-                requires_grad=False,
-            )
-            layer.register_parameter("w2_bias", w2_bias)
-            set_weight_attrs(w2_bias, extra_weight_attrs)
-
-        layer.w13_input_scale = None
-        layer.w2_input_scale = None
-
-        initialize_online_processing(layer)
-
     def _setup_kernel(
         self,
         layer: "FusedMoE",
@@ -430,15 +352,6 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
                 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":
@@ -460,68 +373,9 @@ class _Fp8OnlineMoEBase(FusedMoEMethodBase):
             block_shape=self.weight_block_size,
         )
 
-        # Inject biases into the quant config if the model has them
-        # (e.g. GPT-OSS biased MoE)
-        if quant_config is not None and self.moe.has_bias:
-            w13_bias = getattr(layer, "w13_bias", None)
-            w2_bias = getattr(layer, "w2_bias", None)
-            if w13_bias is not None:
-                quant_config._w1.bias = w13_bias
-            if w2_bias is not None:
-                quant_config._w2.bias = w2_bias
-
+        self._maybe_inject_biases(quant_config, layer)
         return quant_config
 
-    @property
-    def supports_eplb(self) -> bool:
-        return True
-
-    def apply_monolithic(
-        self,
-        layer: "FusedMoE",
-        x: torch.Tensor,
-        router_logits: torch.Tensor,
-    ) -> torch.Tensor | tuple[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 | tuple[torch.Tensor, 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,
-            expert_map=layer.expert_map,
-            apply_router_weight_on_input=layer.apply_router_weight_on_input,
-            shared_experts_input=shared_experts_input,
-        )
-
 
 class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
     """Online tensorwise FP8 MoE quantization.

vllm/model_executor/layers/quantization/online/int8.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch.nn import Module
+
+if TYPE_CHECKING:
+    import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+    from vllm.model_executor.layers.fused_moe import FusedMoE
+    from vllm.model_executor.layers.fused_moe.config import (
+        FusedMoEQuantConfig,
+    )
+
+from vllm.model_executor.layers.fused_moe.oracle.int8 import (
+    make_int8_moe_kernel,
+    make_int8_moe_quant_config,
+    select_int8_moe_backend,
+)
+from vllm.model_executor.layers.quantization.online.moe_base import (
+    OnlineMoEMethodBase,
+)
+from vllm.model_executor.utils import replace_parameter
+
+
+class Int8OnlineMoEMethod(OnlineMoEMethodBase):
+    """Online per-channel INT8 MoE quantization.
+    Loads fp16/bf16 weights and quantizes them per-row to int8 during loading.
+    """
+
+    def __init__(
+        self,
+        *,
+        layer: torch.nn.Module,
+    ):
+        super().__init__(layer.moe_config)
+        self.experts_cls: type[mk.FusedMoEExperts] = select_int8_moe_backend(
+            config=self.moe,
+        )
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        self._quantize_weights(layer)
+        self._setup_kernel(layer)
+
+        layer._already_called_process_weights_after_loading = True
+
+    def _quantize_weights(self, layer: Module) -> None:
+        vmax = torch.iinfo(torch.int8).max
+
+        w13 = torch.empty_like(layer.w13_weight, dtype=torch.int8)
+        w2 = torch.empty_like(layer.w2_weight, dtype=torch.int8)
+        w13_scale = torch.zeros(
+            layer.num_experts,
+            layer.w13_weight.shape[1],
+            device=w13.device,
+            dtype=torch.float32,
+        )
+        w2_scale = torch.zeros(
+            layer.num_experts,
+            layer.w2_weight.shape[1],
+            device=w2.device,
+            dtype=torch.float32,
+        )
+
+        for expert in range(layer.local_num_experts):
+            # w13: per-row quantization over hidden_size dim
+            w = layer.w13_weight[expert, :, :]
+            scales = w.abs().amax(dim=1) / vmax
+            q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
+            w13[expert, :, :] = q.to(torch.int8)
+            w13_scale[expert, :] = scales
+
+            # w2: per-row quantization over intermediate_size dim
+            w = layer.w2_weight[expert, :, :]
+            scales = w.abs().amax(dim=1) / vmax
+            q = w.div(scales.unsqueeze(1)).round().clamp(-vmax, vmax)
+            w2[expert, :, :] = q.to(torch.int8)
+            w2_scale[expert, :] = scales
+
+        replace_parameter(layer, "w13_weight", w13)
+        replace_parameter(layer, "w2_weight", w2)
+        replace_parameter(layer, "w13_scale", w13_scale)
+        replace_parameter(layer, "w2_scale", w2_scale)
+
+    def _setup_kernel(self, layer: "FusedMoE") -> None:
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        assert self.moe_quant_config is not None
+        assert self.experts_cls is not None
+        self.moe_kernel = make_int8_moe_kernel(
+            moe_quant_config=self.moe_quant_config,
+            moe_config=self.moe,
+            experts_cls=self.experts_cls,
+            routing_tables=layer._maybe_init_expert_routing_tables(),
+            shared_experts=layer.shared_experts,
+        )
+
+    def get_fused_moe_quant_config(
+        self, layer: torch.nn.Module
+    ) -> "FusedMoEQuantConfig | None":
+        quant_config = make_int8_moe_quant_config(
+            w1_scale=layer.w13_scale,
+            w2_scale=layer.w2_scale,
+        )
+        self._maybe_inject_biases(quant_config, layer)
+        return quant_config

vllm/model_executor/layers/quantization/online/moe_base.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import abstractmethod
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe import FusedMoEMethodBase
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.model_loader.reload.layerwise import (
+    initialize_online_processing,
+)
+from vllm.model_executor.utils import set_weight_attrs
+
+
+class OnlineMoEMethodBase(FusedMoEMethodBase):
+    """Base for MoE methods that load full-precision weights on meta device
+    and quantize them after loading via the QeRL layerwise processing system.
+    """
+
+    uses_meta_device: bool = 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
+
+        # Fused gate_up_proj (column parallel) — full precision on meta device
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size,
+                device="meta",
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        # down_proj (row parallel) — full precision on meta device
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition,
+                device="meta",
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        # BIASES (for models like GPT-OSS that have biased MoE)
+        if self.moe.has_bias:
+            w13_bias = torch.nn.Parameter(
+                torch.zeros(
+                    num_experts,
+                    2 * intermediate_size_per_partition,
+                    device="meta",
+                    dtype=layer.orig_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,
+                    device="meta",
+                    dtype=layer.orig_dtype,
+                ),
+                requires_grad=False,
+            )
+            layer.register_parameter("w2_bias", w2_bias)
+            set_weight_attrs(w2_bias, extra_weight_attrs)
+
+        layer.w13_input_scale = None
+        layer.w2_input_scale = None
+
+        initialize_online_processing(layer)
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        pass
+
+    def _maybe_inject_biases(
+        self,
+        quant_config: FusedMoEQuantConfig,
+        layer: torch.nn.Module,
+    ) -> None:
+        """Inject biases into the quant config if the model has them
+        (e.g. GPT-OSS biased MoE)."""
+        if self.moe.has_bias:
+            w13_bias = getattr(layer, "w13_bias", None)
+            w2_bias = getattr(layer, "w2_bias", None)
+            if w13_bias is not None:
+                quant_config._w1.bias = w13_bias
+            if w2_bias is not None:
+                quant_config._w2.bias = w2_bias
+
+    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."
+        )
+
+    @property
+    def supports_eplb(self) -> bool:
+        return True
+
+    def apply_monolithic(
+        self,
+        layer: "FusedMoE",  # type: ignore[name-defined] # noqa: F821
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+    ) -> torch.Tensor | tuple[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",  # type: ignore[name-defined] # noqa: F821
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor | tuple[torch.Tensor, 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,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            shared_experts_input=shared_experts_input,
+        )

vllm/model_executor/model_loader/dummy_loader.py

@@ -43,6 +43,13 @@ class DummyModelLoader(BaseModelLoader):
                 # random values to the weights.
                 initialize_dummy_weights(layer, model_config)
 
+                # Some models build derived weights from loaded parameters instead of
+                # storing them in checkpoints. Rebuild those tensors for dummy load.
+                for layer in model.modules():
+                    fuse_indexer_weights = getattr(layer, "fuse_indexer_weights", None)
+                    if callable(fuse_indexer_weights):
+                        fuse_indexer_weights()
+
     def _process_online_quant_layer(
         self,
         layer: nn.Module,

vllm/model_executor/models/deepseek_mtp.py

@@ -30,7 +30,6 @@ from .deepseek_v2 import (
     DeepseekV2DecoderLayer,
     DeepseekV2MixtureOfExperts,
     DeepseekV2MoE,
-    _try_load_fp8_indexer_wk,
     get_spec_layer_idx_from_weight_name,
 )
 from .utils import maybe_prefix
@@ -191,6 +190,10 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
         )
         # Set MoE hyperparameters
         self.set_moe_parameters()
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
 
     def set_moe_parameters(self):
         self.expert_weights = []
@@ -245,12 +248,13 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
             ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
         ]
 
-        # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
-        indexer_fused_mapping = [
-            ("wk_weights_proj", "wk", 0),
-            ("wk_weights_proj", "weights_proj", 1),
-        ]
-        stacked_params_mapping.extend(indexer_fused_mapping)
+        if self.is_fp4_ckpt:
+            # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
+            indexer_fused_mapping = [
+                ("wk_weights_proj", "wk", 0),
+                ("wk_weights_proj", "weights_proj", 1),
+            ]
+            stacked_params_mapping.extend(indexer_fused_mapping)
 
         expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
             self,
@@ -267,7 +271,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
-        _pending_wk_fp8: dict = {}  # FP8 indexer wk dequant buffer
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
@@ -278,12 +281,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
                 rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
             )
             name = self._rewrite_spec_layer_name(spec_layer, name)
-
-            if _try_load_fp8_indexer_wk(
-                name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
-            ):
-                continue
-
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 # Skip non-stacked layers and experts (experts handled below).
                 if weight_name not in name:

vllm/model_executor/models/deepseek_v2.py

@@ -66,10 +66,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8,
 )
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    GroupShape,
-    scaled_dequantize,
-)
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.sparse_attn_indexer import (
     SparseAttnIndexer,
@@ -632,6 +628,10 @@ class Indexer(nn.Module):
         self.vllm_config = vllm_config
         self.config = config
         self.quant_config = quant_config
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
         # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
         self.topk_tokens = config.index_topk
         self.n_head = config.index_n_heads  # 64
@@ -646,16 +646,36 @@ class Indexer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.wq_b",
         )
-        # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
-        # FP8 wk weights are upcasted to BF16 during loading to maintain fusion.
-        self.wk_weights_proj = MergedColumnParallelLinear(
-            hidden_size,
-            [self.head_dim, self.n_head],
-            bias=False,
-            quant_config=None,
-            disable_tp=True,
-            prefix=f"{prefix}.wk_weights_proj",
-        )
+        if self.is_fp4_ckpt:
+            # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
+            # weights_proj does not get quantized,
+            # so we run both with quant_config=None
+            # wk may be upcasted from the default quant;
+            # experiments show fusion is always faster unless WK proj is in FP4,
+            # which is not the case for all known quants.
+            self.wk_weights_proj = MergedColumnParallelLinear(
+                hidden_size,
+                [self.head_dim, self.n_head],
+                bias=False,
+                quant_config=None,
+                disable_tp=True,
+                prefix=f"{prefix}.wk_weights_proj",
+            )
+        else:
+            self.wk = ReplicatedLinear(
+                hidden_size,
+                self.head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.wk",
+            )
+            self.weights_proj = ReplicatedLinear(
+                hidden_size,
+                self.n_head,
+                bias=False,
+                quant_config=None,
+                prefix=f"{prefix}.weights_proj",
+            )
         self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
         self.softmax_scale = self.head_dim**-0.5
 
@@ -696,10 +716,14 @@ class Indexer(nn.Module):
         q_pe, q_nope = torch.split(
             q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
         )
-        # Fused wk + weights_proj: one GEMM, then split
-        kw, _ = self.wk_weights_proj(hidden_states)
-        k = kw[:, : self.head_dim]
-        weights = kw[:, self.head_dim :]
+        if self.is_fp4_ckpt:
+            # Fused wk + weights_proj: one GEMM, then split
+            kw, _ = self.wk_weights_proj(hidden_states)
+            k = kw[:, : self.head_dim]
+            weights = kw[:, self.head_dim :]
+        else:
+            k, _ = self.wk(hidden_states)
+            weights, _ = self.weights_proj(hidden_states)
 
         k = self.k_norm(k)
         k_pe, k_nope = torch.split(
@@ -737,46 +761,6 @@ class Indexer(nn.Module):
         return self.indexer_op(hidden_states, q_fp8, k, weights)
 
 
-def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
-    """
-    We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
-    in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
-    Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK
-    weights and scale, and when both are available, dequantizes to BF16 and stores into
-    the fused wk_weights_proj.weight parameter.
-    """
-    if "indexer.wk." not in name or "wk_weights" in name:
-        return False  # Weight is not an isolated WK weight for the indexer, ignore.
-    is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
-    is_scale = "weight_scale_inv" in name
-    if not is_weight and not is_scale:
-        return False  # WK is not in FP8 format, ignore.
-    # Buffer this tensor (weight or scale) until both have arrived.
-    layer_prefix = name.rsplit(".wk.", 1)[0]  # e.g. "model.layers.0.self_attn.indexer"
-    entry = buf.setdefault(layer_prefix, {})
-    entry["weight" if is_weight else "scale"] = tensor
-    if "weight" not in entry or "scale" not in entry:
-        return True  # still waiting for the other param
-
-    # We have both weight and scale: dequantize FP8 to BF16.
-    weight_fp8, scale_inv = entry["weight"], entry["scale"]
-    del buf[layer_prefix]
-    block_size = weight_fp8.shape[1] // scale_inv.shape[1]
-    weight_bf16 = scaled_dequantize(
-        weight_fp8,
-        scale_inv,
-        group_shape=GroupShape(block_size, block_size),
-        out_dtype=torch.bfloat16,
-    )
-
-    # Load the dequantized weight into shard 0 of the fused buffer.
-    fused_name = f"{layer_prefix}.wk_weights_proj.weight"
-    param = params_dict[fused_name]
-    param.weight_loader(param, weight_bf16, 0)
-    loaded_params.add(fused_name)
-    return True
-
-
 def _min_latency_fused_qkv_a_proj_impl(
     input_: torch.Tensor,
     weight: torch.Tensor,
@@ -1360,6 +1344,10 @@ class DeepseekV2ForCausalLM(
         quant_config = vllm_config.quant_config
         self.config = config
         self.quant_config = quant_config
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
 
         qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
         qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
@@ -1485,13 +1473,13 @@ class DeepseekV2ForCausalLM(
             ("qkv_proj", "k_proj", "k"),
             ("qkv_proj", "v_proj", "v"),
         ]
-        # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
-        _pending_wk_fp8: dict = {}  # When WK is in FP8, we dequant to BF16 for fusion
-        indexer_fused_mapping = [
-            ("wk_weights_proj", "wk", 0),
-            ("wk_weights_proj", "weights_proj", 1),
-        ]
-        stacked_params_mapping.extend(indexer_fused_mapping)
+        if self.is_fp4_ckpt:
+            # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
+            indexer_fused_mapping = [
+                ("wk_weights_proj", "wk", 0),
+                ("wk_weights_proj", "weights_proj", 1),
+            ]
+            stacked_params_mapping.extend(indexer_fused_mapping)
 
         if self.use_mha:
             stacked_params_mapping.extend(mha_params_mapping)
@@ -1528,11 +1516,6 @@ class DeepseekV2ForCausalLM(
                 rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
             )
 
-            if _try_load_fp8_indexer_wk(
-                name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
-            ):
-                continue
-
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 # Skip non-stacked layers and experts (experts handled below).
                 if weight_name not in name:

vllm/model_executor/models/parakeet.py

@@ -99,6 +99,8 @@ class ProjectedParakeet(nn.Module):
             if target is None:
                 target = buffers_dict.get(target_name)
             if target is None:
+                if self._can_skip_missing_named_param(target_name):
+                    continue
                 raise ValueError(f"Unknown weight: {name}")
             weight_loader = getattr(target, "weight_loader", default_weight_loader)
             with torch.no_grad():
@@ -107,6 +109,27 @@ class ProjectedParakeet(nn.Module):
 
         return loaded_params
 
+    def _can_skip_missing_named_param(self, target_name: str) -> bool:
+        if self.config.convolution_bias:
+            return False
+
+        # In transformers v5 (not v4), `convolution_bias=False` is
+        # propagated from parakeet config. If `False`, torch.conv1d will
+        # *skip registering the param*, thus it will be missing in the
+        # module's named params. *If* you happen to also have the bias
+        # tensors in the weights, it will cause a mismatch between the
+        # weights and the params.
+        # This allows us to have `convolution_bias=False` in the sound config,
+        # but still allow for the weights to exist.
+
+        return target_name.endswith(
+            (
+                ".conv.pointwise_conv1.bias",
+                ".conv.depthwise_conv.bias",
+                ".conv.pointwise_conv2.bias",
+            )
+        )
+
 
 EPSILON = 1e-5
 LOG_ZERO_GUARD_VALUE = 2**-24

vllm/model_executor/models/registry.py

@@ -1302,6 +1302,15 @@ ModelRegistry = _ModelRegistry(
     }
 )
 
+if envs.VLLM_USE_SPECIALIZED_MODELS:
+    from vllm.model_executor.specialized_models import get_specialized_models
+
+    for _arch, (_mod, _cls) in get_specialized_models().items():
+        ModelRegistry.models[_arch] = _LazyRegisteredModel(
+            module_name=_mod,
+            class_name=_cls,
+        )
+
 _T = TypeVar("_T")
 
 

vllm/model_executor/specialized_models/README.md

@@ -0,0 +1,19 @@
+# [Experimental] Specialized Models
+
+This directory contains experimental, hand-tuned implementations for a small number of selected models. Each subdirectory targets a specific combination of model architecture (including all tensor shapes), quantization scheme, attention backend, and hardware.
+
+For example, `deepseek_v3_2_nvfp4/` targets `nvidia/DeepSeek-V3.2-NVFP4` with FP8 FlashInfer sparse MLA on Blackwell GPUs.
+
+**To opt in, set `VLLM_USE_SPECIALIZED_MODELS=1`.** When enabled, vLLM will prefer a specialized implementation over the generic one if a match is available.
+
+## Development Philosophy
+
+These implementations prioritize iteration speed and checkpoint-specific performance over broad reuse. They may target a very narrow use case and are not expected to cover the full vLLM feature surface. Known limitations include:
+
+- Parallelism strategy support may be incomplete (e.g. TP only, no EP, or vice versa).
+- `torch.compile` compatibility may be limited or untested.
+- Behavior with checkpoint formats outside the intended target is unsupported.
+
+Also, code duplication across implementations is intentional — each model should be free to evolve and be optimized independently without risk of regressing another.
+
+Code here is experimental and may be short-lived. Generic features and anything intended for long-term support should live in `../models/`.

vllm/model_executor/specialized_models/__init__.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Specialized model implementations.
+
+Each entry maps a vLLM architecture name to a (module_path, class_name)
+tuple, exactly like the main model registry.  When
+``VLLM_USE_SPECIALIZED_MODELS=1`` the main registry merges these entries
+so they take priority over the generic implementations.
+
+To add a new specialized model:
+  1. Create a sub-package under this directory.
+  2. Add the architecture -> (module, class) mapping to ``_MODELS`` below.
+"""
+
+from __future__ import annotations
+
+# ── Model list ───────────────────────────────────────────────────────
+# Maps architecture name -> (fully-qualified module, class name).
+# When the flag is enabled, these override the corresponding entries
+# in the main registry.
+_MODELS: dict[str, tuple[str, str]] = {
+    "DeepseekV32ForCausalLM": (
+        "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4",
+        "DeepseekV32ForCausalLM",
+    ),
+    "DeepSeekMTPModel": (
+        "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4",
+        "DeepSeekMTP",
+    ),
+}
+
+
+def get_specialized_models() -> dict[str, tuple[str, str]]:
+    """Return the specialized model registry."""
+    return _MODELS

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/README.md

@@ -0,0 +1,34 @@
+# nvidia/DeepSeek-V3.2-NVFP4
+
+An optimized implementation for `nvidia/DeepSeek-V3.2-NVFP4` with FP8 FlashInfer MLA on Blackwell GPUs.
+
+The main win comes from aggressively fusing ops in the attention path, across the MLA and sparse-indexer boundary, which is critical for low latency.
+On top of manual fusions, the implementation uses `torch.compile` with vLLM's custom fusion passes to fuse remaining miscellaneous ops.
+It is compatible with piecewise CUDA graphs for prefill and full CUDA graphs for decode.
+
+TP and EP are supported; PP is not.
+MTP is supported.
+
+## Usage
+
+```bash
+export VLLM_USE_SPECIALIZED_MODELS=1
+export VLLM_USE_V2_MODEL_RUNNER=1
+export TRTLLM_ENABLE_PDL=1
+
+NUM_GPUS=4
+
+# With TP
+vllm serve nvidia/DeepSeek-V3.2-NVFP4 \
+    -tp 4 \
+    --compilation-config '{"max_cudagraph_capture_size": 1024}' \
+    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \
+    --kernel-config.enable_flashinfer_autotune=False
+
+# With attention DP + MoE EP
+vllm serve nvidia/DeepSeek-V3.2-NVFP4 \
+    -dp $NUM_GPUS -ep \
+    --compilation-config '{"max_cudagraph_capture_size": 1024}' \
+    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \
+    --kernel-config.enable_flashinfer_autotune=False
+```

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V3.2 model optimized for SM100 (Blackwell)."""
+
+from .model import DeepseekV32ForCausalLM
+from .mtp import DeepSeekMTP
+
+__all__ = ["DeepseekV32ForCausalLM", "DeepSeekMTP"]

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/kernels.py

@@ -0,0 +1,931 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
+
+
+@triton.jit
+def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
+    x = x.to(tl.float32)
+    mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
+    rrms = tl.rsqrt(mean_sq + eps)
+    w = w.to(tl.float32)
+    return (x * rrms) * w
+
+
+@triton.jit
+def _fused_mtp_entry_kernel(
+    inputs_embeds_ptr,
+    inputs_embeds_stride,
+    hidden_states_ptr,
+    hidden_states_stride,
+    positions_ptr,
+    enorm_weight_ptr,
+    hnorm_weight_ptr,
+    out_ptr,
+    out_stride,
+    e_eps,
+    h_eps,
+    HIDDEN_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    tok_idx = tl.program_id(0)
+    which = tl.program_id(1)  # 0: enorm, 1: hnorm
+
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < HIDDEN_SIZE
+
+    if which == 0:
+        position = tl.load(positions_ptr + tok_idx)
+        x = tl.load(
+            inputs_embeds_ptr + tok_idx * inputs_embeds_stride + offs,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        # Mask out inputs_embeds when position == 0 (MTP convention).
+        keep = (position != 0).to(tl.float32)
+        x = x * keep
+        w = tl.load(enorm_weight_ptr + offs, mask=mask).to(tl.float32)
+        mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
+        rrms = tl.rsqrt(mean_sq + e_eps)
+        y = (x * rrms) * w
+        tl.store(
+            out_ptr + tok_idx * out_stride + offs,
+            y,
+            mask=mask,
+        )
+    else:
+        h = tl.load(
+            hidden_states_ptr + tok_idx * hidden_states_stride + offs,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        w = tl.load(hnorm_weight_ptr + offs, mask=mask).to(tl.float32)
+        mean_sq = tl.sum(h * h, axis=0) / HIDDEN_SIZE
+        rrms = tl.rsqrt(mean_sq + h_eps)
+        y = (h * rrms) * w
+        tl.store(
+            out_ptr + tok_idx * out_stride + HIDDEN_SIZE + offs,
+            y,
+            mask=mask,
+        )
+
+
+@triton.jit
+def _fused_mtp_entry_eps_kernel(
+    inputs_embeds_ptr,
+    inputs_embeds_stride,
+    hidden_states_ptr,
+    hidden_states_stride,
+    positions_ptr,
+    enorm_weight_ptr,
+    hnorm_weight_ptr,
+    e_eps_ptr,
+    h_eps_ptr,
+    out_ptr,
+    out_stride,
+    HIDDEN_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Same as _fused_mtp_entry_kernel but reads eps from 0-dim tensors."""
+    tok_idx = tl.program_id(0)
+    which = tl.program_id(1)
+
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < HIDDEN_SIZE
+
+    if which == 0:
+        position = tl.load(positions_ptr + tok_idx)
+        x = tl.load(
+            inputs_embeds_ptr + tok_idx * inputs_embeds_stride + offs,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        keep = (position != 0).to(tl.float32)
+        x = x * keep
+        w = tl.load(enorm_weight_ptr + offs, mask=mask).to(tl.float32)
+        mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
+        e_eps = tl.load(e_eps_ptr)
+        rrms = tl.rsqrt(mean_sq + e_eps)
+        y = (x * rrms) * w
+        tl.store(
+            out_ptr + tok_idx * out_stride + offs,
+            y,
+            mask=mask,
+        )
+    else:
+        h = tl.load(
+            hidden_states_ptr + tok_idx * hidden_states_stride + offs,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        w = tl.load(hnorm_weight_ptr + offs, mask=mask).to(tl.float32)
+        mean_sq = tl.sum(h * h, axis=0) / HIDDEN_SIZE
+        h_eps = tl.load(h_eps_ptr)
+        rrms = tl.rsqrt(mean_sq + h_eps)
+        y = (h * rrms) * w
+        tl.store(
+            out_ptr + tok_idx * out_stride + HIDDEN_SIZE + offs,
+            y,
+            mask=mask,
+        )
+
+
+def _fused_mtp_entry_impl(
+    inputs_embeds: torch.Tensor,
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+    enorm_weight: torch.Tensor,
+    hnorm_weight: torch.Tensor,
+    e_eps: torch.Tensor,
+    h_eps: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    num_tokens, hidden_size = inputs_embeds.shape
+    BLOCK_SIZE = triton.next_power_of_2(hidden_size)
+    _fused_mtp_entry_eps_kernel[(num_tokens, 2)](
+        inputs_embeds,
+        inputs_embeds.stride(0),
+        hidden_states,
+        hidden_states.stride(0),
+        positions,
+        enorm_weight,
+        hnorm_weight,
+        e_eps,
+        h_eps,
+        out,
+        out.stride(0),
+        HIDDEN_SIZE=hidden_size,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=8,
+    )
+    return out
+
+
+def _fused_mtp_entry_fake(
+    inputs_embeds: torch.Tensor,
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+    enorm_weight: torch.Tensor,
+    hnorm_weight: torch.Tensor,
+    e_eps: torch.Tensor,
+    h_eps: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    del (
+        inputs_embeds,
+        hidden_states,
+        positions,
+        enorm_weight,
+        hnorm_weight,
+        e_eps,
+        h_eps,
+    )
+    return out
+
+
+direct_register_custom_op(
+    op_name="fused_mtp_entry",
+    op_func=_fused_mtp_entry_impl,
+    fake_impl=_fused_mtp_entry_fake,
+    mutates_args=["out"],
+    dispatch_key=current_platform.dispatch_key,
+)
+
+
+def fused_mtp_entry(
+    inputs_embeds: torch.Tensor,
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+    enorm_weight: torch.Tensor,
+    hnorm_weight: torch.Tensor,
+    e_eps: torch.Tensor,
+    h_eps: torch.Tensor,
+) -> torch.Tensor:
+    """Fused: mask(pos==0) + enorm(embeds) | hnorm(hidden) -> concat.
+
+    Output is the concatenation [enorm(embeds), hnorm(hidden)] in the
+    last dim, ready to feed into eh_proj. `e_eps`/`h_eps` are 0-dim fp32
+    tensors (not Python floats) so the custom op stays tensor-only.
+    """
+    num_tokens, hidden_size = inputs_embeds.shape
+    out = torch.empty(
+        num_tokens,
+        hidden_size * 2,
+        dtype=inputs_embeds.dtype,
+        device=inputs_embeds.device,
+    )
+    return torch.ops.vllm.fused_mtp_entry(
+        inputs_embeds,
+        hidden_states,
+        positions,
+        enorm_weight,
+        hnorm_weight,
+        e_eps,
+        h_eps,
+        out,
+    )
+
+
+@triton.jit
+def _layer_norm(x, w, b, eps, mask, HIDDEN_SIZE: tl.constexpr):
+    x = x.to(tl.float32)
+    mean = tl.sum(x, axis=0) / HIDDEN_SIZE
+    diff = tl.where(mask, x - mean, 0.0)
+    var = tl.sum(diff * diff, axis=0) / HIDDEN_SIZE
+    rstd = tl.rsqrt(var + eps)
+
+    w = w.to(tl.float32)
+    b = b.to(tl.float32)
+    return (x - mean) * rstd * w + b
+
+
+@triton.jit
+def _rope(
+    base_ptr,
+    head_stride,
+    cos,
+    sin,
+    NUM_HEADS: tl.constexpr,
+    HALF_ROT_DIM: tl.constexpr,
+    START_OFFSET: tl.constexpr,
+    INTERLEAVED: tl.constexpr,
+):
+    head_offset = tl.arange(0, NUM_HEADS)
+    dim_offset = tl.arange(0, HALF_ROT_DIM)
+    base_ptr = base_ptr + head_offset[:, None] * head_stride + START_OFFSET
+    if INTERLEAVED:
+        x1 = tl.load(base_ptr + dim_offset * 2).to(tl.float32)
+        x2 = tl.load(base_ptr + dim_offset * 2 + 1).to(tl.float32)
+        tl.store(base_ptr + dim_offset * 2, x1 * cos - x2 * sin)
+        tl.store(base_ptr + dim_offset * 2 + 1, x2 * cos + x1 * sin)
+    else:
+        x1 = tl.load(base_ptr + dim_offset).to(tl.float32)
+        x2 = tl.load(base_ptr + dim_offset + HALF_ROT_DIM).to(tl.float32)
+        tl.store(base_ptr + dim_offset, x1 * cos - x2 * sin)
+        tl.store(base_ptr + dim_offset + HALF_ROT_DIM, x2 * cos + x1 * sin)
+
+
+@triton.jit
+def _get_cos_sin(
+    cos_sin_cache_ptr,
+    cos_sin_cache_stride,
+    pos,
+    HALF_ROT_DIM: tl.constexpr,
+):
+    block = tl.arange(0, HALF_ROT_DIM)
+    cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block)
+    cos = cos.to(tl.float32)
+    sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM)
+    sin = sin.to(tl.float32)
+    return cos, sin
+
+
+@triton.jit
+def _fp8_ue8m0_quantize(vals):
+    """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
+
+    Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
+    """
+    vals = vals.to(tl.float32)
+    amax = tl.max(tl.abs(vals))
+    scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
+    scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
+    fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv)
+    return fp8_vals, scale
+
+
+@triton.jit
+def _fp8_quant_and_cache_write(
+    vals,
+    mask,
+    slot_idx,
+    kv_cache_ptr,
+    kv_cache_scale_ptr,
+    cache_block_size,
+    cache_stride,
+    offsets,
+    HEAD_DIM: tl.constexpr,
+):
+    k_fp8, scale = _fp8_ue8m0_quantize(vals)
+
+    block_idx = slot_idx // cache_block_size
+    block_offset = slot_idx % cache_block_size
+    block_start = block_idx * cache_block_size * cache_stride
+
+    tl.store(
+        kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
+        k_fp8,
+        mask=mask,
+    )
+    scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
+    tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
+
+
+@triton.jit
+def _fused_norm_rope_kernel(
+    pos_ptr,
+    # Q RMS norm
+    q_c_ptr,
+    q_c_stride,
+    q_rms_norm_w_ptr,
+    q_rms_eps,
+    q_c_out_ptr,
+    q_c_out_stride,
+    Q_DIM: tl.constexpr,
+    Q_BLOCK_SIZE: tl.constexpr,
+    # KV RMS norm
+    kv_ptr,
+    kv_stride,
+    kv_rms_norm_w_ptr,
+    kv_rms_eps,
+    KV_DIM: tl.constexpr,
+    # KV RoPE
+    kpe_ptr,
+    kpe_stride,
+    kpe_rope_cos_sin_cache_ptr,
+    kpe_rope_cos_sin_cache_stride,
+    KPE_HALF_ROT_DIM: tl.constexpr,
+    # Index K layer norm
+    index_k_ptr,
+    index_k_stride,
+    index_k_layer_norm_w_ptr,
+    index_k_layer_norm_bias_ptr,
+    index_k_layer_norm_eps,
+    INDEX_K_DIM: tl.constexpr,
+    INDEX_K_BLOCK_SIZE: tl.constexpr,
+    # Index K RoPE
+    index_k_rope_cos_sin_cache_ptr,
+    index_k_rope_cos_sin_cache_stride,
+    INDEX_K_HALF_ROT_DIM: tl.constexpr,
+    # Index K fp32 scratch buffer for layernorm → RoPE handoff
+    index_k_normed_ptr,
+    # Cache params (shared by indexer K and MLA)
+    slot_mapping_ptr,
+    # Index K FP8 cache
+    indexer_cache_ptr,
+    indexer_cache_scale_ptr,
+    indexer_cache_block_size,
+    indexer_cache_stride,
+    # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
+    mla_cache_ptr,
+    mla_cache_block_stride,
+    mla_cache_entry_stride,
+    MLA_CACHE_FP8: tl.constexpr,
+    mla_cache_scale_ptr,
+    # Top k indices
+    topk_indices_ptr,
+    topk_indices_stride,
+    TOPK: tl.constexpr,
+    TOPK_BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    tok_idx = tl.program_id(1)
+    if pid == 3:
+        # Fill top k indices buffer with -1
+        for i in range(0, TOPK, TOPK_BLOCK_SIZE):
+            offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
+            mask = offset < TOPK
+            tl.store(
+                topk_indices_ptr + tok_idx * topk_indices_stride + offset,
+                -1,
+                mask=mask,
+            )
+        return
+
+    if slot_mapping_ptr is None:
+        # Memory profiling run.
+        return
+    slot_idx = tl.load(slot_mapping_ptr + tok_idx)
+    if slot_idx < 0:
+        # Padding
+        return
+
+    if pid == 2:
+        # Q RMS norm
+        q_block = tl.arange(0, Q_BLOCK_SIZE)
+        q_mask = q_block < Q_DIM
+        q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
+        q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
+        q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
+        tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
+    elif pid == 1:
+        # KV RMS Norm + KV RoPE + MLA concat_and_cache.
+        # Merged so the normed kv_c and RoPE'd k_pe can be written
+        # to the MLA KV cache directly without a separate kernel.
+
+        # KV RMS Norm (result stays in registers for MLA cache write)
+        kv_block = tl.arange(0, KV_DIM)
+        kv_c = tl.load(kv_ptr + tok_idx * kv_stride + kv_block)
+        kv_c_rms_w = tl.load(kv_rms_norm_w_ptr + kv_block)
+        kv_c = _rms_norm(kv_c, kv_c_rms_w, kv_rms_eps, KV_DIM)
+
+        # KV RoPE (interleaved) on k_pe — in registers only.
+        # k_pe is not needed after the cache write (MLA decode reads
+        # from kv_cache), so we skip writing back to kpe_ptr.
+        pos = tl.load(pos_ptr + tok_idx)
+        cos, sin = _get_cos_sin(
+            kpe_rope_cos_sin_cache_ptr,
+            kpe_rope_cos_sin_cache_stride,
+            pos,
+            KPE_HALF_ROT_DIM,
+        )
+        dim_off = tl.arange(0, KPE_HALF_ROT_DIM)
+        kpe_base = kpe_ptr + tok_idx * kpe_stride
+        x1 = tl.load(kpe_base + dim_off * 2).to(tl.float32)
+        x2 = tl.load(kpe_base + dim_off * 2 + 1).to(tl.float32)
+        r1 = x1 * cos - x2 * sin
+        r2 = x2 * cos + x1 * sin
+
+        # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.
+        if mla_cache_entry_stride == 0:
+            return
+
+        mla_block_size = mla_cache_block_stride // mla_cache_entry_stride
+        mla_block_idx = slot_idx // mla_block_size
+        mla_block_off = slot_idx % mla_block_size
+        dst = (
+            mla_cache_ptr
+            + mla_block_idx * mla_cache_block_stride
+            + mla_block_off * mla_cache_entry_stride
+        )
+        # kv_c_normed (KV_DIM elements)
+        if MLA_CACHE_FP8:
+            scale = tl.load(mla_cache_scale_ptr)
+            kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)
+            tl.store(dst + kv_block, kv_c_fp8)
+        else:
+            tl.store(dst + kv_block, kv_c)
+        # k_pe_roped (from registers, interleaved layout)
+        if MLA_CACHE_FP8:
+            tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))
+            tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(tl.float8e4nv))
+        else:
+            tl.store(dst + KV_DIM + dim_off * 2, r1)
+            tl.store(dst + KV_DIM + dim_off * 2 + 1, r2)
+    elif pid == 0:
+        # Fused: Index K LayerNorm + RoPE + FP8 quant + cache write.
+        # Eliminates the separate indexer_k_quant_and_cache kernel launch.
+
+        # 1. LayerNorm → fp32 temp buffer
+        index_k_block = tl.arange(0, INDEX_K_BLOCK_SIZE)
+        index_k_mask = index_k_block < INDEX_K_DIM
+        index_k = tl.load(
+            index_k_ptr + tok_idx * index_k_stride + index_k_block,
+            mask=index_k_mask,
+            other=0.0,
+        )
+        index_k_w = tl.load(index_k_layer_norm_w_ptr + index_k_block, mask=index_k_mask)
+        index_k_b = tl.load(
+            index_k_layer_norm_bias_ptr + index_k_block, mask=index_k_mask
+        )
+        normed = _layer_norm(
+            index_k,
+            index_k_w,
+            index_k_b,
+            index_k_layer_norm_eps,
+            index_k_mask,
+            INDEX_K_DIM,
+        )
+        # Write to a fp32 scratch buffer so RoPE can read the two
+        # halves without Triton pointer-aliasing issues.
+        scratch = index_k_normed_ptr + tok_idx * INDEX_K_DIM
+        tl.store(scratch + index_k_block, normed, mask=index_k_mask)
+
+        # 2. RoPE (neox / non-interleaved) on the full vector.
+        pos = tl.load(pos_ptr + tok_idx)
+        cos_full = tl.load(
+            index_k_rope_cos_sin_cache_ptr
+            + pos * index_k_rope_cos_sin_cache_stride
+            + index_k_block % INDEX_K_HALF_ROT_DIM,
+            mask=index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
+            other=1.0,
+        ).to(tl.float32)
+        sin_full = tl.load(
+            index_k_rope_cos_sin_cache_ptr
+            + pos * index_k_rope_cos_sin_cache_stride
+            + INDEX_K_HALF_ROT_DIM
+            + index_k_block % INDEX_K_HALF_ROT_DIM,
+            mask=index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
+            other=0.0,
+        ).to(tl.float32)
+        # XOR with HALF swaps the first/second half of the rotation
+        # region to get each element's partner.
+        partner_offs = tl.where(
+            index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
+            index_k_block ^ INDEX_K_HALF_ROT_DIM,
+            index_k_block,
+        )
+        full = tl.load(scratch + index_k_block, mask=index_k_mask)
+        # Atomic read for the partner: tl.atomic_add(ptr, 0) returns the
+        # current value with guaranteed store visibility, avoiding the
+        # Triton compiler's aliasing issue with different offset expressions.
+        zeros = tl.zeros([INDEX_K_BLOCK_SIZE], dtype=tl.float32)
+        partner = tl.atomic_add(scratch + partner_offs, zeros, mask=index_k_mask)
+        sign = tl.where(index_k_block < INDEX_K_HALF_ROT_DIM, -1.0, 1.0)
+        roped = full * cos_full + sign * partner * sin_full
+        result = tl.where(index_k_block < 2 * INDEX_K_HALF_ROT_DIM, roped, full)
+
+        # 3. FP8 quantize + cache write from registers.
+        #    No need to write back to index_k_ptr — the only consumer
+        #    (sparse_attn_indexer) reads from the cache, not index_k.
+        _fp8_quant_and_cache_write(
+            result,
+            index_k_mask,
+            slot_idx,
+            indexer_cache_ptr,
+            indexer_cache_scale_ptr,
+            indexer_cache_block_size,
+            indexer_cache_stride,
+            index_k_block,
+            INDEX_K_DIM,
+        )
+
+
+def fused_norm_rope(
+    positions: torch.Tensor,
+    q_c: torch.Tensor,
+    q_rms_norm_w: torch.Tensor,
+    q_rms_eps: float,
+    kv_c: torch.Tensor,
+    kv_rms_norm_w: torch.Tensor,
+    kv_rms_eps: float,
+    k_pe: torch.Tensor,
+    k_rope_cos_sin_cache: torch.Tensor,
+    index_k: torch.Tensor,
+    index_k_layer_norm_w: torch.Tensor,
+    index_k_layer_norm_bias: torch.Tensor,
+    index_k_layer_norm_eps: float,
+    index_k_rope_cos_sin_cache: torch.Tensor,
+    topk_indices_buffer: torch.Tensor,
+    # Cache params for fused writes (single slot_mapping for both caches)
+    slot_mapping: torch.Tensor | None = None,
+    indexer_k_cache: torch.Tensor | None = None,
+    mla_kv_cache: torch.Tensor | None = None,
+    mla_kv_cache_dtype: str = "auto",
+    mla_k_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    assert positions.ndim == 1
+    assert q_c.ndim == 2
+    assert kv_c.ndim == 2
+    assert k_pe.ndim == 2
+    assert index_k.ndim == 2
+    assert topk_indices_buffer.ndim == 2
+
+    num_tokens = positions.shape[0]
+    q_dim = q_c.shape[-1]
+    kv_dim = kv_c.shape[-1]
+    index_k_dim = index_k.shape[-1]
+    topk = topk_indices_buffer.shape[-1]
+    device = positions.device
+
+    # --- Indexer K cache setup ---
+    if indexer_k_cache is not None:
+        assert slot_mapping is not None
+        idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
+        idx_cache_block_size = indexer_k_cache.shape[1]
+        idx_cache_stride = indexer_k_cache.shape[2]
+        if indexer_k_cache.dtype == torch.uint8:
+            indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
+    else:
+        idx_cache_scale_view = torch.empty(0, dtype=torch.float32, device=device)
+        indexer_k_cache = torch.empty(0, dtype=torch.float8_e4m3fn, device=device)
+        slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
+        idx_cache_block_size = 1
+        idx_cache_stride = 1
+
+    # --- MLA KV cache setup ---
+    mla_cache_fp8 = mla_kv_cache_dtype != "auto"
+    if mla_kv_cache is not None:
+        mla_block_stride = mla_kv_cache.stride(0)
+        mla_entry_stride = mla_kv_cache.stride(1)
+        if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
+            mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
+        if mla_k_scale is None:
+            mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
+    else:
+        # Dummy values — pid 2 will skip the MLA cache write because
+        # slot_mapping is all -1.
+        mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)
+        mla_block_stride = 0
+        mla_entry_stride = 0
+        mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
+
+    # fp32 scratch buffer for layernorm output → RoPE handoff.
+    index_k_normed = torch.empty(
+        num_tokens, index_k_dim, dtype=torch.float32, device=device
+    )
+
+    q_c_out = torch.empty_like(q_c)
+    _fused_norm_rope_kernel[(4, num_tokens)](
+        positions,
+        # Q RMS norm
+        q_c,
+        q_c.stride(0),
+        q_rms_norm_w,
+        q_rms_eps,
+        q_c_out,
+        q_c_out.stride(0),
+        q_dim,
+        triton.next_power_of_2(q_dim),
+        # KV RMS norm
+        kv_c,
+        kv_c.stride(0),
+        kv_rms_norm_w,
+        kv_rms_eps,
+        kv_dim,
+        # KV RoPE
+        k_pe,
+        k_pe.stride(0),
+        k_rope_cos_sin_cache,
+        k_rope_cos_sin_cache.stride(0),
+        k_rope_cos_sin_cache.shape[-1] // 2,
+        # Index K layer norm + RoPE + FP8 quant
+        index_k,
+        index_k.stride(0),
+        index_k_layer_norm_w,
+        index_k_layer_norm_bias,
+        index_k_layer_norm_eps,
+        index_k_dim,
+        triton.next_power_of_2(index_k_dim),
+        index_k_rope_cos_sin_cache,
+        index_k_rope_cos_sin_cache.stride(0),
+        index_k_rope_cos_sin_cache.shape[-1] // 2,
+        index_k_normed,
+        # Cache params
+        slot_mapping,
+        indexer_k_cache,
+        idx_cache_scale_view,
+        idx_cache_block_size,
+        idx_cache_stride,
+        # MLA KV cache (uses same slot_mapping)
+        mla_kv_cache,
+        mla_block_stride,
+        mla_entry_stride,
+        mla_cache_fp8,
+        mla_k_scale,
+        # Top k indices buffer
+        topk_indices_buffer,
+        topk_indices_buffer.stride(0),
+        topk,
+        TOPK_BLOCK_SIZE=1024,
+    )
+    return q_c_out
+
+
+@triton.jit
+def _fused_q_kernel(
+    pos_ptr,
+    # MQA query PE: RoPE + FP8 pack into output tail
+    q_pe_ptr,
+    q_pe_stride0,
+    q_pe_stride1,
+    NUM_Q_HEADS: tl.constexpr,
+    q_pe_cos_sin_ptr,
+    q_pe_cos_sin_stride,
+    Q_PE_HALF_ROT_DIM: tl.constexpr,
+    # Index Q RoPE
+    index_q_ptr,
+    index_q_stride0,
+    index_q_stride1,
+    NUM_INDEX_Q_HEADS: tl.constexpr,
+    index_q_cos_sin_ptr,
+    index_q_cos_sin_stride,
+    INDEX_Q_HALF_ROT_DIM: tl.constexpr,
+    # Index Q Quantize
+    index_q_fp8_ptr,
+    index_q_fp8_stride0,
+    index_q_fp8_stride1,
+    INDEX_Q_HEAD_DIM: tl.constexpr,
+    # MQA query pack: quantize ql_nope and RoPE+quantize q_pe into mqa_q_fp8
+    ql_nope_ptr,
+    ql_nope_stride0,
+    ql_nope_stride1,
+    mqa_q_fp8_ptr,
+    mqa_q_fp8_stride0,
+    mqa_q_fp8_stride1,
+    q_scale_ptr,
+    QL_NOPE_DIM: tl.constexpr,
+    QL_NOPE_BLOCK: tl.constexpr,
+    # Index weights
+    index_weights_ptr,
+    index_weights_stride,
+    index_weights_softmax_scale,
+    index_weights_head_scale,
+    index_weights_out_ptr,
+    index_weights_out_stride,
+):
+    pid = tl.program_id(0)
+    tok_idx = tl.program_id(1)
+    head_idx = tl.program_id(2)
+
+    if pid == 2:
+        # ql_nope quantize + pack into the front of mqa_q_fp8.
+        if 2 * head_idx >= NUM_Q_HEADS:
+            return
+
+        scale = tl.load(q_scale_ptr)
+        for local_head in range(2):
+            q_head_idx = head_idx * 2 + local_head
+            if q_head_idx < NUM_Q_HEADS:
+                ql_nope_off = tl.arange(0, QL_NOPE_BLOCK)
+                ql_nope_mask = ql_nope_off < QL_NOPE_DIM
+                ql_nope = tl.load(
+                    ql_nope_ptr
+                    + tok_idx * ql_nope_stride0
+                    + q_head_idx * ql_nope_stride1
+                    + ql_nope_off,
+                    mask=ql_nope_mask,
+                ).to(tl.float32)
+                ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv)
+                tl.store(
+                    mqa_q_fp8_ptr
+                    + tok_idx * mqa_q_fp8_stride0
+                    + q_head_idx * mqa_q_fp8_stride1
+                    + ql_nope_off,
+                    ql_nope_fp8,
+                    mask=ql_nope_mask,
+                )
+        return
+    elif pid == 0:
+        # q_pe RoPE + quantize + pack into the tail of mqa_q_fp8.
+        if 2 * head_idx >= NUM_Q_HEADS:
+            return
+
+        pos = tl.load(pos_ptr + tok_idx)
+        cos, sin = _get_cos_sin(
+            q_pe_cos_sin_ptr,
+            q_pe_cos_sin_stride,
+            pos,
+            Q_PE_HALF_ROT_DIM,
+        )
+
+        scale = tl.load(q_scale_ptr)
+        for local_head in range(2):
+            q_head_idx = head_idx * 2 + local_head
+            if q_head_idx < NUM_Q_HEADS:
+                rot_off = tl.arange(0, Q_PE_HALF_ROT_DIM)
+                x1 = tl.load(
+                    q_pe_ptr
+                    + tok_idx * q_pe_stride0
+                    + q_head_idx * q_pe_stride1
+                    + rot_off * 2,
+                ).to(tl.float32)
+                x2 = tl.load(
+                    q_pe_ptr
+                    + tok_idx * q_pe_stride0
+                    + q_head_idx * q_pe_stride1
+                    + rot_off * 2
+                    + 1
+                ).to(tl.float32)
+                r1 = x1 * cos - x2 * sin
+                r2 = x2 * cos + x1 * sin
+                tl.store(
+                    mqa_q_fp8_ptr
+                    + tok_idx * mqa_q_fp8_stride0
+                    + q_head_idx * mqa_q_fp8_stride1
+                    + QL_NOPE_DIM
+                    + rot_off * 2,
+                    (r1 / scale).to(tl.float8e4nv),
+                )
+                tl.store(
+                    mqa_q_fp8_ptr
+                    + tok_idx * mqa_q_fp8_stride0
+                    + q_head_idx * mqa_q_fp8_stride1
+                    + QL_NOPE_DIM
+                    + rot_off * 2
+                    + 1,
+                    (r2 / scale).to(tl.float8e4nv),
+                )
+        return
+    elif pid == 1:
+        # Index Q RoPE
+        if head_idx >= NUM_INDEX_Q_HEADS:
+            return
+
+        pos = tl.load(pos_ptr + tok_idx)
+        cos, sin = _get_cos_sin(
+            index_q_cos_sin_ptr,
+            index_q_cos_sin_stride,
+            pos,
+            INDEX_Q_HALF_ROT_DIM,
+        )
+        _rope(
+            index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1,
+            0,
+            cos,
+            sin,
+            1,
+            INDEX_Q_HALF_ROT_DIM,
+            0,
+            False,
+        )
+
+        # Index Q Quantize
+        index_q_block = tl.arange(0, INDEX_Q_HEAD_DIM)
+        index_q = tl.load(
+            index_q_ptr
+            + tok_idx * index_q_stride0
+            + head_idx * index_q_stride1
+            + index_q_block
+        )
+
+        index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q)
+        tl.store(
+            index_q_fp8_ptr
+            + tok_idx * index_q_fp8_stride0
+            + head_idx * index_q_fp8_stride1
+            + index_q_block,
+            index_q_fp8,
+        )
+
+        # Index weights update
+        index_weights = tl.load(
+            index_weights_ptr + tok_idx * index_weights_stride + head_idx
+        )
+        index_weights = index_weights.to(tl.float32)
+        index_weights *= index_q_scale
+        index_weights *= index_weights_softmax_scale
+        index_weights *= index_weights_head_scale
+        tl.store(
+            index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
+            index_weights,
+        )
+
+
+def fused_q(
+    positions: torch.Tensor,
+    q_pe: torch.Tensor,
+    q_pe_cos_sin_cache: torch.Tensor,
+    index_q: torch.Tensor,
+    index_q_cos_sin_cache: torch.Tensor,
+    ql_nope: torch.Tensor,
+    q_scale: torch.Tensor,
+    # Index weights
+    index_weights: torch.Tensor,
+    index_weights_softmax_scale: float,
+    index_weights_head_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    assert positions.ndim == 1
+    assert q_pe.ndim == 3
+    assert q_pe_cos_sin_cache.ndim == 2
+    assert index_q.ndim == 3
+    assert index_q_cos_sin_cache.ndim == 2
+
+    num_tokens = positions.shape[0]
+    num_q_heads = q_pe.shape[1]
+    num_index_q_heads = index_q.shape[1]
+    index_q_head_dim = index_q.shape[2]
+    assert ql_nope.ndim == 3
+    assert ql_nope.shape[:2] == q_pe.shape[:2]
+    mqa_q_fp8 = torch.empty(
+        q_pe.shape[0],
+        q_pe.shape[1],
+        ql_nope.shape[2] + q_pe.shape[2],
+        dtype=torch.float8_e4m3fn,
+        device=q_pe.device,
+    )
+
+    index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
+    index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
+    _fused_q_kernel[(3, num_tokens, num_index_q_heads)](
+        positions,
+        q_pe,
+        q_pe.stride(0),
+        q_pe.stride(1),
+        num_q_heads,
+        q_pe_cos_sin_cache,
+        q_pe_cos_sin_cache.stride(0),
+        q_pe_cos_sin_cache.shape[-1] // 2,
+        index_q,
+        index_q.stride(0),
+        index_q.stride(1),
+        num_index_q_heads,
+        index_q_cos_sin_cache,
+        index_q_cos_sin_cache.stride(0),
+        index_q_cos_sin_cache.shape[-1] // 2,
+        index_q_fp8,
+        index_q_fp8.stride(0),
+        index_q_fp8.stride(1),
+        index_q_head_dim,
+        ql_nope,
+        ql_nope.stride(0),
+        ql_nope.stride(1),
+        mqa_q_fp8,
+        mqa_q_fp8.stride(0),
+        mqa_q_fp8.stride(1),
+        q_scale,
+        ql_nope.shape[2],
+        triton.next_power_of_2(ql_nope.shape[2]),
+        index_weights,
+        index_weights.stride(0),
+        index_weights_softmax_scale,
+        index_weights_head_scale,
+        index_weights_out,
+        index_weights_out.stride(0),
+        num_warps=1,  # TODO: Tune this
+    )
+    return index_q_fp8, index_weights_out, mqa_q_fp8

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/layer.py

@@ -0,0 +1,570 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+MLA attention and decoder layer for DeepSeek V3.2 on SM100 (Blackwell).
+
+MLAAttention:
+  KV cache update -> W_UK_T absorption -> sparse attn kernel -> W_UV up-proj
+  MLAAttention kept only as a registration stub for KV cache / backend.
+
+DecoderLayer:
+  Single decoder layer: norm -> attn -> norm -> MoE/MLP.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.attention.mla_attention import MLAAttention
+from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
+from vllm.model_executor.models.deepseek_v2 import (
+    DeepseekV32IndexerCache,
+    yarn_get_mscale,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
+
+from .kernels import fused_norm_rope, fused_q
+from .sparse_indexer import sparse_attn_indexer
+
+
+def dsa(
+    positions: torch.Tensor,
+    q_c: torch.Tensor,
+    kv_c: torch.Tensor,
+    k_pe: torch.Tensor,
+    index_k: torch.Tensor,
+    index_weights: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    layer = get_forward_context().no_compile_layers[layer_name]
+    attn = layer.attn
+    mla = attn.mla_attn
+
+    attn_metadata = get_forward_context().attn_metadata
+    if not isinstance(attn_metadata, dict):
+        output.zero_()
+        return output
+
+    mla_attn_metadata = attn_metadata.get(mla.layer_name)
+    if mla_attn_metadata is None:
+        output.zero_()
+        return output
+
+    num_actual_toks = mla_attn_metadata.num_actual_tokens  # type: ignore[attr-defined]
+    if num_actual_toks == 0:
+        output.zero_()
+        return output
+
+    # Step 2. fused norm + rope + cache writes
+    slot_mapping = None
+    indexer_k_cache = None
+    mla_kv_cache = None
+    mla_k_scale = None
+    idx_meta = attn_metadata.get(attn.indexer_k_cache.prefix)
+    if idx_meta is not None:
+        slot_mapping = idx_meta.slot_mapping  # type: ignore[attr-defined]
+        indexer_k_cache = attn.indexer_k_cache.kv_cache
+        mla_kv_cache = attn.mla_attn.kv_cache
+        mla_k_scale = attn.mla_attn._k_scale
+
+    q_c = fused_norm_rope(
+        positions,
+        q_c,
+        attn.q_a_layernorm_weight,
+        layer.rms_norm_eps,
+        kv_c,
+        attn.kv_a_layernorm_weight,
+        attn.rms_norm_eps,
+        k_pe,
+        attn.rotary_emb.cos_sin_cache,
+        index_k,
+        attn.indexer_k_norm.weight,
+        attn.indexer_k_norm.bias,
+        attn.rms_norm_eps,
+        attn.indexer_rope_emb.cos_sin_cache,
+        attn.topk_indices_buffer,
+        slot_mapping=slot_mapping,
+        indexer_k_cache=indexer_k_cache,
+        mla_kv_cache=mla_kv_cache,
+        mla_kv_cache_dtype=attn.mla_attn.kv_cache_dtype,
+        mla_k_scale=mla_k_scale,
+    )
+
+    # Step 3. q_c -> index_q, q
+    step3_out = torch.mm(q_c, layer._fused_step3_q_w.T)
+    index_q, q = step3_out.split(layer._q_split_sizes, dim=-1)
+    index_q = index_q.view(-1, attn.index_n_heads, attn.index_head_dim)
+    q = q.view(-1, attn.num_local_heads, attn.qk_head_dim)
+
+    # Step 4. Q RoPE + W_UK_T absorption + FP8 packing
+    q_nope, q_pe = q.split(
+        [mla.qk_nope_head_dim, mla.qk_rope_head_dim],
+        dim=-1,
+    )
+    q_nope = q_nope.transpose(0, 1)
+    ql_nope = torch.bmm(q_nope, mla.W_UK_T)
+    ql_nope = ql_nope.transpose(0, 1)
+
+    index_q_fp8, index_weights, mqa_q = fused_q(
+        positions,
+        q_pe,
+        attn.rotary_emb.cos_sin_cache,
+        index_q,
+        attn.indexer_rope_emb.cos_sin_cache,
+        ql_nope,
+        mla._q_scale,
+        index_weights,
+        attn.indexer_softmax_scale,
+        attn.index_n_heads**-0.5,
+    )
+
+    # Steps 5-6. Sparse indexer + MLA sparse decode attention
+    sparse_attn_indexer(
+        attn.indexer_k_cache.prefix,
+        attn.indexer_k_cache.kv_cache,
+        index_q_fp8,
+        index_weights,
+        attn.topk_tokens,
+        attn.index_head_dim,
+        layer.max_model_len,
+        layer.indexer_workspace_size,
+        attn.topk_indices_buffer,
+    )
+
+    mqa_q = mqa_q[:num_actual_toks]
+    kv_cache = mla.kv_cache
+    if mla.kv_cache_dtype.startswith("fp8") and mla.kv_cache_dtype != "fp8_ds_mla":
+        kv_cache = kv_cache.view(torch.float8_e4m3fn)
+    attn_out, _ = mla.impl.forward_mqa(mqa_q, kv_cache, mla_attn_metadata, mla)
+    x = attn_out.view(-1, mla.num_heads, mla.kv_lora_rank).transpose(0, 1)
+
+    out = output[:num_actual_toks].view(-1, mla.num_heads, mla.v_head_dim)
+    out = out.transpose(0, 1)
+    torch.bmm(x, mla.W_UV, out=out)
+    return output
+
+
+def dsa_fake(
+    positions: torch.Tensor,
+    q_c: torch.Tensor,
+    kv_c: torch.Tensor,
+    k_pe: torch.Tensor,
+    index_k: torch.Tensor,
+    index_weights: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    del positions, q_c, kv_c, k_pe, index_k, index_weights, layer_name
+    return output
+
+
+direct_register_custom_op(
+    op_name="monolithic_attn",
+    op_func=dsa,
+    fake_impl=dsa_fake,
+    mutates_args=["output"],
+    dispatch_key=current_platform.dispatch_key,
+)
+
+
+class DeepseekV32DecoderLayer(nn.Module):
+    """
+    Single decoder layer: norm -> attn -> norm -> MoE/MLP.
+    Norms are raw weight + direct kernel call.
+    Gate inlined as raw weight, experts kept as FusedMoE for quantization.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config,
+        layer_idx: int,
+        topk_indices_buffer: torch.Tensor,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+        self.layer_name = prefix
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.rms_norm_eps = config.rms_norm_eps
+        self.q_lora_rank = config.q_lora_rank
+        self.kv_lora_rank = config.kv_lora_rank
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        parallel_config = vllm_config.parallel_config
+        self.indexer_workspace_size = get_max_prefill_buffer_size(vllm_config)
+        self.max_model_len = vllm_config.model_config.max_model_len
+
+        # Use the regular vLLM RMSNorm modules so the compiler sees the
+        # canonical residual-add + RMSNorm pattern.
+        dtype = torch.get_default_dtype()
+        self.input_layernorm = RMSNorm(
+            hidden_size=config.hidden_size,
+            eps=config.rms_norm_eps,
+            dtype=dtype,
+        )
+        self.post_attention_layernorm = RMSNorm(
+            hidden_size=config.hidden_size,
+            eps=config.rms_norm_eps,
+            dtype=dtype,
+        )
+
+        # Fused QKV A-projection lives inside self_attn namespace
+        # for weight loading compatibility with original checkpoint paths
+        from vllm.model_executor.models.deepseek_v2 import (
+            DeepSeekV2FusedQkvAProjLinear,
+        )
+
+        self.self_attn = nn.Module()
+        self.self_attn.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
+            config.hidden_size,
+            [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn.fused_qkv_a_proj",
+        )
+
+        # MLA Attention
+        self.attn = DeepseekV32MLAAttention(
+            vllm_config=vllm_config,
+            config=config,
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            qk_nope_head_dim=config.qk_nope_head_dim,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            v_head_dim=config.v_head_dim,
+            q_lora_rank=self.q_lora_rank,
+            kv_lora_rank=self.kv_lora_rank,
+            max_position_embeddings=getattr(config, "max_position_embeddings", 8192),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            topk_indices_buffer=topk_indices_buffer,
+            prefix=f"{prefix}.self_attn",
+        )
+
+        # MoE or Dense MLP
+        moe_layer_freq = getattr(config, "moe_layer_freq", 1)
+        self.is_moe = (
+            config.n_routed_experts is not None
+            and layer_idx >= config.first_k_dense_replace
+            and layer_idx % moe_layer_freq == 0
+        )
+        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
+
+        from vllm.model_executor.models.deepseek_v2 import (
+            DeepseekV2MLP,
+            DeepseekV2MoE,
+        )
+
+        if self.is_moe:
+            self.mlp = DeepseekV2MoE(
+                config=config,
+                parallel_config=parallel_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            self.mlp = DeepseekV2MLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        # Step 1. hidden_states -> q_c, kv_c, k_pe, index_k, index_weights
+        step1_out = torch.mm(hidden_states, self._fused_step1_hidden_w.T)
+        q_c, kv_c, k_pe, index_k, index_weights = step1_out.split(
+            self._step1_split_sizes,
+            dim=-1,
+        )
+
+        # Steps 2-6. Combined: fused norm/rope + Q projections + sparse MLA.
+        mla = self.attn.mla_attn
+        output_shape = (hidden_states.shape[0], mla.num_heads * mla.v_head_dim)
+        output_dtype = mla.W_UV.dtype
+        attn_out = torch.empty(
+            output_shape,
+            dtype=output_dtype,
+            device=hidden_states.device,
+        )
+        attn_out = torch.ops.vllm.monolithic_attn(
+            positions,
+            q_c,
+            kv_c,
+            k_pe,
+            index_k,
+            index_weights,
+            attn_out,
+            self.layer_name,
+        )
+
+        hidden_states, _ = self.attn.o_proj(attn_out)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+    def fuse_indexer_weights(self) -> None:
+        """Fuse Step 1 and Step 3 BF16 linears used by the inlined path.
+
+        Call after model weights are loaded.
+        """
+        attn = self.attn
+        qkv_a = self.self_attn.fused_qkv_a_proj.weight.data  # [2112, 7168]
+        wk = attn.indexer_wk.weight.data  # [128, 7168]
+        wp = attn.indexer_weights_proj.weight.data  # [64, 7168]
+        if not (qkv_a.dtype == wk.dtype == wp.dtype):
+            raise ValueError(
+                "Cannot fuse Step 1 weights: expected matching dtypes for "
+                "fused_qkv_a_proj, indexer_wk, and indexer_weights_proj."
+            )
+        self._fused_step1_hidden_w = nn.Parameter(
+            torch.cat([qkv_a, wk, wp], dim=0),  # [2304, 7168]
+            requires_grad=False,
+        )
+        self._step1_split_sizes = [
+            self.q_lora_rank,
+            self.kv_lora_rank,
+            self.qk_rope_head_dim,
+            wk.shape[0],
+            wp.shape[0],
+        ]
+
+        wq_b = attn.indexer_wq_b.weight.data
+        q_b = attn.q_b_proj.weight.data
+        if wq_b.dtype != q_b.dtype:
+            raise ValueError(
+                "Cannot fuse Step 3 weights: expected matching dtypes for "
+                "indexer_wq_b and q_b_proj."
+            )
+        self._fused_step3_q_w = nn.Parameter(
+            torch.cat([wq_b, q_b], dim=0),
+            requires_grad=False,
+        )
+        self._q_split_sizes = [wq_b.shape[0], q_b.shape[0]]
+
+
+class DeepseekV32MLAAttention(nn.Module):
+    """
+    MLA attention for DeepSeek V3.2 targeting SM100.
+    MLA forward fully inlined. MLAAttention kept only for KV cache
+    registration and backend/impl initialization.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config,
+        hidden_size: int,
+        num_heads: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        q_lora_rank: int,
+        kv_lora_rank: int,
+        max_position_embeddings: int,
+        cache_config: CacheConfig,
+        quant_config: QuantizationConfig | None,
+        topk_indices_buffer: torch.Tensor,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.num_heads = num_heads
+        self.num_local_heads = num_heads // get_tensor_model_parallel_world_size()
+        self.scaling = self.qk_head_dim**-0.5
+        self.rms_norm_eps = config.rms_norm_eps
+
+        # Q path
+        self.q_a_layernorm_weight = nn.Parameter(
+            torch.ones(q_lora_rank, dtype=torch.get_default_dtype())
+        )
+        self.q_b_proj = ColumnParallelLinear(
+            q_lora_rank,
+            num_heads * self.qk_head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.q_b_proj",
+        )
+
+        # KV path
+        self.kv_a_layernorm_weight = nn.Parameter(
+            torch.ones(kv_lora_rank, dtype=torch.get_default_dtype())
+        )
+        self.kv_b_proj = ColumnParallelLinear(
+            kv_lora_rank,
+            num_heads * (qk_nope_head_dim + v_head_dim),
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_b_proj",
+        )
+
+        # Output projection (TP sync point)
+        self.o_proj = RowParallelLinear(
+            num_heads * v_head_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        # RoPE
+        if config.rope_parameters["rope_type"] != "default":
+            config.rope_parameters["rope_type"] = (
+                "deepseek_yarn"
+                if config.rope_parameters.get("apply_yarn_scaling", True)
+                else "deepseek_llama_scaling"
+            )
+        self.rotary_emb = get_rope(
+            qk_rope_head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=False,
+        )
+        if config.rope_parameters["rope_type"] == "deepseek_yarn":
+            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
+            scaling_factor = config.rope_parameters["factor"]
+            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
+            self.scaling = self.scaling * mscale * mscale
+
+        # V3.2 Sparse Indexer (inlined)
+        self.indexer_rope_emb = get_rope(
+            qk_rope_head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=not getattr(config, "indexer_rope_interleave", False),
+        )
+        self.topk_tokens = config.index_topk
+        self.index_n_heads = config.index_n_heads
+        self.index_head_dim = config.index_head_dim
+        self.indexer_softmax_scale = config.index_head_dim**-0.5
+        self.indexer_quant_block_size = 128
+        self.topk_indices_buffer = topk_indices_buffer
+
+        self.indexer_wq_b = ReplicatedLinear(
+            q_lora_rank,
+            config.index_head_dim * config.index_n_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.indexer.wq_b",
+        )
+        self.indexer_wk = ReplicatedLinear(
+            hidden_size,
+            config.index_head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.indexer.wk",
+        )
+        self.indexer_k_norm = LayerNorm(config.index_head_dim, eps=1e-6)
+        self.indexer_weights_proj = ReplicatedLinear(
+            hidden_size,
+            config.index_n_heads,
+            bias=False,
+            quant_config=None,
+            prefix=f"{prefix}.indexer.weights_proj",
+        )
+
+        idx_dim = config.index_head_dim
+        indexer_cache_head_dim = idx_dim + idx_dim // 128 * 4
+        self.indexer_k_cache = DeepseekV32IndexerCache(
+            head_dim=indexer_cache_head_dim,
+            dtype=torch.uint8,
+            prefix=f"{prefix}.indexer.k_cache",
+            cache_config=cache_config,
+        )
+        self.indexer_op = SparseAttnIndexer(
+            self.indexer_k_cache,
+            self.indexer_quant_block_size,
+            "ue8m0",
+            self.topk_tokens,
+            config.index_head_dim,
+            vllm_config.model_config.max_model_len,
+            get_max_prefill_buffer_size(vllm_config),
+            self.topk_indices_buffer,
+        )
+
+        # MLAAttention stub: only for KV cache registration + backend init.
+        # We never call its forward(); we inline everything below.
+        class _IndexerProxy:
+            def __init__(proxy_self):
+                proxy_self.topk_indices_buffer = topk_indices_buffer
+                proxy_self.indexer_op = self.indexer_op
+
+        self._indexer_proxy = _IndexerProxy()
+        self.mla_attn = MLAAttention(
+            num_heads=self.num_local_heads,
+            scale=self.scaling,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            q_lora_rank=q_lora_rank,
+            kv_lora_rank=kv_lora_rank,
+            kv_b_proj=self.kv_b_proj,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mla_attn",
+            use_sparse=True,
+            indexer=self._indexer_proxy,
+        )
+
+
+def remap_weight_name(name: str) -> str:
+    """Remap checkpoint names that differ from the module layout."""
+    replacements = [
+        (
+            "self_attn.q_a_layernorm.weight",
+            "attn.q_a_layernorm_weight",
+        ),
+        (
+            "self_attn.kv_a_layernorm.weight",
+            "attn.kv_a_layernorm_weight",
+        ),
+        ("self_attn.q_b_proj", "attn.q_b_proj"),
+        ("self_attn.kv_b_proj", "attn.kv_b_proj"),
+        ("self_attn.o_proj", "attn.o_proj"),
+        ("self_attn.indexer.", "attn.indexer_"),
+    ]
+    for old, new in replacements:
+        if old in name:
+            return name.replace(old, new)
+    return name

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/model.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V3.2 NVFP4 model for SM100 (Blackwell)."""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.platforms import current_platform
+
+from .layer import DeepseekV32DecoderLayer, remap_weight_name
+
+logger = init_logger(__name__)
+
+
+@support_torch_compile
+class DeepseekV32Model(nn.Module):
+    fall_back_to_pt_during_load = False
+
+    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.device = current_platform.device_type
+
+        topk_tokens = config.index_topk
+        self.topk_indices_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            topk_tokens,
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.embed_tokens",
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                DeepseekV32DecoderLayer(
+                    vllm_config=vllm_config,
+                    config=config,
+                    layer_idx=i,
+                    topk_indices_buffer=self.topk_indices_buffer,
+                    prefix=f"{prefix}.layers.{i}",
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+
+        self.norm = RMSNorm(
+            hidden_size=config.hidden_size,
+            eps=config.rms_norm_eps,
+            dtype=torch.get_default_dtype(),
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = self.embed_tokens(input_ids)
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(positions, hidden_states, residual)
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+
+class DeepseekV32ForCausalLM(nn.Module):
+    packed_modules_mapping = {
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
+    }
+
+    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.tp_size = get_tensor_model_parallel_world_size()
+
+        self.model = DeepseekV32Model(
+            vllm_config=vllm_config,
+            prefix=f"{prefix}.model" if prefix else "model",
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.lm_head" if prefix else "lm_head",
+        )
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.num_redundant_experts = 0
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors=None,
+        inputs_embeds=None,
+    ) -> torch.Tensor:
+        return self.model(input_ids, positions)
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Delegate to the original DeepSeek V2 weight loader.
+
+        Our module structure matches the original for all weights that
+        need special loading (fused_qkv_a_proj, experts, gate_up_proj).
+        Only layernorm weights and indexer paths differ.
+        """
+        from vllm.model_executor.models.deepseek_v2 import (
+            DeepseekV2ForCausalLM,
+        )
+
+        def _remap_weights():
+            for name, w in weights:
+                yield remap_weight_name(name), w
+
+        self.use_mha = False
+        self.fuse_qkv_a_proj = True
+        self.is_fp4_ckpt = False
+        loaded = DeepseekV2ForCausalLM.load_weights(self, _remap_weights())
+
+        # Fuse indexer linear weights after loading.
+        for layer in self.model.layers:
+            layer.fuse_indexer_weights()
+
+        return loaded

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/mtp.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V3.2 MTP model for SM100 (Blackwell)."""
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepSeekMTPBase
+from vllm.model_executor.models.deepseek_mtp import (
+    DeepSeekMultiTokenPredictor as DeepSeekMultiTokenPredictorBase,
+)
+from vllm.model_executor.models.deepseek_mtp import (
+    DeepSeekMultiTokenPredictorLayer as DeepSeekMultiTokenPredictorLayerBase,
+)
+from vllm.model_executor.models.deepseek_mtp import SharedHead as SharedHeadBase
+from vllm.model_executor.models.deepseek_v2 import DeepseekV2MoE
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+
+from .kernels import fused_mtp_entry
+from .layer import DeepseekV32DecoderLayer
+from .model import remap_weight_name
+
+logger = init_logger(__name__)
+
+
+class SharedHead(SharedHeadBase):
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return rms_norm(hidden_states, self.norm.weight, self.norm.variance_epsilon)
+
+
+class DeepSeekMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayerBase):
+    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
+        nn.Module.__init__(self)
+
+        assert vllm_config.speculative_config is not None
+        config = vllm_config.speculative_config.draft_model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        self.config = config
+        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
+
+        topk_indices_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            config.index_topk,
+            dtype=torch.int32,
+            device=current_platform.device_type,
+        )
+
+        self.shared_head = SharedHead(
+            config=config, prefix=prefix, quant_config=quant_config
+        )
+        self.mtp_block = DeepseekV32DecoderLayer(
+            vllm_config=vllm_config,
+            config=config,
+            layer_idx=int(prefix.rsplit(".", 1)[-1]),
+            topk_indices_buffer=topk_indices_buffer,
+            prefix=prefix,
+        )
+        # Pre-allocated 0-dim eps tensors so fused_mtp_entry can stay
+        # tensor-only (avoids Python-float scalars leaking into the
+        # torch.compile input list).
+        self._e_eps_gpu = torch.full(
+            (),
+            self.enorm.variance_epsilon,
+            dtype=torch.float32,
+            device=current_platform.device_type,
+        )
+        self._h_eps_gpu = torch.full(
+            (),
+            self.hnorm.variance_epsilon,
+            dtype=torch.float32,
+            device=current_platform.device_type,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_index: int = 0,
+    ) -> torch.Tensor:
+        assert inputs_embeds is not None
+        eh_concat = fused_mtp_entry(
+            inputs_embeds,
+            previous_hidden_states,
+            positions,
+            self.enorm.weight,
+            self.hnorm.weight,
+            self._e_eps_gpu,
+            self._h_eps_gpu,
+        )
+        hidden_states = self.eh_proj(eh_concat)
+        hidden_states, residual = self.mtp_block(
+            positions=positions, hidden_states=hidden_states, residual=None
+        )
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class DeepSeekMultiTokenPredictor(DeepSeekMultiTokenPredictorBase):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        config = vllm_config.model_config.hf_config
+        self.mtp_start_layer_idx = config.num_hidden_layers
+        self.num_mtp_layers = config.num_nextn_predict_layers
+
+        self.layers = torch.nn.ModuleDict(
+            {
+                str(idx): DeepSeekMultiTokenPredictorLayer(
+                    vllm_config, f"{prefix}.layers.{idx}"
+                )
+                for idx in range(
+                    self.mtp_start_layer_idx,
+                    self.mtp_start_layer_idx + self.num_mtp_layers,
+                )
+            }
+        )
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+
+
+@support_torch_compile
+class DeepSeekMTP(DeepSeekMTPBase):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        self.config = vllm_config.model_config.hf_config
+        self.quant_config = vllm_config.quant_config
+        assert hasattr(self.config, "index_topk")
+        cache_config = vllm_config.cache_config
+        if cache_config.cache_dtype == "bfloat16":
+            cache_config.cache_dtype = "auto"
+            logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
+        self.model = DeepSeekMultiTokenPredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        self.set_moe_parameters()
+        # Keep the original loader from applying the fused FP4 indexer remap.
+        self.is_fp4_ckpt = False
+
+    def set_moe_parameters(self):
+        self.expert_weights = []
+        self.num_moe_layers = self.config.num_nextn_predict_layers
+        self.num_expert_groups = self.config.n_group
+
+        self.moe_layers = []
+        self.moe_mlp_layers = []
+        example_moe = None
+        for layer in self.model.layers.values():
+            layer = layer.mtp_block
+            assert isinstance(layer, DeepseekV32DecoderLayer)
+            if isinstance(layer.mlp, DeepseekV2MoE):
+                example_moe = layer.mlp
+                self.moe_mlp_layers.append(layer.mlp)
+                self.moe_layers.append(layer.mlp.experts)
+        self.extract_moe_parameters(example_moe)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        del intermediate_tensors
+        return self.model(
+            input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded_params = super().load_weights(weights)
+        for layer in self.model.layers.values():
+            layer.mtp_block.fuse_indexer_weights()
+        return loaded_params
+
+    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
+        name = super()._rewrite_spec_layer_name(spec_layer, name)
+        return remap_weight_name(name)
+
+
+@torch.compile
+def rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
+    orig_dtype = x.dtype
+    x = x.to(torch.float32)
+    mean_sq = (x * x).mean(dim=-1, keepdim=True)
+    rrms = torch.rsqrt(mean_sq + eps)
+    x = x * rrms
+    x = x * w.to(torch.float32)
+    return x.to(orig_dtype)

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/sparse_indexer.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Custom Sparse Attention Indexer layers."""
+
+import torch
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
+from vllm.utils.torch_utils import (
+    LayerNameType,
+    _resolve_layer_name,
+)
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV32IndexerMetadata,
+)
+from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
+from vllm.v1.worker.workspace import current_workspace_manager
+
+logger = init_logger(__name__)
+
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
+
+
+def sparse_attn_indexer(
+    k_cache_prefix: LayerNameType,
+    kv_cache: torch.Tensor,
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    topk_tokens: int,
+    head_dim: int,
+    max_model_len: int,
+    total_seq_lens: int,
+    topk_indices_buffer: torch.Tensor,
+) -> torch.Tensor:
+    # careful! this will be None in dummy run
+    attn_metadata = get_forward_context().attn_metadata
+    fp8_dtype = current_platform.fp8_dtype()
+    k_cache_prefix = _resolve_layer_name(k_cache_prefix)
+
+    # assert isinstance(attn_metadata, dict)
+    if not isinstance(attn_metadata, dict):
+        # Reserve workspace for indexer during profiling run
+        current_workspace_manager().get_simultaneous(
+            ((total_seq_lens, head_dim), torch.float8_e4m3fn),
+            ((total_seq_lens, 4), torch.uint8),
+            ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+        )
+
+        # Dummy allocation to simulate for peak logits tensor memory during inference.
+        # FP8 elements so elements == bytes
+        max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+        _ = torch.empty(max_logits_elems, dtype=torch.uint8, device=q_fp8.device)
+        return None
+
+    attn_metadata = attn_metadata[k_cache_prefix]  # type: ignore[assignment]
+    assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
+    has_decode = attn_metadata.num_decodes > 0
+    has_prefill = attn_metadata.num_prefills > 0
+    num_decode_tokens = attn_metadata.num_decode_tokens
+
+    if has_prefill:
+        prefill_metadata = attn_metadata.prefill
+        assert prefill_metadata is not None
+
+        # Get the full shared workspace buffers once (will allocate on first use)
+        workspace_manager = current_workspace_manager()
+        k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
+            ((total_seq_lens, head_dim), fp8_dtype),
+            ((total_seq_lens, 4), torch.uint8),
+        )
+        for chunk in prefill_metadata.chunks:
+            k_fp8 = k_fp8_full[: chunk.total_seq_lens]
+            k_scale = k_scale_full[: chunk.total_seq_lens]
+
+            if not chunk.skip_kv_gather:
+                ops.cp_gather_indexer_k_quant_cache(
+                    kv_cache,
+                    k_fp8,
+                    k_scale,
+                    chunk.block_table,
+                    chunk.cu_seq_lens,
+                )
+
+            logits = fp8_mqa_logits(
+                q_fp8[chunk.token_start : chunk.token_end],
+                (k_fp8, k_scale.view(torch.float32).flatten()),
+                weights[chunk.token_start : chunk.token_end],
+                chunk.cu_seqlen_ks,
+                chunk.cu_seqlen_ke,
+                clean_logits=False,
+            )
+            num_rows = logits.shape[0]
+
+            topk_indices = topk_indices_buffer[
+                chunk.token_start : chunk.token_end, :topk_tokens
+            ]
+
+            torch.ops._C.top_k_per_row_prefill(
+                logits,
+                chunk.cu_seqlen_ks,
+                chunk.cu_seqlen_ke,
+                topk_indices,
+                num_rows,
+                logits.stride(0),
+                logits.stride(1),
+                topk_tokens,
+            )
+
+    if has_decode:
+        decode_metadata = attn_metadata.decode
+        assert decode_metadata is not None
+        # kv_cache shape [
+        # kv_cache size requirement [num_block, block_size, n_head, head_dim],
+        # we only have [num_block, block_size, head_dim],
+        kv_cache = kv_cache.unsqueeze(-2)
+        decode_lens = decode_metadata.decode_lens
+        if decode_metadata.requires_padding:
+            # pad in edge case where we have short chunked prefill length <
+            # decode_threshold since we unstrictly split
+            # prefill and decode by decode_threshold
+            # (currently set to 1 + speculative tokens)
+            padded_q_fp8_decode_tokens = pack_seq_triton(
+                q_fp8[:num_decode_tokens], decode_lens
+            )
+        else:
+            padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
+                decode_lens.shape[0], -1, *q_fp8.shape[1:]
+            )
+        # TODO: move and optimize below logic with triton kernels
+        batch_size = padded_q_fp8_decode_tokens.shape[0]
+        next_n = padded_q_fp8_decode_tokens.shape[1]
+        num_padded_tokens = batch_size * next_n
+        seq_lens = decode_metadata.seq_lens[:batch_size]
+        # seq_lens is (B, next_n) for native spec decode, (B,) otherwise.
+        # fp8_paged_mqa_logits and all topk kernels accept both shapes.
+        logits = fp8_paged_mqa_logits(
+            padded_q_fp8_decode_tokens,
+            kv_cache,
+            weights[:num_padded_tokens],
+            seq_lens,
+            decode_metadata.block_table,
+            decode_metadata.schedule_metadata,
+            max_model_len=max_model_len,
+            clean_logits=False,
+        )
+        num_rows = logits.shape[0]
+        topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
+
+        workspace_manager = current_workspace_manager()
+        (topk_workspace,) = workspace_manager.get_simultaneous(
+            ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+        )
+        torch.ops._C.persistent_topk(
+            logits,
+            seq_lens,
+            topk_indices,
+            topk_workspace,
+            topk_tokens,
+            attn_metadata.max_seq_len,
+        )
+
+        if decode_metadata.requires_padding:
+            # if padded, we need to unpack
+            # the topk indices removing padded tokens
+            topk_indices = unpack_seq_triton(
+                topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
+                decode_lens,
+            )
+            topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = (
+                topk_indices
+            )

vllm/v1/attention/backends/mla/indexer.py

@@ -8,6 +8,7 @@ import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 from vllm.utils.deep_gemm import (
     get_paged_mqa_logits_metadata,
     has_deep_gemm,
@@ -30,6 +31,40 @@ from vllm.v1.worker.cp_utils import get_total_cp_world_size
 logger = init_logger(__name__)
 
 
+@triton.jit
+def _prepare_uniform_decode_kernel(
+    seq_lens_ptr,
+    decode_seq_lens_ptr,
+    block_table_ptr,
+    block_table_stride,
+    expanded_block_table_ptr,
+    expanded_bt_stride,
+    decode_lens_ptr,
+    max_decode_len,
+    BLOCK_SIZE: tl.constexpr,
+):
+    idx = tl.program_id(0)
+    req_id = idx // max_decode_len
+    local_idx = idx % max_decode_len
+
+    # Compute number of KVs attended to by this token.
+    seq_len = tl.load(seq_lens_ptr + req_id)
+    per_token_seq_len = seq_len - max_decode_len + local_idx + 1
+    tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
+
+    # Copy block table row.
+    src = block_table_ptr + req_id * block_table_stride
+    dst = expanded_block_table_ptr + idx * expanded_bt_stride
+    for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
+        off = i + tl.arange(0, BLOCK_SIZE)
+        mask = off < expanded_bt_stride
+        src_block = tl.load(src + off, mask=mask)
+        tl.store(dst + off, src_block, mask=mask)
+
+    # All reqs now have decode_len = 1.
+    tl.store(decode_lens_ptr + idx, 1)
+
+
 def split_indexer_prefill_chunks(
     seq_lens_cpu: torch.Tensor,
     query_lens_cpu: torch.Tensor,
@@ -405,52 +440,75 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).
         seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, next_n) for native MTP.
         """
+        min_decode_len = int(decode_lens_cpu.min().item())
         if not use_native and max_decode_len > 1:
             assert self.decode_seq_lens_buffer.dim() == 1
-            # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
-            # padding) and decode_lens [3, 1, 4, 0] in the below example comments.
-            # The context lengths are therefore
-            # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
-
-            # 3 + 1 + 4 + 0 = 8
-            actual_expanded = int(decode_lens_cpu.sum().item())
-
-            # Fuse expanded_base and expanded_starts into a single repeat_interleave:
-            # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
-            # where context_start[b] = seq_lens[b] - decode_lens[b].
-            # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
-            # expanded_offsets  = [7, 7, 7, 3, 4, 4, 4, 4]
-            # result            = [8, 9, 10, 7, 9, 10, 11, 12]
-            expanded_offsets = torch.repeat_interleave(
-                seq_lens - decode_lens - query_start_loc,
-                decode_lens,
-                output_size=actual_expanded,
-            )
-
-            # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
-            self.decode_seq_lens_buffer[:actual_expanded] = (
-                expanded_offsets + self.arange_buffer[:actual_expanded] + 1
-            )
-            self.decode_seq_lens_buffer[actual_expanded:] = 0
-            seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
-
-            # Give each of the flattened entries the same block table row as the
-            # original request.
-            self.expanded_block_table_buffer[:actual_expanded] = (
-                torch.repeat_interleave(
-                    block_table, decode_lens, dim=0, output_size=actual_expanded
+            if min_decode_len == max_decode_len:
+                # Uniform decode lengths.
+                num_decode_tokens = num_decodes * max_decode_len
+                _prepare_uniform_decode_kernel[(num_decode_tokens,)](
+                    seq_lens,
+                    self.decode_seq_lens_buffer,
+                    block_table,
+                    block_table.stride(0),
+                    self.expanded_block_table_buffer,
+                    self.expanded_block_table_buffer.stride(0),
+                    self.decode_lens_buffer,
+                    max_decode_len,
+                    BLOCK_SIZE=1024,
                 )
-            )
-            if actual_expanded < num_decode_tokens:
-                self.expanded_block_table_buffer[
-                    actual_expanded:num_decode_tokens, 0
-                ] = 0
-            block_table = self.expanded_block_table_buffer[:num_decode_tokens]
+                self.decode_seq_lens_buffer[num_decode_tokens:] = 0
+                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
+                block_table = self.expanded_block_table_buffer[:num_decode_tokens]
+                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
+                return seq_lens, block_table, decode_lens, num_decode_tokens, False
+            else:
+                # Variable decode lengths.
+                # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
+                # padding) and decode_lens [3, 1, 4, 0] in the below example comments.
+                # The context lengths are therefore
+                # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
 
-            # All reqs now have decode_len=1
-            self.decode_lens_buffer[:num_decode_tokens] = 1
-            decode_lens = self.decode_lens_buffer[:num_decode_tokens]
-            return seq_lens, block_table, decode_lens, num_decode_tokens, False
+                # 3 + 1 + 4 + 0 = 8
+                actual_expanded = int(decode_lens_cpu.sum().item())
+
+                # Fuse expanded_base and expanded_starts into a single
+                # repeat_interleave:
+                # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
+                # where context_start[b] = seq_lens[b] - decode_lens[b].
+                # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
+                # expanded_offsets  = [7, 7, 7, 3, 4, 4, 4, 4]
+                # result            = [8, 9, 10, 7, 9, 10, 11, 12]
+                expanded_offsets = torch.repeat_interleave(
+                    seq_lens - decode_lens - query_start_loc,
+                    decode_lens,
+                    output_size=actual_expanded,
+                )
+
+                # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
+                self.decode_seq_lens_buffer[:actual_expanded] = (
+                    expanded_offsets + self.arange_buffer[:actual_expanded] + 1
+                )
+                self.decode_seq_lens_buffer[actual_expanded:] = 0
+                seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
+
+                # Give each of the flattened entries the same block table row as the
+                # original request.
+                self.expanded_block_table_buffer[:actual_expanded] = (
+                    torch.repeat_interleave(
+                        block_table, decode_lens, dim=0, output_size=actual_expanded
+                    )
+                )
+                if actual_expanded < num_decode_tokens:
+                    self.expanded_block_table_buffer[
+                        actual_expanded:num_decode_tokens, 0
+                    ] = 0
+                block_table = self.expanded_block_table_buffer[:num_decode_tokens]
+
+                # All reqs now have decode_len=1
+                self.decode_lens_buffer[:num_decode_tokens] = 1
+                decode_lens = self.decode_lens_buffer[:num_decode_tokens]
+                return seq_lens, block_table, decode_lens, num_decode_tokens, False
         else:
             # Native path: plain decode (next_n==1) or spec decode
             # with 2D per-token context lengths (next_n > 1).
@@ -459,7 +517,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
             # decode_len < next_n due to padding or short prefills), the simple
             # reshape in sparse_attn_indexer won't work. Use pack_seq_triton
             # (requires_padding) instead.
-            min_decode_len = int(decode_lens_cpu.min().item())
             requires_padding = min_decode_len != max_decode_len
             if use_native and next_n > 1:
                 assert self.decode_seq_lens_buffer.dim() == 2

vllm/v1/worker/gpu/model_runner.py

@@ -363,7 +363,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend(
             self.kv_cache_config, self.vllm_config, self.device
         )
-        initialize_mamba_ssu_backend(self.vllm_config.mamba_config)
+        initialize_mamba_ssu_backend(
+            self.vllm_config.mamba_config, self.kv_cache_config
+        )
         cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
             attn_cg_support.min_cg_support,
             attn_cg_support.min_cg_attn_backend,

vllm/v1/worker/gpu/sample/gumbel.py

@@ -74,6 +74,7 @@ def gumbel_block_argmax(
     temp_ptr,
     seeds_ptr,
     pos_ptr,
+    pos_offset,
     processed_logits_ptr,
     processed_logits_stride,
     APPLY_TEMPERATURE: tl.constexpr,
@@ -98,12 +99,11 @@ def gumbel_block_argmax(
     if temp != 0.0:
         # Calculate the seed for gumbel noise.
         seed = tl.load(seeds_ptr + req_state_idx)
-        pos = tl.load(pos_ptr + token_idx)
+        pos = tl.load(pos_ptr + token_idx) + pos_offset
         gumbel_seed = tl.randint(seed, pos)
 
-        # tl.rand returns fp32, so build a true fp64 uniform from 64 random
-        # bits before applying the double-log transform.
-        u = tl_rand64(gumbel_seed, block, includes_zero=False)
+        # Use FP32 for performance.
+        u = tl.rand(gumbel_seed, block)
         gumbel_noise = -tl.log(-tl.log(u))
 
         # Apply gumbel noise.
@@ -126,6 +126,7 @@ def _gumbel_sample_kernel(
     expanded_idx_mapping_ptr,
     seeds_ptr,
     pos_ptr,
+    pos_offset,
     temp_ptr,
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
@@ -151,6 +152,7 @@ def _gumbel_sample_kernel(
         temp_ptr,
         seeds_ptr,
         pos_ptr,
+        pos_offset,
         processed_logits_ptr,
         processed_logits_stride,
         APPLY_TEMPERATURE=APPLY_TEMPERATURE,
@@ -160,6 +162,33 @@ def _gumbel_sample_kernel(
     tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, value)
 
 
+@triton.jit
+def _gumbel_reduce_kernel(
+    local_argmax_ptr,
+    local_argmax_stride,
+    local_max_ptr,
+    local_max_stride,
+    sampled_ptr,
+    sampled_stride,
+    num_blocks,
+    NUM_BLOCKS_NEXT_POW2: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    offs = tl.arange(0, NUM_BLOCKS_NEXT_POW2)
+    mask = offs < num_blocks
+
+    values = tl.load(
+        local_max_ptr + token_idx * local_max_stride + offs,
+        mask=mask,
+        other=float("-inf"),
+    )
+    _, block_idx = tl.max(values, axis=0, return_indices=True)
+    token_id = tl.load(
+        local_argmax_ptr + token_idx * local_argmax_stride + block_idx,
+    )
+    tl.store(sampled_ptr + token_idx * sampled_stride, token_id)
+
+
 def gumbel_sample(
     logits: torch.Tensor,  # [num_tokens, vocab_size]
     expanded_idx_mapping: torch.Tensor,  # [num_tokens]
@@ -168,6 +197,8 @@ def gumbel_sample(
     pos: torch.Tensor,  # [num_tokens]
     apply_temperature: bool,
     processed_logits_out: torch.Tensor | None = None,  # [num_reqs, vocab_size]
+    out: torch.Tensor | None = None,  # [num_tokens], int64
+    pos_offset: int = 0,
 ) -> torch.Tensor:
     num_tokens, vocab_size = logits.shape
     BLOCK_SIZE = 1024
@@ -186,12 +217,24 @@ def gumbel_sample(
         expanded_idx_mapping,
         seed,
         pos,
+        pos_offset,
         temperature,
         vocab_size,
         BLOCK_SIZE=BLOCK_SIZE,
         APPLY_TEMPERATURE=apply_temperature,
     )
     # NOTE(woosuk): Use int64 for later indexing.
-    max_block_idx = local_max.argmax(dim=-1, keepdim=True)
-    sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
-    return sampled
+    if out is None:
+        out = torch.empty(num_tokens, dtype=torch.int64, device=logits.device)
+    _gumbel_reduce_kernel[(num_tokens,)](
+        local_argmax,
+        local_argmax.stride(0),
+        local_max,
+        local_max.stride(0),
+        out,
+        out.stride(0),
+        num_blocks,
+        NUM_BLOCKS_NEXT_POW2=triton.next_power_of_2(num_blocks),
+        num_warps=1,
+    )
+    return out

vllm/v1/worker/gpu/sample/sampler.py

@@ -39,6 +39,12 @@ class Sampler:
         self.logit_bias_state = LogitBiasState(max_num_reqs, device)
         self.bad_words_state = BadWordsState(req_states)
         self.num_speculative_tokens = num_speculative_tokens
+        # Pre-allocated ones tensor for SamplerOutput.num_sampled (1 per req
+        # in the non-rejection path). Slicing returns a view so downstream
+        # reads see a stable tensor without a kernel launch per call.
+        self._num_sampled_ones = torch.ones(
+            max_num_reqs, dtype=torch.int32, device=device
+        )
 
     def add_request(
         self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
@@ -62,21 +68,36 @@ class Sampler:
         expanded_idx_mapping = input_batch.expanded_idx_mapping
         idx_mapping_np = input_batch.idx_mapping_np
         cu_num_logits_np = input_batch.cu_num_logits_np
-        expanded_local_pos = input_batch.expanded_local_pos
         pos = input_batch.positions[input_batch.logits_indices]
-        input_ids = input_batch.input_ids[input_batch.logits_indices]
 
         # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
         # that num_nans is computed before applying penalties and temperature.
         num_nans = get_num_nans(logits) if self.compute_nans else None
-        sampled, processed_logits = self.sample(
-            logits,
-            expanded_idx_mapping,
-            idx_mapping_np,
-            pos,
-            input_ids,
-            expanded_local_pos,
-        )
+
+        if self._is_sampling_params_noop(idx_mapping_np):
+            # Fast path: no per-request op modifies logits. Skip the fp32 copy,
+            # the input_ids gather (only needed by bias/penalties/bad_words),
+            # and all per-state kernel dispatches.
+            sampled = gumbel_sample(
+                logits,
+                expanded_idx_mapping,
+                self.sampling_states.temperature.gpu,
+                self.sampling_states.seeds.gpu,
+                pos,
+                apply_temperature=True,
+            )
+            processed_logits = logits
+        else:
+            input_ids = input_batch.input_ids[input_batch.logits_indices]
+            expanded_local_pos = input_batch.expanded_local_pos
+            sampled, processed_logits = self.sample(
+                logits,
+                expanded_idx_mapping,
+                idx_mapping_np,
+                pos,
+                input_ids,
+                expanded_local_pos,
+            )
 
         max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
         if max_num_logprobs != NO_LOGPROBS:
@@ -98,7 +119,7 @@ class Sampler:
             sampled_token_ids=sampled.view(-1, 1),
             logprobs_tensors=logprobs_tensors,
             num_nans=num_nans,
-            num_sampled=input_batch.seq_lens.new_ones(input_batch.num_reqs),
+            num_sampled=self._num_sampled_ones[: input_batch.num_reqs],
         )
         return sampler_output
 
@@ -151,6 +172,31 @@ class Sampler:
             logits, expanded_idx_mapping, idx_mapping_np
         )
 
+    def _is_sampling_params_noop(self, idx_mapping_np: np.ndarray) -> bool:
+        """True iff every active request uses pure defaults (argmax / gumbel).
+
+        In that case we can bypass the bf16->fp32 copy and all the per-state
+        kernel dispatches and feed the raw logits directly into
+        gumbel_sample with APPLY_TEMPERATURE=True, which handles both
+        temperature=0 (argmax) and temperature=1 (gumbel noise) without
+        any prior in-place mutation.
+        """
+        states = self.sampling_states
+        temp_np = states.temperature.np[idx_mapping_np]
+        if not np.all((temp_np == 0.0) | (temp_np == 1.0)):
+            return False
+        if np.any(states.min_p.np[idx_mapping_np] != 0.0):
+            return False
+        if np.any(states.top_k.np[idx_mapping_np] != states.vocab_size):
+            return False
+        if np.any(states.top_p.np[idx_mapping_np] != 1.0):
+            return False
+        if np.any(self.penalties_state.use_penalty[idx_mapping_np]):
+            return False
+        if np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]):
+            return False
+        return np.all(self.bad_words_state.num_bad_words.np[idx_mapping_np] == 0)
+
     def sample(
         self,
         logits: torch.Tensor,

vllm/v1/worker/gpu/spec_decode/eagle/speculator.py

@@ -237,23 +237,35 @@ class EagleSpeculator:
             cudagraph_runtime_mode=cudagraph_runtime_mode,
             mm_inputs=mm_inputs,
         )
-        sample_hidden_states = last_hidden_states[last_token_indices]
+        # For MTP, run_model returns the same tensor for both; the two
+        # `[last_token_indices]` gathers below would be redundant, so
+        # write once into self.hidden_states and feed compute_logits from
+        # that view. For eagle3 the two tensors differ, so we still need
+        # both gathers.
+        if last_hidden_states is hidden_states:
+            self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
+            sample_hidden_states = self.hidden_states[:num_reqs]
+        else:
+            sample_hidden_states = last_hidden_states[last_token_indices]
+            self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
         logits = self.model.compute_logits(sample_hidden_states)
 
         # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
-        # used for draft and target sampling.
-        self.draft_tokens[:num_reqs, 0] = gumbel_sample(
+        # used for draft and target sampling. pos_offset=1 folds the +1 into
+        # the kernel itself instead of launching a separate add kernel.
+        gumbel_sample(
             logits,
             idx_mapping,
             self.temperature,
             self.seeds,
-            pos + 1,
+            pos,
             apply_temperature=True,
             processed_logits_out=self.draft_logits[:, 0]
             if self.draft_logits is not None
             else None,
+            out=self.draft_tokens[:num_reqs, 0],
+            pos_offset=1,
         )
-        self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
         self.input_buffers.positions[:num_reqs] = pos
 
     def generate_draft(
@@ -282,19 +294,24 @@ class EagleSpeculator:
             logits = self.model.compute_logits(last_hidden_states)
 
             # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
-            # used for draft and target sampling.
+            # used for draft and target sampling. pos_offset=1 folds the +1
+            # into the kernel instead of launching a separate add.
+            # Write the sampled token directly into the step-th column of
+            # draft_tokens (strided view). update_eagle_inputs below will
+            # re-read it from the same slice with a matching stride.
             draft_tokens = gumbel_sample(
                 logits,
                 idx_mapping,
                 self.temperature,
                 self.seeds,
-                pos + 1,
+                pos,
                 apply_temperature=True,
                 processed_logits_out=self.draft_logits[:, step]
                 if self.draft_logits is not None
                 else None,
+                out=self.draft_tokens[:num_reqs, step],
+                pos_offset=1,
             )
-            self.draft_tokens[:num_reqs, step] = draft_tokens
 
             if step < self.num_speculative_steps - 1:
                 # Update the inputs for the next step.
@@ -760,6 +777,7 @@ def _update_eagle_inputs_kernel(
     seq_lens_ptr,
     max_model_len,
     draft_tokens_ptr,
+    draft_tokens_stride,
     output_hidden_states_ptr,
     output_hidden_states_stride,
     hidden_size,
@@ -768,7 +786,7 @@ def _update_eagle_inputs_kernel(
     req_idx = tl.program_id(0)
 
     # Draft token -> Input ID.
-    draft_token = tl.load(draft_tokens_ptr + req_idx)
+    draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride)
     tl.store(input_ids_ptr + req_idx, draft_token)
 
     # Output hidden states -> Input hidden states.
@@ -813,6 +831,7 @@ def update_eagle_inputs(
         input_buffers.seq_lens,
         max_model_len,
         draft_tokens,
+        draft_tokens.stride(0),
         output_hidden_states,
         output_hidden_states.stride(0),
         hidden_size,

vllm/v1/worker/gpu_model_runner.py

@@ -6738,7 +6738,9 @@ class GPUModelRunner(
         self.may_add_encoder_only_layers_to_kv_cache_config()
         self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
         self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
-        initialize_mamba_ssu_backend(self.vllm_config.mamba_config)
+        initialize_mamba_ssu_backend(
+            self.vllm_config.mamba_config, self.kv_cache_config
+        )
         # The kernel block size for all KV cache groups. For example, if
         # kv_cache_manager uses block_size 256 for a given group, but the attention
         # backends for that group only supports block_size 64, we will return