@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d626108b18

[ROCm][Perf] Fuse DeepSeek-V4 mHC post/pre and RMSNorm with AITER (#52737)

Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

Shanshan Shen committed Aug 20, 2026 at 12:51 UTC d626108b1841888ec90aced33367149a6bbc7e4b

5 files changed +272 -10

vllm/_aiter_ops.py

+96

@@ -3299,6 +3299,8 @@ class rocm_aiter_ops:

3299 hc_sinkhorn_eps: float,

3300 hc_post_mult_value: float,

3301 sinkhorn_repeat: int,

3302 + norm_weight: torch.Tensor | None = None,

3303 + norm_eps: float = 0.0,

3304 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

3305 """

3306 Forward pass for mHC pre block.

@@ -3313,6 +3315,8 @@ class rocm_aiter_ops:

3315 hc_sinkhorn_eps: sinkhorn epsilon

3316 hc_post_mult_value: post-mix multiplier value

3317 sinkhorn_repeat: number of sinkhorn iterations

3318 + norm_weight: optional RMSNorm weight fused into the pre kernel

3319 + norm_eps: epsilon for the fused RMSNorm when norm_weight is set

3320

3321 Returns:

3322 post_mix: shape (..., hc_mult), dtype torch.float32

@@ -3380,6 +3384,8 @@ class rocm_aiter_ops:

3384 hc_sinkhorn_eps,

3385 hc_post_mult_value,

3386 sinkhorn_repeat,

3387 + norm_weight,

3388 + norm_eps,

3389 )

3390 return (

3391 post_mix.view(*outer_shape, hc_mult, 1),

@@ -3465,5 +3471,95 @@ class rocm_aiter_ops:

3471 )

3472 return out.view_as(residual)

3473

3474 + @staticmethod

3475 + def mhc_fused_post_pre(

3476 + x: torch.Tensor,

3477 + residual: torch.Tensor,

3478 + post_layer_mix: torch.Tensor,

3479 + comb_res_mix: torch.Tensor,

3480 + fn: torch.Tensor,

3481 + hc_scale: torch.Tensor,

3482 + hc_base: torch.Tensor,

3483 + rms_eps: float,

3484 + hc_pre_eps: float,

3485 + hc_sinkhorn_eps: float,

3486 + hc_post_mult_value: float,

3487 + sinkhorn_repeat: int,

3488 + norm_weight: torch.Tensor | None = None,

3489 + norm_eps: float = 0.0,

3490 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

3491 + """Fused mHC post + next mHC pre via AITER.

3492 +

3493 + Returns residual_cur, post_mix, comb_mix, layer_input in vLLM order

3494 + (AITER returns post_mix, comb_mix, layer_input, next_residual).

3495 + """

3496 + from aiter.ops.mhc import mhc_fused_post_pre

3497 +

3498 + assert x.dtype == torch.bfloat16

3499 + assert residual.dtype == torch.bfloat16

3500 + assert fn.dtype == torch.float32

3501 + assert hc_scale.dtype == torch.float32

3502 + assert hc_base.dtype == torch.float32

3503 +

3504 + hc_mult = residual.shape[-2]

3505 + hidden_size = residual.shape[-1]

3506 + outer_shape = residual.shape[:-2]

3507 +

3508 + residual_flat = residual.view(-1, hc_mult, hidden_size)

3509 + num_tokens = residual_flat.shape[0]

3510 + x_flat = x.view(num_tokens, hidden_size)

3511 + post_flat = post_layer_mix.view(num_tokens, hc_mult, 1)

3512 + comb_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)

3513 +

3514 + if num_tokens == 0:

3515 + return (

3516 + torch.empty_like(residual_flat).view_as(residual),

3517 + torch.empty(

3518 + *outer_shape,

3519 + hc_mult,

3520 + 1,

3521 + dtype=torch.float32,

3522 + device=residual.device,

3523 + ),

3524 + torch.empty(

3525 + *outer_shape,

3526 + hc_mult,

3527 + hc_mult,

3528 + dtype=torch.float32,

3529 + device=residual.device,

3530 + ),

3531 + torch.empty(

3532 + *outer_shape,

3533 + hidden_size,

3534 + dtype=torch.bfloat16,

3535 + device=residual.device,

3536 + ),

3537 + )

3538 +

3539 + with torch.device(residual_flat.device):

3540 + post_mix, comb_mix, layer_input, next_residual = mhc_fused_post_pre(

3541 + x_flat,

3542 + residual_flat,

3543 + post_flat,

3544 + comb_flat,

3545 + fn,

3546 + hc_scale,

3547 + hc_base,

3548 + rms_eps,

3549 + hc_pre_eps,

3550 + hc_sinkhorn_eps,

3551 + hc_post_mult_value,

3552 + sinkhorn_repeat,

3553 + norm_weight,

3554 + norm_eps,

3555 + )

3556 +

3557 + return (

3558 + next_residual.view_as(residual),

3559 + post_mix.view(*outer_shape, hc_mult, 1),

3560 + comb_mix.view(*outer_shape, hc_mult, hc_mult),

3561 + layer_input.view(*outer_shape, hidden_size),

3562 + )

3563 +

3564

3565 rocm_aiter_ops.register_ops_once()

vllm/model_executor/kernels/mhc/aiter.py

+102

@@ -16,6 +16,8 @@ def mhc_pre_aiter(

16 hc_post_mult_value: float,

17 sinkhorn_repeat: int,

18 n_splits: int = 1,

19 + norm_weight: torch.Tensor | None = None,

20 + norm_eps: float = 0.0,

21 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

22 """

23 Forward pass for mHC pre block.

@@ -31,6 +33,8 @@ def mhc_pre_aiter(

33 hc_post_mult_value: post-mix multiplier value

34 sinkhorn_repeat: number of sinkhorn iterations

35 n_splits: split-k factor;

36 + norm_weight: optional RMSNorm weight fused into the pre kernel

37 + norm_eps: epsilon for the fused RMSNorm when norm_weight is set

39 Returns:

40 post_mix: shape (..., hc_mult), dtype torch.float32

@@ -52,6 +56,8 @@ def mhc_pre_aiter(

56 hc_sinkhorn_eps,

57 hc_post_mult_value,

58 sinkhorn_repeat,

59 + norm_weight,

60 + norm_eps,

61 )

@@ -66,6 +72,8 @@ def _mhc_pre_aiter_fake(

72 hc_post_mult_value: float,

73 sinkhorn_repeat: int,

74 n_splits: int = 1,

75 + norm_weight: torch.Tensor | None = None,

76 + norm_eps: float = 0.0,

77 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

78 hc_mult = residual.shape[-2]

79 hidden_size = residual.shape[-1]

@@ -124,6 +132,94 @@ def _mhc_post_aiter_fake(

132 return torch.empty_like(residual)

133

134

135 +def mhc_fused_post_pre_aiter(

136 + x: torch.Tensor,

137 + residual: torch.Tensor,

138 + post_layer_mix: torch.Tensor,

139 + comb_res_mix: torch.Tensor,

140 + fn: torch.Tensor,

141 + hc_scale: torch.Tensor,

142 + hc_base: torch.Tensor,

143 + rms_eps: float,

144 + hc_pre_eps: float,

145 + hc_sinkhorn_eps: float,

146 + hc_post_mult_value: float,

147 + sinkhorn_repeat: int,

148 + n_splits: int = 1,

149 + tile_n: int = 1,

150 + norm_weight: torch.Tensor | None = None,

151 + norm_eps: float = 0.0,

152 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

153 + """Fused mHC post + next mHC pre on ROCm via AITER.

154 +

155 + Returns residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur.

156 + """

157 + hidden_size = residual.shape[-1]

158 + assert hidden_size % 256 == 0

159 + from vllm._aiter_ops import rocm_aiter_ops

160 +

161 + return rocm_aiter_ops.mhc_fused_post_pre(

162 + x,

163 + residual,

164 + post_layer_mix,

165 + comb_res_mix,

166 + fn,

167 + hc_scale,

168 + hc_base,

169 + rms_eps,

170 + hc_pre_eps,

171 + hc_sinkhorn_eps,

172 + hc_post_mult_value,

173 + sinkhorn_repeat,

174 + norm_weight,

175 + norm_eps,

176 + )

177 +

178 +

179 +def _mhc_fused_post_pre_aiter_fake(

180 + x: torch.Tensor,

181 + residual: torch.Tensor,

182 + post_layer_mix: torch.Tensor,

183 + comb_res_mix: torch.Tensor,

184 + fn: torch.Tensor,

185 + hc_scale: torch.Tensor,

186 + hc_base: torch.Tensor,

187 + rms_eps: float,

188 + hc_pre_eps: float,

189 + hc_sinkhorn_eps: float,

190 + hc_post_mult_value: float,

191 + sinkhorn_repeat: int,

192 + n_splits: int = 1,

193 + tile_n: int = 1,

194 + norm_weight: torch.Tensor | None = None,

195 + norm_eps: float = 0.0,

196 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

197 + hc_mult = residual.shape[-2]

198 + hidden_size = residual.shape[-1]

199 + outer_shape = residual.shape[:-2]

200 + post_mix = torch.empty(

201 + *outer_shape,

202 + hc_mult,

203 + 1,

204 + dtype=torch.float32,

205 + device=residual.device,

206 + )

207 + comb_mix = torch.empty(

208 + *outer_shape,

209 + hc_mult,

210 + hc_mult,

211 + dtype=torch.float32,

212 + device=residual.device,

213 + )

214 + layer_input = torch.empty(

215 + *outer_shape,

216 + hidden_size,

217 + dtype=torch.bfloat16,

218 + device=residual.device,

219 + )

220 + return torch.empty_like(residual), post_mix, comb_mix, layer_input

221 +

222 +

223 direct_register_custom_op(

224 op_name="mhc_pre_aiter",

225 op_func=mhc_pre_aiter,

@@ -136,3 +232,9 @@ direct_register_custom_op(

232 mutates_args=[],

233 fake_impl=_mhc_post_aiter_fake,

234 )

235 +direct_register_custom_op(

236 + op_name="mhc_fused_post_pre_aiter",

237 + op_func=mhc_fused_post_pre_aiter,

238 + mutates_args=[],

239 + fake_impl=_mhc_fused_post_pre_aiter_fake,

240 +)

vllm/model_executor/layers/mhc.py

+28 -2

@@ -90,7 +90,8 @@ class MHCPreOp(CustomOp):

90 norm_eps: float = 0.0,

91 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

92 hidden_size = residual.shape[-1]

93 - if HAS_AITER_MHC and hidden_size % 256 == 0:

93 + hc_mult = residual.shape[-2]

94 + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:

95 return torch.ops.vllm.mhc_pre_aiter(

96 residual,

97 fn,

@@ -101,6 +102,9 @@ class MHCPreOp(CustomOp):

102 hc_sinkhorn_eps,

103 hc_post_mult_value,

104 sinkhorn_repeat,

105 + n_splits,

106 + norm_weight,

107 + norm_eps,

108 )

109 elif HAS_TILELANG_MHC:

110 return torch.ops.vllm.mhc_pre_tilelang(

@@ -222,7 +226,8 @@ class MHCPostOp(CustomOp):

226 comb_res_mix: torch.Tensor,

227 ) -> torch.Tensor:

228 hidden_size = residual.shape[-1]

225 - if HAS_AITER_MHC and hidden_size % 256 == 0:

229 + hc_mult = residual.shape[-2]

230 + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:

231 return torch.ops.vllm.mhc_post_aiter(

232 x,

233 residual,

@@ -444,6 +449,27 @@ class MHCFusedPostPreOp(CustomOp):

449 norm_weight: torch.Tensor | None = None,

450 norm_eps: float = 0.0,

451 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

452 + hidden_size = residual.shape[-1]

453 + hc_mult = residual.shape[-2]

454 + if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:

455 + return torch.ops.vllm.mhc_fused_post_pre_aiter(

456 + x,

457 + residual,

458 + post_layer_mix,

459 + comb_res_mix,

460 + fn,

461 + hc_scale,

462 + hc_base,

463 + rms_eps,

464 + hc_pre_eps,

465 + hc_sinkhorn_eps,

466 + hc_post_mult_value,

467 + sinkhorn_repeat,

468 + n_splits,

469 + tile_n,

470 + norm_weight,

471 + norm_eps,

472 + )

473 if HAS_TILELANG_MHC:

474 return torch.ops.vllm.mhc_fused_post_pre_tilelang(

475 x,

vllm/models/deepseek_v4/amd/dspark.py

+2 -2

@@ -9,8 +9,8 @@ ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for

9 attention + MHC CustomOp path) instead of the nvidia one;

10 * route the MHC head through the ``HCHeadOp`` CustomOp dispatcher (aiter /

11 tilelang / triton / torch) instead of calling the tilelang kernels directly,

12 - and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (False on the aiter

13 - path, where the decoder layer already applies hc_post in-layer);

12 + and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER

13 + or TileLang fused MHC is available; False only on the torch fallback);

14 * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` /

15 ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py).

vllm/models/deepseek_v4/amd/model.py

+44 -6

@@ -305,6 +305,10 @@ class DeepseekV4MoE(nn.Module):

305 return final_hidden_states.view(org_shape)

306

307

308 +# Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm.

309 +_AITER_MHC_FUSED_RMSNORM_SIZES = frozenset({1280, 2560, 4096, 7168})

310 +

311 +

312 class DeepseekV4DecoderLayer(nn.Module):

313 def __init__(

314 self,

@@ -384,9 +388,18 @@ class DeepseekV4DecoderLayer(nn.Module):

388 self.mhc_pre = MHCPreOp()

389 self.mhc_post = MHCPostOp()

390 self.mhc_fused_post_pre = MHCFusedPostPreOp()

387 - self.use_fused_mhc = HAS_TILELANG_MHC and not (

388 - HAS_AITER_MHC and self.hidden_size % 256 == 0

391 + # AITER mhc kernels (pre/post/fused) require hc_mult == 4.

392 + use_aiter_mhc = (

393 + HAS_AITER_MHC and self.hidden_size % 256 == 0 and self.hc_mult == 4

394 )

395 + # Prefer AITER fused post+pre when eligible; otherwise TileLang.

396 + self.use_fused_mhc = use_aiter_mhc or HAS_TILELANG_MHC

397 + # Fold attn/ffn RMSNorm into MHC only when the active backend's

398 + # fused-rmsnorm path supports this hidden size.

399 + if use_aiter_mhc:

400 + self.fuse_mhc_rmsnorm = self.hidden_size in _AITER_MHC_FUSED_RMSNORM_SIZES

401 + else:

402 + self.fuse_mhc_rmsnorm = HAS_TILELANG_MHC and self.use_fused_mhc

403

404 def hc_pre(

405 self,

@@ -394,7 +407,13 @@ class DeepseekV4DecoderLayer(nn.Module):

407 hc_fn: torch.Tensor,

408 hc_scale: torch.Tensor,

409 hc_base: torch.Tensor,

410 + norm_weight: torch.Tensor | None = None,

411 + norm_eps: float = 0.0,

412 ):

413 + """Reduce HC residual streams to the next sub-layer input.

414 +

415 + When ``norm_weight`` is set, RMSNorm is fused into the pre kernel.

416 + """

417 post_mix, res_mix, layer_input = self.mhc_pre(

418 residual=x,

419 fn=hc_fn,

@@ -405,6 +424,8 @@ class DeepseekV4DecoderLayer(nn.Module):

424 hc_sinkhorn_eps=self.hc_eps,

425 hc_post_mult_value=self.hc_post_alpha,

426 sinkhorn_repeat=self.hc_sinkhorn_iters,

427 + norm_weight=norm_weight,

428 + norm_eps=norm_eps,

429 )

430 return layer_input, post_mix, res_mix

431

@@ -426,11 +447,20 @@ class DeepseekV4DecoderLayer(nn.Module):

447 res_mix: torch.Tensor | None = None,

448 residual: torch.Tensor | None = None,

449 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

450 + attn_norm_weight = self.attn_norm.weight if self.fuse_mhc_rmsnorm else None

451 + attn_norm_eps = (

452 + self.attn_norm.variance_epsilon if self.fuse_mhc_rmsnorm else 0.0

453 + )

454 if residual is None:

455 # Run standalone hc_pre on first layer

456 residual = x

457 x, post_mix, res_mix = self.hc_pre(

433 - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base

458 + x,

459 + self.hc_attn_fn,

460 + self.hc_attn_scale,

461 + self.hc_attn_base,

462 + norm_weight=attn_norm_weight,

463 + norm_eps=attn_norm_eps,

464 )

465 else:

466 residual, post_mix, res_mix, x = self.mhc_fused_post_pre(

@@ -446,11 +476,16 @@ class DeepseekV4DecoderLayer(nn.Module):

476 self.hc_eps,

477 self.hc_post_alpha,

478 self.hc_sinkhorn_iters,

479 + norm_weight=attn_norm_weight,

480 + norm_eps=attn_norm_eps,

481 )

482

451 - x = self.attn_norm(x)

483 + if not self.fuse_mhc_rmsnorm:

484 + x = self.attn_norm(x)

485 x = self.attn(positions, x, None)

486

487 + ffn_norm_weight = self.ffn_norm.weight if self.fuse_mhc_rmsnorm else None

488 + ffn_norm_eps = self.ffn_norm.variance_epsilon if self.fuse_mhc_rmsnorm else 0.0

489 residual, post_mix, res_mix, x = self.mhc_fused_post_pre(

490 x,

491 residual,

@@ -464,8 +499,11 @@ class DeepseekV4DecoderLayer(nn.Module):

499 self.hc_eps,

500 self.hc_post_alpha,

501 self.hc_sinkhorn_iters,

502 + norm_weight=ffn_norm_weight,

503 + norm_eps=ffn_norm_eps,

504 )

468 - x = self.ffn_norm(x)

505 + if not self.fuse_mhc_rmsnorm:

506 + x = self.ffn_norm(x)

507 x = self.ffn(x, input_ids)

508 return x, residual, post_mix, res_mix

509

@@ -678,7 +716,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

716 residual,

717 )

718 if (idx + 1) in self.aux_hidden_state_layers:

681 - # On the unfused (aiter) path the layer already applied hc_post,

719 + # On the unfused path the layer already applied hc_post,

720 # so hidden_states is the reconstructed stream; on the fused

721 # path reconstruct it via hc_post before averaging.

722 if layer.use_fused_mhc: