@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9b49f92344

[Perf][DSV4.1] Overlap mHC coefficients for small TP batches (#57603)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Codex <noreply@openai.com>

Woosuk Kwon committed Sep 20, 2026 at 17:33 UTC 9b49f92344312c41ad61e05282c8e6a2d9bafb7f

6 files changed +357 -46

tests/kernels/test_mhc_kernels.py

+129 -12

@@ -266,6 +266,77 @@ def test_deepseek_v41_mhc_pre_delayed(

266 torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2)

267

268

269 +@pytest.mark.skipif(

270 + not HAS_TILELANG_MHC or not current_platform.is_device_capability_family(100),

271 + reason="SM100 TileLang mHC required",

272 +)

273 +@pytest.mark.parametrize("num_tokens", [0, 1, 8, 9, 32, 128, 256])

274 +@pytest.mark.parametrize("entry", ["broadcast", "identity", "carried"])

275 +@pytest.mark.parametrize("hidden_size,hc_mult", [(4096, 2), (5120, 4), (7168, 4)])

276 +def test_deepseek_v41_mhc_overlap_preserves_outputs_on_replay(

277 + num_tokens, entry, hidden_size, hc_mult

278 +):

279 + """The input is usable before joining; coefficients stay exact on graph replay."""

280 + from vllm.models.deepseek_v41.nvidia.ops.mhc import mhc_pre_delayed_overlap

281 + from vllm.utils.deep_gemm import is_deep_gemm_supported

282 +

283 + if not is_deep_gemm_supported():

284 + pytest.skip("DeepGEMM required")

285 + set_random_seed(42)

286 + mix_size = hc_mult * (hc_mult + 2)

287 + residual = torch.randn(

288 + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE

289 + )

290 + fn = torch.randn(mix_size, hc_mult, hidden_size, device=DEVICE) * 0.02

291 + x = None

292 + pre_mix = None

293 + if entry == "broadcast":

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

295 + residual = x.unsqueeze(1).expand(-1, hc_mult, -1).contiguous()

296 + fn = fn.sum(1)

297 + else:

298 + fn = fn.flatten(1)

299 + if entry == "carried":

300 + pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE)

301 + scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE)

302 + base = torch.randn(mix_size, device=DEVICE)

303 + weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE).uniform_(

304 + 0.5, 1.5

305 + )

306 + args = (residual, fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20)

307 + kwargs = dict(pre_mix=pre_mix, x=x, norm_weight=weight, norm_eps=1e-20)

308 + side = torch.cuda.Stream()

309 +

310 + def run():

311 + outputs = mhc_pre_delayed_overlap(*args, **kwargs, stream=side)

312 + consumed_input = outputs[2].clone()

313 + torch.cuda.current_stream().wait_stream(side)

314 + return (*outputs, consumed_input)

315 +

316 + def check(actual):

317 + expected = mhc_pre_delayed_tilelang(*args, **kwargs)

318 + for a, b in zip(actual[:4], expected, strict=True):

319 + torch.testing.assert_close(a, b, atol=0, rtol=0)

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

321 +

322 + check(run())

323 + if num_tokens == 0:

324 + return

325 + graph = torch.cuda.CUDAGraph()

326 + with torch.cuda.graph(graph):

327 + captured = run()

328 + for _ in range(5):

329 + if x is not None:

330 + x.normal_()

331 + residual.copy_(x.unsqueeze(1).expand_as(residual))

332 + else:

333 + residual.normal_()

334 + if pre_mix is not None:

335 + pre_mix.uniform_()

336 + graph.replay()

337 + check(captured)

338 +

339 +

340 def mhc_fused_post_pre_delayed_ref(

341 x: torch.Tensor,

342 residual: torch.Tensor,

@@ -433,10 +504,16 @@ def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_

504

505 @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")

506 @pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"])

507 +@pytest.mark.parametrize(

508 + "mhc_mode", ["disabled", "overlap", "piecewise", "large_batch"]

509 +)

510 def test_deepseek_v41_decoder_mixes_match_torch(

437 - entry, monkeypatch, default_vllm_config

511 + entry, mhc_mode, monkeypatch, default_vllm_config

512 ):

513 """Preserve carried pre-mixes and Engram ordering at every decoder entry."""

514 + from vllm.models.deepseek_v41.nvidia.ops.mhc import MHC_OVERLAP_MAX_TOKENS

515 +

516 + num_tokens = MHC_OVERLAP_MAX_TOKENS + 1 if mhc_mode == "large_batch" else 3

517 set_random_seed(0)

518 decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer)

519 nn.Module.__init__(decoder)

@@ -446,6 +523,29 @@ def test_deepseek_v41_decoder_mixes_match_torch(

523 decoder.rms_norm_eps = 1e-20

524 decoder.hc_post_alpha = 2.0

525 decoder.use_sequence_parallel = False

526 + decoder.mhc_stream = None

527 + if mhc_mode != "disabled":

528 + from vllm.utils.deep_gemm import is_deep_gemm_supported

529 +

530 + if (

531 + not current_platform.is_device_capability_family(100)

532 + or not is_deep_gemm_supported()

533 + ):

534 + pytest.skip("SM100 DeepGEMM required for overlap")

535 + decoder.mhc_stream = torch.cuda.Stream()

536 + if mhc_mode in ("piecewise", "large_batch"):

537 +

538 + def unexpected_overlap(*args, **kwargs):

539 + pytest.fail("unsupported execution must retain native mHC")

540 +

541 + monkeypatch.setattr(

542 + "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_overlap",

543 + unexpected_overlap,

544 + )

545 + monkeypatch.setattr(

546 + "vllm.models.deepseek_v41.nvidia.ops.mhc.mhc_pre_delayed_overlap",

547 + unexpected_overlap,

548 + )

549 decoder.engram = None

550 from vllm.model_executor.layers.layernorm import RMSNorm

551

@@ -461,18 +561,18 @@ def test_deepseek_v41_decoder_mixes_match_torch(

561 decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3)

562 decoder.hc_attn_base = torch.randn(24)

563 decoder.hc_ffn_base = torch.randn(24)

464 - x = torch.randn(3, 5120, dtype=torch.bfloat16)

465 - positions = torch.arange(3)

564 + x = torch.randn(num_tokens, 5120, dtype=torch.bfloat16)

565 + positions = torch.arange(num_tokens)

566 kwargs = {}

567 if entry != "broadcast":

468 - kwargs["pre_mix"] = torch.rand(3, 4)

568 + kwargs["pre_mix"] = torch.rand(num_tokens, 4)

569 if entry == "pipeline":

470 - x = torch.randn(3, 4, 5120, dtype=torch.bfloat16)

570 + x = torch.randn(num_tokens, 4, 5120, dtype=torch.bfloat16)

571 if entry in ("residual", "engram"):

572 kwargs.update(

473 - residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16),

474 - post_mix=torch.rand(3, 4, 1),

475 - res_mix=torch.rand(3, 4, 4),

573 + residual=torch.randn(num_tokens, 4, 5120, dtype=torch.bfloat16),

574 + post_mix=torch.rand(num_tokens, 4, 1),

575 + res_mix=torch.rand(num_tokens, 4, 4),

576 )

577 if entry == "engram":

578

@@ -483,15 +583,22 @@ def test_deepseek_v41_decoder_mixes_match_torch(

583 return residual + 0.125

584

585 decoder.engram = FakeEngram()

486 - kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32)

586 + kwargs["engram_hashes"] = torch.zeros(num_tokens, 1, 1, dtype=torch.int32)

587

488 - actual = decoder(x, positions, None, **kwargs)

588 + from vllm.config import CUDAGraphMode

589 + from vllm.forward_context import set_forward_context

590 +

591 + mode = CUDAGraphMode.PIECEWISE if mhc_mode == "piecewise" else CUDAGraphMode.NONE

592 + with set_forward_context(None, default_vllm_config, cudagraph_runtime_mode=mode):

593 + actual = decoder(x, positions, None, **kwargs)

594

595 def reference(*args, norm_weight, norm_eps, **kwargs):

596 post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs)

597 return post, res, decoder.attn_norm(collapsed), pre

598

494 - def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, **kw):

599 + def fused_reference(

600 + x, residual, post_mix, res_mix, *args, capture_aux=False, stream=None, **kw

601 + ):

602 residual = mhc_post_torch(x, residual, post_mix, res_mix)

603 aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0)

604 return residual, *reference(residual, *args, **kw), aux

@@ -500,6 +607,7 @@ def test_deepseek_v41_decoder_mixes_match_torch(

607 "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang",

608 reference,

609 )

610 + decoder.mhc_stream = None

611 monkeypatch.setattr(

612 "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre",

613 fused_reference,

@@ -530,6 +638,7 @@ def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_conf

638 decoder.rms_norm_eps = 1e-6

639 decoder.hc_post_alpha = 2.0

640 decoder.use_sequence_parallel = False

641 + decoder.mhc_stream = None

642 decoder.engram = None

643 from vllm.model_executor.layers.layernorm import RMSNorm

644

@@ -772,7 +881,8 @@ def hc_head_ref(

881 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])

882 @pytest.mark.parametrize("hidden_size", [4096, 7168])

883 @pytest.mark.parametrize("hc_mult", [4])

775 -def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):

884 +@pytest.mark.parametrize("fused_norm", [False, True])

885 +def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult, fused_norm):

886 torch.set_default_device(DEVICE)

887 set_random_seed(0)

888

@@ -802,6 +912,12 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):

912 hc_post_alpha,

913 sinkhorn_repeat,

914 )

915 + norm_kwargs = {}

916 + if fused_norm:

917 + # Non-shifted mHC must still use its freshly computed pre-mix.

918 + weight = torch.ones(hidden_size, dtype=torch.bfloat16)

919 + norm_kwargs = dict(norm_weight=weight, norm_eps=rms_eps)

920 + ref = (*ref[:2], F.rms_norm(ref[2], (hidden_size,), weight, rms_eps))

921 out = torch.ops.vllm.mhc_pre_tilelang(

922 residual,

923 fn,

@@ -812,6 +928,7 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):

928 hc_sinkhorn_eps,

929 hc_post_alpha,

930 sinkhorn_repeat,

931 + **norm_kwargs,

932 )

933

934 for actual, expected in zip(out, ref, strict=True):

vllm/model_executor/kernels/mhc/tilelang_kernels.py

+21 -17

@@ -305,8 +305,12 @@ def mhc_pre_big_fuse_with_norm_tilelang(

305 use_pre_mix_in: bool = False,

306 save_pre_mix: bool = False,

307 rms_numel: int = 0,

308 + split_mode: str = "fused",

309 write_aux: bool = False,

310 ):

311 + # Split modes require shifted mHC: input collapse uses a carried pre-mix.

312 + assert split_mode in ("fused", "stats", "input")

313 + assert split_mode == "fused" or save_pre_mix

314 num_tokens = T.dynamic("num_tokens")

315 hc_mult3 = hc_mult * (2 + hc_mult)

316 if gemm_last_dim < 0:

@@ -330,26 +334,26 @@ def mhc_pre_big_fuse_with_norm_tilelang(

334 aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]

335

336 with T.Kernel(num_tokens, threads=96) as i:

333 - rms = T.alloc_fragment(1, T.float32)

334 - mixes = T.alloc_fragment(hc_mult3, T.float32)

335 - T.clear(mixes)

336 - rms[0] = 0

337 -

337 if ENABLE_PDL:

338 T.pdl_sync()

340 -

341 - for i_split in T.serial(n_splits):

342 - rms[0] += gemm_out_sqrsum[i_split, i]

343 - rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps)

344 - for j in T.Parallel(hc_mult3):

345 - mixes[j] = 0

346 - for i_split in T.serial(n_splits):

347 - mixes[j] += gemm_out_mul[i_split, i, j]

348 - mixes[j] *= rms[0]

339 mixes_shared = T.alloc_shared(hc_mult3, T.float32)

350 - T.copy(mixes, mixes_shared)

340 + if split_mode != "input":

341 + rms = T.alloc_fragment(1, T.float32)

342 + mixes = T.alloc_fragment(hc_mult3, T.float32)

343 + T.clear(mixes)

344 + rms[0] = 0

345

352 - if T.get_thread_binding() < 32:

346 + for i_split in T.serial(n_splits):

347 + rms[0] += gemm_out_sqrsum[i_split, i]

348 + rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps)

349 + for j in T.Parallel(hc_mult3):

350 + mixes[j] = 0

351 + for i_split in T.serial(n_splits):

352 + mixes[j] += gemm_out_mul[i_split, i, j]

353 + mixes[j] *= rms[0]

354 + T.copy(mixes, mixes_shared)

355 +

356 + if split_mode != "input" and T.get_thread_binding() < 32:

357 cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)

358 for j in T.Parallel(hc_mult):

359 if save_pre_mix:

@@ -395,7 +399,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(

399

400 for j, k in T.Parallel(hc_mult, hc_mult):

401 comb_mix[i, j * hc_mult + k] = cm[j, k]

398 - else:

402 + elif split_mode != "stats" and T.get_thread_binding() >= 32:

403 pre_mix_shared = T.alloc_shared(hc_mult, T.float32)

404 for j in T.Parallel(hc_mult):

405 if use_pre_mix_in:

vllm/model_executor/kernels/mhc/warmup.py

@@ -36,6 +36,7 @@ class MHCPreNormKernel(VllmJitKernel["MHCPreNormKernel.CompileKey"]):

36 use_pre_mix_in: bool

37 rms_numel: int

38 save_pre_mix: bool = True

39 + split_mode: str = "fused"

40 write_aux: bool = False

42 kernel: Any = staticmethod(mhc_pre_big_fuse_with_norm_tilelang)

vllm/models/deepseek_v41/nvidia/model.py

+60 -16

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import typing

4 from collections.abc import Callable, Iterable

5 +from functools import partial

6 from itertools import islice

8 import regex as re

@@ -17,7 +18,11 @@ from vllm.distributed import (

18 get_tensor_model_parallel_rank,

19 get_tensor_model_parallel_world_size,

20 )

20 -from vllm.forward_context import get_forward_context, is_forward_context_available

21 +from vllm.forward_context import (

22 + get_forward_context,

23 + in_piecewise_cudagraph,

24 + is_forward_context_available,

25 +)

26 from vllm.logger import init_logger

27 from vllm.model_executor.kernels.mhc.tilelang import (

28 mhc_post_tilelang,

@@ -84,6 +89,11 @@ from ..common.engram import EngramLayout, NgramHashState

89 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask

90 from .engram import Engram, gather_engram_hashes

91 from .ops.mega_mhc import mhc_shifted_post_pre

92 +from .ops.mhc import (

93 + MHC_OVERLAP_MAX_TOKENS,

94 + mhc_pre_delayed_overlap,

95 + supports_mhc_overlap,

96 +)

98 if typing.TYPE_CHECKING:

99 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata

@@ -188,11 +198,13 @@ class DeepseekV4DecoderLayer(nn.Module):

198 aux_stream_list: list[torch.cuda.Stream] | None = None,

199 candidate_block_buffer: torch.Tensor | None = None,

200 engram_layout: EngramLayout | None = None,

201 + mhc_stream: torch.cuda.Stream | None = None,

202 ):

203 super().__init__()

204

205 config = vllm_config.model_config.hf_config

206 self.hidden_size = config.hidden_size

207 + self.mhc_stream = mhc_stream

208 self.use_sequence_parallel = _use_sequence_parallel(vllm_config)

209

210 self.engram: Engram | None = None

@@ -269,19 +281,30 @@ class DeepseekV4DecoderLayer(nn.Module):

281 variants.append(

282 {**hc_stream, "use_pre_mix_in": True, "write_aux": True}

283 )

284 + # Shifted mHC splits the epilogue in two: the input collapse on the

285 + # caller stream and coefficient generation on a side stream.

286 + split_modes = (

287 + ("fused", "stats", "input") if mhc_stream is not None else ("fused",)

288 + )

289 for variant in variants:

273 - MHC_PRE_NORM_KERNEL.register_warmup(

274 - max_tokens=max_tokens,

275 - hidden_size=self.hidden_size,

276 - rms_eps=self.rms_norm_eps,

277 - hc_pre_eps=self.hc_eps,

278 - hc_sinkhorn_eps=self.hc_eps,

279 - hc_post_mult_value=self.hc_post_alpha,

280 - sinkhorn_repeat=self.hc_sinkhorn_iters,

281 - norm_eps=self.rms_norm_eps,

282 - hc_mult=self.hc_mult,

283 - **variant,

284 - )

290 + for split_mode in split_modes:

291 + MHC_PRE_NORM_KERNEL.register_warmup(

292 + split_mode=split_mode,

293 + max_tokens=(

294 + max_tokens

295 + if split_mode == "fused"

296 + else min(max_tokens, MHC_OVERLAP_MAX_TOKENS)

297 + ),

298 + hidden_size=self.hidden_size,

299 + rms_eps=self.rms_norm_eps,

300 + hc_pre_eps=self.hc_eps,

301 + hc_sinkhorn_eps=self.hc_eps,

302 + hc_post_mult_value=self.hc_post_alpha,

303 + sinkhorn_repeat=self.hc_sinkhorn_iters,

304 + norm_eps=self.rms_norm_eps,

305 + hc_mult=self.hc_mult,

306 + **variant,

307 + )

308 mix_hc = (2 + self.hc_mult) * self.hc_mult

309 hc_dim = self.hc_mult * self.hidden_size

310 self.hc_attn_fn = nn.Parameter(

@@ -351,6 +374,18 @@ class DeepseekV4DecoderLayer(nn.Module):

374 torch.Tensor | None,

375 ]:

376 previous_aux: torch.Tensor | None = None

377 + mhc_stream = self.mhc_stream

378 + if mhc_stream is not None and (

379 + in_piecewise_cudagraph()

380 + or not 0 < positions.shape[0] <= MHC_OVERLAP_MAX_TOKENS

381 + ):

382 + # A side stream cannot remain unjoined across breakable graph segments.

383 + mhc_stream = None

384 + mhc_pre = (

385 + partial(mhc_pre_delayed_overlap, stream=mhc_stream)

386 + if mhc_stream is not None

387 + else mhc_pre_delayed_tilelang

388 + )

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

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

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

@@ -360,7 +395,7 @@ class DeepseekV4DecoderLayer(nn.Module):

395 # copies and the identity pre-mix selects copy 0.

396 assert self.hc_attn_fn_broadcast is not None

397 residual = x.unsqueeze(1).expand(-1, self.hc_mult, -1).contiguous()

363 - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(

398 + post_mix, res_mix, x, attn_pre = mhc_pre(

399 residual,

400 self.hc_attn_fn_broadcast,

401 self.hc_attn_scale,

@@ -376,7 +411,7 @@ class DeepseekV4DecoderLayer(nn.Module):

411 )

412 else:

413 residual = x

379 - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(

414 + post_mix, res_mix, x, attn_pre = mhc_pre(

415 residual,

416 self.hc_attn_fn,

417 self.hc_attn_scale,

@@ -403,7 +438,7 @@ class DeepseekV4DecoderLayer(nn.Module):

438 engram_hashes[:, self.engram.layer_hash_index],

439 engram_mask,

440 )

406 - post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(

441 + post_mix, res_mix, x, attn_pre = mhc_pre(

442 residual,

443 self.hc_attn_fn,

444 self.hc_attn_scale,

@@ -437,6 +472,7 @@ class DeepseekV4DecoderLayer(nn.Module):

472 norm_weight=self.attn_norm.weight,

473 norm_eps=self.attn_norm.variance_epsilon,

474 capture_aux=capture_previous_aux,

475 + stream=mhc_stream,

476 )

477 if capture_previous_aux:

478 previous_aux = aux

@@ -448,6 +484,8 @@ class DeepseekV4DecoderLayer(nn.Module):

484 if self.use_sequence_parallel:

485 x = sp_reduce_scatter(x)

486

487 + if mhc_stream is not None:

488 + torch.cuda.current_stream().wait_stream(mhc_stream)

489 residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre(

490 x,

491 residual,

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

502 pre_mix=attn_pre,

503 norm_weight=self.ffn_norm.weight,

504 norm_eps=self.ffn_norm.variance_epsilon,

505 + stream=mhc_stream,

506 )

507 x = self.ffn(x, input_ids, mega_gate_metadata)

508 + if mhc_stream is not None:

509 + torch.cuda.current_stream().wait_stream(mhc_stream)

510 return x, residual, post_mix, res_mix, ffn_pre, previous_aux

511

512

@@ -497,6 +538,8 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

538 # (compressor kv_score, indexer.weights_proj). fused_wqa_wkv stays on

539 # the default stream.

540 aux_stream_list = [torch.cuda.Stream() for _ in range(3)]

541 + # Keep mHC independent of the streams used inside attention.

542 + mhc_stream = torch.cuda.Stream() if supports_mhc_overlap(vllm_config) else None

543

544 # Reserved topk indices buffer for all Indexer layers to reuse.

545 self.topk_indices_buffer = torch.empty(

@@ -541,6 +584,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

584 aux_stream_list=aux_stream_list,

585 candidate_block_buffer=self.candidate_block_buffer,

586 engram_layout=self.engram_layout,

587 + mhc_stream=mhc_stream,

588 ),

589 prefix=f"{prefix}.layers",

590 )

vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py

+29 -1

@@ -7,6 +7,7 @@ import torch

8 from vllm.model_executor.kernels.mhc.tilelang import (

9 mhc_fused_post_pre_delayed_tilelang,

10 + mhc_post_tilelang,

11 )

12 from vllm.platforms import current_platform

13 from vllm.utils.deep_gemm import (

@@ -106,10 +107,37 @@ def mhc_shifted_post_pre(

107 norm_weight: torch.Tensor | None = None,

108 norm_eps: float = 1e-6,

109 capture_aux: bool = False,

110 + *,

111 + stream: torch.cuda.Stream | None = None,

112 ) -> tuple[

113 torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor

114 ]:

112 - """Use Mega-mHC for carried mixing, retaining TileLang for aux capture."""

115 + """Dispatch shifted post/pre to overlap, Mega-mHC, or fused TileLang.

116 +

117 + When stream is supplied, join it before consuming the returned coefficients.

118 + """

119 + if stream is not None:

120 + from .mhc import mhc_pre_delayed_overlap

121 +

122 + residual = mhc_post_tilelang(x, residual, post_layer_mix, comb_res_mix)

123 + aux = residual.mean(dim=1) if capture_aux else x.new_empty(0, x.shape[1])

124 + pre_outputs = mhc_pre_delayed_overlap(

125 + residual,

126 + fn,

127 + hc_scale,

128 + hc_base,

129 + rms_eps,

130 + hc_pre_eps,

131 + hc_sinkhorn_eps,

132 + hc_post_mult_value,

133 + sinkhorn_repeat,

134 + pre_mix=pre_mix,

135 + norm_weight=norm_weight,

136 + norm_eps=norm_eps,

137 + stream=stream,

138 + )

139 + return residual, *pre_outputs, aux

140 +

141 if (

142 pre_mix is not None

143 and norm_weight is not None

vllm/models/deepseek_v41/nvidia/ops/mhc.py new

+117

@@ -0,0 +1,117 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Split shifted mHC pre so coefficient generation can overlap the sublayer."""

4 +

5 +from functools import partial

6 +

7 +import torch

8 +

9 +from vllm.config import VllmConfig

10 +from vllm.platforms import current_platform

11 +from vllm.utils.deep_gemm import is_deep_gemm_supported

12 +

13 +# GB200 TP4/FlashInfer improves through 16 tokens; larger screens tie or regress.

14 +MHC_OVERLAP_MAX_TOKENS = 16

15 +

16 +

17 +def supports_mhc_overlap(vllm_config: VllmConfig) -> bool:

18 + """Check kernel requirements and safety of sharing the coefficient stream."""

19 + config = vllm_config.model_config.hf_config

20 + # DeepGEMM's prenorm kernel requires K % 64 == 0, N % 8 == 0, and N <= 32.

21 + mix_size = config.hc_mult * (config.hc_mult + 2)

22 + return (

23 + current_platform.is_device_capability_family(100)

24 + and is_deep_gemm_supported()

25 + and config.hidden_size % 64 == 0

26 + and 0 < mix_size <= 32

27 + and mix_size % 8 == 0

28 + and not vllm_config.parallel_config.use_ubatching

29 + )

30 +

31 +

32 +def mhc_pre_delayed_overlap(

33 + residual: torch.Tensor,

34 + fn: torch.Tensor,

35 + hc_scale: torch.Tensor,

36 + hc_base: torch.Tensor,

37 + rms_eps: float,

38 + hc_pre_eps: float,

39 + hc_sinkhorn_eps: float,

40 + hc_post_mult_value: float,

41 + sinkhorn_repeat: int,

42 + pre_mix: torch.Tensor | None = None,

43 + x: torch.Tensor | None = None,

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

45 + norm_eps: float = 1e-6,

46 + *,

47 + stream: torch.cuda.Stream,

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

49 + """Prepare the input on the caller stream and coefficients on another stream.

50 +

51 + Returns post mix, residual mix, normalized input, and next pre-mix. Only the

52 + input is ready on the caller stream; join stream before using coefficients.

53 + The caller must retain the weights and coefficient outputs until the join.

54 + """

55 + from vllm.model_executor.kernels.mhc.warmup import (

56 + MHC_PRE_NORM_KERNEL,

57 + compute_mhc_pre_num_splits,

58 + )

59 + from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm

60 +

61 + n, hc, hidden = residual.shape

62 + assert residual.is_contiguous() and residual.dtype == torch.bfloat16

63 + assert norm_weight is not None

64 + if x is None:

65 + x = residual.view(n, hc * hidden)

66 + assert x.is_contiguous()

67 + post = torch.empty((n, hc), device=residual.device, dtype=torch.float32)

68 + comb = torch.empty((n, hc * hc), device=residual.device, dtype=torch.float32)

69 + next_pre = torch.empty_like(post)

70 + layer_input = torch.empty((n, hidden), device=residual.device, dtype=torch.bfloat16)

71 + outputs = post.unsqueeze(-1), comb.view(n, hc, hc), layer_input, next_pre

72 + if n == 0:

73 + return outputs

74 + splits = compute_mhc_pre_num_splits(x.shape[1], n)

75 + mix = torch.empty(

76 + (splits, n, hc * (hc + 2)), device=residual.device, dtype=torch.float32

77 + )

78 + sqr = torch.empty((splits, n), device=residual.device, dtype=torch.float32)

79 + epilogue = partial(

80 + MHC_PRE_NORM_KERNEL,

81 + mix,

82 + sqr,

83 + hc_scale,

84 + hc_base,

85 + residual,

86 + post,

87 + comb,

88 + layer_input,

89 + norm_weight,

90 + pre_mix if pre_mix is not None else post,

91 + next_pre,

92 + layer_input, # Unused aux output in split modes.

93 + hidden_size=hidden,

94 + rms_eps=rms_eps,

95 + hc_pre_eps=hc_pre_eps,

96 + hc_sinkhorn_eps=hc_sinkhorn_eps,

97 + hc_post_mult_value=hc_post_mult_value,

98 + sinkhorn_repeat=sinkhorn_repeat,

99 + norm_eps=norm_eps,

100 + hc_mult=hc,

101 + use_pre_mix_in=pre_mix is not None,

102 + save_pre_mix=True,

103 + rms_numel=x.shape[1],

104 + )

105 + main = torch.cuda.current_stream()

106 + # Prioritize input readiness for small batches before releasing statistics.

107 + if n <= 8:

108 + epilogue(split_mode="input")

109 + stream.wait_stream(main)

110 + with torch.cuda.stream(stream):

111 + tf32_hc_prenorm_gemm(x, fn, mix, sqr, splits)

112 + epilogue(split_mode="stats")

113 + for tensor in (x, mix, sqr):

114 + tensor.record_stream(stream)

115 + if n > 8:

116 + epilogue(split_mode="input")

117 + return outputs