perf(deepseek): Multi-stream overlap of indexer wk+weights_proj with QKV-A for DeepSeek V3.2 NSA

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

.buildkite/test_areas/distributed.yaml

@@ -196,8 +196,6 @@ 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,9 +44,8 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/pooling_params.py @noooop @DarkLight1337
 /vllm/tokenizers @DarkLight1337 @njhill
 /vllm/renderers @DarkLight1337 @njhill
-/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33
-/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33
-/vllm/parser @aarnphm @chaunceyjiang @sfeng33
+/vllm/reasoning @aarnphm @chaunceyjiang
+/vllm/tool_parsers @aarnphm @chaunceyjiang
 
 # vLLM V1
 /vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -92,10 +91,7 @@ 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/reasoning @aarnphm @chaunceyjiang @sfeng33
-/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33
-/tests/tool_use @aarnphm @chaunceyjiang @sfeng33
+/tests/v1/determinism @yewentao256 
 
 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor

docs/governance/committers.md

@@ -31,7 +31,6 @@ 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
@@ -42,7 +41,6 @@ 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
@@ -88,7 +86,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, @MatthewBonanni
+    - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg
 - 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
@@ -107,7 +105,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, @MatthewBonanni
+    - @WoosukKwon, @LucasWilkinson, @heheda12345
 - FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB
     - @tlrmchlsmth
 - Quantization: Various quantization config, weight loading, and kernel.
@@ -121,7 +119,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, @sfeng33
+    - @chaunceyjiang, @aarnphm
 
 ### Entrypoints
 
@@ -135,7 +133,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, @MatthewBonanni
+    - @WoosukKwon, @benchislett, @luccafong
 - Structured Output: The structured output implementation
     - @russellb, @aarnphm
 - RL: The RL related features such as collective rpc, sleep mode, etc.
@@ -155,8 +153,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 
 ### External Kernels Integration
 
-- FlashAttention: @LucasWilkinson, @MatthewBonanni
-- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni
+- FlashAttention: @LucasWilkinson
+- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon
 - 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="nvidia/DeepSeek-V3.2-NVFP4", enforce_eager=True, tensor_parallel_size=4, kernel_config={"enable_flashinfer_autotune": False})
+    llm = LLM(model="facebook/opt-125m")
     # 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,46 +2,31 @@
 # 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 = [
-    [0] * 10_000,
-    [1] * 10_000,
-    [2] * 10_000,
-    [3] * 10_000,
-    [4] * 10_000,
-    [5] * 10_000,
-    [6] * 10_000,
-    [7] * 10_000,
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
 ]
 # Create a sampling params object.
-sampling_params = SamplingParams(temperature=0.0)
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
 
 def main():
     # Create an LLM.
     llm = LLM(
-        model="nvidia/DeepSeek-V3.2-NVFP4",
-        tensor_parallel_size=4,
-        kernel_config={"enable_flashinfer_autotune": False},
+        model="facebook/opt-125m",
+        tensor_parallel_size=1,
         profiler_config={
             "profiler": "torch",
-            "torch_profiler_dir": f"./vllm_profile/bsz{len(prompts)}/",
+            "torch_profiler_dir": "./vllm_profile",
         },
-        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

@@ -1,126 +0,0 @@
-# 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,7 +41,6 @@ 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
 
 
@@ -284,7 +283,6 @@ 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)
 
@@ -668,7 +666,6 @@ 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,11 +13,6 @@ 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
@@ -27,40 +22,22 @@ 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(), _kv_cache_config_with_ssu())
+    initialize_mamba_ssu_backend(MambaConfig())
     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), _kv_cache_config_with_ssu()
-    )
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
     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), _kv_cache_config_with_ssu()
-    )
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER))
     backend = get_mamba_ssu_backend()
     assert isinstance(backend, FlashInferSSUBackend)
     assert backend.name == "flashinfer"
@@ -76,25 +53,6 @@ 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"):
@@ -103,9 +61,7 @@ def test_flashinfer_import_error():
 
 def test_triton_basic_call():
     set_random_seed(0)
-    initialize_mamba_ssu_backend(
-        MambaConfig(backend=MambaBackendEnum.TRITON), _kv_cache_config_with_ssu()
-    )
+    initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON))
     device = "cuda"
     batch_size = 2
     dim = 64

tests/kernels/moe/test_moe_layer.py

@@ -465,14 +465,6 @@ 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,3 +140,88 @@ 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,5 +38,6 @@ 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,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """The request function for API endpoints."""
 
-import codecs
 import io
 import json
 import os
@@ -26,12 +25,11 @@ 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 = self._decoder.decode(chunk_bytes)
+        chunk_str = chunk_bytes.decode("utf-8")
         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 TYPE_CHECKING, ClassVar
+from typing import ClassVar
 
 from typing_extensions import Self, override
 
@@ -17,8 +17,20 @@ from vllm.utils.import_utils import PlaceholderModule
 
 from .utils import sanitize_filename
 
-if TYPE_CHECKING:
+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:
+    seaborn = PlaceholderModule("seaborn")
 
 
 @dataclass
@@ -253,20 +265,6 @@ 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 TYPE_CHECKING, ClassVar
+from typing import ClassVar
 
 from vllm.utils.collection_utils import full_groupby
 from vllm.utils.import_utils import PlaceholderModule
@@ -14,8 +14,20 @@ from vllm.utils.import_utils import PlaceholderModule
 from .plot import DummyExecutor, _json_load_bytes
 from .utils import sanitize_filename
 
-if TYPE_CHECKING:
+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:
+    seaborn = PlaceholderModule("seaborn")
 
 
 def _first_present(run_data: dict[str, object], keys: list[str]):
@@ -183,20 +195,6 @@ 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/compilation/cuda_graph.py

@@ -290,9 +290,14 @@ class CUDAGraphWrapper:
                     # across layers will make the cudagraph capture very slow.
                     # therefore, we only run gc for the first graph,
                     # and disable gc for the rest of the graphs.
-                    stack.enter_context(patch("gc.collect", lambda: None))
                     stack.enter_context(
-                        patch("torch.accelerator.empty_cache", lambda: None)
+                        patch("gc.collect", lambda *args, **kwargs: None)
+                    )
+                    stack.enter_context(
+                        patch(
+                            "torch.accelerator.empty_cache",
+                            lambda *args, **kwargs: None,
+                        )
                     )
 
                 if self.graph_pool is not None:

vllm/config/compilation.py

@@ -737,8 +737,6 @@ 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,10 +19,6 @@ 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,18 +1550,6 @@ 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,2 +1,19 @@
 # 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,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
-import sys
 import typing
 
 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
@@ -15,17 +14,6 @@ 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."""
 
@@ -50,28 +38,18 @@ class BenchmarkSubcommand(CLISubcommand):
         )
         bench_subparsers = bench_parser.add_subparsers(required=True, dest="bench_type")
 
-        # 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}"
-                )
+        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,9 +100,10 @@ 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 and
-# https://github.com/vllm-project/vllm/issues/10619.
+# see https://github.com/vllm-project/vllm/issues/10480
 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,7 +216,6 @@ 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
@@ -242,6 +241,7 @@ if TYPE_CHECKING:
     VLLM_DEBUG_WORKSPACE: bool = False
     VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
+    VLLM_DISABLE_INDEXER_STREAM: bool = False
     VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
     VLLM_USE_V2_MODEL_RUNNER: bool = False
     VLLM_LOG_MODEL_INSPECTION: bool = False
@@ -1521,10 +1521,6 @@ 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"))
@@ -1634,6 +1630,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int(
         int(os.getenv("VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256))
     ),
+    # Disables parallel execution of indexer q_b_proj via separate cuda stream
+    "VLLM_DISABLE_INDEXER_STREAM": lambda: bool(
+        int(os.getenv("VLLM_DISABLE_INDEXER_STREAM", "0"))
+    ),
     # Format for saving torch.compile cache artifacts
     # - "binary": saves as binary file
     #     Safe for multiple vllm serve processes accessing the same torch compile cache.

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

@@ -1,84 +0,0 @@
-# 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,7 +15,6 @@ 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__)
 
@@ -189,22 +188,12 @@ _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
 _mamba_ssu_backend: MambaSSUBackend | None = None
 
 
-def initialize_mamba_ssu_backend(
-    mamba_config: MambaConfig,
-    kv_cache_config: KVCacheConfig,
-) -> None:
+def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None:
     """Initialize the global Mamba SSU backend.
 
-    No-op if `kv_cache_config` contains no specs that call
-    selective_state_update.
+    Args:
+        mamba_config: Mamba configuration.
     """
-    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
@@ -214,11 +203,7 @@ def initialize_mamba_ssu_backend(
             f"Valid options: {list(_BACKEND_REGISTRY.keys())}"
         )
 
-    backend_cls = _BACKEND_REGISTRY[backend]
-    if isinstance(_mamba_ssu_backend, backend_cls):
-        return
-
-    _mamba_ssu_backend = backend_cls(mamba_config)
+    _mamba_ssu_backend = _BACKEND_REGISTRY[backend](mamba_config)
     logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
 
 

vllm/model_executor/layers/mla.py

@@ -4,10 +4,20 @@ from dataclasses import dataclass
 
 import torch
 
-from vllm.config import CacheConfig
+import vllm.envs as envs
+from vllm.config import CacheConfig, get_current_vllm_config
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
 from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.attention import MLAAttention
 from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.utils.torch_utils import current_stream, direct_register_custom_op
+
+logger = init_logger(__name__)
+
+# Token threshold for multi-stream indexer overlap.
+# Disables multi-stream for batches > 1024 to avoid SM contention.
+_INDEXER_STREAM_TOKEN_THRESHOLD = 1024
 
 
 @dataclass
@@ -27,6 +37,229 @@ class MLAModules:
     is_sparse: bool
     topk_indices_buffer: torch.Tensor | None
     indexer_rotary_emb: torch.nn.Module | None = None
+    alt_stream: torch.cuda.Stream | None = None
+
+
+class _WkForkModule(torch.nn.Module):
+    """Compiled module for wk_weights_proj+k_norm on alt_stream.
+
+    Wraps the indexer's fused wk_weights_proj and k_norm into a single
+    compilation unit.  When compiled with torch.compile the operations
+    benefit from Inductor optimizations:
+      - wk_weights_proj:  single fused GEMM for wk + weights_proj
+      - k_norm:  operator fusion with surrounding ops
+
+    The compiled module is called inside the mla_wk_fork custom op,
+    which runs it on alt_stream concurrent with QKV-A on the main
+    stream.
+
+    Returns a concatenated ``[k, raw_weights]`` tensor; the join
+    caller splits it back using known ``wk_dim`` and ``weights_dim``.
+
+    Sub-modules are stored via ``object.__setattr__`` so they do NOT
+    appear in ``_modules`` / ``state_dict()``.  This prevents:
+      1. Duplicate parameter entries (they are shared with Indexer).
+      2. State-dict key mismatches during weight loading.
+      3. ``isinstance`` false-positives when tests use MagicMock.
+    """
+
+    def __init__(self, wk_weights_proj, k_norm, head_dim):
+        super().__init__()
+        object.__setattr__(self, "wk_weights_proj", wk_weights_proj)
+        object.__setattr__(self, "k_norm", k_norm)
+        object.__setattr__(self, "head_dim", head_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        kw, _ = self.wk_weights_proj(hidden_states)
+        k = kw[:, : self.head_dim]
+        raw_weights = kw[:, self.head_dim :]
+        k = self.k_norm(k)
+        return torch.cat([k, raw_weights], dim=-1)
+
+
+# ---- Multi-Stream wk_weights_proj Overlap Custom Ops ----
+#
+# Two minimal custom ops overlap wk_weights_proj+k_norm with QKV-A:
+#   mla_wk_fork: launches a COMPILED wk_weights_proj+k_norm module on
+#                alt_stream (concurrent with QKV-A on main)
+#   mla_wk_join: waits for alt_stream, returns pre-computed
+#                [k | raw_weights] concatenated tensor
+#
+# CRITICAL DESIGN PRINCIPLES:
+#   1. ALL indexer GEMMs (wq_b) and q_b_proj MUST stay inside the
+#      main torch.compile graph.
+#   2. Fork operations MUST ALSO be compiled — running them eagerly
+#      loses operator fusion and kernel selection overhead.
+#   3. The fix: a SEPARATELY torch.compile'd _WkForkModule wraps
+#      wk_weights_proj+k_norm.  The compiled module is called inside
+#      the fork custom op on alt_stream.
+#
+# WHY wk_weights_proj+k_norm:
+#   The fused wk_weights_proj GEMM depends ONLY on hidden_states (the
+#   layer input).  It can start at the VERY BEGINNING of the forward
+#   pass, concurrent with the QKV-A GEMM on the main stream.
+#
+#   Alt stream (compiled):  wk_weights_proj fused GEMM + k_norm
+#   Main stream (compiled): QKV-A + Q-A LN + Q-B proj + kv preprocess
+#                           + RoPE
+#   Alt < Main → fork is completely hidden!
+#
+# The indexer call stays INLINE in forward() (traced by torch.compile).
+# Indexer.forward() receives pre-computed k via precomputed_k and raw
+# weights via precomputed_weights, skipping its own wk_weights_proj
+# and k_norm.  The remaining indexer GEMM (wq_b) and
+# sparse_attn_indexer stay in the compiled graph.
+#
+# Pattern EXTENDS MoE shared expert streaming (default_moe_runner.py):
+#   1. Register the layer in static_forward_context during __init__
+#   2. Custom ops retrieve the layer by name from forward_context
+#   3. Stream fork/join happens inside the custom ops (opaque)
+#   4. Fake implementations provide output shape for symbolic execution
+#   5. NOT in _attention_ops — opaque nodes inside compiled region
+#   6. tags=(torch.Tag.needs_fixed_stride_order,) prevents Inductor
+#      stride conversion overhead
+#
+# DIFFERENCE from MoE: the MoE shared expert runs EAGERLY inside its
+# custom op.  Here we add a SEPARATE torch.compile unit (_WkForkModule)
+# for the fork operations.  This is a novel extension; a graceful
+# fallback to eager is included in case torch.compile fails.
+#
+# Fork/Join symmetry:
+#   The fork sets wrapper._wk_forked = True when multi-stream is used.
+#   The join checks this flag to decide whether to wait_stream.  This
+#   ensures fork and join ALWAYS agree on whether multi-stream is active.
+
+
+def _mla_wk_fork(
+    hidden_states: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    """Launch compiled wk_weights_proj+k_norm on alt_stream.
+
+    Returns a clone of hidden_states to establish a data dependency
+    while obeying PyTorch's custom-op contract: outputs MUST NOT alias
+    inputs.  Returning the input directly caused undefined behaviour in
+    the Inductor buffer-assignment pass (incorrect buffer reuse in the
+    generated code) and triggered a CUDA-graph capture error.
+
+    Stores the concatenated [k, raw_weights] result in
+    ``wrapper._fork_result`` for the join op.
+
+    The fork calls ``wrapper._compiled_fork_ops`` — a separately
+    torch.compile'd _WkForkModule — so that the operations benefit
+    from Inductor optimisations (operator fusion, kernel selection)
+    even when running on the alt_stream.
+    """
+    wrapper = get_forward_context().no_compile_layers[layer_name]
+    indexer = wrapper.indexer
+
+    if indexer is None or not wrapper.is_sparse:
+        wrapper._wk_forked = False
+        # Clone to satisfy the custom-op no-alias contract.
+        # The fake impl returns torch.empty_like (new tensor),
+        # so the real impl must also return a non-aliasing tensor.
+        return hidden_states.clone()
+
+    use_multi_stream = (
+        wrapper.alt_stream is not None
+        and not envs.VLLM_DISABLE_INDEXER_STREAM
+        and hidden_states.shape[0] <= _INDEXER_STREAM_TOKEN_THRESHOLD
+    )
+
+    fork_ops = wrapper._compiled_fork_ops
+
+    if use_multi_stream:
+        main_stream = current_stream()
+        alt_stream = wrapper.alt_stream
+
+        # Prevent GC from freeing hidden_states while alt_stream reads it.
+        hidden_states.record_stream(alt_stream)
+
+        # alt_stream waits for hidden_states to be ready on main.
+        alt_stream.wait_stream(main_stream)
+
+        # Launch compiled wk_weights_proj+k_norm on alt_stream
+        # (concurrent with QKV-A on main).
+        with torch.cuda.stream(alt_stream):
+            wrapper._fork_result = fork_ops(hidden_states)
+
+        wrapper._wk_forked = True
+    else:
+        # Sequential: run compiled fork ops on main stream.
+        wrapper._fork_result = fork_ops(hidden_states)
+        wrapper._wk_forked = False
+
+    # Clone to satisfy the custom-op no-alias contract.
+    # The clone is a lightweight memcpy (e.g. ~14 KB for decode
+    # batch_size=1 with hidden_size=7168 in bf16).  Both streams
+    # read the original hidden_states concurrently; the clone
+    # provides a separate buffer for downstream compiled code
+    # (QKV-A on main stream) so the Inductor's buffer-liveness
+    # analysis stays correct.
+    return hidden_states.clone()
+
+
+def _mla_wk_fork_fake(
+    hidden_states: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    return torch.empty_like(hidden_states)
+
+
+direct_register_custom_op(
+    op_name="mla_wk_fork",
+    op_func=_mla_wk_fork,
+    mutates_args=[],
+    fake_impl=_mla_wk_fork_fake,
+    tags=(torch.Tag.needs_fixed_stride_order,),
+)
+
+
+def _mla_wk_join(
+    hidden_states: torch.Tensor,
+    layer_name: str,
+    join_dim: int,
+) -> torch.Tensor:
+    """Get pre-computed [k, raw_weights], waiting for alt_stream if needed.
+
+    Returns the concatenated tensor stored by ``_mla_wk_fork``.
+    Shape: ``[num_tokens, join_dim]`` where ``join_dim = wk_dim + weights_dim``.
+    Only waits if the fork op set ``wrapper._wk_forked = True``,
+    ensuring symmetric fork/join behaviour.
+    """
+    wrapper = get_forward_context().no_compile_layers[layer_name]
+
+    # Check the flag set by fork — guarantees fork/join symmetry.
+    if getattr(wrapper, "_wk_forked", False):
+        main_stream = current_stream()
+        main_stream.wait_stream(wrapper.alt_stream)
+        wrapper._wk_forked = False
+
+    # Return the concatenated [k, raw_weights] produced by the fork.
+    # The caller splits using known wk_dim and weights_dim.
+    return wrapper._fork_result
+
+
+def _mla_wk_join_fake(
+    hidden_states: torch.Tensor,
+    layer_name: str,
+    join_dim: int,
+) -> torch.Tensor:
+    return torch.empty(
+        hidden_states.shape[0],
+        join_dim,
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+
+direct_register_custom_op(
+    op_name="mla_wk_join",
+    op_func=_mla_wk_join,
+    mutates_args=[],
+    fake_impl=_mla_wk_join_fake,
+    tags=(torch.Tag.needs_fixed_stride_order,),
+)
 
 
 # --8<-- [start:multi_head_latent_attention]
@@ -86,11 +319,53 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
         self.indexer = mla_modules.indexer
         self.indexer_rope_emb = mla_modules.indexer_rotary_emb
         self.is_sparse = mla_modules.is_sparse
+        self.alt_stream = mla_modules.alt_stream
+        # Flag for symmetric fork/join. Set by _mla_wk_fork, checked by
+        # _mla_wk_join. Ensures join only waits when fork actually
+        # launched work on alt_stream.
+        self._wk_forked = False
 
         if self.indexer is not None:
             assert hasattr(self.indexer, "topk_tokens")
             self.topk_tokens = self.indexer.topk_tokens
             self.topk_indices_buffer = mla_modules.topk_indices_buffer
+            # Store dimensions for the fork/join custom ops.
+            # wk_dim: output dimension of indexer wk (head_dim=128)
+            # weights_dim: output dimension of indexer weights_proj (n_head=64)
+            # join_dim: total concatenated dim returned by mla_wk_join
+            self.wk_dim = self.indexer.head_dim
+            self.weights_dim = self.indexer.n_head
+            self.join_dim = self.wk_dim + self.weights_dim
+
+            # Compile wk_weights_proj+k_norm as a SEPARATE torch.compile
+            # unit.  The compiled module runs on alt_stream inside the
+            # mla_wk_fork custom op, concurrent with QKV-A on main.
+            # Uses object.__setattr__ to avoid registering as a sub-module
+            # (prevents state_dict / weight-loading duplication).
+            #
+            # NOTE: This EXTENDS the MoE shared-expert streaming pattern
+            # (default_moe_runner.py) — the MoE pattern runs shared experts
+            # EAGERLY, while we add a separate torch.compile unit for the
+            # fork ops.  Graceful fallback to eager if compilation fails.
+            _fork_mod = _WkForkModule(
+                self.indexer.wk_weights_proj,
+                self.indexer.k_norm,
+                self.indexer.head_dim,
+            )
+            try:
+                _compiled = torch.compile(_fork_mod, dynamic=True)
+            except Exception:
+                logger.warning(
+                    "Failed to compile MLA fork ops for layer %s, "
+                    "falling back to eager execution.",
+                    prefix,
+                )
+                _compiled = _fork_mod
+            object.__setattr__(
+                self,
+                "_compiled_fork_ops",
+                _compiled,
+            )
 
         self.mla_attn = MLAAttention(
             num_heads=self.num_heads,
@@ -110,6 +385,13 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
 
         self.prefix = prefix
 
+        # Register in static_forward_context so the fork/join custom ops
+        # (mla_wk_fork, mla_wk_join) can retrieve this wrapper.
+        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
+
     def forward(
         self,
         positions: torch.Tensor,
@@ -130,12 +412,24 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
                 "q_b_proj is required when q_lora_rank is not None"
             )
 
+            # Fork: launch wk_weights_proj+k_norm on alt_stream,
+            # concurrent with QKV-A. Opaque to torch.compile.
+            # Fused GEMM hidden behind QKV-A+Q-A LN+Q-B on main.
+            # All other GEMMs stay INSIDE torch.compile scope.
+            hidden_states = torch.ops.vllm.mla_wk_fork(
+                hidden_states,
+                self.prefix,
+            )
+
+            # QKV-A GEMM on main stream — COMPILED, concurrent with wk.
             qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
             q_c, kv_lora = qkv_lora.split(
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                 dim=-1,
             )
             q_c = self.q_a_layernorm(q_c)
+
+            # q_b_proj on main stream — INSIDE torch.compile scope.
             q = self.q_b_proj(q_c)[0]
         else:
             assert self.kv_a_proj_with_mqa is not None, (
@@ -159,9 +453,28 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
                 positions, q[..., self.qk_nope_head_dim :], k_pe
             )
 
-        if self.indexer and self.is_sparse:
-            _topk_indices = self.indexer(
-                hidden_states, q_c, positions, self.indexer_rope_emb
+        # Join wk_weights_proj + run indexer INLINE (COMPILED on main).
+        # The indexer GEMM (wq_b) stays in torch.compile scope.
+        # wk_weights_proj+k_norm run on alt_stream (hidden behind QKV-A).
+        # sparse_attn_indexer remains a PIECEWISE split point (as original).
+        if self.indexer is not None and self.is_sparse:
+            k_weights = torch.ops.vllm.mla_wk_join(
+                hidden_states,
+                self.prefix,
+                self.join_dim,
+            )
+            # Split the concatenated join result into k and raw_weights.
+            k_pre, weights_pre = k_weights.split(
+                [self.wk_dim, self.weights_dim],
+                dim=-1,
+            )
+            self.indexer(
+                hidden_states,
+                q_c,
+                positions,
+                self.indexer_rope_emb,
+                precomputed_k=k_pre,
+                precomputed_weights=weights_pre,
             )
 
         if llama_4_scaling is not None:
@@ -171,7 +484,10 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
             q,
             kv_c_normed,
             k_pe,
-            output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim),
+            output_shape=(
+                hidden_states.shape[0],
+                self.num_heads * self.v_head_dim,
+            ),
         )
 
         return self.o_proj(attn_out)[0]

vllm/model_executor/layers/quantization/__init__.py

@@ -40,7 +40,6 @@ 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))
 
@@ -48,6 +47,7 @@ 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,25 +5,27 @@ from typing import Any
 
 import torch
 
-from vllm.model_executor.layers.fused_moe import FusedMoE
+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.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.layers.quantization.online.int8 import (
-    Int8OnlineMoEMethod,
-)
+from vllm.model_executor.utils import set_weight_attrs
 
 
 class ExpertsInt8Config(QuantizationConfig):
-    """Online int8 quantization for MoE expert weights.
-    Linear layers are left unquantized.
-
-    Backward-compatible config for ``--quantization experts_int8``.
-    Prefer ``--quantization int8_per_channel``
-    """
+    """Config class for Int8 experts quantization."""
 
     def __init__(self) -> None:
         super().__init__()
@@ -54,5 +56,149 @@ class ExpertsInt8Config(QuantizationConfig):
         if isinstance(layer, LinearBase):
             return UnquantizedLinearMethod()
         elif isinstance(layer, FusedMoE):
-            return Int8OnlineMoEMethod(layer=layer)
+            return ExpertsInt8MoEMethod(self, layer.moe_config)
         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,7 +9,6 @@ from vllm.config.quantization import (
     OnlineQuantizationConfigArgs,
     OnlineQuantScheme,
 )
-from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
 )
@@ -34,11 +33,6 @@ 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):
@@ -102,13 +96,7 @@ class OnlineQuantizationConfig(QuantizationConfig):
                 return UnquantizedLinearMethod()
 
             linear_scheme = self.args.linear_scheme_override or self.args.global_scheme
-            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:
+            if linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
                 return Fp8PerBlockOnlineLinearMethod()
             else:
                 return Fp8PerTensorOnlineLinearMethod()
@@ -121,9 +109,7 @@ class OnlineQuantizationConfig(QuantizationConfig):
                 return UnquantizedFusedMoEMethod(layer.moe_config)
 
             moe_scheme = self.args.moe_scheme_override or self.args.global_scheme
-            if moe_scheme == OnlineQuantScheme.INT8_PER_CHANNEL_WEIGHT_ONLY:
-                return Int8OnlineMoEMethod(layer=layer)
-            elif moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
+            if moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
                 return Fp8PerBlockOnlineMoEMethod(layer=layer)
             else:
                 return Fp8PerTensorOnlineMoEMethod(layer=layer)

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

@@ -10,6 +10,7 @@ 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
@@ -18,15 +19,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,
@@ -43,7 +44,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
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import per_block_cast_to_fp8
 
@@ -267,15 +268,21 @@ class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
 # ---------------------------------------------------------------------------
 
 
-class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
+class _Fp8OnlineMoEBase(FusedMoEMethodBase):
     """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,
@@ -306,6 +313,77 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
             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",
@@ -352,6 +430,15 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
                 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":
@@ -373,9 +460,68 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
             block_shape=self.weight_block_size,
         )
 
-        self._maybe_inject_biases(quant_config, layer)
+        # 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
+
         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

@@ -1,109 +0,0 @@
-# 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

@@ -1,172 +0,0 @@
-# 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,13 +43,6 @@ 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,6 +30,7 @@ from .deepseek_v2 import (
     DeepseekV2DecoderLayer,
     DeepseekV2MixtureOfExperts,
     DeepseekV2MoE,
+    _try_load_fp8_indexer_wk,
     get_spec_layer_idx_from_weight_name,
 )
 from .utils import maybe_prefix
@@ -190,10 +191,6 @@ 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 = []
@@ -248,13 +245,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
             ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
         ]
 
-        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)
+        # 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,
@@ -271,6 +267,7 @@ 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
@@ -281,6 +278,12 @@ 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,6 +66,10 @@ 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,
@@ -628,10 +632,6 @@ 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,36 +646,16 @@ class Indexer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.wq_b",
         )
-        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",
-            )
+        # 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",
+        )
         self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
         self.softmax_scale = self.head_dim**-0.5
 
@@ -709,23 +689,31 @@ class Indexer(nn.Module):
         )
 
     def forward(
-        self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
+        self,
+        hidden_states: torch.Tensor,
+        qr: torch.Tensor,
+        positions,
+        rotary_emb,
+        precomputed_k: torch.Tensor | None = None,
+        precomputed_weights: torch.Tensor | None = None,
     ) -> torch.Tensor:
         q, _ = self.wq_b(qr)
         q = q.view(-1, self.n_head, self.head_dim)
         q_pe, q_nope = torch.split(
             q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
         )
-        if self.is_fp4_ckpt:
+
+        # Use pre-computed k and weights from mla_wk_fork/join when available.
+        # wk_weights_proj+k_norm were already computed on alt_stream.
+        if precomputed_k is not None:
+            k = precomputed_k
+            weights = precomputed_weights
+        else:
             # 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 = self.k_norm(k)
         k_pe, k_nope = torch.split(
             k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
         )
@@ -761,6 +749,46 @@ 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,
@@ -872,6 +900,7 @@ class DeepseekV2MLAAttention(nn.Module):
         prefix: str = "",
         topk_indices_buffer: torch.Tensor | None = None,
         input_size: int | None = None,
+        alt_stream: torch.cuda.Stream | None = None,
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
@@ -1008,6 +1037,7 @@ class DeepseekV2MLAAttention(nn.Module):
             indexer_rotary_emb=self.indexer_rope_emb,
             is_sparse=self.is_v32,
             topk_indices_buffer=topk_indices_buffer,
+            alt_stream=alt_stream,
         )
 
         self.mla_attn = MultiHeadLatentAttentionWrapper(
@@ -1041,6 +1071,7 @@ class DeepseekV2DecoderLayer(nn.Module):
         prefix: str,
         config: DeepseekV2Config | None = None,
         topk_indices_buffer: torch.Tensor | None = None,
+        alt_stream: torch.cuda.Stream | None = None,
     ) -> None:
         super().__init__()
 
@@ -1091,6 +1122,7 @@ class DeepseekV2DecoderLayer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.self_attn",
             topk_indices_buffer=topk_indices_buffer,
+            **({"alt_stream": alt_stream} if alt_stream is not None else {}),
         )
 
         if (
@@ -1193,6 +1225,17 @@ class DeepseekV2Model(nn.Module):
         else:
             topk_indices_buffer = None
 
+        # Create alt_stream for multi-stream indexer parallelism.
+        # Single stream shared across ALL layers. Matches SGLang design.
+        if (
+            self.is_v32
+            and current_platform.is_cuda_alike()
+            and vllm_config.model_config.use_mla
+        ):
+            self.alt_stream = torch.cuda.Stream()
+        else:
+            self.alt_stream = None
+
         if get_pp_group().is_first_rank:
             self.embed_tokens = VocabParallelEmbedding(
                 config.vocab_size,
@@ -1208,6 +1251,7 @@ class DeepseekV2Model(nn.Module):
                 vllm_config,
                 prefix,
                 topk_indices_buffer=topk_indices_buffer,
+                alt_stream=self.alt_stream,
             ),
             prefix=f"{prefix}.layers",
         )
@@ -1344,10 +1388,6 @@ 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)
@@ -1473,13 +1513,13 @@ class DeepseekV2ForCausalLM(
             ("qkv_proj", "k_proj", "k"),
             ("qkv_proj", "v_proj", "v"),
         ]
-        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)
+        # 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.use_mha:
             stacked_params_mapping.extend(mha_params_mapping)
@@ -1516,6 +1556,11 @@ 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,8 +99,6 @@ 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():
@@ -109,27 +107,6 @@ 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,15 +1302,6 @@ 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

@@ -1,19 +0,0 @@
-# [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

@@ -1,36 +0,0 @@
-# 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

@@ -1,34 +0,0 @@
-# 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

@@ -1,8 +0,0 @@
-# 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

@@ -1,931 +0,0 @@
-# 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

@@ -1,570 +0,0 @@
-# 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

@@ -1,151 +0,0 @@
-# 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

@@ -1,209 +0,0 @@
-# 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

@@ -1,175 +0,0 @@
-# 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,7 +8,6 @@ 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,
@@ -31,40 +30,6 @@ 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,
@@ -440,75 +405,52 @@ 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
-            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,
+            # 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
                 )
-                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].
+            )
+            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]
 
-                # 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
+            # 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).
@@ -517,6 +459,7 @@ 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,9 +363,7 @@ 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, self.kv_cache_config
-        )
+        initialize_mamba_ssu_backend(self.vllm_config.mamba_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,7 +74,6 @@ def gumbel_block_argmax(
     temp_ptr,
     seeds_ptr,
     pos_ptr,
-    pos_offset,
     processed_logits_ptr,
     processed_logits_stride,
     APPLY_TEMPERATURE: tl.constexpr,
@@ -99,11 +98,12 @@ 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_offset
+        pos = tl.load(pos_ptr + token_idx)
         gumbel_seed = tl.randint(seed, pos)
 
-        # Use FP32 for performance.
-        u = tl.rand(gumbel_seed, block)
+        # 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)
         gumbel_noise = -tl.log(-tl.log(u))
 
         # Apply gumbel noise.
@@ -126,7 +126,6 @@ def _gumbel_sample_kernel(
     expanded_idx_mapping_ptr,
     seeds_ptr,
     pos_ptr,
-    pos_offset,
     temp_ptr,
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
@@ -152,7 +151,6 @@ def _gumbel_sample_kernel(
         temp_ptr,
         seeds_ptr,
         pos_ptr,
-        pos_offset,
         processed_logits_ptr,
         processed_logits_stride,
         APPLY_TEMPERATURE=APPLY_TEMPERATURE,
@@ -162,33 +160,6 @@ 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]
@@ -197,8 +168,6 @@ 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
@@ -217,24 +186,12 @@ 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.
-    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
+    max_block_idx = local_max.argmax(dim=-1, keepdim=True)
+    sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
+    return sampled

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

@@ -39,12 +39,6 @@ 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
@@ -68,36 +62,21 @@ 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
-
-        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,
-            )
+        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:
@@ -119,7 +98,7 @@ class Sampler:
             sampled_token_ids=sampled.view(-1, 1),
             logprobs_tensors=logprobs_tensors,
             num_nans=num_nans,
-            num_sampled=self._num_sampled_ones[: input_batch.num_reqs],
+            num_sampled=input_batch.seq_lens.new_ones(input_batch.num_reqs),
         )
         return sampler_output
 
@@ -172,31 +151,6 @@ 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,35 +237,23 @@ class EagleSpeculator:
             cudagraph_runtime_mode=cudagraph_runtime_mode,
             mm_inputs=mm_inputs,
         )
-        # 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]
+        sample_hidden_states = last_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. pos_offset=1 folds the +1 into
-        # the kernel itself instead of launching a separate add kernel.
-        gumbel_sample(
+        # used for draft and target sampling.
+        self.draft_tokens[:num_reqs, 0] = gumbel_sample(
             logits,
             idx_mapping,
             self.temperature,
             self.seeds,
-            pos,
+            pos + 1,
             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(
@@ -294,24 +282,19 @@ 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. 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.
+            # used for draft and target sampling.
             draft_tokens = gumbel_sample(
                 logits,
                 idx_mapping,
                 self.temperature,
                 self.seeds,
-                pos,
+                pos + 1,
                 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.
@@ -777,7 +760,6 @@ 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,
@@ -786,7 +768,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_tokens_stride)
+    draft_token = tl.load(draft_tokens_ptr + req_idx)
     tl.store(input_ids_ptr + req_idx, draft_token)
 
     # Output hidden states -> Input hidden states.
@@ -831,7 +813,6 @@ 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,9 +6738,7 @@ 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, self.kv_cache_config
-        )
+        initialize_mamba_ssu_backend(self.vllm_config.mamba_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