@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0af34418e9

[ROCm][DSv4.1][Perf] Run the delayed mHC seams through aiter's fused Triton kernel (#58655)

Signed-off-by: Ahmed <Muhammad.Ahmed@amd.com>

Muhammad Ahmed committed Sep 28, 2026 at 22:19 UTC 0af34418e99b972029b53130c9a9665b4e692536

6 files changed +324 -16

tests/kernels/mhc/test_mhc_kernels.py

+116 -8

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.mhc import (

27 HAS_AITER_MHC,

28 HAS_AITER_MHC_FUSED,

29 HAS_AITER_MHC_FUSED_NORM,

30 + HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM,

31 HAS_AITER_MHC_PRE_NORM,

32 HAS_TILELANG_MHC,

33 MHCFusedPostPreOp,

@@ -1300,18 +1301,109 @@ def test_mhc_pre_delayed_rocm_aiter(num_tokens, carried):

1301 torch.testing.assert_close(actual[2], expected[2], atol=0, rtol=0)

1302

1303

1304 +@pytest.mark.skipif(

1305 + not (current_platform.is_rocm() and HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM),

1306 + reason="AITER fused delayed-seam kernel required (gfx950)",

1307 +)

1308 +@pytest.mark.parametrize("num_tokens", [1, 2, 7, 128, 1024])

1309 +@pytest.mark.parametrize("seam", ["post", "no_post", "identity_pre"])

1310 +def test_mhc_pre_delayed_rocm_aiter_fused_rms_norm(num_tokens, seam):

1311 + """One fused kernel must match the whole reference seam: the post block

1312 + (the previous seam's post/comb gates fold ``sublayer_out`` into the

1313 + streams), the delayed pre block (projection + Sinkhorn producing the next

1314 + seam's post/comb/pre gates, plus the collapse with the carried pre-mix),

1315 + and the RMSNorm folded onto the collapse.

1316 +

1317 + ``num_tokens`` spans the kernel's split-K config buckets. ``seam`` covers

1318 + the three call shapes the model makes: a regular seam (post folded in), an

1319 + Engram seam (post already applied), and the draft entry seam

1320 + (``pre_mix=None`` = identity pre-mix).

1321 + """

1322 + set_random_seed(0)

1323 + hc_mult, hidden_size = 4, 5120

1324 + residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs(

1325 + num_tokens=num_tokens, hidden_size=hidden_size, hc_mult=hc_mult

1326 + )

1327 + rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6

1328 + args = (

1329 + residual,

1330 + fn,

1331 + hc_scale,

1332 + hc_base,

1333 + rms_eps,

1334 + hc_pre_eps,

1335 + hc_sinkhorn_eps,

1336 + 1.0,

1337 + 20,

1338 + )

1339 +

1340 + pre_mix = (

1341 + None

1342 + if seam == "identity_pre"

1343 + else torch.rand(num_tokens, hc_mult, dtype=torch.float32, device=DEVICE) + 0.5

1344 + )

1345 + post_kwargs = {}

1346 + residual_ref = residual

1347 + if seam == "post":

1348 + sublayer_out = torch.randn(

1349 + num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE

1350 + )

1351 + post_layer_mix = 2 * torch.rand(

1352 + num_tokens, hc_mult, 1, dtype=torch.float32, device=DEVICE

1353 + )

1354 + comb_res_mix = torch.softmax(

1355 + torch.randn(num_tokens, hc_mult, hc_mult, device=DEVICE), dim=-1

1356 + )

1357 + post_kwargs = dict(

1358 + sublayer_out=sublayer_out,

1359 + post_layer_mix=post_layer_mix,

1360 + comb_res_mix=comb_res_mix,

1361 + )

1362 + residual_ref = mhc_post_torch(

1363 + sublayer_out, residual, post_layer_mix, comb_res_mix

1364 + )

1365 +

1366 + expected = mhc_pre_delayed_torch(residual_ref, *args[1:], pre_mix=pre_mix)

1367 + expected_layer_input = F.rms_norm(

1368 + expected[2], (hidden_size,), norm_weight, norm_eps

1369 + )

1370 +

1371 + residual_out, *actual = object.__new__(MHCPreDelayedOp).forward_hip(

1372 + *args,

1373 + pre_mix=pre_mix,

1374 + norm_weight=norm_weight,

1375 + norm_eps=norm_eps,

1376 + **post_kwargs,

1377 + )

1378 +

1379 + if seam == "post":

1380 + torch.testing.assert_close(residual_out, residual_ref)

1381 + else:

1382 + # No post was requested, so the residual comes straight back.

1383 + assert residual_out is residual

1384 + # Gates: post_mix, comb_mix, next_pre_mix (bf16 MFMA vs the fp32 reference).

1385 + for i in (0, 1, 3):

1386 + torch.testing.assert_close(actual[i], expected[i], atol=5e-4, rtol=1e-3)

1387 + # The collapse is bf16 and the kernel projects the once-rounded R', so

1388 + # near-zero values need an absolute tolerance (same bound as aiter's UT).

1389 + torch.testing.assert_close(actual[2], expected_layer_input, atol=2e-2, rtol=1e-2)

1390 +

1391 +

1392 @pytest.mark.skipif(

1393 not (current_platform.is_rocm() and HAS_AITER_MHC),

1394 reason="AITER mHC required",

1395 )

1396 def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch):

1308 - """The broadcast seam and a fused norm must not take the AITER path.

1309 -

1310 - Neither is expressible with AITER's pre kernels: the broadcast projects a

1311 - narrower ``x``, and ``mhc_pre_gemm_sqrsum`` folds no RMSNorm. What is

1312 - asserted here is the routing decision, which is what this gate owns; the

1313 - numerics of whichever fallback it lands on are covered by

1314 - ``test_deepseek_v41_mhc_pre_delayed``.

1397 + """The broadcast seam must not take the AITER path; a fused norm must not

1398 + take the unfused AITER op.

1399 +

1400 + The broadcast projects a narrower ``x``, which no AITER pre kernel takes;

1401 + ``mhc_pre_gemm_sqrsum`` folds no RMSNorm, so a fused norm routes to the

1402 + fused seam kernel where it exists (gfx950) and away from AITER otherwise.

1403 + What is asserted here is the routing decision, which is what this gate

1404 + owns; the numerics of the paths are covered by

1405 + ``test_deepseek_v41_mhc_pre_delayed`` and

1406 + ``test_mhc_pre_delayed_rocm_aiter_fused_rms_norm``.

1407 """

1408 set_random_seed(0)

1409 hc_mult, hidden_size = 4, 5120

@@ -1331,6 +1423,18 @@ def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch):

1423

1424 monkeypatch.setattr(torch.ops.vllm, "mhc_pre_delayed_aiter", spy)

1425

1426 + fused_op = torch.ops.vllm.mhc_fused_post_pre_delayed_rms_norm_aiter

1427 + took_fused = False

1428 +

1429 + def fused_spy(*spy_args, **spy_kwargs):

1430 + nonlocal took_fused

1431 + took_fused = True

1432 + return fused_op(*spy_args, **spy_kwargs)

1433 +

1434 + monkeypatch.setattr(

1435 + torch.ops.vllm, "mhc_fused_post_pre_delayed_rms_norm_aiter", fused_spy

1436 + )

1437 +

1438 x = residual[:, 0].contiguous()

1439 broadcast_fn = fn.view(-1, hc_mult, hidden_size).sum(1)

1440 broadcast_residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()

@@ -1342,11 +1446,15 @@ def test_mhc_pre_delayed_rocm_aiter_declines_unsupported(monkeypatch):

1446 )

1447 assert residual_out is broadcast_residual

1448 assert not took_aiter, "the broadcast seam must not reach AITER"

1449 + assert not took_fused, "the broadcast seam must not reach AITER"

1450 for i in range(4):

1451 torch.testing.assert_close(actual[i], expected[i], atol=1e-4, rtol=1e-3)

1452

1453 + # A fused norm must not reach the unfused AITER op; where the fused seam

1454 + # kernel exists (gfx950) it routes there instead.

1455 op.forward_hip(*args, norm_weight=norm_weight, norm_eps=1e-6)

1349 - assert not took_aiter, "a fused norm must not reach AITER"

1456 + assert not took_aiter, "a fused norm must not reach the unfused AITER op"

1457 + assert took_fused == mhc_layers.HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM

1458

1459 # Positive control: the same inputs without a norm do take the AITER path,

1460 # so the two declines above are the gate discriminating rather than the

vllm/_aiter_ops.py

+51

@@ -4235,6 +4235,57 @@ class rocm_aiter_ops:

4235 next_pre_mix.view(*outer_shape, hc_mult),

4236 )

4237

4238 + @staticmethod

4239 + def mhc_fused_post_pre_delayed_rms_norm(

4240 + residual: torch.Tensor,

4241 + fn: torch.Tensor,

4242 + hc_scale: torch.Tensor,

4243 + hc_base: torch.Tensor,

4244 + rms_eps: float,

4245 + hc_pre_eps: float,

4246 + hc_sinkhorn_eps: float,

4247 + hc_post_mult_value: float,

4248 + sinkhorn_repeat: int,

4249 + pre_mix: torch.Tensor | None,

4250 + sublayer_out: torch.Tensor | None,

4251 + post_layer_mix: torch.Tensor | None,

4252 + comb_res_mix: torch.Tensor | None,

4253 + norm_weight: torch.Tensor,

4254 + norm_eps: float,

4255 + residual_out: torch.Tensor | None,

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

4257 + """:meth:`mhc_pre_delayed` with the RMSNorm of the collapse folded in.

4258 +

4259 + Runs aiter's fused Triton kernel ``mhc_fused_post_pre_delayed_rmsnorm``;

4260 + the returned ``layer_input`` is already normalised with ``norm_weight`` /

4261 + ``norm_eps``.

4262 + """

4263 + from aiter.ops.triton.fusions.mhc_fused_post_pre_delayed_rmsnorm import (

4264 + mhc_fused_post_pre_delayed_rmsnorm,

4265 + )

4266 +

4267 + _, post_mix, comb_mix, layer_input, next_pre_mix = (

4268 + mhc_fused_post_pre_delayed_rmsnorm(

4269 + residual,

4270 + fn,

4271 + hc_scale,

4272 + hc_base,

4273 + rms_eps,

4274 + hc_pre_eps,

4275 + hc_sinkhorn_eps,

4276 + hc_post_mult_value,

4277 + sinkhorn_repeat,

4278 + pre_mix,

4279 + sublayer_out,

4280 + post_layer_mix,

4281 + comb_res_mix,

4282 + norm_weight,

4283 + norm_eps,

4284 + residual_out=residual_out,

4285 + )

4286 + )

4287 + return post_mix, comb_mix, layer_input, next_pre_mix

4288 +

4289 @staticmethod

4290 def mhc_fused_post_pre_delayed_prefers_unfused(num_tokens: int) -> bool:

4291 """True when AITER's own heuristic favours separate post and pre.

vllm/model_executor/kernels/mhc/__init__.py

@@ -13,6 +13,7 @@ __all__ = [

13 "hc_head_fused_kernel_cuda",

14 "mhc_pre_aiter",

15 "mhc_pre_delayed_aiter",

16 + "mhc_fused_post_pre_delayed_rms_norm_aiter",

17 "mhc_post_aiter",

18 "mhc_fused_post_pre_aiter",

19 "hc_head_fused_aiter",

vllm/model_executor/kernels/mhc/aiter.py

+79

@@ -207,6 +207,79 @@ def _mhc_pre_delayed_aiter_fake(

207 )

208

209

210 +def mhc_fused_post_pre_delayed_rms_norm_aiter(

211 + residual: torch.Tensor,

212 + fn: torch.Tensor,

213 + hc_scale: torch.Tensor,

214 + hc_base: torch.Tensor,

215 + rms_eps: float,

216 + hc_pre_eps: float,

217 + hc_sinkhorn_eps: float,

218 + hc_post_mult_value: float,

219 + sinkhorn_repeat: int,

220 + pre_mix: torch.Tensor | None,

221 + sublayer_out: torch.Tensor | None,

222 + post_layer_mix: torch.Tensor | None,

223 + comb_res_mix: torch.Tensor | None,

224 + norm_weight: torch.Tensor,

225 + norm_eps: float,

226 + residual_out: torch.Tensor | None,

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

228 + """:func:`mhc_pre_delayed_aiter` with the RMSNorm of the collapse folded in:

229 + ``layer_input`` comes back normalised with ``norm_weight`` / ``norm_eps``."""

230 + from vllm._aiter_ops import rocm_aiter_ops

231 +

232 + return rocm_aiter_ops.mhc_fused_post_pre_delayed_rms_norm(

233 + residual,

234 + fn,

235 + hc_scale,

236 + hc_base,

237 + rms_eps,

238 + hc_pre_eps,

239 + hc_sinkhorn_eps,

240 + hc_post_mult_value,

241 + sinkhorn_repeat,

242 + pre_mix,

243 + sublayer_out,

244 + post_layer_mix,

245 + comb_res_mix,

246 + norm_weight,

247 + norm_eps,

248 + residual_out,

249 + )

250 +

251 +

252 +def _mhc_fused_post_pre_delayed_rms_norm_aiter_fake(

253 + residual: torch.Tensor,

254 + fn: torch.Tensor,

255 + hc_scale: torch.Tensor,

256 + hc_base: torch.Tensor,

257 + rms_eps: float,

258 + hc_pre_eps: float,

259 + hc_sinkhorn_eps: float,

260 + hc_post_mult_value: float,

261 + sinkhorn_repeat: int,

262 + pre_mix: torch.Tensor | None,

263 + sublayer_out: torch.Tensor | None,

264 + post_layer_mix: torch.Tensor | None,

265 + comb_res_mix: torch.Tensor | None,

266 + norm_weight: torch.Tensor,

267 + norm_eps: float,

268 + residual_out: torch.Tensor | None,

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

270 + return _mhc_pre_delayed_aiter_fake(

271 + residual,

272 + fn,

273 + hc_scale,

274 + hc_base,

275 + rms_eps,

276 + hc_pre_eps,

277 + hc_sinkhorn_eps,

278 + hc_post_mult_value,

279 + sinkhorn_repeat,

280 + )

281 +

282 +

283 def mhc_post_aiter(

284 x: torch.Tensor,

285 residual: torch.Tensor,

@@ -335,6 +408,12 @@ direct_register_custom_op(

408 mutates_args=["residual_out"],

409 fake_impl=_mhc_pre_delayed_aiter_fake,

410 )

411 +direct_register_custom_op(

412 + op_name="mhc_fused_post_pre_delayed_rms_norm_aiter",

413 + op_func=mhc_fused_post_pre_delayed_rms_norm_aiter,

414 + mutates_args=["residual_out"],

415 + fake_impl=_mhc_fused_post_pre_delayed_rms_norm_aiter_fake,

416 +)

417 direct_register_custom_op(

418 op_name="mhc_post_aiter",

419 op_func=mhc_post_aiter,

vllm/model_executor/layers/mhc.py

+48 -4

@@ -59,9 +59,21 @@ def _aiter_mhc_op_accepts_norm(op_name: str) -> bool:

59 return False

62 +def _has_aiter_mhc_fused_post_pre_delayed_rms_norm() -> bool:

63 + if not HAS_AITER_MHC:

64 + return False

65 + from vllm.platforms.rocm import on_gfx950

66 +

67 + # The fused delayed seam kernel is only validated on gfx950.

68 + return on_gfx950()

69 +

70 +

71 HAS_AITER_MHC_FUSED = _has_aiter_mhc_fused()

72 HAS_AITER_MHC_PRE_NORM = _aiter_mhc_op_accepts_norm("mhc_pre")

73 HAS_AITER_MHC_FUSED_NORM = _aiter_mhc_op_accepts_norm("mhc_fused_post_pre")

74 +HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM = (

75 + _has_aiter_mhc_fused_post_pre_delayed_rms_norm()

76 +)

79 def _aiter_mhc_supported(

@@ -388,10 +400,42 @@ class MHCPreDelayedOp(CustomOp):

400 post_layer_mix: torch.Tensor | None = None,

401 comb_res_mix: torch.Tensor | None = None,

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

391 - # The AITER delayed path drives mhc_pre_gemm_sqrsum against `residual`

392 - # and folds no RMSNorm, so it cannot serve the model-entry broadcast

393 - # (which projects a narrower `x`) or a fused norm. Both are handled by

394 - # TileLang, or by the reference when TileLang is unavailable.

403 + # aiter's fused delayed seam: post-mix, gate projection, and the

404 + # collapse with the carried pre-mix and its RMSNorm. It always applies

405 + # the norm and projects `residual` itself.

406 + if (

407 + HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM

408 + and x is None

409 + and norm_weight is not None

410 + and _aiter_mhc_supported(residual, norm_weight, supports_norm=True)

411 + ):

412 + residual_out = (

413 + torch.empty_like(residual) if sublayer_out is not None else None

414 + )

415 + rest = torch.ops.vllm.mhc_fused_post_pre_delayed_rms_norm_aiter(

416 + residual,

417 + fn,

418 + hc_scale,

419 + hc_base,

420 + rms_eps,

421 + hc_pre_eps,

422 + hc_sinkhorn_eps,

423 + hc_post_mult_value,

424 + sinkhorn_repeat,

425 + pre_mix,

426 + sublayer_out,

427 + post_layer_mix,

428 + comb_res_mix,

429 + norm_weight,

430 + norm_eps,

431 + residual_out,

432 + )

433 + return (residual if residual_out is None else residual_out), *rest

434 + # The unfused AITER delayed path drives mhc_pre_gemm_sqrsum against

435 + # `residual` and folds no RMSNorm, so it cannot serve the model-entry

436 + # broadcast (which projects a narrower `x`) or a fused norm the branch

437 + # above did not take. Both are handled by TileLang, or by the reference

438 + # when TileLang is unavailable.

439 if (

440 x is None

441 and norm_weight is None

vllm/models/deepseek_v41/amd/model.py

+29 -4

@@ -22,7 +22,11 @@ from vllm.model_executor.layers.fused_moe import (

22 )

23 from vllm.model_executor.layers.layernorm import RMSNorm

24 from vllm.model_executor.layers.logits_processor import LogitsProcessor

25 -from vllm.model_executor.layers.mhc import MHCPostOp, MHCPreDelayedOp

25 +from vllm.model_executor.layers.mhc import (

26 + HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM,

27 + MHCPostOp,

28 + MHCPreDelayedOp,

29 +)

30 from vllm.model_executor.layers.quantization import QuantizationConfig

31 from vllm.model_executor.layers.vocab_parallel_embedding import (

32 ParallelLMHead,

@@ -251,6 +255,10 @@ class DeepseekV4DecoderLayer(nn.Module):

255 )

256 self.mhc_pre_delayed = MHCPreDelayedOp()

257 self.mhc_post = MHCPostOp()

258 + # Where aiter's fused seam kernel runs (gfx950), it folds the following

259 + # attn_norm / ffn_norm into its collapse, so the separate norms are

260 + # skipped for the seams it takes.

261 + self.fuse_seam_norm = HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM

262

263 @staticmethod

264 def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor:

@@ -269,6 +277,11 @@ class DeepseekV4DecoderLayer(nn.Module):

277 engram_hashes: torch.Tensor | None = None,

278 engram_mask: torch.Tensor | None = None,

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

280 + # Layer 0's attention seam projects the 2-D embedding with the folded

281 + # hc_attn_fn_broadcast instead of the 4-stream residual with hc_attn_fn.

282 + # The fused kernel only takes the latter, so that seam keeps the

283 + # separate attn_norm.

284 + fuse_attn_norm = self.fuse_seam_norm and not (residual is None and x.dim() == 2)

285 # The reference collapses each sublayer's input with the *previous*

286 # sublayer's pre-mix: attention uses the pre-mix carried in (identity

287 # for the first layer), the FFN uses this layer's attention pre-mix.

@@ -303,6 +316,8 @@ class DeepseekV4DecoderLayer(nn.Module):

316 self.hc_post_alpha,

317 self.hc_sinkhorn_iters,

318 pre_mix=pre_mix,

319 + norm_weight=self.attn_norm.weight if self.fuse_seam_norm else None,

320 + norm_eps=self.attn_norm.variance_epsilon,

321 )

322 else:

323 pre_args = (

@@ -328,7 +343,11 @@ class DeepseekV4DecoderLayer(nn.Module):

343 engram_mask,

344 )

345 residual, post_mix, res_mix, x, attn_pre = self.mhc_pre_delayed(

331 - residual, *pre_args, pre_mix=pre_mix

346 + residual,

347 + *pre_args,

348 + pre_mix=pre_mix,

349 + norm_weight=self.attn_norm.weight if self.fuse_seam_norm else None,

350 + norm_eps=self.attn_norm.variance_epsilon,

351 )

352 else:

353 (

@@ -344,8 +363,11 @@ class DeepseekV4DecoderLayer(nn.Module):

363 sublayer_out=x,

364 post_layer_mix=post_mix,

365 comb_res_mix=res_mix,

366 + norm_weight=self.attn_norm.weight if self.fuse_seam_norm else None,

367 + norm_eps=self.attn_norm.variance_epsilon,

368 )

348 - x = self.attn_norm(x)

369 + if not fuse_attn_norm:

370 + x = self.attn_norm(x)

371

372 if self.use_sequence_parallel:

373 x = sp_all_gather(x)[: positions.shape[0]]

@@ -368,8 +390,11 @@ class DeepseekV4DecoderLayer(nn.Module):

390 sublayer_out=x,

391 post_layer_mix=post_mix,

392 comb_res_mix=res_mix,

393 + norm_weight=self.ffn_norm.weight if self.fuse_seam_norm else None,

394 + norm_eps=self.ffn_norm.variance_epsilon,

395 )

372 - x = self.ffn_norm(x)

396 + if not self.fuse_seam_norm:

397 + x = self.ffn_norm(x)

398 x = self.ffn(x, input_ids)

399 return x, residual, post_mix, res_mix, ffn_pre

400