[XPU]feat: add DeepSeek-V4 XPU attention decode path (#42953)

Signed-off-by: Ma Jian <jian1.ma@intel.com>

vllm/model_executor/kernels/linear/scaled_mm/xpu.py

@@ -100,6 +100,21 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
             "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale"
         )
         scale = getattr(layer, scale_attr)
+        # Models with scale_fmt=ue8m0 (e.g. DeepSeek-V4) store weight scales
+        # as float8_e8m0fnu. The oneDNN fp8_gemm kernel dispatches to its
+        # "block quant" path only when NEITHER scale is e8m0:
+        #
+        #   is_block_quant = (m1_sc != e8m0) && (m2_sc != e8m0) && ...
+        #
+        # Since activation scales are always float32 (use_ue8m0=False on XPU,
+        # DeepGEMM requires Hopper/Blackwell), an e8m0 weight scale causes
+        # is_block_quant=false and falls into the wrong per-channel path,
+        # producing NaN. Converting e8m0→float32 here at load time (one-time,
+        # negligible overhead for small scale tensors) ensures the kernel sees
+        # matching dtypes and correctly enters the block-quant path with the
+        # actual group_size derived from scale tensor shapes.
+        if scale.dtype == torch.float8_e8m0fnu:
+            scale = scale.to(torch.float32)
         replace_parameter(layer, scale_attr, scale.data.t().contiguous())
 
     def apply_block_scaled_mm(

vllm/model_executor/layers/mhc.py

@@ -147,6 +147,36 @@ class MHCPreOp(CustomOp):
             sinkhorn_repeat,
         )
 
+    def forward_xpu(
+        self,
+        residual: torch.Tensor,
+        fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_eps: float,
+        hc_pre_eps: float,
+        hc_sinkhorn_eps: float,
+        hc_post_mult_value: float,
+        sinkhorn_repeat: int,
+        n_splits: int = 1,
+        norm_weight: torch.Tensor | None = None,
+        norm_eps: float = 0.0,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.forward_native(
+            residual,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            n_splits,
+            norm_weight,
+            norm_eps,
+        )
+
 
 # --8<-- [start:mhc_post]
 @CustomOp.register("mhc_post")
@@ -215,6 +245,20 @@ class MHCPostOp(CustomOp):
             comb_res_mix,
         )
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        post_layer_mix: torch.Tensor,
+        comb_res_mix: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(
+            x,
+            residual,
+            post_layer_mix,
+            comb_res_mix,
+        )
+
 
 # --8<-- [start:hc_head]
 @CustomOp.register("hc_head")
@@ -300,6 +344,36 @@ class HCHeadOp(CustomOp):
     def forward_native(self, *args, **kwargs):
         raise NotImplementedError("Native implementation of hc_head is not available")
 
+    def forward_xpu(
+        self,
+        hidden_states: torch.Tensor,
+        hc_fn: torch.Tensor,
+        hc_scale: torch.Tensor,
+        hc_base: torch.Tensor,
+        rms_norm_eps: float,
+        hc_eps: float,
+    ) -> torch.Tensor:
+        hc_mult, hidden_size = hidden_states.shape[-2:]
+        outer_shape = hidden_states.shape[:-2]
+        hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
+        num_tokens = hs_flat.shape[0]
+
+        out = torch.empty(
+            num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
+        )
+        torch.ops.vllm.hc_head_triton(
+            hs_flat,
+            hc_fn,
+            hc_scale,
+            hc_base,
+            out,
+            hidden_size,
+            rms_norm_eps,
+            hc_eps,
+            hc_mult,
+        )
+        return out.view(*outer_shape, hidden_size)
+
 
 # --8<-- [start:mhc_fused_post_pre]
 @CustomOp.register("mhc_fused_post_pre")

vllm/models/deepseek_v4/__init__.py

@@ -7,21 +7,22 @@ picks the right one for the current platform and re-exports the public
 classes used by the model registry and quantization config lookup.
 """
 
-from typing import TYPE_CHECKING
-
 from vllm.platforms import current_platform
 
 from .quant_config import DeepseekV4FP8Config
 
 # Pick the per-platform implementation. The NVIDIA branch is the static
-# default that mypy sees; the ROCm branch overrides it at runtime and is
+# default that mypy sees; the ROCm/XPU branches override at runtime and are
 # kept type-compatible via ``# type: ignore[assignment]``.
-if TYPE_CHECKING or not current_platform.is_rocm():
-    from .nvidia.model import DeepseekV4ForCausalLM
-    from .nvidia.mtp import DeepSeekV4MTP
+if current_platform.is_rocm():
+    from .amd.model import DeepseekV4ForCausalLM
+    from .amd.mtp import DeepSeekV4MTP
+elif current_platform.is_xpu():
+    from .xpu.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
+    from .xpu.mtp import DeepSeekV4MTP  # type: ignore[assignment]
 else:
-    from .amd.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
-    from .amd.mtp import DeepSeekV4MTP  # type: ignore[assignment]
+    from .nvidia.model import DeepseekV4ForCausalLM  # type: ignore[assignment]
+    from .nvidia.mtp import DeepSeekV4MTP  # type: ignore[assignment]
 
 __all__ = [
     "DeepSeekV4MTP",

vllm/models/deepseek_v4/compressor.py

@@ -368,7 +368,7 @@ class DeepseekCompressor(nn.Module):
                 fp8_scale=fp8_scale,
             )
         else:
-            # Indexer path (head_dim == 128) or AMD: triton, legacy UE8M0 only.
+            # Indexer path (head_dim == 128) or non-CUDA GPUs (AMD, XPU, etc.).
             compress_norm_rope_store_fn = compress_norm_rope_store_triton
             extra_kwargs = {}
 

vllm/models/deepseek_v4/xpu/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/deepseek_v4/xpu/mtp.py

@@ -0,0 +1,511 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MTP draft model for DeepSeek V4 (internal codename: DeepseekV4).
+
+Split from ``deepseek_mtp.py`` because the V4 architecture introduces several
+pieces that have no analogue in V3/V32:
+  * separate ``e_proj`` / ``h_proj`` with fp8 linear quantization (instead of
+    the fused ``eh_proj``);
+  * ``hc_head`` hypercompressed vocab projection applied in ``compute_logits``;
+  * ``DeepseekV4DecoderLayer`` with its own aux-stream management;
+  * V4-specific checkpoint weight-name remapping in ``load_weights``.
+"""
+
+import typing
+from collections.abc import Callable, Iterable
+
+import regex as re
+import torch
+import torch.nn as nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mhc import HCHeadOp
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.deepseek_mtp import SharedHead
+from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+
+from .model import (
+    DeepseekV4DecoderLayer,
+    make_deepseek_v4_expert_params_mapping,
+)
+
+logger = init_logger(__name__)
+
+# MoE expert scales are fused into per-layer w13/w2 tensors. The exact
+# parameter suffix depends on which FusedMoE method handles the experts:
+# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``;
+# - fp8 experts (Fp8MoEMethod with block_quant=True) register
+#   ``w{1,2,3}_weight_scale_inv``.
+# Other FP8 linear scales (including shared experts) always use
+# ``.weight_scale_inv``. Mirrors the per-instance mapper built by
+# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py.
+_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
+
+
+class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        topk_indices_buffer: torch.Tensor,
+        prefix: str,
+        aux_stream_list: list | None = None,
+    ) -> None:
+        super().__init__()
+
+        assert vllm_config.speculative_config is not None
+        config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.config = config
+        quant_config = vllm_config.quant_config
+        self.rms_norm_eps = config.rms_norm_eps
+
+        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # V4 keeps e_ and h_ proj separate (with fp8 linear quant) rather than
+        # fusing them the way V3 does with eh_proj.
+        self.e_proj = ReplicatedLinear(
+            config.hidden_size,
+            config.hidden_size,
+            bias=False,
+            return_bias=False,
+            quant_config=quant_config,
+        )
+        self.h_proj = ReplicatedLinear(
+            config.hidden_size,
+            config.hidden_size,
+            bias=False,
+            return_bias=False,
+            quant_config=quant_config,
+        )
+
+        self.hc_eps = config.hc_eps
+        self.hc_mult = config.hc_mult
+        self.hc_dim = self.hc_mult * config.hidden_size
+        self.hc_head_fn = nn.Parameter(
+            torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32),
+            requires_grad=False,
+        )
+        self.hc_head_base = nn.Parameter(
+            torch.empty(self.hc_mult, dtype=torch.float32),
+            requires_grad=False,
+        )
+        self.hc_head_scale = nn.Parameter(
+            torch.empty(1, dtype=torch.float32),
+            requires_grad=False,
+        )
+
+        self.shared_head = SharedHead(
+            config=config, prefix=prefix, quant_config=quant_config
+        )
+        self.mtp_block = DeepseekV4DecoderLayer(
+            vllm_config,
+            prefix,
+            topk_indices_buffer=topk_indices_buffer,
+            aux_stream_list=aux_stream_list,
+        )
+
+        self.hc_head_op = HCHeadOp()
+
+    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
+        # masking inputs at position 0, as not needed by MTP
+        inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
+        inputs_embeds = self.enorm(inputs_embeds)
+
+        # Target stashes pre-hc_head residual as flat (T, hc_mult * D);
+        # reshape to (T, hc_mult, D) — the training-time layout.
+        previous_hidden_states = previous_hidden_states.view(
+            -1, self.hc_mult, self.config.hidden_size
+        )
+        previous_hidden_states = self.hnorm(previous_hidden_states)
+        hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
+            inputs_embeds
+        ).unsqueeze(-2)
+        hidden_states, residual, post_mix, res_mix = self.mtp_block(
+            positions=positions, x=hidden_states, input_ids=None
+        )
+        # Return the flat pre-hc_head residual so it can be re-fed as the
+        # next spec step's `previous_hidden_states` when
+        # num_speculative_tokens > 1. hc_head is deferred to compute_logits.
+        return hidden_states.flatten(1)
+
+
+class DeepSeekV4MultiTokenPredictor(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        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.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,
+        )
+
+        # Disable aux streams on XPU — no multi-stream overlap support.
+        aux_stream_list = None
+
+        # to map the exact layer index from weights
+        self.layers = torch.nn.ModuleDict(
+            {
+                str(idx): DeepSeekV4MultiTokenPredictorLayer(
+                    vllm_config,
+                    self.topk_indices_buffer,
+                    f"{prefix}.layers.{idx}",
+                    aux_stream_list=aux_stream_list,
+                )
+                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)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
+            input_ids,
+            positions,
+            previous_hidden_states,
+            inputs_embeds,
+            current_step_idx,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
+        # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply
+        # hc_head here so logits are computed from the dense hidden state.
+        hidden_states = hidden_states.view(
+            -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size
+        )
+        hidden_states = mtp_layer.hc_head_op(
+            hidden_states,
+            mtp_layer.hc_head_fn,
+            mtp_layer.hc_head_scale,
+            mtp_layer.hc_head_base,
+            mtp_layer.rms_norm_eps,
+            mtp_layer.hc_eps,
+        )
+        logits = self.logits_processor(
+            mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
+        )
+        return logits
+
+
+@support_torch_compile
+class DeepSeekV4MTP(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.config = vllm_config.model_config.hf_config
+        self.quant_config = vllm_config.quant_config
+        self.model = DeepSeekV4MultiTokenPredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | 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:
+        hidden_states = self.model(
+            input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
+        )
+        return hidden_states
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor | None:
+        return self.model.compute_logits(hidden_states, spec_step_idx)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Weight name remapping for checkpoint compatibility.
+        # Maps checkpoint weight paths to model parameter paths.
+        WEIGHT_NAME_REMAPPING: dict[str, str] = {
+            ".emb.tok_emb.weight": ".embed_tokens.weight",
+            ".head.weight": ".shared_head.head.weight",
+            ".norm.weight": ".shared_head.norm.weight",
+            # Pre-MoE norm + gate are now owned by
+            # ``DeepseekV4MoE.norm_gate`` (see NormGatedLinear).
+            ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
+            ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
+            ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
+        }
+
+        def _remap_weight_name(name: str) -> str:
+            """Remap checkpoint weight names to model parameter names."""
+            for old_pattern, new_pattern in WEIGHT_NAME_REMAPPING.items():
+                if old_pattern in name:
+                    name = name.replace(old_pattern, new_pattern)
+            return name
+
+        def _find_mtp_layer_idx(name: str) -> int:
+            subnames = name.split(".")
+            for subname in subnames:
+                try:
+                    # we return the first encountered integer
+                    return int(subname)
+                except ValueError:
+                    continue
+            return 0
+
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("gate_up_proj", "w1", 0),
+            ("gate_up_proj", "w3", 1),
+            ("attn.fused_wqa_wkv", "attn.wq_a", 0),
+            ("attn.fused_wqa_wkv", "attn.wkv", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        # TP for attention
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
+        n_head = self.config.num_attention_heads
+        n_local_head = n_head // tp_size
+        head_rank_start = n_local_head * tp_rank
+        head_rank_end = n_local_head * (tp_rank + 1)
+
+        # Pre-compute expert mapping ONCE.
+        first_layer = next(iter(self.model.layers.values()))
+        if first_layer.mtp_block.ffn.use_mega_moe:
+            expert_mapping = make_deepseek_v4_expert_params_mapping(
+                self.config.n_routed_experts
+            )
+        else:
+            expert_mapping = FusedMoE.make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="w1",
+                ckpt_down_proj_name="w2",
+                ckpt_up_proj_name="w3",
+                num_experts=self.config.n_routed_experts,
+            )
+
+        # FP8 experts register ``..._weight_scale_inv`` (block_quant) while
+        # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix
+        # for the rename below based on the model's expert dtype.
+        expert_scale_suffix = (
+            ".weight_scale"
+            if getattr(self.config, "expert_dtype", "fp4") == "fp4"
+            else ".weight_scale_inv"
+        )
+
+        for name, loaded_weight in weights:
+            mtp_layer_idx = _find_mtp_layer_idx(name)
+            # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to
+            # `model.layers.{num_hidden_layers + i}.*` so that
+            # get_spec_layer_idx_from_weight_name can identify them.
+            name = name.replace(
+                f"mtp.{mtp_layer_idx}.",
+                f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.",
+            )
+
+            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
+            if spec_layer is None:
+                continue
+
+            name = _remap_weight_name(name)
+            name = self._rewrite_spec_layer_name(spec_layer, name)
+
+            if spec_layer != self.model.mtp_start_layer_idx and ".layers" not in name:
+                continue
+            if name.endswith(".scale"):
+                suffix = (
+                    expert_scale_suffix
+                    if _EXPERT_SCALE_RE.search(name)
+                    else ".weight_scale_inv"
+                )
+                name = name.removesuffix(".scale") + suffix
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                # Skip non-stacked layers and experts (experts handled below).
+                if ".experts." in name:
+                    continue
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(name)
+                break
+            else:
+                if ".experts." in name:
+                    # Reinterpret E8M0 scales as uint8 to preserve raw
+                    # exponent bytes; numeric copy_() would zero them.
+                    # Mirrors the main DeepseekV4 loader.
+                    if (
+                        "weight_scale" in name
+                        and loaded_weight.dtype == torch.float8_e8m0fnu
+                    ):
+                        loaded_weight = loaded_weight.view(torch.uint8)
+                    for mapping in expert_mapping:
+                        param_name, weight_name, expert_id, expert_shard_id = mapping
+                        if weight_name not in name:
+                            continue
+                        name_mapped = name.replace(weight_name, param_name)
+                        param = params_dict[name_mapped]
+                        # We should ask the weight loader to return success or not
+                        # here since otherwise we may skip experts with other
+                        # available replicas.
+                        weight_loader = typing.cast(
+                            Callable[..., bool], param.weight_loader
+                        )
+                        success = weight_loader(
+                            param,
+                            loaded_weight,
+                            name_mapped,
+                            shard_id=expert_shard_id,
+                            expert_id=expert_id,
+                            return_success=True,
+                        )
+                        if success:
+                            name = name_mapped
+                            loaded_params.add(name_mapped)
+                            break
+                    continue
+                elif "attn_sink" in name:
+                    narrow_weight = loaded_weight[head_rank_start:head_rank_end]
+                    n = narrow_weight.shape[0]
+                    params_dict[name][:n].copy_(narrow_weight)
+                    loaded_params.add(name)
+                    continue
+                else:
+                    if ".shared_experts.w2" in name:
+                        name = name.replace(
+                            ".shared_experts.w2", ".shared_experts.down_proj"
+                        )
+                    if name.endswith(".ffn.gate.bias"):
+                        # ``e_score_correction_bias`` lives on
+                        # ``norm_gate`` directly (not on the inner gate).
+                        name = name.replace(
+                            ".ffn.gate.bias",
+                            ".ffn.norm_gate.e_score_correction_bias",
+                        )
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add(name)
+                    continue
+
+        loaded_layers: set[int] = set()
+        for param_name in loaded_params:
+            spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
+            if spec_layer is not None:
+                loaded_layers.add(spec_layer)
+        for layer_idx in range(
+            self.model.mtp_start_layer_idx,
+            self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
+        ):
+            if layer_idx not in loaded_layers:
+                raise ValueError(
+                    f"MTP speculative decoding layer {layer_idx} weights "
+                    f"missing from checkpoint. The checkpoint may have "
+                    f"been quantized without including the MTP layers. "
+                    f"Use a checkpoint that includes MTP layer weights, "
+                    f"or disable speculative decoding."
+                )
+        self.finalize_mega_moe_weights()
+        logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
+        return loaded_params
+
+    def finalize_mega_moe_weights(self) -> None:
+        for layer in self.model.layers.values():
+            layer.mtp_block.ffn.finalize_mega_moe_weights()
+
+    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
+        """
+        Rewrite the weight name to match the format of the original model.
+        Add .mtp_block for modules in transformer layer block for spec layer
+        and rename shared layer weights to be top level.
+        """
+        spec_layer_weight_names = [
+            "embed_tokens",
+            "enorm",
+            "hnorm",
+            "h_proj",
+            "e_proj",
+            "shared_head",
+            "hc_head_fn",
+            "hc_head_base",
+            "hc_head_scale",
+        ]
+        shared_weight_names = ["embed_tokens"]
+        spec_layer_weight = False
+        shared_weight = False
+        for weight_name in spec_layer_weight_names:
+            if weight_name in name:
+                spec_layer_weight = True
+                if weight_name in shared_weight_names:
+                    shared_weight = True
+                break
+        if not spec_layer_weight:
+            # treat rest weights as weights for transformer layer block
+            name = name.replace(
+                f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
+            )
+        elif shared_weight:
+            # treat shared weights as top level weights
+            name = name.replace(f"model.layers.{spec_layer}.", "model.")
+        return name

vllm/models/deepseek_v4/xpu/xpu_qnorm_rope_kv_fp8_insert.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""XPU Triton replacement for fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.
+
+Does: Q per-head RMSNorm + GPT-J RoPE, KV GPT-J RoPE + UE8M0 FP8 quant + insert.
+Uses the existing quantize_and_insert_k_cache for the FP8 portion.
+"""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+
+HEAD_DIM = 512
+ROPE_DIM = 64
+NOPE_DIM = HEAD_DIM - ROPE_DIM
+HALF_ROPE = ROPE_DIM // 2
+
+
+@triton.jit
+def _xpu_qnorm_rope_kernel(
+    q_ptr,  # [num_tokens, num_heads, HEAD_DIM]
+    kv_ptr,  # [num_tokens, HEAD_DIM]
+    kv_out_ptr,  # [num_tokens, HEAD_DIM] bf16 (RoPE-applied kv for cache insert)
+    position_ids_ptr,
+    cos_sin_cache_ptr,
+    eps: tl.constexpr,
+    num_tokens,
+    num_heads: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    ROPE_DIM: tl.constexpr,
+    NOPE_DIM: tl.constexpr,
+    HALF_ROPE: tl.constexpr,
+):
+    """Apply per-head RMSNorm + GPT-J RoPE on Q, GPT-J RoPE on KV.
+
+    GPT-J interleaved format: pairs are (data[2i], data[2i+1]).
+    cos_sin_cache layout: [max_pos, ROPE_DIM] with first HALF_ROPE=cos,
+    second HALF_ROPE=sin.
+    """
+    token_idx = tl.program_id(0)
+    head_idx = tl.program_id(1)
+
+    if token_idx >= num_tokens:
+        return
+
+    pos = tl.load(position_ids_ptr + token_idx).to(tl.int64)
+
+    # Load cos/sin for this position
+    rope_pair_idx = tl.arange(0, HALF_ROPE)
+    cos_val = tl.load(cos_sin_cache_ptr + pos * ROPE_DIM + rope_pair_idx).to(tl.float32)
+    sin_val = tl.load(
+        cos_sin_cache_ptr + pos * ROPE_DIM + HALF_ROPE + rope_pair_idx
+    ).to(tl.float32)
+
+    if head_idx < num_heads:
+        # ========== Q: per-head RMSNorm + GPT-J RoPE ==========
+        q_base = q_ptr + token_idx * num_heads * HEAD_DIM + head_idx * HEAD_DIM
+
+        # Load full head
+        offs = tl.arange(0, HEAD_DIM)
+        q_vals = tl.load(q_base + offs).to(tl.float32)
+
+        # RMSNorm (no weight)
+        sq_sum = tl.sum(q_vals * q_vals, axis=0)
+        rms = tl.rsqrt(sq_sum / HEAD_DIM + eps)
+        q_vals = q_vals * rms
+
+        # Store ONLY the NoPE portion (positions 0..NOPE_DIM-1)
+        nope_mask = offs < NOPE_DIM
+        tl.store(q_base + offs, q_vals.to(q_ptr.type.element_ty), mask=nope_mask)
+
+        # GPT-J interleaved RoPE on the last ROPE_DIM dimensions:
+        even_offs = NOPE_DIM + rope_pair_idx * 2
+        odd_offs = NOPE_DIM + rope_pair_idx * 2 + 1
+
+        # Re-load original values at rope positions and normalize
+        q_even = tl.load(q_base + even_offs).to(tl.float32) * rms
+        q_odd = tl.load(q_base + odd_offs).to(tl.float32) * rms
+
+        new_even = q_even * cos_val - q_odd * sin_val
+        new_odd = q_even * sin_val + q_odd * cos_val
+
+        # Store rotated RoPE values
+        tl.store(q_base + even_offs, new_even.to(q_ptr.type.element_ty))
+        tl.store(q_base + odd_offs, new_odd.to(q_ptr.type.element_ty))
+    else:
+        # ========== KV: GPT-J RoPE only ==========
+        kv_base = kv_ptr + token_idx * HEAD_DIM
+        kv_out_base = kv_out_ptr + token_idx * HEAD_DIM
+
+        # Copy full KV unchanged first
+        offs = tl.arange(0, HEAD_DIM)
+        kv_full = tl.load(kv_base + offs)
+        tl.store(kv_out_base + offs, kv_full)
+
+        # GPT-J interleaved RoPE on the last ROPE_DIM dimensions
+        even_offs = NOPE_DIM + rope_pair_idx * 2
+        odd_offs = NOPE_DIM + rope_pair_idx * 2 + 1
+
+        kv_even = tl.load(kv_base + even_offs).to(tl.float32)
+        kv_odd = tl.load(kv_base + odd_offs).to(tl.float32)
+
+        new_even = kv_even * cos_val - kv_odd * sin_val
+        new_odd = kv_even * sin_val + kv_odd * cos_val
+
+        tl.store(kv_out_base + even_offs, new_even.to(kv_out_ptr.type.element_ty))
+        tl.store(kv_out_base + odd_offs, new_odd.to(kv_out_ptr.type.element_ty))
+
+
+def xpu_qnorm_rope_kv_fp8_insert(
+    q: torch.Tensor,  # [num_tokens, num_heads, HEAD_DIM] bf16, in-place
+    kv: torch.Tensor,  # [num_tokens, HEAD_DIM] bf16
+    swa_kv_cache: torch.Tensor,  # [num_blocks, block_size, 584] or flat uint8
+    slot_mapping: torch.Tensor,  # [num_tokens] int64
+    positions: torch.Tensor,  # [num_tokens] int64
+    cos_sin_cache: torch.Tensor,  # [max_pos, ROPE_DIM]
+    eps: float,
+    block_size: int,
+):
+    """XPU Triton: qnorm+rope on Q, rope on KV, then FP8 UE8M0 quant+insert."""
+    from vllm.models.deepseek_v4.common.ops.cache_utils import (
+        quantize_and_insert_k_cache,
+    )
+
+    num_tokens = q.shape[0]
+    num_heads = q.shape[1]
+
+    # Allocate temp buffer for RoPE-applied KV
+    kv_roped = torch.empty_like(kv)
+
+    # Grid: one program per (token, head_or_kv)
+    # head_idx < num_heads: process Q head
+    # head_idx == num_heads: process KV
+    grid = (num_tokens, num_heads + 1)
+    _xpu_qnorm_rope_kernel[grid](
+        q,
+        kv,
+        kv_roped,
+        positions,
+        cos_sin_cache,
+        eps,
+        num_tokens,
+        num_heads=num_heads,
+        HEAD_DIM=HEAD_DIM,
+        ROPE_DIM=ROPE_DIM,
+        NOPE_DIM=NOPE_DIM,
+        HALF_ROPE=HALF_ROPE,
+    )
+
+    # FP8 UE8M0 quant + paged insert (reuse existing Triton kernel)
+    # swa_kv_cache may be [num_blocks, block_size, 584] or [num_blocks, flat]
+    # quantize_and_insert_k_cache expects [num_blocks, block_bytes] uint8
+    cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
+    quantize_and_insert_k_cache(
+        kv_roped,
+        cache_2d,
+        slot_mapping,
+        block_size=block_size,
+    )

vllm/models/deepseek_v4/xpu/xpu_sparse.py

@@ -0,0 +1,350 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""XPU DeepSeek-V4 attention subclass.
+
+Subclasses the shared ``DeepseekV4Attention`` ABC and provides XPU-native
+Triton kernels for decode (FP8 dequant + BF16 attention) and prefill
+(BF16 gathered KV + sparse attention).
+"""
+
+from typing import TYPE_CHECKING, cast
+
+import torch
+
+from vllm.forward_context import get_forward_context
+from vllm.models.deepseek_v4.attention import DeepseekV4Attention
+from vllm.models.deepseek_v4.common.ops import (
+    combine_topk_swa_indices,
+    compute_global_topk_indices_and_lens,
+    dequantize_and_gather_k_cache,
+)
+from vllm.models.deepseek_v4.sparse_mla import (
+    DeepseekV4FlashMLABackend,
+    DeepseekV4FlashMLAMetadata,
+)
+from vllm.models.deepseek_v4.xpu.xpu_sparse_decode_fp8 import (
+    xpu_sparse_decode_fp8,
+)
+from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
+from vllm.v1.worker.workspace import current_workspace_manager
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
+
+
+class DeepseekV4XPUSparseBackend(DeepseekV4FlashMLABackend):
+    @staticmethod
+    def get_name() -> str:
+        return "XPU_V4_MLA_SPARSE"
+
+
+class DeepseekV4XPUAttention(DeepseekV4Attention):
+    """XPU sparse MLA attention layer for DeepSeek V4."""
+
+    backend_cls = DeepseekV4XPUSparseBackend
+    use_flashmla_fp8_layout = True
+
+    def __init__(self, *args, **kwargs) -> None:
+        # torch.cuda.Event() raises RuntimeError on XPU ("dummy base class").
+        # The Base and DeepseekV4Indexer both create cuda Events in __init__, so
+        # we temporarily redirect torch.cuda.Event → torch.xpu.Event.
+        _orig_event = torch.cuda.Event
+        torch.cuda.Event = torch.xpu.Event  # type: ignore[misc]
+        try:
+            super().__init__(*args, **kwargs)
+        finally:
+            torch.cuda.Event = _orig_event  # type: ignore[misc]
+
+    def _fused_qnorm_rope_kv_insert(self, q, kv, positions, attn_metadata):
+        from typing import cast
+
+        if not isinstance(attn_metadata, dict):
+            # Profile run: no-op, just return q (no padding needed on XPU).
+            return q
+
+        swa_metadata = cast(
+            "DeepseekSparseSWAMetadata | None",
+            attn_metadata.get(self.swa_cache_layer.prefix),
+        )
+        assert swa_metadata is not None
+
+        from vllm.models.deepseek_v4.xpu.xpu_qnorm_rope_kv_fp8_insert import (
+            xpu_qnorm_rope_kv_fp8_insert,
+        )
+
+        xpu_qnorm_rope_kv_fp8_insert(
+            q,
+            kv,
+            self.swa_cache_layer.kv_cache,
+            swa_metadata.slot_mapping,
+            positions,
+            self.rotary_emb.cos_sin_cache,
+            self.eps,
+            swa_metadata.block_size,
+        )
+        return q
+
+    @classmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        return num_heads
+
+    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        # XPU uses BF16 reference wo_a path (same as ROCm).
+        from vllm.models.deepseek_v4.amd.rocm import rocm_inv_rope_einsum
+
+        z = rocm_inv_rope_einsum(
+            self.rotary_emb,
+            o,
+            positions,
+            self.rope_head_dim,
+            self.n_local_groups,
+            self.o_lora_rank,
+            self.wo_a,
+        )
+        return self.wo_b(z.flatten(1))
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        positions: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        assert output.shape == q.shape, (
+            f"output buffer shape {output.shape} must match q shape {q.shape}"
+        )
+        assert output.dtype == q.dtype, (
+            f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
+        )
+
+        forward_context = get_forward_context()
+        attn_metadata = forward_context.attn_metadata
+
+        if attn_metadata is None:
+            # Warmup dummy run: reserve workspace, skip actual kernels.
+            swa_only = self.compress_ratio <= 1
+            N = (
+                0
+                if swa_only
+                else (self.max_model_len + self.compress_ratio - 1)
+                // self.compress_ratio
+            )
+            M = N + self.window_size + self.max_num_batched_tokens
+            current_workspace_manager().get_simultaneous(
+                ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
+            )
+            output.zero_()
+            return
+
+        assert isinstance(attn_metadata, dict)
+        flashmla_metadata = cast(
+            DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)
+        )
+        swa_metadata = cast(
+            "DeepseekSparseSWAMetadata | None",
+            attn_metadata.get(self.swa_cache_layer.prefix),
+        )
+        assert swa_metadata is not None
+
+        swa_only = self.compress_ratio <= 1
+        self_kv_cache = self.kv_cache if not swa_only else None
+        swa_kv_cache = self.swa_cache_layer.kv_cache
+
+        # Split prefill and decode
+        num_decodes = swa_metadata.num_decodes
+        num_prefills = swa_metadata.num_prefills
+        num_decode_tokens = swa_metadata.num_decode_tokens
+
+        if num_prefills > 0:
+            self._forward_prefill(
+                q=q[num_decode_tokens:],
+                positions=positions[num_decode_tokens:],
+                compressed_k_cache=self_kv_cache,
+                swa_k_cache=swa_kv_cache,
+                output=output[num_decode_tokens:],
+                attn_metadata=flashmla_metadata,
+                swa_metadata=swa_metadata,
+            )
+        if num_decodes > 0:
+            self._forward_decode(
+                q=q[:num_decode_tokens],
+                kv_cache=self_kv_cache,
+                swa_metadata=swa_metadata,
+                attn_metadata=flashmla_metadata,
+                swa_only=swa_only,
+                output=output[:num_decode_tokens],
+            )
+
+    def _forward_decode(
+        self,
+        q: torch.Tensor,
+        kv_cache: torch.Tensor | None,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_only: bool,
+        output: torch.Tensor,
+    ) -> None:
+        num_decodes = swa_metadata.num_decodes
+        num_decode_tokens = swa_metadata.num_decode_tokens
+
+        topk_indices = None
+        topk_lens = None
+        if not swa_only:
+            assert attn_metadata is not None
+            assert swa_metadata.is_valid_token is not None
+            block_size = attn_metadata.block_size // self.compress_ratio
+            is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
+            if self.compress_ratio == 4:
+                # C4A: local indices differ per layer (filled by Indexer).
+                assert self.topk_indices_buffer is not None
+                global_indices, topk_lens = compute_global_topk_indices_and_lens(
+                    self.topk_indices_buffer[:num_decode_tokens],
+                    swa_metadata.token_to_req_indices,
+                    attn_metadata.block_table[:num_decodes],
+                    block_size,
+                    is_valid,
+                )
+                topk_indices = global_indices.view(num_decode_tokens, 1, -1)
+            else:
+                # C128A: pre-computed during metadata build.
+                topk_indices = attn_metadata.c128a_global_decode_topk_indices
+                topk_lens = attn_metadata.c128a_decode_topk_lens
+
+        swa_indices = swa_metadata.decode_swa_indices
+        swa_lens = swa_metadata.decode_swa_lens
+
+        assert swa_indices is not None and swa_lens is not None
+        xpu_sparse_decode_fp8(
+            q=q,
+            kv_cache=kv_cache,
+            swa_kv_cache=self.swa_cache_layer.kv_cache,
+            swa_only=swa_only,
+            topk_indices=topk_indices,
+            topk_lens=topk_lens,
+            swa_indices=swa_indices,
+            swa_lens=swa_lens,
+            attn_sink=self.attn_sink,
+            softmax_scale=self.scale,
+            head_dim=self.head_dim,
+            nope_head_dim=self.nope_head_dim,
+            rope_head_dim=self.rope_head_dim,
+            out=output,
+        )
+
+    def _forward_prefill(
+        self,
+        q: torch.Tensor,
+        positions: torch.Tensor,
+        compressed_k_cache: torch.Tensor | None,
+        swa_k_cache: torch.Tensor,
+        output: torch.Tensor,
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_metadata: "DeepseekSparseSWAMetadata",
+    ) -> None:
+        swa_only = attn_metadata is None
+
+        num_prefills = swa_metadata.num_prefills
+        num_prefill_tokens = swa_metadata.num_prefill_tokens
+        num_decodes = swa_metadata.num_decodes
+        num_decode_tokens = swa_metadata.num_decode_tokens
+
+        # Use pre-computed prefill metadata.
+        seq_lens = swa_metadata.prefill_seq_lens
+        gather_lens = swa_metadata.prefill_gather_lens
+        assert seq_lens is not None
+        assert gather_lens is not None
+
+        # Derive prefill-local token offsets from the full query_start_loc_cpu.
+        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
+        query_start_loc = swa_metadata.query_start_loc
+        assert query_start_loc_cpu is not None
+        assert query_start_loc is not None
+        prefill_token_base = query_start_loc_cpu[num_decodes]
+
+        if not swa_only:
+            if self.compress_ratio == 4:
+                assert self.topk_indices_buffer is not None
+                topk_indices = self.topk_indices_buffer[num_decode_tokens:]
+                topk_indices = topk_indices[:num_prefill_tokens]
+            else:
+                # C128A: pre-computed during metadata build.
+                assert attn_metadata is not None
+                topk_indices = attn_metadata.c128a_prefill_topk_indices
+            top_k = topk_indices.shape[-1]
+            N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
+        else:
+            assert self.topk_indices_buffer is not None
+            topk_indices = self.topk_indices_buffer[num_decode_tokens:]
+            top_k = 0
+            N = 0
+
+        M = N + self.window_size + self.max_num_batched_tokens
+        chunk_size_const = self.PREFILL_CHUNK_SIZE
+        num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const
+
+        workspace_manager = current_workspace_manager()
+        kv = workspace_manager.get_simultaneous(
+            ((chunk_size_const, M, q.shape[-1]), torch.bfloat16),
+        )[0]
+        for chunk_idx in range(num_chunks):
+            chunk_start = chunk_idx * chunk_size_const
+            chunk_end = min(chunk_start + chunk_size_const, num_prefills)
+            chunk_size = chunk_end - chunk_start
+            if not swa_only:
+                # Gather compressed KV
+                assert attn_metadata is not None
+                block_table = attn_metadata.block_table[num_decodes:]
+                dequantize_and_gather_k_cache(
+                    kv[:chunk_size],
+                    compressed_k_cache,
+                    seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
+                    gather_lens=None,
+                    block_table=block_table[chunk_start:chunk_end],
+                    block_size=attn_metadata.block_size // self.compress_ratio,
+                    offset=0,
+                )
+
+            # Gather SWA KV
+            swa_block_table = swa_metadata.block_table[num_decodes:]
+            dequantize_and_gather_k_cache(
+                kv[:chunk_size],
+                swa_k_cache,
+                seq_lens=seq_lens[chunk_start:chunk_end],
+                gather_lens=gather_lens[chunk_start:chunk_end],
+                block_table=swa_block_table[chunk_start:chunk_end],
+                block_size=swa_metadata.block_size,
+                offset=N,
+            )
+
+            # Combine the topk indices and SWA indices for gathered KV cache
+            query_start = (
+                query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
+            )
+            query_end = (
+                query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
+            )
+
+            combined_indices, combined_lens = combine_topk_swa_indices(
+                topk_indices[query_start:query_end],
+                query_start_loc[
+                    num_decodes + chunk_start : num_decodes + chunk_end + 1
+                ],
+                seq_lens[chunk_start:chunk_end],
+                gather_lens[chunk_start:chunk_end],
+                self.window_size,
+                self.compress_ratio,
+                top_k,
+                M,
+                N,
+            )
+
+            kv_ws = kv[:chunk_size].reshape(-1, 1, q.shape[-1])
+            out, _, _ = triton_bf16_mla_sparse_interface(
+                q=q[query_start:query_end],
+                kv=kv_ws,
+                indices=combined_indices.unsqueeze(1),
+                sm_scale=self.scale,
+                d_v=q.shape[-1],
+                block_dpe=0,
+            )
+            output[query_start:query_end] = out

vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py

@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""XPU sparse decode for DeepSeek V4 with FP8 KV cache.
+
+Strategy: dequantize FP8 UE8M0 KV cache pages to BF16 on the fly,
+then reuse the BF16 sparse MLA attention kernel (xpu_sparse_mla_bf16).
+This keeps the external KV cache layout identical to CUDA/ROCm.
+"""
+
+import torch
+
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.ops.xpu_mla_sparse import (
+    triton_bf16_mla_sparse_interface,
+)
+
+# FP8 DS MLA cache layout constants
+TOKEN_FP8_DIM = 448  # NoPE portion in FP8
+TOKEN_BF16_DIM = 64  # RoPE portion in BF16
+TOKEN_SCALE_DIM = 8  # UE8M0 scales per token
+QUANT_BLOCK_SIZE = 64  # Elements per quant block
+OUTPUT_DIM = 512  # = TOKEN_FP8_DIM + TOKEN_BF16_DIM after dequant
+TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2  # 576 bytes per token
+
+
+@triton.jit
+def _dequant_gather_slots_kernel(
+    # Output workspace: [total_slots, OUTPUT_DIM] bf16
+    out_ptr,
+    # FP8 paged cache base pointer (uint8 flat)
+    cache_ptr,
+    # Global slot indices: [total_slots] int32
+    indices_ptr,
+    # Cache geometry
+    cache_block_size: tl.constexpr,
+    token_data_size: tl.constexpr,  # 576
+    block_stride: tl.int64,  # k_cache.stride(0) — total uint8 per block
+    fp8_dim: tl.constexpr,  # 448
+    bf16_dim: tl.constexpr,  # 64
+    scale_dim: tl.constexpr,  # 8
+    quant_block: tl.constexpr,  # 64
+    output_dim: tl.constexpr,  # 512
+    n_quant_blocks: tl.constexpr,  # 7
+):
+    """Dequantize scattered FP8 slots into a flat BF16 workspace.
+
+    Grid: [total_slots] — one program per slot to gather.
+
+    Cache block layout (block_size tokens):
+      [0, block_size*576): Token data, each token 448 FP8 + 128 BF16
+      [block_size*576, block_size*576 + block_size*8): Scales
+    """
+    pid = tl.program_id(0)
+
+    # Load global slot index
+    slot_idx = tl.load(indices_ptr + pid).to(tl.int64)
+
+    # Output pointer for this slot
+    out_row_ptr = out_ptr + pid * output_dim
+
+    # Handle invalid slots (index < 0): write zeros
+    if slot_idx < 0:
+        zero = tl.zeros([quant_block], dtype=tl.bfloat16)
+        for i in tl.static_range(0, 512, 64):
+            offsets = i + tl.arange(0, quant_block)
+            mask = offsets < output_dim
+            tl.store(out_row_ptr + offsets, zero, mask=mask)
+        return
+
+    # Compute block and position within block
+    block_idx = slot_idx // cache_block_size
+    pos_in_block = slot_idx % cache_block_size
+
+    # Block base pointer
+    block_base = cache_ptr + block_idx * block_stride
+
+    # Token data: at offset pos_in_block * token_data_size within block
+    token_data_ptr = block_base + pos_in_block * token_data_size
+
+    # Scale: after all token data, at offset
+    # cache_block_size * token_data_size + pos_in_block * scale_dim
+    scale_region_offset = tl.cast(cache_block_size, tl.int64) * token_data_size
+    token_scale_ptr = block_base + scale_region_offset + pos_in_block * scale_dim
+
+    # ========== Dequantize FP8 portion (448 elements) ==========
+    for qblock_idx in tl.static_range(n_quant_blocks):
+        qblock_start = qblock_idx * quant_block
+        offsets = qblock_start + tl.arange(0, quant_block)
+        mask = offsets < fp8_dim
+
+        # Load FP8 as uint8 and bitcast
+        x_uint8 = tl.load(token_data_ptr + offsets, mask=mask, other=0)
+        x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
+        x_float = x_fp8.to(tl.float32)
+
+        # Load UE8M0 scale: scale = 2^(stored_value - 127)
+        encoded_scale = tl.load(token_scale_ptr + qblock_idx)
+        exponent = encoded_scale.to(tl.float32) - 127.0
+        scale = tl.exp2(exponent)
+
+        # Dequantize and store as bf16
+        x_dequant = x_float * scale
+        tl.store(out_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask)
+
+    # ========== Copy BF16 portion (64 elements) directly ==========
+    bf16_src_ptr = (token_data_ptr + fp8_dim).to(tl.pointer_type(tl.bfloat16))
+    bf16_out_ptr = (out_row_ptr + fp8_dim).to(tl.pointer_type(tl.bfloat16))
+
+    for j in tl.static_range(bf16_dim // 16):
+        chunk_offsets = j * 16 + tl.arange(0, 16)
+        bf16_vals = tl.load(bf16_src_ptr + chunk_offsets)
+        tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)
+
+
+def dequant_gather_slots(
+    out: torch.Tensor,  # [total_slots, 512] bf16, pre-allocated
+    cache: torch.Tensor,  # [num_blocks, block_size, head_bytes] uint8
+    indices: torch.Tensor,  # [total_slots] int32, global slot IDs
+    cache_block_size: int,  # block_size for this cache
+) -> None:
+    """Dequantize FP8 UE8M0 pages at scattered slot indices into bf16."""
+    total_slots = indices.shape[0]
+    if total_slots == 0:
+        return
+
+    block_stride = cache.stride(0)
+
+    _dequant_gather_slots_kernel[(total_slots,)](
+        out,
+        cache,
+        indices,
+        cache_block_size=cache_block_size,
+        token_data_size=TOKEN_DATA_SIZE,
+        block_stride=block_stride,
+        fp8_dim=TOKEN_FP8_DIM,
+        bf16_dim=TOKEN_BF16_DIM,
+        scale_dim=TOKEN_SCALE_DIM,
+        quant_block=QUANT_BLOCK_SIZE,
+        output_dim=OUTPUT_DIM,
+        n_quant_blocks=7,
+    )
+
+
+def xpu_sparse_decode_fp8(
+    q: torch.Tensor,  # [num_tokens, num_heads, head_dim]
+    kv_cache: torch.Tensor | None,  # [num_blocks, block_size, head_bytes] uint8
+    swa_kv_cache: torch.Tensor,  # [num_blocks, swa_block_size, head_bytes] uint8
+    swa_only: bool,
+    topk_indices: torch.Tensor | None,  # [num_tokens, 1, topk] global slot IDs
+    topk_lens: torch.Tensor | None,
+    swa_indices: torch.Tensor,  # [num_tokens, 1, swa_k] global slot IDs
+    swa_lens: torch.Tensor,
+    attn_sink: torch.Tensor,
+    softmax_scale: float,
+    head_dim: int,
+    nope_head_dim: int,
+    rope_head_dim: int,
+    out: torch.Tensor,  # [num_tokens, num_heads, head_dim]
+) -> None:
+    """XPU decode: dequant FP8 pages to BF16, then BF16 sparse MLA attention.
+
+    Keeps external FP8 KV cache layout identical to CUDA/ROCm.
+    Performance is slower due to on-the-fly dequant, but correctness is
+    guaranteed by reusing the validated BF16 attention kernel.
+    """
+    num_tokens = q.shape[0]
+    device = q.device
+
+    # Determine max topk and swa widths
+    if not swa_only and topk_indices is not None:
+        topk_idx_2d = (
+            topk_indices.squeeze(1) if topk_indices.dim() == 3 else topk_indices
+        )
+        max_topk = topk_idx_2d.shape[1]
+    else:
+        topk_idx_2d = None
+        max_topk = 0
+
+    swa_idx_2d = swa_indices.squeeze(1) if swa_indices.dim() == 3 else swa_indices
+    max_swa = swa_idx_2d.shape[1]
+
+    K_total = max_topk + max_swa
+
+    # Allocate flat workspace: [num_tokens * K_total, 512] bf16
+    workspace = torch.empty(
+        (num_tokens * K_total, OUTPUT_DIM), dtype=torch.bfloat16, device=device
+    )
+    ws_3d = workspace.view(num_tokens, K_total, OUTPUT_DIM)
+
+    # Dequant+gather topk slots from compressed cache
+    if not swa_only and topk_idx_2d is not None and kv_cache is not None:
+        topk_flat = topk_idx_2d.reshape(-1).to(torch.int32)
+        topk_buf = torch.empty(
+            (num_tokens * max_topk, OUTPUT_DIM), dtype=torch.bfloat16, device=device
+        )
+        compressed_block_size = kv_cache.shape[1]
+        dequant_gather_slots(topk_buf, kv_cache, topk_flat, compressed_block_size)
+        ws_3d[:, :max_topk, :] = topk_buf.view(num_tokens, max_topk, OUTPUT_DIM)
+
+    # Dequant+gather SWA slots
+    swa_flat = swa_idx_2d.reshape(-1).to(torch.int32)
+    swa_buf = torch.empty(
+        (num_tokens * max_swa, OUTPUT_DIM), dtype=torch.bfloat16, device=device
+    )
+    swa_block_size = swa_kv_cache.shape[1]
+    dequant_gather_slots(swa_buf, swa_kv_cache, swa_flat, swa_block_size)
+
+    ws_3d[:, max_topk:, :] = swa_buf.view(num_tokens, max_swa, OUTPUT_DIM)
+
+    # Build combined indices into the flat workspace and combined lengths.
+    # Workspace layout per token t: [topk_0..topk_{max_topk-1}, swa_0..swa_{max_swa-1}]
+    # Flat index for token t, position p = t * K_total + p
+    #
+    # IMPORTANT: The attention kernel uses combined_lens as a position cutoff —
+    # it only reads indices[0..combined_lens-1]. So indices must be PACKED
+    # contiguously: [valid_topk_indices..., valid_swa_indices..., -1 padding...]
+    if not swa_only and topk_lens is not None:
+        combined_lens = (topk_lens + swa_lens).to(torch.int32)
+    else:
+        combined_lens = swa_lens.to(torch.int32)
+
+    max_combined = int(combined_lens.max().item()) if combined_lens.numel() > 0 else 0
+    # Round up to BLOCK_N=16 alignment for kernel efficiency
+    _BLOCK_N = 16
+    max_combined_padded = ((max_combined + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N
+
+    # Build packed index table: [num_tokens, max_combined_padded]
+    # Each token t: [topk_0..topk_{tlen-1}, swa_0..swa_{slen-1}, -1 padding]
+    # Vectorized: for each token, topk indices are t*K_total + 0..tlen-1,
+    #             swa indices are t*K_total + max_topk + 0..slen-1
+    combined_indices = torch.full(
+        (num_tokens, max_combined_padded),
+        fill_value=-1,
+        dtype=torch.int32,
+        device=device,
+    )
+
+    token_offsets = (
+        torch.arange(num_tokens, device=device, dtype=torch.int32) * K_total
+    )  # [B]
+
+    if not swa_only and topk_lens is not None:
+        # Pack topk: for each token, write t*K_total + 0..tlen-1 at positions 0..tlen-1
+        max_tlen = int(topk_lens.max().item())
+        topk_range = torch.arange(max_tlen, device=device, dtype=torch.int32).unsqueeze(
+            0
+        )
+        topk_valid = topk_range < topk_lens.unsqueeze(1)
+        topk_ws_indices = token_offsets.unsqueeze(1) + topk_range
+        combined_indices[:, :max_tlen] = torch.where(
+            topk_valid,
+            topk_ws_indices,
+            torch.tensor(-1, dtype=torch.int32, device=device),
+        )
+        # Pack swa after topk: positions tlen..tlen+slen-1
+        # Since tlen varies per token, we need per-token offset
+        swa_range = torch.arange(max_swa, device=device, dtype=torch.int32).unsqueeze(0)
+        swa_valid = swa_range < swa_lens.unsqueeze(1)
+        swa_ws_indices = token_offsets.unsqueeze(1) + max_topk + swa_range
+        # Write at position topk_lens[t] + swa_pos for each token
+        for t_idx in range(num_tokens):
+            tlen = int(topk_lens[t_idx].item())
+            slen = int(swa_lens[t_idx].item())
+            combined_indices[t_idx, tlen : tlen + slen] = swa_ws_indices[t_idx, :slen]
+    else:
+        # SWA-only: pack swa indices at positions 0..slen-1
+        # Use min(max_swa, max_combined_padded) because combined_indices only
+        # has max_combined_padded columns, and all valid entries fit within it.
+        effective_swa = min(max_swa, max_combined_padded)
+        swa_range = torch.arange(
+            effective_swa, device=device, dtype=torch.int32
+        ).unsqueeze(0)
+        swa_valid = swa_range < swa_lens.unsqueeze(1)
+        swa_ws_indices = token_offsets.unsqueeze(1) + swa_range  # max_topk=0
+        combined_indices[:, :effective_swa] = torch.where(
+            swa_valid,
+            swa_ws_indices,
+            torch.tensor(-1, dtype=torch.int32, device=device),
+        )
+
+    # Call BF16 sparse MLA kernel
+    out_attn, _, _ = triton_bf16_mla_sparse_interface(
+        q=q,
+        kv=workspace.unsqueeze(1),
+        indices=combined_indices.unsqueeze(1),
+        sm_scale=softmax_scale,
+        d_v=q.shape[-1],
+        block_dpe=0,
+    )
+    out.copy_(out_attn)

vllm/v1/attention/ops/xpu_mla_sparse.py

@@ -180,11 +180,17 @@ def triton_bf16_mla_sparse_interface(
     indices: torch.Tensor,  # [num_tokens, num_heads_kv, topk]
     sm_scale: float,
     d_v: int = 512,
+    block_dpe: int = 64,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     out : [num_tokens, num_heads_q, d_v]
     max_logits : [num_tokens, num_heads_q]
     lse : logsumexp, [num_tokens, num_heads_q]
+
+    Args:
+        block_dpe: Size of positional embedding portion of dim_qk.
+            Set to 0 when q/kv contain only the nope latent (e.g. DSv4
+            prefill where RoPE is not split out).
     """
     num_tokens, num_heads_q, dim_qk = q.shape
     _, num_heads_kv, _ = kv.shape
@@ -194,8 +200,8 @@ def triton_bf16_mla_sparse_interface(
     _, _, index_topk = indices.shape
 
     BLOCK_H = 16
-    BLOCK_DMODEL = 512
-    BLOCK_DPE = 64
+    BLOCK_DPE = block_dpe
+    BLOCK_DMODEL = dim_qk - BLOCK_DPE
     BLOCK_M = 32
     BLOCK_N = 16
     BLOCK_DV = 512