A high-throughput and memory-efficient inference and serving engine for LLMs
[DSV4][XPU] Add MHC fused_post_pre support (#44144)
Signed-off-by: Ma Jian <jian1.ma@intel.com>
Ma Jian committed
Jun 9, 2026 at 17:23 UTC
70db1488c5d5f428b23b86223ec6bd4ec89b210f
2 files changed
+112
-17
vllm/model_executor/layers/mhc.py
+72
-3
@@ -466,7 +466,76 @@ class MHCFusedPostPreOp(CustomOp):
466
norm_eps,
467
)
468
469
- def forward_native(self, *args, **kwargs):
470
- raise NotImplementedError(
471
- "Native implementation of mhc_fused_post_pre is not available"
469
+ def forward_native(
470
+ self,
471
+ x: torch.Tensor,
472
+ residual: torch.Tensor,
473
+ post_layer_mix: torch.Tensor,
474
+ comb_res_mix: torch.Tensor,
475
+ fn: torch.Tensor,
476
+ hc_scale: torch.Tensor,
477
+ hc_base: torch.Tensor,
478
+ rms_eps: float,
479
+ hc_pre_eps: float,
480
+ hc_sinkhorn_eps: float,
481
+ hc_post_mult_value: float,
482
+ sinkhorn_repeat: int,
483
+ n_splits: int = 1,
484
+ tile_n: int = 1,
485
+ norm_weight: torch.Tensor | None = None,
486
+ norm_eps: float = 0.0,
487
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
488
+ # Decompose into post + pre (no fused kernel available).
489
+ residual_cur = mhc_kernels.mhc_post_torch(
490
+ x, residual, post_layer_mix, comb_res_mix
491
+ )
492
+ post_mix_cur, comb_mix_cur, layer_input_cur = mhc_kernels.mhc_pre_torch(
493
+ residual_cur,
494
+ fn,
495
+ hc_scale,
496
+ hc_base,
497
+ rms_eps,
498
+ hc_pre_eps,
499
+ hc_sinkhorn_eps,
500
+ hc_post_mult_value,
501
+ sinkhorn_repeat,
502
+ )
503
+ return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
504
+
505
+ def forward_xpu(
506
+ self,
507
+ x: torch.Tensor,
508
+ residual: torch.Tensor,
509
+ post_layer_mix: torch.Tensor,
510
+ comb_res_mix: torch.Tensor,
511
+ fn: torch.Tensor,
512
+ hc_scale: torch.Tensor,
513
+ hc_base: torch.Tensor,
514
+ rms_eps: float,
515
+ hc_pre_eps: float,
516
+ hc_sinkhorn_eps: float,
517
+ hc_post_mult_value: float,
518
+ sinkhorn_repeat: int,
519
+ n_splits: int = 1,
520
+ tile_n: int = 1,
521
+ norm_weight: torch.Tensor | None = None,
522
+ norm_eps: float = 0.0,
523
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
524
+ return self.forward_native(
525
+ x,
526
+ residual,
527
+ post_layer_mix,
528
+ comb_res_mix,
529
+ fn,
530
+ hc_scale,
531
+ hc_base,
532
+ rms_eps,
533
+ hc_pre_eps,
534
+ hc_sinkhorn_eps,
535
+ hc_post_mult_value,
536
+ sinkhorn_repeat,
537
+ n_splits,
538
+ tile_n,
539
+ norm_weight,
540
+ norm_eps,
541
)
vllm/models/deepseek_v4/xpu/model.py
+40
-14
@@ -930,22 +930,48 @@ class DeepseekV4DecoderLayer(nn.Module):
930
) -> tuple[
931
torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
932
]:
933
- residual = x
934
- x, post, comb = self.hc_pre(
935
- x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
936
- )
933
+ if residual is None:
934
+ # First layer: run standalone hc_pre
935
+ residual = x
936
+ x, post_mix, res_mix = self.hc_pre(
937
+ x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
938
+ )
939
+ else:
940
+ residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
941
+ x,
942
+ residual,
943
+ post_mix,
944
+ res_mix,
945
+ self.hc_attn_fn,
946
+ self.hc_attn_scale,
947
+ self.hc_attn_base,
948
+ self.rms_norm_eps,
949
+ self.hc_eps,
950
+ self.hc_eps,
951
+ self.hc_post_alpha,
952
+ self.hc_sinkhorn_iters,
953
+ )
954
+
955
x = self.attn_norm(x)
956
x = self.attn(positions, x, None)
939
- x = self.hc_post(x, residual, post, comb)
957
941
- residual = x
942
- x, post, comb = self.hc_pre(
943
- x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
958
+ residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
959
+ x,
960
+ residual,
961
+ post_mix,
962
+ res_mix,
963
+ self.hc_ffn_fn,
964
+ self.hc_ffn_scale,
965
+ self.hc_ffn_base,
966
+ self.rms_norm_eps,
967
+ self.hc_eps,
968
+ self.hc_eps,
969
+ self.hc_post_alpha,
970
+ self.hc_sinkhorn_iters,
971
)
972
x = self.ffn_norm(x)
973
x = self.ffn(x, input_ids)
947
- x = self.hc_post(x, residual, post, comb)
948
- return x, None, None, None
974
+ return x, residual, post_mix, res_mix
975
976
977
@support_torch_compile
@@ -1096,6 +1122,10 @@ class DeepseekV4Model(nn.Module):
1122
res_mix,
1123
residual,
1124
)
1125
+ # The fused path defers the final hc_post to the next layer's
1126
+ # fused_post_pre. After the last layer we must apply it explicitly.
1127
+ if layer is not None:
1128
+ hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
1129
1130
if not get_pp_group().is_last_rank:
1131
return IntermediateTensors({"hidden_states": hidden_states})
@@ -1196,8 +1226,6 @@ class DeepseekV4Model(nn.Module):
1226
elif "attn_sink" in name:
1227
if is_pp_missing_parameter(name, self):
1228
continue
1199
- if name not in params_dict:
1200
- continue
1229
narrow_weight = loaded_weight[head_rank_start:head_rank_end]
1230
n = narrow_weight.shape[0]
1231
params_dict[name][:n].copy_(narrow_weight)
@@ -1206,8 +1234,6 @@ class DeepseekV4Model(nn.Module):
1234
else:
1235
if is_pp_missing_parameter(name, self):
1236
continue
1209
- if name not in params_dict:
1210
- continue
1237
param = params_dict[name]
1238
weight_loader = getattr(
1239
param, "weight_loader", default_weight_loader