A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 08:45 UTC
[Perf] Intergrate Tile Kernels `head_compute_mix_kernel` for Deepseek-V4 (#41255)
Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Isotr0py committed
May 2, 2026 at 03:01 UTC
a9484dac7b734096ed26db4902454da7e497d2c3
2 files changed
+153
-9
vllm/model_executor/layers/mhc.py
+134
@@ -448,3 +448,137 @@ direct_register_custom_op(
448
mutates_args=[],
449
fake_impl=_mhc_post_fake,
450
)
451
+
452
+
453
+@tilelang.jit(
454
+ pass_configs={
455
+ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
456
+ tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
457
+ tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
458
+ },
459
+)
460
+def hc_head_fuse_tilelang(
461
+ residual,
462
+ fn,
463
+ hc_scale,
464
+ hc_base,
465
+ out,
466
+ hidden_size: int,
467
+ rms_eps: float,
468
+ hc_eps: float,
469
+ hc_mult: int = 4,
470
+ n_thr: int = 128,
471
+ h_blk: int = 1024,
472
+):
473
+ """Two-pass fused kernel for hc_head.
474
+
475
+ Pass 1: accumulate per-token squared sum and hc_mult dot-products
476
+ (projections onto fn rows) using cross-thread reducers.
477
+ Pass 2: apply sigmoid-gated weighted sum of residual channels to output.
478
+
479
+ Avoids materialising mixes / rsqrt / pre tensors to global memory.
480
+ """
481
+ num_tokens = T.dynamic("num_tokens")
482
+ hc_dim = hc_mult * hidden_size
483
+ h_block = math.gcd(h_blk, hidden_size)
484
+ n_h = hidden_size // h_block
485
+
486
+ residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type]
487
+ fn: T.Tensor[[hc_mult, hc_dim], T.float32] # type: ignore[no-redef,valid-type]
488
+ hc_scale: T.Tensor[[1], T.float32] # type: ignore[no-redef,valid-type]
489
+ hc_base: T.Tensor[[hc_mult], T.float32] # type: ignore[no-redef,valid-type]
490
+ out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type]
491
+
492
+ with T.Kernel(num_tokens, threads=n_thr) as i:
493
+ T.pdl_sync()
494
+
495
+ # ------------------------------------------------------------------
496
+ # Pass 1 – for each residual channel m_c and h_block:
497
+ # • accumulate squared sum (for RMS norm denominator)
498
+ # • accumulate hc_mult dot-products with fn rows
499
+ # ------------------------------------------------------------------
500
+ sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all")
501
+ mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all")
502
+ T.fill(sqrsum_r, 0.0)
503
+ T.fill(mixes_r, 0.0)
504
+
505
+ for m_c in T.serial(hc_mult):
506
+ for i_h in T.serial(n_h):
507
+ x_local = T.alloc_fragment(h_block, T.float32)
508
+ T.copy(residual[i, m_c, i_h * h_block], x_local)
509
+
510
+ for k in T.Parallel(h_block):
511
+ sqrsum_r[0] += x_local[k] * x_local[k]
512
+
513
+ for m_m in T.unroll(hc_mult):
514
+ fn_local = T.alloc_fragment(h_block, T.float32)
515
+ T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local)
516
+ for k in T.Parallel(h_block):
517
+ mixes_r[m_m] += x_local[k] * fn_local[k]
518
+
519
+ T.finalize_reducer(sqrsum_r)
520
+ T.finalize_reducer(mixes_r)
521
+
522
+ # ------------------------------------------------------------------
523
+ # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps
524
+ # ------------------------------------------------------------------
525
+ pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
526
+ rsqrt_val = T.alloc_fragment(1, T.float32)
527
+ rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps)
528
+ for m in T.Parallel(hc_mult):
529
+ pre_mix_shared[m] = (
530
+ T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps
531
+ )
532
+
533
+ # ------------------------------------------------------------------
534
+ # Pass 2 – apply_mix: pipelined weighted sum over residual channels
535
+ # ------------------------------------------------------------------
536
+ for i0_h in T.Pipelined(n_h, num_stages=2):
537
+ xs = T.alloc_shared((hc_mult, h_block), T.bfloat16)
538
+ xl = T.alloc_fragment((hc_mult, h_block), T.float32)
539
+ T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True)
540
+ T.copy(xs, xl)
541
+
542
+ ol = T.alloc_fragment(h_block, T.float32)
543
+ T.clear(ol)
544
+ for i_hc in T.serial(hc_mult):
545
+ pre = pre_mix_shared[i_hc]
546
+ for i1_h in T.Parallel(h_block):
547
+ ol[i1_h] += pre * xl[i_hc, i1_h]
548
+
549
+ T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
550
+
551
+ T.pdl_trigger()
552
+
553
+
554
+def _hc_head_fused_kernel(
555
+ hs_flat: torch.Tensor,
556
+ fn: torch.Tensor,
557
+ hc_scale: torch.Tensor,
558
+ hc_base: torch.Tensor,
559
+ out: torch.Tensor,
560
+ hidden_size: int,
561
+ rms_eps: float,
562
+ hc_eps: float,
563
+ hc_mult: int,
564
+) -> None:
565
+ """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
566
+ if hs_flat.shape[0] > 0:
567
+ hc_head_fuse_tilelang(
568
+ hs_flat,
569
+ fn,
570
+ hc_scale,
571
+ hc_base,
572
+ out,
573
+ hidden_size,
574
+ rms_eps,
575
+ hc_eps,
576
+ hc_mult,
577
+ )
578
+
579
+
580
+direct_register_custom_op(
581
+ op_name="hc_head_fused_kernel",
582
+ op_func=_hc_head_fused_kernel,
583
+ mutates_args=["out"],
584
+)
vllm/model_executor/models/deepseek_v4.py
+19
-9
@@ -7,7 +7,6 @@ from itertools import islice
7
import regex as re
8
import torch
9
import torch.nn as nn
10
-import torch.nn.functional as F
10
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig, get_current_vllm_config
@@ -1456,14 +1455,25 @@ def hc_head(
1455
rms_norm_eps: float,
1456
hc_eps: float,
1457
) -> torch.Tensor:
1459
- x = hidden_states
1460
- shape, dtype = x.size(), x.dtype
1461
- x = x.flatten(1).float()
1462
- rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + rms_norm_eps)
1463
- mixes = F.linear(x, hc_fn) * rsqrt
1464
- pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
1465
- y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1)
1466
- return y.to(dtype)
1458
+ hc_mult, hidden_size = hidden_states.shape[-2:]
1459
+ outer_shape = hidden_states.shape[:-2]
1460
+ hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
1461
+ num_tokens = hs_flat.shape[0]
1462
+ out = torch.empty(
1463
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
1464
+ )
1465
+ torch.ops.vllm.hc_head_fused_kernel(
1466
+ hs_flat,
1467
+ hc_fn,
1468
+ hc_scale,
1469
+ hc_base,
1470
+ out,
1471
+ hidden_size,
1472
+ rms_norm_eps,
1473
+ hc_eps,
1474
+ hc_mult,
1475
+ )
1476
+ return out.view(*outer_shape, hidden_size)
1477
1478
1479
def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: