A high-throughput and memory-efficient inference and serving engine for LLMs
[XPU]feat: add DeepSeek-V4 XPU attention decode path (#42953)
Signed-off-by: Ma Jian <jian1.ma@intel.com>
Ma Jian committed
Jun 8, 2026 at 13:27 UTC
eebce65756f0beb08547439728da98b5e1a7119c
11 files changed
+2759
-11
vllm/model_executor/kernels/linear/scaled_mm/xpu.py
+15
@@ -100,6 +100,21 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
100
"weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale"
101
)
102
scale = getattr(layer, scale_attr)
103
+ # Models with scale_fmt=ue8m0 (e.g. DeepSeek-V4) store weight scales
104
+ # as float8_e8m0fnu. The oneDNN fp8_gemm kernel dispatches to its
105
+ # "block quant" path only when NEITHER scale is e8m0:
106
+ #
107
+ # is_block_quant = (m1_sc != e8m0) && (m2_sc != e8m0) && ...
108
+ #
109
+ # Since activation scales are always float32 (use_ue8m0=False on XPU,
110
+ # DeepGEMM requires Hopper/Blackwell), an e8m0 weight scale causes
111
+ # is_block_quant=false and falls into the wrong per-channel path,
112
+ # producing NaN. Converting e8m0→float32 here at load time (one-time,
113
+ # negligible overhead for small scale tensors) ensures the kernel sees
114
+ # matching dtypes and correctly enters the block-quant path with the
115
+ # actual group_size derived from scale tensor shapes.
116
+ if scale.dtype == torch.float8_e8m0fnu:
117
+ scale = scale.to(torch.float32)
118
replace_parameter(layer, scale_attr, scale.data.t().contiguous())
119
120
def apply_block_scaled_mm(
vllm/model_executor/layers/mhc.py
+74
@@ -147,6 +147,36 @@ class MHCPreOp(CustomOp):
147
sinkhorn_repeat,
148
)
149
150
+ def forward_xpu(
151
+ self,
152
+ residual: torch.Tensor,
153
+ fn: torch.Tensor,
154
+ hc_scale: torch.Tensor,
155
+ hc_base: torch.Tensor,
156
+ rms_eps: float,
157
+ hc_pre_eps: float,
158
+ hc_sinkhorn_eps: float,
159
+ hc_post_mult_value: float,
160
+ sinkhorn_repeat: int,
161
+ n_splits: int = 1,
162
+ norm_weight: torch.Tensor | None = None,
163
+ norm_eps: float = 0.0,
164
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
165
+ return self.forward_native(
166
+ residual,
167
+ fn,
168
+ hc_scale,
169
+ hc_base,
170
+ rms_eps,
171
+ hc_pre_eps,
172
+ hc_sinkhorn_eps,
173
+ hc_post_mult_value,
174
+ sinkhorn_repeat,
175
+ n_splits,
176
+ norm_weight,
177
+ norm_eps,
178
+ )
179
+
180
181
# --8<-- [start:mhc_post]
182
@CustomOp.register("mhc_post")
@@ -215,6 +245,20 @@ class MHCPostOp(CustomOp):
245
comb_res_mix,
246
)
247
248
+ def forward_xpu(
249
+ self,
250
+ x: torch.Tensor,
251
+ residual: torch.Tensor,
252
+ post_layer_mix: torch.Tensor,
253
+ comb_res_mix: torch.Tensor,
254
+ ) -> torch.Tensor:
255
+ return self.forward_native(
256
+ x,
257
+ residual,
258
+ post_layer_mix,
259
+ comb_res_mix,
260
+ )
261
+
262
263
# --8<-- [start:hc_head]
264
@CustomOp.register("hc_head")
@@ -300,6 +344,36 @@ class HCHeadOp(CustomOp):
344
def forward_native(self, *args, **kwargs):
345
raise NotImplementedError("Native implementation of hc_head is not available")
346
347
+ def forward_xpu(
348
+ self,
349
+ hidden_states: torch.Tensor,
350
+ hc_fn: torch.Tensor,
351
+ hc_scale: torch.Tensor,
352
+ hc_base: torch.Tensor,
353
+ rms_norm_eps: float,
354
+ hc_eps: float,
355
+ ) -> torch.Tensor:
356
+ hc_mult, hidden_size = hidden_states.shape[-2:]
357
+ outer_shape = hidden_states.shape[:-2]
358
+ hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
359
+ num_tokens = hs_flat.shape[0]
360
+
361
+ out = torch.empty(
362
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
363
+ )
364
+ torch.ops.vllm.hc_head_triton(
365
+ hs_flat,
366
+ hc_fn,
367
+ hc_scale,
368
+ hc_base,
369
+ out,
370
+ hidden_size,
371
+ rms_norm_eps,
372
+ hc_eps,
373
+ hc_mult,
374
+ )
375
+ return out.view(*outer_shape, hidden_size)
376
+
377
378
# --8<-- [start:mhc_fused_post_pre]
379
@CustomOp.register("mhc_fused_post_pre")
vllm/models/deepseek_v4/__init__.py
+9
-8
@@ -7,21 +7,22 @@ picks the right one for the current platform and re-exports the public
7
classes used by the model registry and quantization config lookup.
8
"""
9
10
-from typing import TYPE_CHECKING
11
-
10
from vllm.platforms import current_platform
11
12
from .quant_config import DeepseekV4FP8Config
13
14
# Pick the per-platform implementation. The NVIDIA branch is the static
17
-# default that mypy sees; the ROCm branch overrides it at runtime and is
15
+# default that mypy sees; the ROCm/XPU branches override at runtime and are
16
# kept type-compatible via ``# type: ignore[assignment]``.
19
-if TYPE_CHECKING or not current_platform.is_rocm():
20
- from .nvidia.model import DeepseekV4ForCausalLM
21
- from .nvidia.mtp import DeepSeekV4MTP
17
+if current_platform.is_rocm():
18
+ from .amd.model import DeepseekV4ForCausalLM
19
+ from .amd.mtp import DeepSeekV4MTP
20
+elif current_platform.is_xpu():
21
+ from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment]
22
+ from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment]
23
else:
23
- from .amd.model import DeepseekV4ForCausalLM # type: ignore[assignment]
24
- from .amd.mtp import DeepSeekV4MTP # type: ignore[assignment]
24
+ from .nvidia.model import DeepseekV4ForCausalLM # type: ignore[assignment]
25
+ from .nvidia.mtp import DeepSeekV4MTP # type: ignore[assignment]
26
27
__all__ = [
28
"DeepSeekV4MTP",
vllm/models/deepseek_v4/compressor.py
+1
-1
@@ -368,7 +368,7 @@ class DeepseekCompressor(nn.Module):
368
fp8_scale=fp8_scale,
369
)
370
else:
371
- # Indexer path (head_dim == 128) or AMD: triton, legacy UE8M0 only.
371
+ # Indexer path (head_dim == 128) or non-CUDA GPUs (AMD, XPU, etc.).
372
compress_norm_rope_store_fn = compress_norm_rope_store_triton
373
extra_kwargs = {}
374
vllm/models/deepseek_v4/xpu/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v4/xpu/model.py
new
+1340
@@ -0,0 +1,1340 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import typing
4
+from collections.abc import Callable, Iterable
5
+from itertools import islice
6
+
7
+import regex as re
8
+import torch
9
+import torch.nn as nn
10
+
11
+from vllm.compilation.decorators import support_torch_compile
12
+from vllm.config import VllmConfig
13
+from vllm.distributed import (
14
+ get_ep_group,
15
+ get_pp_group,
16
+ get_tensor_model_parallel_rank,
17
+ get_tensor_model_parallel_world_size,
18
+)
19
+from vllm.forward_context import get_forward_context
20
+from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
21
+from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
22
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
23
+ fused_topk_bias,
24
+)
25
+from vllm.model_executor.layers.layernorm import RMSNorm
26
+from vllm.model_executor.layers.linear import (
27
+ MergedColumnParallelLinear,
28
+ RowParallelLinear,
29
+)
30
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
31
+from vllm.model_executor.layers.mhc import (
32
+ HCHeadOp,
33
+ MHCFusedPostPreOp,
34
+ MHCPostOp,
35
+ MHCPreOp,
36
+)
37
+from vllm.model_executor.layers.quantization import QuantizationConfig
38
+from vllm.model_executor.layers.vocab_parallel_embedding import (
39
+ ParallelLMHead,
40
+ VocabParallelEmbedding,
41
+)
42
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43
+from vllm.model_executor.models.interfaces import SupportsPP
44
+from vllm.model_executor.models.utils import (
45
+ AutoWeightsLoader,
46
+ PPMissingLayer,
47
+ WeightsMapper,
48
+ extract_layer_index,
49
+ is_pp_missing_parameter,
50
+ make_layers,
51
+ maybe_prefix,
52
+)
53
+from vllm.model_executor.utils import set_weight_attrs
54
+from vllm.models.deepseek_v4.xpu.xpu_sparse import DeepseekV4XPUAttention
55
+from vllm.platforms import current_platform
56
+from vllm.sequence import IntermediateTensors
57
+from vllm.triton_utils import tl, triton
58
+from vllm.utils.torch_utils import direct_register_custom_op
59
+
60
+
61
+class DeepseekV4MLP(nn.Module):
62
+ def __init__(
63
+ self,
64
+ hidden_size: int,
65
+ intermediate_size: int,
66
+ hidden_act: str,
67
+ swiglu_limit: float | None = None,
68
+ quant_config: QuantizationConfig | None = None,
69
+ reduce_results: bool = True,
70
+ is_sequence_parallel: bool = False,
71
+ prefix: str = "",
72
+ ) -> None:
73
+ super().__init__()
74
+
75
+ # If is_sequence_parallel, the input and output tensors are sharded
76
+ # across the ranks within the tp_group. In this case the weights are
77
+ # replicated and no collective ops are needed.
78
+ # Otherwise we use standard TP with an allreduce at the end.
79
+ self.gate_up_proj = MergedColumnParallelLinear(
80
+ hidden_size,
81
+ [intermediate_size] * 2,
82
+ bias=False,
83
+ quant_config=quant_config,
84
+ disable_tp=is_sequence_parallel,
85
+ prefix=f"{prefix}.gate_up_proj",
86
+ )
87
+ self.down_proj = RowParallelLinear(
88
+ intermediate_size,
89
+ hidden_size,
90
+ bias=False,
91
+ quant_config=quant_config,
92
+ reduce_results=reduce_results,
93
+ disable_tp=is_sequence_parallel,
94
+ prefix=f"{prefix}.down_proj",
95
+ )
96
+ if hidden_act != "silu":
97
+ raise ValueError(
98
+ f"Unsupported activation: {hidden_act}. Only silu is supported for now."
99
+ )
100
+ if swiglu_limit is not None:
101
+ self.act_fn = SiluAndMulWithClamp(swiglu_limit)
102
+ else:
103
+ self.act_fn = SiluAndMul()
104
+
105
+ def forward(self, x):
106
+ gate_up, _ = self.gate_up_proj(x)
107
+ x = self.act_fn(gate_up)
108
+ x, _ = self.down_proj(x)
109
+ return x
110
+
111
+
112
+@triton.jit
113
+def _deepseek_v4_stage_mega_moe_inputs_kernel(
114
+ hidden_states,
115
+ x_fp8,
116
+ x_sf,
117
+ topk_ids,
118
+ topk_weights,
119
+ topk_idx_out,
120
+ topk_weights_out,
121
+ hidden_stride_m: tl.constexpr,
122
+ hidden_stride_k: tl.constexpr,
123
+ x_stride_m: tl.constexpr,
124
+ x_stride_k: tl.constexpr,
125
+ x_sf_stride_m: tl.constexpr,
126
+ x_sf_stride_k: tl.constexpr,
127
+ topk_ids_stride_m: tl.constexpr,
128
+ topk_ids_stride_k: tl.constexpr,
129
+ topk_weights_stride_m: tl.constexpr,
130
+ topk_weights_stride_k: tl.constexpr,
131
+ topk_idx_stride_m: tl.constexpr,
132
+ topk_idx_stride_k: tl.constexpr,
133
+ topk_weights_out_stride_m: tl.constexpr,
134
+ topk_weights_out_stride_k: tl.constexpr,
135
+ hidden_size: tl.constexpr,
136
+ top_k: tl.constexpr,
137
+ BLOCK_K: tl.constexpr,
138
+ GROUP_K: tl.constexpr,
139
+ BLOCK_TOPK: tl.constexpr,
140
+) -> None:
141
+ token_id = tl.program_id(0)
142
+ k_block_id = tl.program_id(1)
143
+
144
+ k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
145
+ k_mask = k_offsets < hidden_size
146
+ hidden = tl.load(
147
+ hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
148
+ mask=k_mask,
149
+ other=0.0,
150
+ ).to(tl.float32)
151
+
152
+ num_groups: tl.constexpr = BLOCK_K // GROUP_K
153
+ hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
154
+ amax = tl.max(hidden_groups, axis=1)
155
+ amax = tl.maximum(amax, 1.0e-4)
156
+
157
+ scale = amax / 448.0
158
+ scale_bits = scale.to(tl.uint32, bitcast=True)
159
+ scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
160
+ tl.uint32
161
+ )
162
+ scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
163
+ rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
164
+
165
+ hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
166
+ scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
167
+ scaled = tl.reshape(scaled, [BLOCK_K])
168
+ fp8 = scaled.to(tl.float8e4nv)
169
+ tl.store(
170
+ x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
171
+ fp8,
172
+ mask=k_mask,
173
+ )
174
+
175
+ scale_offsets = tl.arange(0, num_groups)
176
+ packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
177
+ tl.store(
178
+ x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
179
+ packed_scale,
180
+ )
181
+
182
+ if k_block_id == 0:
183
+ topk_offsets = tl.arange(0, BLOCK_TOPK)
184
+ topk_mask = topk_offsets < top_k
185
+
186
+ ids = tl.load(
187
+ topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
188
+ mask=topk_mask,
189
+ other=0,
190
+ ).to(tl.int64)
191
+ tl.store(
192
+ topk_idx_out
193
+ + token_id * topk_idx_stride_m
194
+ + topk_offsets * topk_idx_stride_k,
195
+ ids,
196
+ mask=topk_mask,
197
+ )
198
+
199
+ weights = tl.load(
200
+ topk_weights
201
+ + token_id * topk_weights_stride_m
202
+ + topk_offsets * topk_weights_stride_k,
203
+ mask=topk_mask,
204
+ other=0.0,
205
+ )
206
+ tl.store(
207
+ topk_weights_out
208
+ + token_id * topk_weights_out_stride_m
209
+ + topk_offsets * topk_weights_out_stride_k,
210
+ weights,
211
+ mask=topk_mask,
212
+ )
213
+
214
+
215
+def _stage_deepseek_v4_mega_moe_inputs(
216
+ hidden_states: torch.Tensor,
217
+ topk_weights: torch.Tensor,
218
+ topk_ids: torch.Tensor,
219
+ x_fp8: torch.Tensor,
220
+ x_sf: torch.Tensor,
221
+ topk_idx_out: torch.Tensor,
222
+ topk_weights_out: torch.Tensor,
223
+) -> None:
224
+ num_tokens, hidden_size = hidden_states.shape
225
+ if num_tokens == 0:
226
+ return
227
+ if hidden_size % 128 != 0:
228
+ raise ValueError(
229
+ "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
230
+ "a multiple of 128."
231
+ )
232
+ top_k = topk_ids.shape[1]
233
+ if topk_weights.shape != topk_ids.shape:
234
+ raise ValueError(
235
+ "DeepSeek V4 MegaMoE input staging requires topk_weights and "
236
+ "topk_ids to have the same shape."
237
+ )
238
+
239
+ block_k = 128
240
+ grid = (num_tokens, triton.cdiv(hidden_size, block_k))
241
+ block_topk = triton.next_power_of_2(top_k)
242
+ _deepseek_v4_stage_mega_moe_inputs_kernel[grid](
243
+ hidden_states,
244
+ x_fp8,
245
+ x_sf,
246
+ topk_ids,
247
+ topk_weights,
248
+ topk_idx_out,
249
+ topk_weights_out,
250
+ hidden_states.stride(0),
251
+ hidden_states.stride(1),
252
+ x_fp8.stride(0),
253
+ x_fp8.stride(1),
254
+ x_sf.stride(0),
255
+ x_sf.stride(1),
256
+ topk_ids.stride(0),
257
+ topk_ids.stride(1),
258
+ topk_weights.stride(0),
259
+ topk_weights.stride(1),
260
+ topk_idx_out.stride(0),
261
+ topk_idx_out.stride(1),
262
+ topk_weights_out.stride(0),
263
+ topk_weights_out.stride(1),
264
+ hidden_size,
265
+ top_k,
266
+ BLOCK_K=block_k,
267
+ GROUP_K=32,
268
+ BLOCK_TOPK=block_topk,
269
+ num_warps=4,
270
+ )
271
+
272
+
273
+def make_deepseek_v4_expert_params_mapping(
274
+ num_experts: int,
275
+) -> list[tuple[str, str, int, str]]:
276
+ return [
277
+ (
278
+ "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_",
279
+ f"experts.{expert_id}.{weight_name}.",
280
+ expert_id,
281
+ shard_id,
282
+ )
283
+ for expert_id in range(num_experts)
284
+ for shard_id, weight_name in [
285
+ ("w1", "w1"),
286
+ ("w2", "w2"),
287
+ ("w3", "w3"),
288
+ ]
289
+ ]
290
+
291
+
292
+class DeepseekV4MegaMoEExperts(nn.Module):
293
+ _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
294
+
295
+ def __init__(
296
+ self,
297
+ vllm_config: VllmConfig,
298
+ *,
299
+ num_experts: int,
300
+ num_local_experts: int,
301
+ experts_start_idx: int,
302
+ top_k: int,
303
+ hidden_size: int,
304
+ intermediate_size: int,
305
+ prefix: str = "",
306
+ ):
307
+ super().__init__()
308
+ self.prefix = prefix
309
+ self.num_experts = num_experts
310
+ self.num_local_experts = num_local_experts
311
+ self.experts_start_idx = experts_start_idx
312
+ self.experts_end_idx = experts_start_idx + num_local_experts
313
+ self.top_k = top_k
314
+ self.hidden_size = hidden_size
315
+ self.intermediate_size = intermediate_size
316
+ self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
317
+
318
+ weight_attrs = {"weight_loader": self.weight_loader}
319
+ self.w13_weight = nn.Parameter(
320
+ torch.zeros(
321
+ num_local_experts,
322
+ 2 * intermediate_size,
323
+ hidden_size // 2,
324
+ dtype=torch.uint8,
325
+ ),
326
+ requires_grad=False,
327
+ )
328
+ set_weight_attrs(self.w13_weight, weight_attrs)
329
+
330
+ self.w13_weight_scale = nn.Parameter(
331
+ torch.zeros(
332
+ num_local_experts,
333
+ 2 * intermediate_size,
334
+ hidden_size // 32,
335
+ dtype=torch.uint8,
336
+ ),
337
+ requires_grad=False,
338
+ )
339
+ set_weight_attrs(self.w13_weight_scale, weight_attrs)
340
+ self.w13_weight_scale.quant_method = "block"
341
+
342
+ self.w2_weight = nn.Parameter(
343
+ torch.zeros(
344
+ num_local_experts,
345
+ hidden_size,
346
+ intermediate_size // 2,
347
+ dtype=torch.uint8,
348
+ ),
349
+ requires_grad=False,
350
+ )
351
+ set_weight_attrs(self.w2_weight, weight_attrs)
352
+
353
+ self.w2_weight_scale = nn.Parameter(
354
+ torch.zeros(
355
+ num_local_experts,
356
+ hidden_size,
357
+ intermediate_size // 32,
358
+ dtype=torch.uint8,
359
+ ),
360
+ requires_grad=False,
361
+ )
362
+ set_weight_attrs(self.w2_weight_scale, weight_attrs)
363
+ self.w2_weight_scale.quant_method = "block"
364
+
365
+ self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
366
+ self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
367
+
368
+ # Register in the static forward context so the custom-op wrapper
369
+ # can look up this module by name from within a torch.compile graph.
370
+ compilation_config = vllm_config.compilation_config
371
+ if prefix in compilation_config.static_forward_context:
372
+ raise ValueError(f"Duplicate layer name: {prefix}")
373
+ compilation_config.static_forward_context[prefix] = self
374
+
375
+ def _map_global_expert_id(self, expert_id: int) -> int:
376
+ if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:
377
+ return -1
378
+ return expert_id - self.experts_start_idx
379
+
380
+ def weight_loader(
381
+ self,
382
+ param: nn.Parameter,
383
+ loaded_weight: torch.Tensor,
384
+ weight_name: str,
385
+ shard_id: str,
386
+ expert_id: int,
387
+ return_success: bool = False,
388
+ ) -> bool | None:
389
+ local_expert_id = self._map_global_expert_id(expert_id)
390
+ if local_expert_id == -1:
391
+ return False if return_success else None
392
+
393
+ expert_data = param.data[local_expert_id]
394
+ if shard_id in ("w1", "w3"):
395
+ if "w13_" not in weight_name:
396
+ return False if return_success else None
397
+ shard_offset = 0 if shard_id == "w1" else self.intermediate_size
398
+ expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size)
399
+ elif shard_id == "w2":
400
+ if "w2_" not in weight_name:
401
+ return False if return_success else None
402
+ else:
403
+ raise ValueError(f"Unsupported expert shard id: {shard_id}")
404
+
405
+ if expert_data.shape != loaded_weight.shape:
406
+ raise ValueError(
407
+ f"DeepSeek V4 MegaMoE expert weight shape mismatch for "
408
+ f"{weight_name}: parameter shard {tuple(expert_data.shape)} "
409
+ f"vs checkpoint {tuple(loaded_weight.shape)}"
410
+ )
411
+ expert_data.copy_(loaded_weight)
412
+ return True if return_success else None
413
+
414
+ @staticmethod
415
+ def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor:
416
+ return (sf.to(torch.int32) << 23).view(torch.float32)
417
+
418
+ def _check_runtime_supported(self) -> None:
419
+ raise NotImplementedError("DeepSeek V4 MegaMoE is not supported on XPU.")
420
+
421
+ def finalize_weights(self) -> None:
422
+ if self._transformed_l1_weights is not None:
423
+ return
424
+
425
+ self._check_runtime_supported()
426
+ import vllm.third_party.deep_gemm as deep_gemm
427
+
428
+ w13_scale = deep_gemm.transform_sf_into_required_layout(
429
+ self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
430
+ 2 * self.intermediate_size,
431
+ self.hidden_size,
432
+ (1, 32),
433
+ self.num_local_experts,
434
+ )
435
+ w2_scale = deep_gemm.transform_sf_into_required_layout(
436
+ self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
437
+ self.hidden_size,
438
+ self.intermediate_size,
439
+ (1, 32),
440
+ self.num_local_experts,
441
+ )
442
+ self._transformed_l1_weights, self._transformed_l2_weights = (
443
+ deep_gemm.transform_weights_for_mega_moe(
444
+ (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
445
+ (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
446
+ )
447
+ )
448
+ # Drop the original loader-side parameters: the MegaMoE kernels only
449
+ # consume the transformed views above. transform_weights_for_mega_moe
450
+ # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
451
+ # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
452
+ # aliases the original storage, and _transformed_l2_weights still holds
453
+ # it, so the storage stays live after we drop the Parameter.
454
+ self.w13_weight = None
455
+ self.w13_weight_scale = None
456
+ self.w2_weight = None
457
+ self.w2_weight_scale = None
458
+
459
+ def get_symm_buffer(self):
460
+ import vllm.third_party.deep_gemm as deep_gemm
461
+
462
+ group = get_ep_group().device_group
463
+ device = torch.accelerator.current_device_index()
464
+ key = (
465
+ id(group),
466
+ device,
467
+ self.num_experts,
468
+ self.max_num_tokens,
469
+ self.top_k,
470
+ self.hidden_size,
471
+ self.intermediate_size,
472
+ )
473
+ symm_buffer = self._symm_buffer_cache.get(key)
474
+ if symm_buffer is None:
475
+ symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
476
+ group,
477
+ self.num_experts,
478
+ self.max_num_tokens,
479
+ self.top_k,
480
+ self.hidden_size,
481
+ self.intermediate_size,
482
+ )
483
+ self._symm_buffer_cache[key] = symm_buffer
484
+ return symm_buffer
485
+
486
+ def forward(
487
+ self,
488
+ hidden_states: torch.Tensor,
489
+ topk_weights: torch.Tensor,
490
+ topk_ids: torch.Tensor,
491
+ *,
492
+ activation_clamp: float | None,
493
+ fast_math: bool = True,
494
+ ) -> torch.Tensor:
495
+ if hidden_states.shape[0] > self.max_num_tokens:
496
+ raise ValueError(
497
+ f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
498
+ f"but the symmetric buffer was sized for {self.max_num_tokens}."
499
+ )
500
+ y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
501
+ torch.ops.vllm.deepseek_v4_mega_moe_experts(
502
+ hidden_states,
503
+ topk_weights,
504
+ topk_ids,
505
+ y,
506
+ self.prefix,
507
+ activation_clamp,
508
+ fast_math,
509
+ )
510
+ return y
511
+
512
+ def _run_mega_moe(
513
+ self,
514
+ hidden_states: torch.Tensor,
515
+ topk_weights: torch.Tensor,
516
+ topk_ids: torch.Tensor,
517
+ y: torch.Tensor,
518
+ activation_clamp: float | None,
519
+ fast_math: bool,
520
+ ) -> None:
521
+ import vllm.third_party.deep_gemm as deep_gemm
522
+
523
+ symm_buffer = self.get_symm_buffer()
524
+ num_tokens = hidden_states.shape[0]
525
+ _stage_deepseek_v4_mega_moe_inputs(
526
+ hidden_states,
527
+ topk_weights,
528
+ topk_ids,
529
+ symm_buffer.x[:num_tokens],
530
+ symm_buffer.x_sf[:num_tokens],
531
+ symm_buffer.topk_idx[:num_tokens],
532
+ symm_buffer.topk_weights[:num_tokens],
533
+ )
534
+
535
+ # This method must have been already called during the weight loading phase.
536
+ # We call it again here to cover the dummy weight loading case.
537
+ self.finalize_weights()
538
+
539
+ assert self._transformed_l1_weights is not None
540
+ assert self._transformed_l2_weights is not None
541
+ deep_gemm.fp8_fp4_mega_moe(
542
+ y,
543
+ self._transformed_l1_weights,
544
+ self._transformed_l2_weights,
545
+ symm_buffer,
546
+ activation_clamp=activation_clamp,
547
+ fast_math=fast_math,
548
+ )
549
+
550
+
551
+DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
552
+
553
+
554
+def _deepseek_v4_mega_moe_experts_op(
555
+ hidden_states: torch.Tensor,
556
+ topk_weights: torch.Tensor,
557
+ topk_ids: torch.Tensor,
558
+ out: torch.Tensor,
559
+ layer_name: str,
560
+ activation_clamp: float | None,
561
+ fast_math: bool,
562
+) -> None:
563
+ self = get_forward_context().no_compile_layers[layer_name]
564
+ self._run_mega_moe(
565
+ hidden_states,
566
+ topk_weights,
567
+ topk_ids,
568
+ out,
569
+ activation_clamp,
570
+ fast_math,
571
+ )
572
+
573
+
574
+def _deepseek_v4_mega_moe_experts_op_fake(
575
+ hidden_states: torch.Tensor,
576
+ topk_weights: torch.Tensor,
577
+ topk_ids: torch.Tensor,
578
+ out: torch.Tensor,
579
+ layer_name: str,
580
+ activation_clamp: float | None,
581
+ fast_math: bool,
582
+) -> None:
583
+ return None
584
+
585
+
586
+direct_register_custom_op(
587
+ op_name="deepseek_v4_mega_moe_experts",
588
+ op_func=_deepseek_v4_mega_moe_experts_op,
589
+ mutates_args=["out"],
590
+ fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
591
+)
592
+
593
+
594
+class DeepseekV4MoE(nn.Module):
595
+ def __init__(
596
+ self,
597
+ vllm_config: VllmConfig,
598
+ prefix: str = "",
599
+ ):
600
+ super().__init__()
601
+
602
+ self.tp_size = get_tensor_model_parallel_world_size()
603
+ config = vllm_config.model_config.hf_config
604
+ quant_config = vllm_config.quant_config
605
+ self.prefix = prefix
606
+ self.use_mega_moe = (
607
+ vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
608
+ )
609
+ if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
610
+ raise NotImplementedError(
611
+ "DeepSeek V4 MegaMoE currently requires expert parallel. "
612
+ "Enable it with --enable-expert-parallel, or pick a different "
613
+ "moe backend."
614
+ )
615
+
616
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
617
+ self.hidden_size = config.hidden_size
618
+
619
+ self.n_routed_experts = config.n_routed_experts
620
+ self.n_activated_experts = config.num_experts_per_tok
621
+ self.moe_intermediate_size = config.moe_intermediate_size
622
+ self.swiglu_limit = config.swiglu_limit
623
+ self.renormalize = config.norm_topk_prob
624
+ self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
625
+ if self.use_mega_moe and self.scoring_func != "sqrtsoftplus":
626
+ raise NotImplementedError(
627
+ "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
628
+ )
629
+ if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
630
+ raise NotImplementedError(
631
+ "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
632
+ f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
633
+ "deep_gemm_mega_moe for this checkpoint."
634
+ )
635
+
636
+ self.gate = GateLinear(
637
+ input_size=config.hidden_size,
638
+ output_size=config.n_routed_experts,
639
+ bias=False,
640
+ out_dtype=torch.float32,
641
+ prefix=f"{prefix}.gate",
642
+ )
643
+
644
+ self.gate.e_score_correction_bias = None
645
+ self.gate.tid2eid = None
646
+ is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
647
+ self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
648
+ if is_hash_moe:
649
+ # hash MoE doesn't use e_score_correction_bias
650
+ # Use randint instead of empty to avoid garbage values causing
651
+ # invalid memory access in dummy mode (--load-format="dummy")
652
+ self.gate.tid2eid = nn.Parameter(
653
+ torch.randint(
654
+ 0,
655
+ config.n_routed_experts,
656
+ (config.vocab_size, config.num_experts_per_tok),
657
+ dtype=self.hash_indices_dtype,
658
+ ),
659
+ requires_grad=False,
660
+ )
661
+ elif getattr(config, "topk_method", None) == "noaux_tc":
662
+ self.gate.e_score_correction_bias = nn.Parameter(
663
+ torch.empty(config.n_routed_experts, dtype=torch.float32),
664
+ requires_grad=False,
665
+ )
666
+
667
+ if config.n_shared_experts is None:
668
+ self.shared_experts = None
669
+ else:
670
+ intermediate_size = config.moe_intermediate_size * config.n_shared_experts
671
+
672
+ self.shared_experts = DeepseekV4MLP(
673
+ hidden_size=config.hidden_size,
674
+ intermediate_size=intermediate_size,
675
+ hidden_act=config.hidden_act,
676
+ swiglu_limit=self.swiglu_limit,
677
+ quant_config=quant_config,
678
+ reduce_results=self.use_mega_moe,
679
+ prefix=f"{prefix}.shared_experts",
680
+ )
681
+
682
+ if self.use_mega_moe:
683
+ self._init_mega_moe_experts(vllm_config, config, prefix)
684
+ else:
685
+ self._init_fused_moe_experts(config, quant_config, prefix)
686
+
687
+ def _init_mega_moe_experts(
688
+ self,
689
+ vllm_config: VllmConfig,
690
+ config,
691
+ prefix: str,
692
+ ) -> None:
693
+ self.ep_group = get_ep_group()
694
+ self.ep_size = self.ep_group.world_size
695
+ self.ep_rank = self.ep_group.rank_in_group
696
+ assert config.n_routed_experts % self.ep_size == 0
697
+
698
+ self.n_local_experts = config.n_routed_experts // self.ep_size
699
+ self.experts_start_idx = self.ep_rank * self.n_local_experts
700
+ self.experts_end_idx = self.experts_start_idx + self.n_local_experts
701
+
702
+ self.experts = DeepseekV4MegaMoEExperts(
703
+ vllm_config,
704
+ num_experts=config.n_routed_experts,
705
+ num_local_experts=self.n_local_experts,
706
+ experts_start_idx=self.experts_start_idx,
707
+ top_k=config.num_experts_per_tok,
708
+ hidden_size=config.hidden_size,
709
+ intermediate_size=config.moe_intermediate_size,
710
+ prefix=f"{prefix}.experts",
711
+ )
712
+
713
+ def _init_fused_moe_experts(
714
+ self,
715
+ config,
716
+ quant_config,
717
+ prefix: str,
718
+ ) -> None:
719
+ self.tp_rank = get_tensor_model_parallel_rank()
720
+ assert config.n_routed_experts % self.tp_size == 0
721
+
722
+ self.n_local_experts = config.n_routed_experts // self.tp_size
723
+ self.experts_start_idx = self.tp_rank * self.n_local_experts
724
+ self.experts_end_idx = self.experts_start_idx + self.n_local_experts
725
+ self.experts = FusedMoE(
726
+ shared_experts=self.shared_experts,
727
+ gate=self.gate,
728
+ num_experts=config.n_routed_experts,
729
+ top_k=config.num_experts_per_tok,
730
+ hidden_size=config.hidden_size,
731
+ intermediate_size=config.moe_intermediate_size,
732
+ renormalize=config.norm_topk_prob,
733
+ quant_config=quant_config,
734
+ prefix=f"{prefix}.experts",
735
+ scoring_func=self.scoring_func,
736
+ routed_scaling_factor=self.routed_scaling_factor,
737
+ e_score_correction_bias=self.gate.e_score_correction_bias,
738
+ hash_indices_table=self.gate.tid2eid,
739
+ swiglu_limit=self.swiglu_limit,
740
+ router_logits_dtype=torch.float32,
741
+ )
742
+
743
+ def forward(
744
+ self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
745
+ ) -> torch.Tensor:
746
+ if self.gate.tid2eid is not None and input_ids is None:
747
+ raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
748
+
749
+ if not self.use_mega_moe:
750
+ return self._forward_fused_moe(hidden_states, input_ids)
751
+
752
+ org_shape = hidden_states.shape
753
+ router_logits, _ = self.gate(hidden_states)
754
+ topk_weights, topk_ids = fused_topk_bias(
755
+ hidden_states=hidden_states,
756
+ gating_output=router_logits,
757
+ scoring_func=self.scoring_func,
758
+ e_score_correction_bias=self.gate.e_score_correction_bias.data
759
+ if self.gate.e_score_correction_bias is not None
760
+ else None,
761
+ topk=self.n_activated_experts,
762
+ renormalize=self.renormalize,
763
+ indices_type=self.hash_indices_dtype,
764
+ input_tokens=input_ids,
765
+ hash_indices_table=self.gate.tid2eid,
766
+ routed_scaling_factor=self.routed_scaling_factor,
767
+ )
768
+ activation_clamp = (
769
+ float(self.swiglu_limit) if self.swiglu_limit is not None else None
770
+ )
771
+ final_hidden_states = self.experts(
772
+ hidden_states,
773
+ topk_weights,
774
+ topk_ids,
775
+ activation_clamp=activation_clamp,
776
+ )
777
+
778
+ if self.shared_experts is not None:
779
+ shared_output = self.shared_experts(hidden_states)
780
+ final_hidden_states += shared_output
781
+
782
+ return final_hidden_states.view(org_shape)
783
+
784
+ def _forward_fused_moe(
785
+ self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
786
+ ) -> torch.Tensor:
787
+ org_shape = hidden_states.shape
788
+ if self.experts.is_internal_router:
789
+ final_hidden_states = self.experts(
790
+ hidden_states=hidden_states,
791
+ router_logits=hidden_states,
792
+ input_ids=input_ids,
793
+ )
794
+ else:
795
+ router_logits, _ = self.gate(hidden_states)
796
+ final_hidden_states = self.experts(
797
+ hidden_states=hidden_states,
798
+ router_logits=router_logits,
799
+ input_ids=input_ids,
800
+ )
801
+
802
+ return final_hidden_states.view(org_shape)
803
+
804
+ def finalize_mega_moe_weights(self) -> None:
805
+ if self.use_mega_moe:
806
+ self.experts.finalize_weights()
807
+
808
+
809
+class DeepseekV4DecoderLayer(nn.Module):
810
+ def __init__(
811
+ self,
812
+ vllm_config,
813
+ prefix,
814
+ topk_indices_buffer: torch.Tensor | None = None,
815
+ aux_stream_list: list | None = None,
816
+ ):
817
+ super().__init__()
818
+
819
+ # Lazy import to avoid top-level tilelang dependency.
820
+ # Registers both torch.ops.vllm.mhc_pre and mhc_post
821
+ import vllm.model_executor.layers.mhc # noqa: F401
822
+
823
+ config = vllm_config.model_config.hf_config
824
+ self.hidden_size = config.hidden_size
825
+
826
+ self.rms_norm_eps = config.rms_norm_eps
827
+ self.attn = DeepseekV4XPUAttention(
828
+ vllm_config,
829
+ prefix=f"{prefix}.attn",
830
+ topk_indices_buffer=topk_indices_buffer,
831
+ aux_stream_list=aux_stream_list,
832
+ )
833
+ self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
834
+
835
+ self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
836
+ self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
837
+ self.hc_mult = config.hc_mult
838
+ self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
839
+ self.hc_eps = config.hc_eps
840
+ self.hc_post_alpha = 2.0
841
+ mix_hc = (2 + self.hc_mult) * self.hc_mult
842
+ hc_dim = self.hc_mult * self.hidden_size
843
+ self.hc_attn_fn = nn.Parameter(
844
+ torch.empty(
845
+ (mix_hc, hc_dim),
846
+ dtype=torch.float32,
847
+ ),
848
+ requires_grad=False,
849
+ )
850
+ self.hc_ffn_fn = nn.Parameter(
851
+ torch.empty(
852
+ (mix_hc, hc_dim),
853
+ dtype=torch.float32,
854
+ ),
855
+ requires_grad=False,
856
+ )
857
+ self.hc_attn_base = nn.Parameter(
858
+ torch.empty(
859
+ mix_hc,
860
+ dtype=torch.float32,
861
+ ),
862
+ requires_grad=False,
863
+ )
864
+ self.hc_ffn_base = nn.Parameter(
865
+ torch.empty(
866
+ mix_hc,
867
+ dtype=torch.float32,
868
+ ),
869
+ requires_grad=False,
870
+ )
871
+ self.hc_attn_scale = nn.Parameter(
872
+ torch.empty(
873
+ 3,
874
+ dtype=torch.float32,
875
+ ),
876
+ requires_grad=False,
877
+ )
878
+ self.hc_ffn_scale = nn.Parameter(
879
+ torch.empty(
880
+ 3,
881
+ dtype=torch.float32,
882
+ ),
883
+ requires_grad=False,
884
+ )
885
+ self.mhc_pre = MHCPreOp()
886
+ self.mhc_post = MHCPostOp()
887
+ self.mhc_fused_post_pre = MHCFusedPostPreOp()
888
+
889
+ def hc_pre(
890
+ self,
891
+ x: torch.Tensor,
892
+ hc_fn: torch.Tensor,
893
+ hc_scale: torch.Tensor,
894
+ hc_base: torch.Tensor,
895
+ ):
896
+ post_mix, res_mix, layer_input = self.mhc_pre(
897
+ residual=x,
898
+ fn=hc_fn,
899
+ hc_scale=hc_scale,
900
+ hc_base=hc_base,
901
+ rms_eps=self.rms_norm_eps,
902
+ hc_pre_eps=self.hc_eps,
903
+ hc_sinkhorn_eps=self.hc_eps,
904
+ hc_post_mult_value=self.hc_post_alpha,
905
+ sinkhorn_repeat=self.hc_sinkhorn_iters,
906
+ )
907
+ return layer_input, post_mix, res_mix
908
+
909
+ def hc_post(
910
+ self,
911
+ x: torch.Tensor,
912
+ residual: torch.Tensor,
913
+ post: torch.Tensor,
914
+ comb: torch.Tensor,
915
+ ):
916
+ return self.mhc_post(x, residual, post, comb)
917
+
918
+ def forward(
919
+ self,
920
+ x: torch.Tensor,
921
+ positions: torch.Tensor,
922
+ input_ids: torch.Tensor | None,
923
+ post_mix: torch.Tensor | None = None,
924
+ res_mix: torch.Tensor | None = None,
925
+ residual: torch.Tensor | None = None,
926
+ ) -> tuple[
927
+ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
928
+ ]:
929
+ residual = x
930
+ x, post, comb = self.hc_pre(
931
+ x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
932
+ )
933
+ x = self.attn_norm(x)
934
+ x = self.attn(positions, x, None)
935
+ x = self.hc_post(x, residual, post, comb)
936
+
937
+ residual = x
938
+ x, post, comb = self.hc_pre(
939
+ x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
940
+ )
941
+ x = self.ffn_norm(x)
942
+ x = self.ffn(x, input_ids)
943
+ x = self.hc_post(x, residual, post, comb)
944
+ return x, None, None, None
945
+
946
+
947
+@support_torch_compile
948
+class DeepseekV4Model(nn.Module):
949
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
950
+ super().__init__()
951
+
952
+ config = vllm_config.model_config.hf_config
953
+ quant_config = vllm_config.quant_config
954
+ self.config = config
955
+ self.use_mega_moe = (
956
+ vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
957
+ )
958
+ if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
959
+ raise NotImplementedError(
960
+ "DeepSeek V4 MegaMoE currently requires expert parallel. "
961
+ "Enable it with --enable-expert-parallel, or pick a different "
962
+ "moe backend."
963
+ )
964
+ self.vocab_size = config.vocab_size
965
+ self.hc_eps = config.hc_eps
966
+ self.hc_mult = config.hc_mult
967
+ self.hc_dim = self.hc_mult * config.hidden_size
968
+ self.rms_norm_eps = config.rms_norm_eps
969
+
970
+ # Disable aux streams on XPU — no multi-stream overlap support.
971
+ aux_stream_list = None
972
+
973
+ self.device = current_platform.device_type
974
+ # Reserved topk indices buffer for all Indexer layers to reuse.
975
+ self.topk_indices_buffer = torch.empty(
976
+ vllm_config.scheduler_config.max_num_batched_tokens,
977
+ config.index_topk,
978
+ dtype=torch.int32,
979
+ device=self.device,
980
+ )
981
+
982
+ if get_pp_group().is_first_rank:
983
+ self.embed_tokens = VocabParallelEmbedding(
984
+ config.vocab_size,
985
+ config.hidden_size,
986
+ quant_config=quant_config,
987
+ prefix=f"{prefix}.embed_tokens",
988
+ )
989
+ else:
990
+ self.embed_tokens = PPMissingLayer()
991
+
992
+ self.start_layer, self.end_layer, self.layers = make_layers(
993
+ config.num_hidden_layers,
994
+ lambda prefix: DeepseekV4DecoderLayer(
995
+ vllm_config,
996
+ prefix=prefix,
997
+ topk_indices_buffer=self.topk_indices_buffer,
998
+ aux_stream_list=aux_stream_list,
999
+ ),
1000
+ prefix=f"{prefix}.layers",
1001
+ )
1002
+
1003
+ if get_pp_group().is_last_rank:
1004
+ self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
1005
+ else:
1006
+ self.norm = PPMissingLayer()
1007
+
1008
+ self.hc_head_fn = nn.Parameter(
1009
+ torch.empty(
1010
+ self.hc_mult,
1011
+ self.hc_dim,
1012
+ dtype=torch.float32,
1013
+ ),
1014
+ requires_grad=False,
1015
+ )
1016
+ self.hc_head_base = nn.Parameter(
1017
+ torch.empty(
1018
+ self.hc_mult,
1019
+ dtype=torch.float32,
1020
+ ),
1021
+ requires_grad=False,
1022
+ )
1023
+ self.hc_head_scale = nn.Parameter(
1024
+ torch.empty(1, dtype=torch.float32),
1025
+ requires_grad=False,
1026
+ )
1027
+ self.hc_head_op = HCHeadOp()
1028
+ # Pre-hc_head residual stream buffer for the MTP draft. Stable
1029
+ # address so the copy_ in forward() refreshes it correctly across
1030
+ # captured shapes. Only allocated on the last PP rank — that's
1031
+ # where MTP target hidden states are produced.
1032
+ if get_pp_group().is_last_rank:
1033
+ self._mtp_hidden_buffer = torch.empty(
1034
+ vllm_config.scheduler_config.max_num_batched_tokens,
1035
+ self.hc_dim,
1036
+ dtype=vllm_config.model_config.dtype,
1037
+ device=self.device,
1038
+ )
1039
+ else:
1040
+ self._mtp_hidden_buffer = None
1041
+
1042
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1043
+ return self.embed_tokens(input_ids)
1044
+
1045
+ def make_empty_intermediate_tensors(
1046
+ self,
1047
+ batch_size: int,
1048
+ dtype: torch.dtype,
1049
+ device: torch.device,
1050
+ ) -> IntermediateTensors:
1051
+ # PP intermediate tensors carry the multi-stream hidden_states
1052
+ # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the
1053
+ # token embedding to hc_mult streams before the first decoder
1054
+ # layer and keeps that shape until hc_head() collapses it.
1055
+ return IntermediateTensors(
1056
+ {
1057
+ "hidden_states": torch.zeros(
1058
+ (batch_size, self.hc_mult, self.config.hidden_size),
1059
+ dtype=dtype,
1060
+ device=device,
1061
+ ),
1062
+ }
1063
+ )
1064
+
1065
+ def forward(
1066
+ self,
1067
+ input_ids: torch.Tensor,
1068
+ positions: torch.Tensor,
1069
+ intermediate_tensors: IntermediateTensors | None,
1070
+ inputs_embeds: torch.Tensor | None = None,
1071
+ ) -> torch.Tensor | IntermediateTensors:
1072
+ if get_pp_group().is_first_rank:
1073
+ if inputs_embeds is not None:
1074
+ hidden_states = inputs_embeds
1075
+ else:
1076
+ hidden_states = self.embed_input_ids(input_ids)
1077
+ hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
1078
+ else:
1079
+ assert intermediate_tensors is not None
1080
+ hidden_states = intermediate_tensors["hidden_states"]
1081
+
1082
+ if self.use_mega_moe:
1083
+ input_ids = input_ids.to(torch.int64)
1084
+
1085
+ residual, post_mix, res_mix = None, None, None
1086
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
1087
+ hidden_states, residual, post_mix, res_mix = layer(
1088
+ hidden_states,
1089
+ positions,
1090
+ input_ids,
1091
+ post_mix,
1092
+ res_mix,
1093
+ residual,
1094
+ )
1095
+
1096
+ if not get_pp_group().is_last_rank:
1097
+ return IntermediateTensors({"hidden_states": hidden_states})
1098
+
1099
+ # Stash pre-hc_head residual for the MTP draft (captured copy_).
1100
+ num_tokens = hidden_states.shape[0]
1101
+ self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1102
+
1103
+ hidden_states = self.hc_head_op(
1104
+ hidden_states,
1105
+ self.hc_head_fn,
1106
+ self.hc_head_scale,
1107
+ self.hc_head_base,
1108
+ self.rms_norm_eps,
1109
+ self.hc_eps,
1110
+ )
1111
+ hidden_states = self.norm(hidden_states)
1112
+ return hidden_states
1113
+
1114
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1115
+ stacked_params_mapping = [
1116
+ # (param_name, shard_name, shard_id)
1117
+ ("gate_up_proj", "w1", 0),
1118
+ ("gate_up_proj", "w3", 1),
1119
+ ("attn.fused_wqa_wkv", "attn.wq_a", 0),
1120
+ ("attn.fused_wqa_wkv", "attn.wkv", 1),
1121
+ ("compressor.fused_wkv_wgate", "compressor.wkv", 0),
1122
+ ("compressor.fused_wkv_wgate", "compressor.wgate", 1),
1123
+ ]
1124
+ params_dict = dict(self.named_parameters())
1125
+ loaded_params: set[str] = set()
1126
+
1127
+ # TP for attention
1128
+ tp_size = get_tensor_model_parallel_world_size()
1129
+ tp_rank = get_tensor_model_parallel_rank()
1130
+ n_head = self.config.num_attention_heads
1131
+ n_local_head = n_head // tp_size
1132
+ head_rank_start = n_local_head * tp_rank
1133
+ head_rank_end = n_local_head * (tp_rank + 1)
1134
+
1135
+ # Pre-compute expert mapping ONCE.
1136
+ expert_mapping = self.get_expert_mapping()
1137
+
1138
+ for name, loaded_weight in weights:
1139
+ for param_name, weight_name, shard_id in stacked_params_mapping:
1140
+ # Skip non-stacked layers and experts (experts handled below).
1141
+ if ".experts." in name:
1142
+ continue
1143
+ if weight_name not in name:
1144
+ continue
1145
+ name = name.replace(weight_name, param_name)
1146
+
1147
+ if is_pp_missing_parameter(name, self):
1148
+ break
1149
+ param = params_dict[name]
1150
+ weight_loader = param.weight_loader
1151
+ weight_loader(param, loaded_weight, shard_id)
1152
+ loaded_params.add(name)
1153
+ break
1154
+ else:
1155
+ if ".experts." in name:
1156
+ # E8M0 scales are stored as float8_e8m0fnu in
1157
+ # checkpoints but the MoE param is uint8. copy_()
1158
+ # would do a numeric conversion (e.g. 2^-7 → 0),
1159
+ # destroying the raw exponent bytes.
1160
+ if (
1161
+ "weight_scale" in name
1162
+ and loaded_weight.dtype == torch.float8_e8m0fnu
1163
+ ):
1164
+ loaded_weight = loaded_weight.view(torch.uint8)
1165
+ for mapping in expert_mapping:
1166
+ param_name, weight_name, expert_id, expert_shard_id = mapping
1167
+ if weight_name not in name:
1168
+ continue
1169
+ name_mapped = name.replace(weight_name, param_name)
1170
+ if is_pp_missing_parameter(name_mapped, self):
1171
+ continue
1172
+ param = params_dict[name_mapped]
1173
+ # We should ask the weight loader to return success or not
1174
+ # here since otherwise we may skip experts with other
1175
+ # available replicas.
1176
+ weight_loader = typing.cast(
1177
+ Callable[..., bool], param.weight_loader
1178
+ )
1179
+ success = weight_loader(
1180
+ param,
1181
+ loaded_weight,
1182
+ name_mapped,
1183
+ shard_id=expert_shard_id,
1184
+ expert_id=expert_id,
1185
+ return_success=True,
1186
+ )
1187
+ if success:
1188
+ name = name_mapped
1189
+ break
1190
+ loaded_params.add(name_mapped)
1191
+ continue
1192
+ elif "attn_sink" in name:
1193
+ if is_pp_missing_parameter(name, self):
1194
+ continue
1195
+ if name not in params_dict:
1196
+ continue
1197
+ narrow_weight = loaded_weight[head_rank_start:head_rank_end]
1198
+ n = narrow_weight.shape[0]
1199
+ params_dict[name][:n].copy_(narrow_weight)
1200
+ loaded_params.add(name)
1201
+ continue
1202
+ else:
1203
+ if is_pp_missing_parameter(name, self):
1204
+ continue
1205
+ if name not in params_dict:
1206
+ continue
1207
+ param = params_dict[name]
1208
+ weight_loader = getattr(
1209
+ param, "weight_loader", default_weight_loader
1210
+ )
1211
+ weight_loader(param, loaded_weight)
1212
+ loaded_params.add(name)
1213
+ continue
1214
+
1215
+ return loaded_params
1216
+
1217
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1218
+ first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
1219
+ if first_layer.ffn.use_mega_moe:
1220
+ return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
1221
+ # Params for weights, fp8 weight scales, fp8 activation scales
1222
+ # (param_name, weight_name, expert_id, shard_id)
1223
+ return FusedMoE.make_expert_params_mapping(
1224
+ self,
1225
+ ckpt_gate_proj_name="w1",
1226
+ ckpt_down_proj_name="w2",
1227
+ ckpt_up_proj_name="w3",
1228
+ num_experts=self.config.n_routed_experts,
1229
+ )
1230
+
1231
+ def finalize_mega_moe_weights(self) -> None:
1232
+ for layer in islice(self.layers, self.start_layer, self.end_layer):
1233
+ layer.ffn.finalize_mega_moe_weights()
1234
+
1235
+
1236
+def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1237
+ if expert_dtype == "fp4":
1238
+ # MXFP4 experts use Mxfp4MoEMethod, which registers scales as
1239
+ # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and
1240
+ # shared experts use Fp8LinearMethod's block scales, which
1241
+ # register as ``weight_scale_inv``.
1242
+ scale_regex = {
1243
+ re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
1244
+ re.compile(r"\.scale$"): ".weight_scale_inv",
1245
+ }
1246
+ else:
1247
+ # FP8 experts use Fp8MoEMethod (block_quant=True), which registers
1248
+ # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
1249
+ # there.
1250
+ scale_regex = {
1251
+ re.compile(r"\.scale$"): ".weight_scale_inv",
1252
+ }
1253
+ return WeightsMapper(
1254
+ orig_to_new_prefix={
1255
+ "layers.": "model.layers.",
1256
+ "embed.": "model.embed.",
1257
+ "norm.": "model.norm.",
1258
+ "hc_head": "model.hc_head",
1259
+ "mtp.": "model.mtp.",
1260
+ },
1261
+ orig_to_new_regex=scale_regex,
1262
+ orig_to_new_suffix={
1263
+ "head.weight": "lm_head.weight",
1264
+ "embed.weight": "embed_tokens.weight",
1265
+ ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
1266
+ },
1267
+ orig_to_new_substr={
1268
+ ".shared_experts.w2": ".shared_experts.down_proj",
1269
+ },
1270
+ )
1271
+
1272
+
1273
+class DeepseekV4ForCausalLM(nn.Module, SupportsPP):
1274
+ model_cls = DeepseekV4Model
1275
+
1276
+ # Default mapper assumes the original FP4-expert checkpoint layout.
1277
+ # Overridden per-instance in __init__ when expert_dtype != "fp4".
1278
+ hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
1279
+
1280
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1281
+ super().__init__()
1282
+
1283
+ config = vllm_config.model_config.hf_config
1284
+ self.config = config
1285
+ expert_dtype = getattr(config, "expert_dtype", "fp4")
1286
+ if expert_dtype != "fp4":
1287
+ self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype)
1288
+
1289
+ self.model = self.model_cls(
1290
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
1291
+ )
1292
+ if get_pp_group().is_last_rank:
1293
+ self.lm_head = ParallelLMHead(
1294
+ config.vocab_size,
1295
+ config.hidden_size,
1296
+ prefix=maybe_prefix(prefix, "lm_head"),
1297
+ )
1298
+ else:
1299
+ self.lm_head = PPMissingLayer()
1300
+ self.logits_processor = LogitsProcessor(config.vocab_size)
1301
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
1302
+ self.model.make_empty_intermediate_tensors
1303
+ )
1304
+
1305
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1306
+ return self.model.embed_input_ids(input_ids)
1307
+
1308
+ def compute_logits(
1309
+ self,
1310
+ hidden_states: torch.Tensor,
1311
+ ) -> torch.Tensor | None:
1312
+ logits = self.logits_processor(self.lm_head, hidden_states)
1313
+ return logits
1314
+
1315
+ def forward(
1316
+ self,
1317
+ input_ids: torch.Tensor,
1318
+ positions: torch.Tensor,
1319
+ intermediate_tensors: IntermediateTensors | None = None,
1320
+ inputs_embeds: torch.Tensor | None = None,
1321
+ ) -> torch.Tensor | IntermediateTensors:
1322
+ hidden_states = self.model(
1323
+ input_ids, positions, intermediate_tensors, inputs_embeds
1324
+ )
1325
+ return hidden_states
1326
+
1327
+ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
1328
+ """Pre-hc_head residual stream buffer (max_num_batched_tokens,
1329
+ hc_mult * hidden_size) for the MTP draft model. Populated by
1330
+ forward(); valid after each target step."""
1331
+ return getattr(self.model, "_mtp_hidden_buffer", None)
1332
+
1333
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1334
+ loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
1335
+ loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1336
+ self.model.finalize_mega_moe_weights()
1337
+ return loaded_params
1338
+
1339
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1340
+ return self.model.get_expert_mapping()
vllm/models/deepseek_v4/xpu/mtp.py
new
+511
@@ -0,0 +1,511 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""MTP draft model for DeepSeek V4 (internal codename: DeepseekV4).
4
+
5
+Split from ``deepseek_mtp.py`` because the V4 architecture introduces several
6
+pieces that have no analogue in V3/V32:
7
+ * separate ``e_proj`` / ``h_proj`` with fp8 linear quantization (instead of
8
+ the fused ``eh_proj``);
9
+ * ``hc_head`` hypercompressed vocab projection applied in ``compute_logits``;
10
+ * ``DeepseekV4DecoderLayer`` with its own aux-stream management;
11
+ * V4-specific checkpoint weight-name remapping in ``load_weights``.
12
+"""
13
+
14
+import typing
15
+from collections.abc import Callable, Iterable
16
+
17
+import regex as re
18
+import torch
19
+import torch.nn as nn
20
+
21
+from vllm.compilation.decorators import support_torch_compile
22
+from vllm.config import VllmConfig
23
+from vllm.distributed import (
24
+ get_tensor_model_parallel_rank,
25
+ get_tensor_model_parallel_world_size,
26
+)
27
+from vllm.logger import init_logger
28
+from vllm.model_executor.layers.fused_moe import FusedMoE
29
+from vllm.model_executor.layers.layernorm import RMSNorm
30
+from vllm.model_executor.layers.linear import ReplicatedLinear
31
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
32
+from vllm.model_executor.layers.mhc import HCHeadOp
33
+from vllm.model_executor.layers.vocab_parallel_embedding import (
34
+ VocabParallelEmbedding,
35
+)
36
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37
+from vllm.model_executor.models.deepseek_mtp import SharedHead
38
+from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
39
+from vllm.model_executor.models.utils import maybe_prefix
40
+from vllm.platforms import current_platform
41
+from vllm.sequence import IntermediateTensors
42
+
43
+from .model import (
44
+ DeepseekV4DecoderLayer,
45
+ make_deepseek_v4_expert_params_mapping,
46
+)
47
+
48
+logger = init_logger(__name__)
49
+
50
+# MoE expert scales are fused into per-layer w13/w2 tensors. The exact
51
+# parameter suffix depends on which FusedMoE method handles the experts:
52
+# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``;
53
+# - fp8 experts (Fp8MoEMethod with block_quant=True) register
54
+# ``w{1,2,3}_weight_scale_inv``.
55
+# Other FP8 linear scales (including shared experts) always use
56
+# ``.weight_scale_inv``. Mirrors the per-instance mapper built by
57
+# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py.
58
+_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")
59
+
60
+
61
+class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
62
+ def __init__(
63
+ self,
64
+ vllm_config: VllmConfig,
65
+ topk_indices_buffer: torch.Tensor,
66
+ prefix: str,
67
+ aux_stream_list: list | None = None,
68
+ ) -> None:
69
+ super().__init__()
70
+
71
+ assert vllm_config.speculative_config is not None
72
+ config = vllm_config.speculative_config.draft_model_config.hf_config
73
+ self.config = config
74
+ quant_config = vllm_config.quant_config
75
+ self.rms_norm_eps = config.rms_norm_eps
76
+
77
+ self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
78
+ self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
79
+
80
+ # V4 keeps e_ and h_ proj separate (with fp8 linear quant) rather than
81
+ # fusing them the way V3 does with eh_proj.
82
+ self.e_proj = ReplicatedLinear(
83
+ config.hidden_size,
84
+ config.hidden_size,
85
+ bias=False,
86
+ return_bias=False,
87
+ quant_config=quant_config,
88
+ )
89
+ self.h_proj = ReplicatedLinear(
90
+ config.hidden_size,
91
+ config.hidden_size,
92
+ bias=False,
93
+ return_bias=False,
94
+ quant_config=quant_config,
95
+ )
96
+
97
+ self.hc_eps = config.hc_eps
98
+ self.hc_mult = config.hc_mult
99
+ self.hc_dim = self.hc_mult * config.hidden_size
100
+ self.hc_head_fn = nn.Parameter(
101
+ torch.empty(self.hc_mult, self.hc_dim, dtype=torch.float32),
102
+ requires_grad=False,
103
+ )
104
+ self.hc_head_base = nn.Parameter(
105
+ torch.empty(self.hc_mult, dtype=torch.float32),
106
+ requires_grad=False,
107
+ )
108
+ self.hc_head_scale = nn.Parameter(
109
+ torch.empty(1, dtype=torch.float32),
110
+ requires_grad=False,
111
+ )
112
+
113
+ self.shared_head = SharedHead(
114
+ config=config, prefix=prefix, quant_config=quant_config
115
+ )
116
+ self.mtp_block = DeepseekV4DecoderLayer(
117
+ vllm_config,
118
+ prefix,
119
+ topk_indices_buffer=topk_indices_buffer,
120
+ aux_stream_list=aux_stream_list,
121
+ )
122
+
123
+ self.hc_head_op = HCHeadOp()
124
+
125
+ def forward(
126
+ self,
127
+ input_ids: torch.Tensor,
128
+ positions: torch.Tensor,
129
+ previous_hidden_states: torch.Tensor,
130
+ inputs_embeds: torch.Tensor | None = None,
131
+ spec_step_index: int = 0,
132
+ ) -> torch.Tensor:
133
+ assert inputs_embeds is not None
134
+ # masking inputs at position 0, as not needed by MTP
135
+ inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
136
+ inputs_embeds = self.enorm(inputs_embeds)
137
+
138
+ # Target stashes pre-hc_head residual as flat (T, hc_mult * D);
139
+ # reshape to (T, hc_mult, D) — the training-time layout.
140
+ previous_hidden_states = previous_hidden_states.view(
141
+ -1, self.hc_mult, self.config.hidden_size
142
+ )
143
+ previous_hidden_states = self.hnorm(previous_hidden_states)
144
+ hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
145
+ inputs_embeds
146
+ ).unsqueeze(-2)
147
+ hidden_states, residual, post_mix, res_mix = self.mtp_block(
148
+ positions=positions, x=hidden_states, input_ids=None
149
+ )
150
+ # Return the flat pre-hc_head residual so it can be re-fed as the
151
+ # next spec step's `previous_hidden_states` when
152
+ # num_speculative_tokens > 1. hc_head is deferred to compute_logits.
153
+ return hidden_states.flatten(1)
154
+
155
+
156
+class DeepSeekV4MultiTokenPredictor(nn.Module):
157
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
158
+ super().__init__()
159
+ config = vllm_config.model_config.hf_config
160
+ self.mtp_start_layer_idx = config.num_hidden_layers
161
+ self.num_mtp_layers = config.num_nextn_predict_layers
162
+ self.device = current_platform.device_type
163
+
164
+ topk_tokens = config.index_topk
165
+ self.topk_indices_buffer = torch.empty(
166
+ vllm_config.scheduler_config.max_num_batched_tokens,
167
+ topk_tokens,
168
+ dtype=torch.int32,
169
+ device=self.device,
170
+ )
171
+
172
+ # Disable aux streams on XPU — no multi-stream overlap support.
173
+ aux_stream_list = None
174
+
175
+ # to map the exact layer index from weights
176
+ self.layers = torch.nn.ModuleDict(
177
+ {
178
+ str(idx): DeepSeekV4MultiTokenPredictorLayer(
179
+ vllm_config,
180
+ self.topk_indices_buffer,
181
+ f"{prefix}.layers.{idx}",
182
+ aux_stream_list=aux_stream_list,
183
+ )
184
+ for idx in range(
185
+ self.mtp_start_layer_idx,
186
+ self.mtp_start_layer_idx + self.num_mtp_layers,
187
+ )
188
+ }
189
+ )
190
+ self.embed_tokens = VocabParallelEmbedding(
191
+ config.vocab_size,
192
+ config.hidden_size,
193
+ prefix=maybe_prefix(prefix, "embed_tokens"),
194
+ )
195
+ self.logits_processor = LogitsProcessor(config.vocab_size)
196
+
197
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
198
+ return self.embed_tokens(input_ids)
199
+
200
+ def forward(
201
+ self,
202
+ input_ids: torch.Tensor,
203
+ positions: torch.Tensor,
204
+ previous_hidden_states: torch.Tensor,
205
+ inputs_embeds: torch.Tensor | None = None,
206
+ spec_step_idx: int = 0,
207
+ ) -> torch.Tensor:
208
+ if inputs_embeds is None:
209
+ inputs_embeds = self.embed_tokens(input_ids)
210
+ current_step_idx = spec_step_idx % self.num_mtp_layers
211
+ return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
212
+ input_ids,
213
+ positions,
214
+ previous_hidden_states,
215
+ inputs_embeds,
216
+ current_step_idx,
217
+ )
218
+
219
+ def compute_logits(
220
+ self,
221
+ hidden_states: torch.Tensor,
222
+ spec_step_idx: int = 0,
223
+ ) -> torch.Tensor:
224
+ current_step_idx = spec_step_idx % self.num_mtp_layers
225
+ mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
226
+ # MTP forward returns the pre-hc_head residual (T, hc_mult * D); apply
227
+ # hc_head here so logits are computed from the dense hidden state.
228
+ hidden_states = hidden_states.view(
229
+ -1, mtp_layer.hc_mult, mtp_layer.config.hidden_size
230
+ )
231
+ hidden_states = mtp_layer.hc_head_op(
232
+ hidden_states,
233
+ mtp_layer.hc_head_fn,
234
+ mtp_layer.hc_head_scale,
235
+ mtp_layer.hc_head_base,
236
+ mtp_layer.rms_norm_eps,
237
+ mtp_layer.hc_eps,
238
+ )
239
+ logits = self.logits_processor(
240
+ mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
241
+ )
242
+ return logits
243
+
244
+
245
+@support_torch_compile
246
+class DeepSeekV4MTP(nn.Module):
247
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
248
+ super().__init__()
249
+ self.config = vllm_config.model_config.hf_config
250
+ self.quant_config = vllm_config.quant_config
251
+ self.model = DeepSeekV4MultiTokenPredictor(
252
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
253
+ )
254
+
255
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
256
+ return self.model.embed_input_ids(input_ids)
257
+
258
+ def forward(
259
+ self,
260
+ input_ids: torch.Tensor | None,
261
+ positions: torch.Tensor,
262
+ hidden_states: torch.Tensor,
263
+ intermediate_tensors: IntermediateTensors | None = None,
264
+ inputs_embeds: torch.Tensor | None = None,
265
+ spec_step_idx: int = 0,
266
+ ) -> torch.Tensor:
267
+ hidden_states = self.model(
268
+ input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
269
+ )
270
+ return hidden_states
271
+
272
+ def compute_logits(
273
+ self,
274
+ hidden_states: torch.Tensor,
275
+ spec_step_idx: int = 0,
276
+ ) -> torch.Tensor | None:
277
+ return self.model.compute_logits(hidden_states, spec_step_idx)
278
+
279
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
280
+ # Weight name remapping for checkpoint compatibility.
281
+ # Maps checkpoint weight paths to model parameter paths.
282
+ WEIGHT_NAME_REMAPPING: dict[str, str] = {
283
+ ".emb.tok_emb.weight": ".embed_tokens.weight",
284
+ ".head.weight": ".shared_head.head.weight",
285
+ ".norm.weight": ".shared_head.norm.weight",
286
+ # Pre-MoE norm + gate are now owned by
287
+ # ``DeepseekV4MoE.norm_gate`` (see NormGatedLinear).
288
+ ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
289
+ ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
290
+ ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
291
+ }
292
+
293
+ def _remap_weight_name(name: str) -> str:
294
+ """Remap checkpoint weight names to model parameter names."""
295
+ for old_pattern, new_pattern in WEIGHT_NAME_REMAPPING.items():
296
+ if old_pattern in name:
297
+ name = name.replace(old_pattern, new_pattern)
298
+ return name
299
+
300
+ def _find_mtp_layer_idx(name: str) -> int:
301
+ subnames = name.split(".")
302
+ for subname in subnames:
303
+ try:
304
+ # we return the first encountered integer
305
+ return int(subname)
306
+ except ValueError:
307
+ continue
308
+ return 0
309
+
310
+ stacked_params_mapping = [
311
+ # (param_name, shard_name, shard_id)
312
+ ("gate_up_proj", "w1", 0),
313
+ ("gate_up_proj", "w3", 1),
314
+ ("attn.fused_wqa_wkv", "attn.wq_a", 0),
315
+ ("attn.fused_wqa_wkv", "attn.wkv", 1),
316
+ ]
317
+ params_dict = dict(self.named_parameters())
318
+ loaded_params: set[str] = set()
319
+
320
+ # TP for attention
321
+ tp_size = get_tensor_model_parallel_world_size()
322
+ tp_rank = get_tensor_model_parallel_rank()
323
+ n_head = self.config.num_attention_heads
324
+ n_local_head = n_head // tp_size
325
+ head_rank_start = n_local_head * tp_rank
326
+ head_rank_end = n_local_head * (tp_rank + 1)
327
+
328
+ # Pre-compute expert mapping ONCE.
329
+ first_layer = next(iter(self.model.layers.values()))
330
+ if first_layer.mtp_block.ffn.use_mega_moe:
331
+ expert_mapping = make_deepseek_v4_expert_params_mapping(
332
+ self.config.n_routed_experts
333
+ )
334
+ else:
335
+ expert_mapping = FusedMoE.make_expert_params_mapping(
336
+ self,
337
+ ckpt_gate_proj_name="w1",
338
+ ckpt_down_proj_name="w2",
339
+ ckpt_up_proj_name="w3",
340
+ num_experts=self.config.n_routed_experts,
341
+ )
342
+
343
+ # FP8 experts register ``..._weight_scale_inv`` (block_quant) while
344
+ # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix
345
+ # for the rename below based on the model's expert dtype.
346
+ expert_scale_suffix = (
347
+ ".weight_scale"
348
+ if getattr(self.config, "expert_dtype", "fp4") == "fp4"
349
+ else ".weight_scale_inv"
350
+ )
351
+
352
+ for name, loaded_weight in weights:
353
+ mtp_layer_idx = _find_mtp_layer_idx(name)
354
+ # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to
355
+ # `model.layers.{num_hidden_layers + i}.*` so that
356
+ # get_spec_layer_idx_from_weight_name can identify them.
357
+ name = name.replace(
358
+ f"mtp.{mtp_layer_idx}.",
359
+ f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.",
360
+ )
361
+
362
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
363
+ if spec_layer is None:
364
+ continue
365
+
366
+ name = _remap_weight_name(name)
367
+ name = self._rewrite_spec_layer_name(spec_layer, name)
368
+
369
+ if spec_layer != self.model.mtp_start_layer_idx and ".layers" not in name:
370
+ continue
371
+ if name.endswith(".scale"):
372
+ suffix = (
373
+ expert_scale_suffix
374
+ if _EXPERT_SCALE_RE.search(name)
375
+ else ".weight_scale_inv"
376
+ )
377
+ name = name.removesuffix(".scale") + suffix
378
+ for param_name, weight_name, shard_id in stacked_params_mapping:
379
+ # Skip non-stacked layers and experts (experts handled below).
380
+ if ".experts." in name:
381
+ continue
382
+ if weight_name not in name:
383
+ continue
384
+ name = name.replace(weight_name, param_name)
385
+
386
+ param = params_dict[name]
387
+ weight_loader = param.weight_loader
388
+ weight_loader(param, loaded_weight, shard_id)
389
+ loaded_params.add(name)
390
+ break
391
+ else:
392
+ if ".experts." in name:
393
+ # Reinterpret E8M0 scales as uint8 to preserve raw
394
+ # exponent bytes; numeric copy_() would zero them.
395
+ # Mirrors the main DeepseekV4 loader.
396
+ if (
397
+ "weight_scale" in name
398
+ and loaded_weight.dtype == torch.float8_e8m0fnu
399
+ ):
400
+ loaded_weight = loaded_weight.view(torch.uint8)
401
+ for mapping in expert_mapping:
402
+ param_name, weight_name, expert_id, expert_shard_id = mapping
403
+ if weight_name not in name:
404
+ continue
405
+ name_mapped = name.replace(weight_name, param_name)
406
+ param = params_dict[name_mapped]
407
+ # We should ask the weight loader to return success or not
408
+ # here since otherwise we may skip experts with other
409
+ # available replicas.
410
+ weight_loader = typing.cast(
411
+ Callable[..., bool], param.weight_loader
412
+ )
413
+ success = weight_loader(
414
+ param,
415
+ loaded_weight,
416
+ name_mapped,
417
+ shard_id=expert_shard_id,
418
+ expert_id=expert_id,
419
+ return_success=True,
420
+ )
421
+ if success:
422
+ name = name_mapped
423
+ loaded_params.add(name_mapped)
424
+ break
425
+ continue
426
+ elif "attn_sink" in name:
427
+ narrow_weight = loaded_weight[head_rank_start:head_rank_end]
428
+ n = narrow_weight.shape[0]
429
+ params_dict[name][:n].copy_(narrow_weight)
430
+ loaded_params.add(name)
431
+ continue
432
+ else:
433
+ if ".shared_experts.w2" in name:
434
+ name = name.replace(
435
+ ".shared_experts.w2", ".shared_experts.down_proj"
436
+ )
437
+ if name.endswith(".ffn.gate.bias"):
438
+ # ``e_score_correction_bias`` lives on
439
+ # ``norm_gate`` directly (not on the inner gate).
440
+ name = name.replace(
441
+ ".ffn.gate.bias",
442
+ ".ffn.norm_gate.e_score_correction_bias",
443
+ )
444
+ param = params_dict[name]
445
+ weight_loader = getattr(
446
+ param, "weight_loader", default_weight_loader
447
+ )
448
+ weight_loader(param, loaded_weight)
449
+ loaded_params.add(name)
450
+ continue
451
+
452
+ loaded_layers: set[int] = set()
453
+ for param_name in loaded_params:
454
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
455
+ if spec_layer is not None:
456
+ loaded_layers.add(spec_layer)
457
+ for layer_idx in range(
458
+ self.model.mtp_start_layer_idx,
459
+ self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
460
+ ):
461
+ if layer_idx not in loaded_layers:
462
+ raise ValueError(
463
+ f"MTP speculative decoding layer {layer_idx} weights "
464
+ f"missing from checkpoint. The checkpoint may have "
465
+ f"been quantized without including the MTP layers. "
466
+ f"Use a checkpoint that includes MTP layer weights, "
467
+ f"or disable speculative decoding."
468
+ )
469
+ self.finalize_mega_moe_weights()
470
+ logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
471
+ return loaded_params
472
+
473
+ def finalize_mega_moe_weights(self) -> None:
474
+ for layer in self.model.layers.values():
475
+ layer.mtp_block.ffn.finalize_mega_moe_weights()
476
+
477
+ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
478
+ """
479
+ Rewrite the weight name to match the format of the original model.
480
+ Add .mtp_block for modules in transformer layer block for spec layer
481
+ and rename shared layer weights to be top level.
482
+ """
483
+ spec_layer_weight_names = [
484
+ "embed_tokens",
485
+ "enorm",
486
+ "hnorm",
487
+ "h_proj",
488
+ "e_proj",
489
+ "shared_head",
490
+ "hc_head_fn",
491
+ "hc_head_base",
492
+ "hc_head_scale",
493
+ ]
494
+ shared_weight_names = ["embed_tokens"]
495
+ spec_layer_weight = False
496
+ shared_weight = False
497
+ for weight_name in spec_layer_weight_names:
498
+ if weight_name in name:
499
+ spec_layer_weight = True
500
+ if weight_name in shared_weight_names:
501
+ shared_weight = True
502
+ break
503
+ if not spec_layer_weight:
504
+ # treat rest weights as weights for transformer layer block
505
+ name = name.replace(
506
+ f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
507
+ )
508
+ elif shared_weight:
509
+ # treat shared weights as top level weights
510
+ name = name.replace(f"model.layers.{spec_layer}.", "model.")
511
+ return name
vllm/models/deepseek_v4/xpu/xpu_qnorm_rope_kv_fp8_insert.py
new
+159
@@ -0,0 +1,159 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""XPU Triton replacement for fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.
4
+
5
+Does: Q per-head RMSNorm + GPT-J RoPE, KV GPT-J RoPE + UE8M0 FP8 quant + insert.
6
+Uses the existing quantize_and_insert_k_cache for the FP8 portion.
7
+"""
8
+
9
+import torch
10
+
11
+from vllm.triton_utils import tl, triton
12
+
13
+HEAD_DIM = 512
14
+ROPE_DIM = 64
15
+NOPE_DIM = HEAD_DIM - ROPE_DIM
16
+HALF_ROPE = ROPE_DIM // 2
17
+
18
+
19
+@triton.jit
20
+def _xpu_qnorm_rope_kernel(
21
+ q_ptr, # [num_tokens, num_heads, HEAD_DIM]
22
+ kv_ptr, # [num_tokens, HEAD_DIM]
23
+ kv_out_ptr, # [num_tokens, HEAD_DIM] bf16 (RoPE-applied kv for cache insert)
24
+ position_ids_ptr,
25
+ cos_sin_cache_ptr,
26
+ eps: tl.constexpr,
27
+ num_tokens,
28
+ num_heads: tl.constexpr,
29
+ HEAD_DIM: tl.constexpr,
30
+ ROPE_DIM: tl.constexpr,
31
+ NOPE_DIM: tl.constexpr,
32
+ HALF_ROPE: tl.constexpr,
33
+):
34
+ """Apply per-head RMSNorm + GPT-J RoPE on Q, GPT-J RoPE on KV.
35
+
36
+ GPT-J interleaved format: pairs are (data[2i], data[2i+1]).
37
+ cos_sin_cache layout: [max_pos, ROPE_DIM] with first HALF_ROPE=cos,
38
+ second HALF_ROPE=sin.
39
+ """
40
+ token_idx = tl.program_id(0)
41
+ head_idx = tl.program_id(1)
42
+
43
+ if token_idx >= num_tokens:
44
+ return
45
+
46
+ pos = tl.load(position_ids_ptr + token_idx).to(tl.int64)
47
+
48
+ # Load cos/sin for this position
49
+ rope_pair_idx = tl.arange(0, HALF_ROPE)
50
+ cos_val = tl.load(cos_sin_cache_ptr + pos * ROPE_DIM + rope_pair_idx).to(tl.float32)
51
+ sin_val = tl.load(
52
+ cos_sin_cache_ptr + pos * ROPE_DIM + HALF_ROPE + rope_pair_idx
53
+ ).to(tl.float32)
54
+
55
+ if head_idx < num_heads:
56
+ # ========== Q: per-head RMSNorm + GPT-J RoPE ==========
57
+ q_base = q_ptr + token_idx * num_heads * HEAD_DIM + head_idx * HEAD_DIM
58
+
59
+ # Load full head
60
+ offs = tl.arange(0, HEAD_DIM)
61
+ q_vals = tl.load(q_base + offs).to(tl.float32)
62
+
63
+ # RMSNorm (no weight)
64
+ sq_sum = tl.sum(q_vals * q_vals, axis=0)
65
+ rms = tl.rsqrt(sq_sum / HEAD_DIM + eps)
66
+ q_vals = q_vals * rms
67
+
68
+ # Store ONLY the NoPE portion (positions 0..NOPE_DIM-1)
69
+ nope_mask = offs < NOPE_DIM
70
+ tl.store(q_base + offs, q_vals.to(q_ptr.type.element_ty), mask=nope_mask)
71
+
72
+ # GPT-J interleaved RoPE on the last ROPE_DIM dimensions:
73
+ even_offs = NOPE_DIM + rope_pair_idx * 2
74
+ odd_offs = NOPE_DIM + rope_pair_idx * 2 + 1
75
+
76
+ # Re-load original values at rope positions and normalize
77
+ q_even = tl.load(q_base + even_offs).to(tl.float32) * rms
78
+ q_odd = tl.load(q_base + odd_offs).to(tl.float32) * rms
79
+
80
+ new_even = q_even * cos_val - q_odd * sin_val
81
+ new_odd = q_even * sin_val + q_odd * cos_val
82
+
83
+ # Store rotated RoPE values
84
+ tl.store(q_base + even_offs, new_even.to(q_ptr.type.element_ty))
85
+ tl.store(q_base + odd_offs, new_odd.to(q_ptr.type.element_ty))
86
+ else:
87
+ # ========== KV: GPT-J RoPE only ==========
88
+ kv_base = kv_ptr + token_idx * HEAD_DIM
89
+ kv_out_base = kv_out_ptr + token_idx * HEAD_DIM
90
+
91
+ # Copy full KV unchanged first
92
+ offs = tl.arange(0, HEAD_DIM)
93
+ kv_full = tl.load(kv_base + offs)
94
+ tl.store(kv_out_base + offs, kv_full)
95
+
96
+ # GPT-J interleaved RoPE on the last ROPE_DIM dimensions
97
+ even_offs = NOPE_DIM + rope_pair_idx * 2
98
+ odd_offs = NOPE_DIM + rope_pair_idx * 2 + 1
99
+
100
+ kv_even = tl.load(kv_base + even_offs).to(tl.float32)
101
+ kv_odd = tl.load(kv_base + odd_offs).to(tl.float32)
102
+
103
+ new_even = kv_even * cos_val - kv_odd * sin_val
104
+ new_odd = kv_even * sin_val + kv_odd * cos_val
105
+
106
+ tl.store(kv_out_base + even_offs, new_even.to(kv_out_ptr.type.element_ty))
107
+ tl.store(kv_out_base + odd_offs, new_odd.to(kv_out_ptr.type.element_ty))
108
+
109
+
110
+def xpu_qnorm_rope_kv_fp8_insert(
111
+ q: torch.Tensor, # [num_tokens, num_heads, HEAD_DIM] bf16, in-place
112
+ kv: torch.Tensor, # [num_tokens, HEAD_DIM] bf16
113
+ swa_kv_cache: torch.Tensor, # [num_blocks, block_size, 584] or flat uint8
114
+ slot_mapping: torch.Tensor, # [num_tokens] int64
115
+ positions: torch.Tensor, # [num_tokens] int64
116
+ cos_sin_cache: torch.Tensor, # [max_pos, ROPE_DIM]
117
+ eps: float,
118
+ block_size: int,
119
+):
120
+ """XPU Triton: qnorm+rope on Q, rope on KV, then FP8 UE8M0 quant+insert."""
121
+ from vllm.models.deepseek_v4.common.ops.cache_utils import (
122
+ quantize_and_insert_k_cache,
123
+ )
124
+
125
+ num_tokens = q.shape[0]
126
+ num_heads = q.shape[1]
127
+
128
+ # Allocate temp buffer for RoPE-applied KV
129
+ kv_roped = torch.empty_like(kv)
130
+
131
+ # Grid: one program per (token, head_or_kv)
132
+ # head_idx < num_heads: process Q head
133
+ # head_idx == num_heads: process KV
134
+ grid = (num_tokens, num_heads + 1)
135
+ _xpu_qnorm_rope_kernel[grid](
136
+ q,
137
+ kv,
138
+ kv_roped,
139
+ positions,
140
+ cos_sin_cache,
141
+ eps,
142
+ num_tokens,
143
+ num_heads=num_heads,
144
+ HEAD_DIM=HEAD_DIM,
145
+ ROPE_DIM=ROPE_DIM,
146
+ NOPE_DIM=NOPE_DIM,
147
+ HALF_ROPE=HALF_ROPE,
148
+ )
149
+
150
+ # FP8 UE8M0 quant + paged insert (reuse existing Triton kernel)
151
+ # swa_kv_cache may be [num_blocks, block_size, 584] or [num_blocks, flat]
152
+ # quantize_and_insert_k_cache expects [num_blocks, block_bytes] uint8
153
+ cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
154
+ quantize_and_insert_k_cache(
155
+ kv_roped,
156
+ cache_2d,
157
+ slot_mapping,
158
+ block_size=block_size,
159
+ )
vllm/models/deepseek_v4/xpu/xpu_sparse.py
new
+350
@@ -0,0 +1,350 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""XPU DeepSeek-V4 attention subclass.
4
+
5
+Subclasses the shared ``DeepseekV4Attention`` ABC and provides XPU-native
6
+Triton kernels for decode (FP8 dequant + BF16 attention) and prefill
7
+(BF16 gathered KV + sparse attention).
8
+"""
9
+
10
+from typing import TYPE_CHECKING, cast
11
+
12
+import torch
13
+
14
+from vllm.forward_context import get_forward_context
15
+from vllm.models.deepseek_v4.attention import DeepseekV4Attention
16
+from vllm.models.deepseek_v4.common.ops import (
17
+ combine_topk_swa_indices,
18
+ compute_global_topk_indices_and_lens,
19
+ dequantize_and_gather_k_cache,
20
+)
21
+from vllm.models.deepseek_v4.sparse_mla import (
22
+ DeepseekV4FlashMLABackend,
23
+ DeepseekV4FlashMLAMetadata,
24
+)
25
+from vllm.models.deepseek_v4.xpu.xpu_sparse_decode_fp8 import (
26
+ xpu_sparse_decode_fp8,
27
+)
28
+from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
29
+from vllm.v1.worker.workspace import current_workspace_manager
30
+
31
+if TYPE_CHECKING:
32
+ from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
33
+
34
+
35
+class DeepseekV4XPUSparseBackend(DeepseekV4FlashMLABackend):
36
+ @staticmethod
37
+ def get_name() -> str:
38
+ return "XPU_V4_MLA_SPARSE"
39
+
40
+
41
+class DeepseekV4XPUAttention(DeepseekV4Attention):
42
+ """XPU sparse MLA attention layer for DeepSeek V4."""
43
+
44
+ backend_cls = DeepseekV4XPUSparseBackend
45
+ use_flashmla_fp8_layout = True
46
+
47
+ def __init__(self, *args, **kwargs) -> None:
48
+ # torch.cuda.Event() raises RuntimeError on XPU ("dummy base class").
49
+ # The Base and DeepseekV4Indexer both create cuda Events in __init__, so
50
+ # we temporarily redirect torch.cuda.Event → torch.xpu.Event.
51
+ _orig_event = torch.cuda.Event
52
+ torch.cuda.Event = torch.xpu.Event # type: ignore[misc]
53
+ try:
54
+ super().__init__(*args, **kwargs)
55
+ finally:
56
+ torch.cuda.Event = _orig_event # type: ignore[misc]
57
+
58
+ def _fused_qnorm_rope_kv_insert(self, q, kv, positions, attn_metadata):
59
+ from typing import cast
60
+
61
+ if not isinstance(attn_metadata, dict):
62
+ # Profile run: no-op, just return q (no padding needed on XPU).
63
+ return q
64
+
65
+ swa_metadata = cast(
66
+ "DeepseekSparseSWAMetadata | None",
67
+ attn_metadata.get(self.swa_cache_layer.prefix),
68
+ )
69
+ assert swa_metadata is not None
70
+
71
+ from vllm.models.deepseek_v4.xpu.xpu_qnorm_rope_kv_fp8_insert import (
72
+ xpu_qnorm_rope_kv_fp8_insert,
73
+ )
74
+
75
+ xpu_qnorm_rope_kv_fp8_insert(
76
+ q,
77
+ kv,
78
+ self.swa_cache_layer.kv_cache,
79
+ swa_metadata.slot_mapping,
80
+ positions,
81
+ self.rotary_emb.cos_sin_cache,
82
+ self.eps,
83
+ swa_metadata.block_size,
84
+ )
85
+ return q
86
+
87
+ @classmethod
88
+ def get_padded_num_q_heads(cls, num_heads: int) -> int:
89
+ return num_heads
90
+
91
+ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
92
+ # XPU uses BF16 reference wo_a path (same as ROCm).
93
+ from vllm.models.deepseek_v4.amd.rocm import rocm_inv_rope_einsum
94
+
95
+ z = rocm_inv_rope_einsum(
96
+ self.rotary_emb,
97
+ o,
98
+ positions,
99
+ self.rope_head_dim,
100
+ self.n_local_groups,
101
+ self.o_lora_rank,
102
+ self.wo_a,
103
+ )
104
+ return self.wo_b(z.flatten(1))
105
+
106
+ def forward_mqa(
107
+ self,
108
+ q: torch.Tensor,
109
+ kv: torch.Tensor,
110
+ positions: torch.Tensor,
111
+ output: torch.Tensor,
112
+ ) -> None:
113
+ assert output.shape == q.shape, (
114
+ f"output buffer shape {output.shape} must match q shape {q.shape}"
115
+ )
116
+ assert output.dtype == q.dtype, (
117
+ f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
118
+ )
119
+
120
+ forward_context = get_forward_context()
121
+ attn_metadata = forward_context.attn_metadata
122
+
123
+ if attn_metadata is None:
124
+ # Warmup dummy run: reserve workspace, skip actual kernels.
125
+ swa_only = self.compress_ratio <= 1
126
+ N = (
127
+ 0
128
+ if swa_only
129
+ else (self.max_model_len + self.compress_ratio - 1)
130
+ // self.compress_ratio
131
+ )
132
+ M = N + self.window_size + self.max_num_batched_tokens
133
+ current_workspace_manager().get_simultaneous(
134
+ ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
135
+ )
136
+ output.zero_()
137
+ return
138
+
139
+ assert isinstance(attn_metadata, dict)
140
+ flashmla_metadata = cast(
141
+ DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)
142
+ )
143
+ swa_metadata = cast(
144
+ "DeepseekSparseSWAMetadata | None",
145
+ attn_metadata.get(self.swa_cache_layer.prefix),
146
+ )
147
+ assert swa_metadata is not None
148
+
149
+ swa_only = self.compress_ratio <= 1
150
+ self_kv_cache = self.kv_cache if not swa_only else None
151
+ swa_kv_cache = self.swa_cache_layer.kv_cache
152
+
153
+ # Split prefill and decode
154
+ num_decodes = swa_metadata.num_decodes
155
+ num_prefills = swa_metadata.num_prefills
156
+ num_decode_tokens = swa_metadata.num_decode_tokens
157
+
158
+ if num_prefills > 0:
159
+ self._forward_prefill(
160
+ q=q[num_decode_tokens:],
161
+ positions=positions[num_decode_tokens:],
162
+ compressed_k_cache=self_kv_cache,
163
+ swa_k_cache=swa_kv_cache,
164
+ output=output[num_decode_tokens:],
165
+ attn_metadata=flashmla_metadata,
166
+ swa_metadata=swa_metadata,
167
+ )
168
+ if num_decodes > 0:
169
+ self._forward_decode(
170
+ q=q[:num_decode_tokens],
171
+ kv_cache=self_kv_cache,
172
+ swa_metadata=swa_metadata,
173
+ attn_metadata=flashmla_metadata,
174
+ swa_only=swa_only,
175
+ output=output[:num_decode_tokens],
176
+ )
177
+
178
+ def _forward_decode(
179
+ self,
180
+ q: torch.Tensor,
181
+ kv_cache: torch.Tensor | None,
182
+ swa_metadata: "DeepseekSparseSWAMetadata",
183
+ attn_metadata: DeepseekV4FlashMLAMetadata | None,
184
+ swa_only: bool,
185
+ output: torch.Tensor,
186
+ ) -> None:
187
+ num_decodes = swa_metadata.num_decodes
188
+ num_decode_tokens = swa_metadata.num_decode_tokens
189
+
190
+ topk_indices = None
191
+ topk_lens = None
192
+ if not swa_only:
193
+ assert attn_metadata is not None
194
+ assert swa_metadata.is_valid_token is not None
195
+ block_size = attn_metadata.block_size // self.compress_ratio
196
+ is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
197
+ if self.compress_ratio == 4:
198
+ # C4A: local indices differ per layer (filled by Indexer).
199
+ assert self.topk_indices_buffer is not None
200
+ global_indices, topk_lens = compute_global_topk_indices_and_lens(
201
+ self.topk_indices_buffer[:num_decode_tokens],
202
+ swa_metadata.token_to_req_indices,
203
+ attn_metadata.block_table[:num_decodes],
204
+ block_size,
205
+ is_valid,
206
+ )
207
+ topk_indices = global_indices.view(num_decode_tokens, 1, -1)
208
+ else:
209
+ # C128A: pre-computed during metadata build.
210
+ topk_indices = attn_metadata.c128a_global_decode_topk_indices
211
+ topk_lens = attn_metadata.c128a_decode_topk_lens
212
+
213
+ swa_indices = swa_metadata.decode_swa_indices
214
+ swa_lens = swa_metadata.decode_swa_lens
215
+
216
+ assert swa_indices is not None and swa_lens is not None
217
+ xpu_sparse_decode_fp8(
218
+ q=q,
219
+ kv_cache=kv_cache,
220
+ swa_kv_cache=self.swa_cache_layer.kv_cache,
221
+ swa_only=swa_only,
222
+ topk_indices=topk_indices,
223
+ topk_lens=topk_lens,
224
+ swa_indices=swa_indices,
225
+ swa_lens=swa_lens,
226
+ attn_sink=self.attn_sink,
227
+ softmax_scale=self.scale,
228
+ head_dim=self.head_dim,
229
+ nope_head_dim=self.nope_head_dim,
230
+ rope_head_dim=self.rope_head_dim,
231
+ out=output,
232
+ )
233
+
234
+ def _forward_prefill(
235
+ self,
236
+ q: torch.Tensor,
237
+ positions: torch.Tensor,
238
+ compressed_k_cache: torch.Tensor | None,
239
+ swa_k_cache: torch.Tensor,
240
+ output: torch.Tensor,
241
+ attn_metadata: DeepseekV4FlashMLAMetadata | None,
242
+ swa_metadata: "DeepseekSparseSWAMetadata",
243
+ ) -> None:
244
+ swa_only = attn_metadata is None
245
+
246
+ num_prefills = swa_metadata.num_prefills
247
+ num_prefill_tokens = swa_metadata.num_prefill_tokens
248
+ num_decodes = swa_metadata.num_decodes
249
+ num_decode_tokens = swa_metadata.num_decode_tokens
250
+
251
+ # Use pre-computed prefill metadata.
252
+ seq_lens = swa_metadata.prefill_seq_lens
253
+ gather_lens = swa_metadata.prefill_gather_lens
254
+ assert seq_lens is not None
255
+ assert gather_lens is not None
256
+
257
+ # Derive prefill-local token offsets from the full query_start_loc_cpu.
258
+ query_start_loc_cpu = swa_metadata.query_start_loc_cpu
259
+ query_start_loc = swa_metadata.query_start_loc
260
+ assert query_start_loc_cpu is not None
261
+ assert query_start_loc is not None
262
+ prefill_token_base = query_start_loc_cpu[num_decodes]
263
+
264
+ if not swa_only:
265
+ if self.compress_ratio == 4:
266
+ assert self.topk_indices_buffer is not None
267
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:]
268
+ topk_indices = topk_indices[:num_prefill_tokens]
269
+ else:
270
+ # C128A: pre-computed during metadata build.
271
+ assert attn_metadata is not None
272
+ topk_indices = attn_metadata.c128a_prefill_topk_indices
273
+ top_k = topk_indices.shape[-1]
274
+ N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
275
+ else:
276
+ assert self.topk_indices_buffer is not None
277
+ topk_indices = self.topk_indices_buffer[num_decode_tokens:]
278
+ top_k = 0
279
+ N = 0
280
+
281
+ M = N + self.window_size + self.max_num_batched_tokens
282
+ chunk_size_const = self.PREFILL_CHUNK_SIZE
283
+ num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const
284
+
285
+ workspace_manager = current_workspace_manager()
286
+ kv = workspace_manager.get_simultaneous(
287
+ ((chunk_size_const, M, q.shape[-1]), torch.bfloat16),
288
+ )[0]
289
+ for chunk_idx in range(num_chunks):
290
+ chunk_start = chunk_idx * chunk_size_const
291
+ chunk_end = min(chunk_start + chunk_size_const, num_prefills)
292
+ chunk_size = chunk_end - chunk_start
293
+ if not swa_only:
294
+ # Gather compressed KV
295
+ assert attn_metadata is not None
296
+ block_table = attn_metadata.block_table[num_decodes:]
297
+ dequantize_and_gather_k_cache(
298
+ kv[:chunk_size],
299
+ compressed_k_cache,
300
+ seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio,
301
+ gather_lens=None,
302
+ block_table=block_table[chunk_start:chunk_end],
303
+ block_size=attn_metadata.block_size // self.compress_ratio,
304
+ offset=0,
305
+ )
306
+
307
+ # Gather SWA KV
308
+ swa_block_table = swa_metadata.block_table[num_decodes:]
309
+ dequantize_and_gather_k_cache(
310
+ kv[:chunk_size],
311
+ swa_k_cache,
312
+ seq_lens=seq_lens[chunk_start:chunk_end],
313
+ gather_lens=gather_lens[chunk_start:chunk_end],
314
+ block_table=swa_block_table[chunk_start:chunk_end],
315
+ block_size=swa_metadata.block_size,
316
+ offset=N,
317
+ )
318
+
319
+ # Combine the topk indices and SWA indices for gathered KV cache
320
+ query_start = (
321
+ query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
322
+ )
323
+ query_end = (
324
+ query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
325
+ )
326
+
327
+ combined_indices, combined_lens = combine_topk_swa_indices(
328
+ topk_indices[query_start:query_end],
329
+ query_start_loc[
330
+ num_decodes + chunk_start : num_decodes + chunk_end + 1
331
+ ],
332
+ seq_lens[chunk_start:chunk_end],
333
+ gather_lens[chunk_start:chunk_end],
334
+ self.window_size,
335
+ self.compress_ratio,
336
+ top_k,
337
+ M,
338
+ N,
339
+ )
340
+
341
+ kv_ws = kv[:chunk_size].reshape(-1, 1, q.shape[-1])
342
+ out, _, _ = triton_bf16_mla_sparse_interface(
343
+ q=q[query_start:query_end],
344
+ kv=kv_ws,
345
+ indices=combined_indices.unsqueeze(1),
346
+ sm_scale=self.scale,
347
+ d_v=q.shape[-1],
348
+ block_dpe=0,
349
+ )
350
+ output[query_start:query_end] = out
vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py
new
+290
@@ -0,0 +1,290 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""XPU sparse decode for DeepSeek V4 with FP8 KV cache.
4
+
5
+Strategy: dequantize FP8 UE8M0 KV cache pages to BF16 on the fly,
6
+then reuse the BF16 sparse MLA attention kernel (xpu_sparse_mla_bf16).
7
+This keeps the external KV cache layout identical to CUDA/ROCm.
8
+"""
9
+
10
+import torch
11
+
12
+from vllm.triton_utils import tl, triton
13
+from vllm.v1.attention.ops.xpu_mla_sparse import (
14
+ triton_bf16_mla_sparse_interface,
15
+)
16
+
17
+# FP8 DS MLA cache layout constants
18
+TOKEN_FP8_DIM = 448 # NoPE portion in FP8
19
+TOKEN_BF16_DIM = 64 # RoPE portion in BF16
20
+TOKEN_SCALE_DIM = 8 # UE8M0 scales per token
21
+QUANT_BLOCK_SIZE = 64 # Elements per quant block
22
+OUTPUT_DIM = 512 # = TOKEN_FP8_DIM + TOKEN_BF16_DIM after dequant
23
+TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 # 576 bytes per token
24
+
25
+
26
+@triton.jit
27
+def _dequant_gather_slots_kernel(
28
+ # Output workspace: [total_slots, OUTPUT_DIM] bf16
29
+ out_ptr,
30
+ # FP8 paged cache base pointer (uint8 flat)
31
+ cache_ptr,
32
+ # Global slot indices: [total_slots] int32
33
+ indices_ptr,
34
+ # Cache geometry
35
+ cache_block_size: tl.constexpr,
36
+ token_data_size: tl.constexpr, # 576
37
+ block_stride: tl.int64, # k_cache.stride(0) — total uint8 per block
38
+ fp8_dim: tl.constexpr, # 448
39
+ bf16_dim: tl.constexpr, # 64
40
+ scale_dim: tl.constexpr, # 8
41
+ quant_block: tl.constexpr, # 64
42
+ output_dim: tl.constexpr, # 512
43
+ n_quant_blocks: tl.constexpr, # 7
44
+):
45
+ """Dequantize scattered FP8 slots into a flat BF16 workspace.
46
+
47
+ Grid: [total_slots] — one program per slot to gather.
48
+
49
+ Cache block layout (block_size tokens):
50
+ [0, block_size*576): Token data, each token 448 FP8 + 128 BF16
51
+ [block_size*576, block_size*576 + block_size*8): Scales
52
+ """
53
+ pid = tl.program_id(0)
54
+
55
+ # Load global slot index
56
+ slot_idx = tl.load(indices_ptr + pid).to(tl.int64)
57
+
58
+ # Output pointer for this slot
59
+ out_row_ptr = out_ptr + pid * output_dim
60
+
61
+ # Handle invalid slots (index < 0): write zeros
62
+ if slot_idx < 0:
63
+ zero = tl.zeros([quant_block], dtype=tl.bfloat16)
64
+ for i in tl.static_range(0, 512, 64):
65
+ offsets = i + tl.arange(0, quant_block)
66
+ mask = offsets < output_dim
67
+ tl.store(out_row_ptr + offsets, zero, mask=mask)
68
+ return
69
+
70
+ # Compute block and position within block
71
+ block_idx = slot_idx // cache_block_size
72
+ pos_in_block = slot_idx % cache_block_size
73
+
74
+ # Block base pointer
75
+ block_base = cache_ptr + block_idx * block_stride
76
+
77
+ # Token data: at offset pos_in_block * token_data_size within block
78
+ token_data_ptr = block_base + pos_in_block * token_data_size
79
+
80
+ # Scale: after all token data, at offset
81
+ # cache_block_size * token_data_size + pos_in_block * scale_dim
82
+ scale_region_offset = tl.cast(cache_block_size, tl.int64) * token_data_size
83
+ token_scale_ptr = block_base + scale_region_offset + pos_in_block * scale_dim
84
+
85
+ # ========== Dequantize FP8 portion (448 elements) ==========
86
+ for qblock_idx in tl.static_range(n_quant_blocks):
87
+ qblock_start = qblock_idx * quant_block
88
+ offsets = qblock_start + tl.arange(0, quant_block)
89
+ mask = offsets < fp8_dim
90
+
91
+ # Load FP8 as uint8 and bitcast
92
+ x_uint8 = tl.load(token_data_ptr + offsets, mask=mask, other=0)
93
+ x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
94
+ x_float = x_fp8.to(tl.float32)
95
+
96
+ # Load UE8M0 scale: scale = 2^(stored_value - 127)
97
+ encoded_scale = tl.load(token_scale_ptr + qblock_idx)
98
+ exponent = encoded_scale.to(tl.float32) - 127.0
99
+ scale = tl.exp2(exponent)
100
+
101
+ # Dequantize and store as bf16
102
+ x_dequant = x_float * scale
103
+ tl.store(out_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask)
104
+
105
+ # ========== Copy BF16 portion (64 elements) directly ==========
106
+ bf16_src_ptr = (token_data_ptr + fp8_dim).to(tl.pointer_type(tl.bfloat16))
107
+ bf16_out_ptr = (out_row_ptr + fp8_dim).to(tl.pointer_type(tl.bfloat16))
108
+
109
+ for j in tl.static_range(bf16_dim // 16):
110
+ chunk_offsets = j * 16 + tl.arange(0, 16)
111
+ bf16_vals = tl.load(bf16_src_ptr + chunk_offsets)
112
+ tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)
113
+
114
+
115
+def dequant_gather_slots(
116
+ out: torch.Tensor, # [total_slots, 512] bf16, pre-allocated
117
+ cache: torch.Tensor, # [num_blocks, block_size, head_bytes] uint8
118
+ indices: torch.Tensor, # [total_slots] int32, global slot IDs
119
+ cache_block_size: int, # block_size for this cache
120
+) -> None:
121
+ """Dequantize FP8 UE8M0 pages at scattered slot indices into bf16."""
122
+ total_slots = indices.shape[0]
123
+ if total_slots == 0:
124
+ return
125
+
126
+ block_stride = cache.stride(0)
127
+
128
+ _dequant_gather_slots_kernel[(total_slots,)](
129
+ out,
130
+ cache,
131
+ indices,
132
+ cache_block_size=cache_block_size,
133
+ token_data_size=TOKEN_DATA_SIZE,
134
+ block_stride=block_stride,
135
+ fp8_dim=TOKEN_FP8_DIM,
136
+ bf16_dim=TOKEN_BF16_DIM,
137
+ scale_dim=TOKEN_SCALE_DIM,
138
+ quant_block=QUANT_BLOCK_SIZE,
139
+ output_dim=OUTPUT_DIM,
140
+ n_quant_blocks=7,
141
+ )
142
+
143
+
144
+def xpu_sparse_decode_fp8(
145
+ q: torch.Tensor, # [num_tokens, num_heads, head_dim]
146
+ kv_cache: torch.Tensor | None, # [num_blocks, block_size, head_bytes] uint8
147
+ swa_kv_cache: torch.Tensor, # [num_blocks, swa_block_size, head_bytes] uint8
148
+ swa_only: bool,
149
+ topk_indices: torch.Tensor | None, # [num_tokens, 1, topk] global slot IDs
150
+ topk_lens: torch.Tensor | None,
151
+ swa_indices: torch.Tensor, # [num_tokens, 1, swa_k] global slot IDs
152
+ swa_lens: torch.Tensor,
153
+ attn_sink: torch.Tensor,
154
+ softmax_scale: float,
155
+ head_dim: int,
156
+ nope_head_dim: int,
157
+ rope_head_dim: int,
158
+ out: torch.Tensor, # [num_tokens, num_heads, head_dim]
159
+) -> None:
160
+ """XPU decode: dequant FP8 pages to BF16, then BF16 sparse MLA attention.
161
+
162
+ Keeps external FP8 KV cache layout identical to CUDA/ROCm.
163
+ Performance is slower due to on-the-fly dequant, but correctness is
164
+ guaranteed by reusing the validated BF16 attention kernel.
165
+ """
166
+ num_tokens = q.shape[0]
167
+ device = q.device
168
+
169
+ # Determine max topk and swa widths
170
+ if not swa_only and topk_indices is not None:
171
+ topk_idx_2d = (
172
+ topk_indices.squeeze(1) if topk_indices.dim() == 3 else topk_indices
173
+ )
174
+ max_topk = topk_idx_2d.shape[1]
175
+ else:
176
+ topk_idx_2d = None
177
+ max_topk = 0
178
+
179
+ swa_idx_2d = swa_indices.squeeze(1) if swa_indices.dim() == 3 else swa_indices
180
+ max_swa = swa_idx_2d.shape[1]
181
+
182
+ K_total = max_topk + max_swa
183
+
184
+ # Allocate flat workspace: [num_tokens * K_total, 512] bf16
185
+ workspace = torch.empty(
186
+ (num_tokens * K_total, OUTPUT_DIM), dtype=torch.bfloat16, device=device
187
+ )
188
+ ws_3d = workspace.view(num_tokens, K_total, OUTPUT_DIM)
189
+
190
+ # Dequant+gather topk slots from compressed cache
191
+ if not swa_only and topk_idx_2d is not None and kv_cache is not None:
192
+ topk_flat = topk_idx_2d.reshape(-1).to(torch.int32)
193
+ topk_buf = torch.empty(
194
+ (num_tokens * max_topk, OUTPUT_DIM), dtype=torch.bfloat16, device=device
195
+ )
196
+ compressed_block_size = kv_cache.shape[1]
197
+ dequant_gather_slots(topk_buf, kv_cache, topk_flat, compressed_block_size)
198
+ ws_3d[:, :max_topk, :] = topk_buf.view(num_tokens, max_topk, OUTPUT_DIM)
199
+
200
+ # Dequant+gather SWA slots
201
+ swa_flat = swa_idx_2d.reshape(-1).to(torch.int32)
202
+ swa_buf = torch.empty(
203
+ (num_tokens * max_swa, OUTPUT_DIM), dtype=torch.bfloat16, device=device
204
+ )
205
+ swa_block_size = swa_kv_cache.shape[1]
206
+ dequant_gather_slots(swa_buf, swa_kv_cache, swa_flat, swa_block_size)
207
+
208
+ ws_3d[:, max_topk:, :] = swa_buf.view(num_tokens, max_swa, OUTPUT_DIM)
209
+
210
+ # Build combined indices into the flat workspace and combined lengths.
211
+ # Workspace layout per token t: [topk_0..topk_{max_topk-1}, swa_0..swa_{max_swa-1}]
212
+ # Flat index for token t, position p = t * K_total + p
213
+ #
214
+ # IMPORTANT: The attention kernel uses combined_lens as a position cutoff —
215
+ # it only reads indices[0..combined_lens-1]. So indices must be PACKED
216
+ # contiguously: [valid_topk_indices..., valid_swa_indices..., -1 padding...]
217
+ if not swa_only and topk_lens is not None:
218
+ combined_lens = (topk_lens + swa_lens).to(torch.int32)
219
+ else:
220
+ combined_lens = swa_lens.to(torch.int32)
221
+
222
+ max_combined = int(combined_lens.max().item()) if combined_lens.numel() > 0 else 0
223
+ # Round up to BLOCK_N=16 alignment for kernel efficiency
224
+ _BLOCK_N = 16
225
+ max_combined_padded = ((max_combined + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N
226
+
227
+ # Build packed index table: [num_tokens, max_combined_padded]
228
+ # Each token t: [topk_0..topk_{tlen-1}, swa_0..swa_{slen-1}, -1 padding]
229
+ # Vectorized: for each token, topk indices are t*K_total + 0..tlen-1,
230
+ # swa indices are t*K_total + max_topk + 0..slen-1
231
+ combined_indices = torch.full(
232
+ (num_tokens, max_combined_padded),
233
+ fill_value=-1,
234
+ dtype=torch.int32,
235
+ device=device,
236
+ )
237
+
238
+ token_offsets = (
239
+ torch.arange(num_tokens, device=device, dtype=torch.int32) * K_total
240
+ ) # [B]
241
+
242
+ if not swa_only and topk_lens is not None:
243
+ # Pack topk: for each token, write t*K_total + 0..tlen-1 at positions 0..tlen-1
244
+ max_tlen = int(topk_lens.max().item())
245
+ topk_range = torch.arange(max_tlen, device=device, dtype=torch.int32).unsqueeze(
246
+ 0
247
+ )
248
+ topk_valid = topk_range < topk_lens.unsqueeze(1)
249
+ topk_ws_indices = token_offsets.unsqueeze(1) + topk_range
250
+ combined_indices[:, :max_tlen] = torch.where(
251
+ topk_valid,
252
+ topk_ws_indices,
253
+ torch.tensor(-1, dtype=torch.int32, device=device),
254
+ )
255
+ # Pack swa after topk: positions tlen..tlen+slen-1
256
+ # Since tlen varies per token, we need per-token offset
257
+ swa_range = torch.arange(max_swa, device=device, dtype=torch.int32).unsqueeze(0)
258
+ swa_valid = swa_range < swa_lens.unsqueeze(1)
259
+ swa_ws_indices = token_offsets.unsqueeze(1) + max_topk + swa_range
260
+ # Write at position topk_lens[t] + swa_pos for each token
261
+ for t_idx in range(num_tokens):
262
+ tlen = int(topk_lens[t_idx].item())
263
+ slen = int(swa_lens[t_idx].item())
264
+ combined_indices[t_idx, tlen : tlen + slen] = swa_ws_indices[t_idx, :slen]
265
+ else:
266
+ # SWA-only: pack swa indices at positions 0..slen-1
267
+ # Use min(max_swa, max_combined_padded) because combined_indices only
268
+ # has max_combined_padded columns, and all valid entries fit within it.
269
+ effective_swa = min(max_swa, max_combined_padded)
270
+ swa_range = torch.arange(
271
+ effective_swa, device=device, dtype=torch.int32
272
+ ).unsqueeze(0)
273
+ swa_valid = swa_range < swa_lens.unsqueeze(1)
274
+ swa_ws_indices = token_offsets.unsqueeze(1) + swa_range # max_topk=0
275
+ combined_indices[:, :effective_swa] = torch.where(
276
+ swa_valid,
277
+ swa_ws_indices,
278
+ torch.tensor(-1, dtype=torch.int32, device=device),
279
+ )
280
+
281
+ # Call BF16 sparse MLA kernel
282
+ out_attn, _, _ = triton_bf16_mla_sparse_interface(
283
+ q=q,
284
+ kv=workspace.unsqueeze(1),
285
+ indices=combined_indices.unsqueeze(1),
286
+ sm_scale=softmax_scale,
287
+ d_v=q.shape[-1],
288
+ block_dpe=0,
289
+ )
290
+ out.copy_(out_attn)
vllm/v1/attention/ops/xpu_mla_sparse.py
+8
-2
@@ -180,11 +180,17 @@ def triton_bf16_mla_sparse_interface(
180
indices: torch.Tensor, # [num_tokens, num_heads_kv, topk]
181
sm_scale: float,
182
d_v: int = 512,
183
+ block_dpe: int = 64,
184
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
185
"""
186
out : [num_tokens, num_heads_q, d_v]
187
max_logits : [num_tokens, num_heads_q]
188
lse : logsumexp, [num_tokens, num_heads_q]
189
+
190
+ Args:
191
+ block_dpe: Size of positional embedding portion of dim_qk.
192
+ Set to 0 when q/kv contain only the nope latent (e.g. DSv4
193
+ prefill where RoPE is not split out).
194
"""
195
num_tokens, num_heads_q, dim_qk = q.shape
196
_, num_heads_kv, _ = kv.shape
@@ -194,8 +200,8 @@ def triton_bf16_mla_sparse_interface(
200
_, _, index_topk = indices.shape
201
202
BLOCK_H = 16
197
- BLOCK_DMODEL = 512
198
- BLOCK_DPE = 64
203
+ BLOCK_DPE = block_dpe
204
+ BLOCK_DMODEL = dim_qk - BLOCK_DPE
205
BLOCK_M = 32
206
BLOCK_N = 16
207
BLOCK_DV = 512