fuse more

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

examples/basic/offline_inference/basic.py

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

examples/offline_inference/simple_profiling.py

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

vllm/config/compilation.py

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

vllm/envs.py

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

vllm/model_executor/model_loader/dummy_loader.py

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

vllm/model_executor/models/deepseek_mtp.py

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

vllm/model_executor/models/deepseek_v2.py

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

vllm/model_executor/models/registry.py

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

vllm/model_executor/specialized_models/README.md

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

vllm/model_executor/specialized_models/__init__.py

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

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/README.md

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

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/__init__.py

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

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/kernels.py

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

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/layer.py

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

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/model.py

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

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/mtp.py

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

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/sparse_indexer.py

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

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

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

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

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

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

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