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 2, 12:15 UTC
[DSv4] Move mHC tilelang kernels & Don't use CustomOP in dsv4/nvidia (#43905)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 28, 2026 at 19:25 UTC
7bd45da5857db1d83dbf224203f4936264af8ded
6 files changed
+72
-102
tests/kernels/test_mhc_kernels.py
+3
-8
@@ -340,22 +340,17 @@ def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
340
hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
341
rms_eps = hc_eps = 1e-6
342
343
- out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16)
344
- out.fill_(float("nan"))
345
-
346
- result = torch.ops.vllm.hc_head_fused_kernel_tilelang(
343
+ out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
344
residual,
345
fn,
346
hc_scale,
347
hc_base,
351
- out,
352
- hidden_size,
348
rms_eps,
349
hc_eps,
355
- hc_mult,
350
)
351
358
- assert result is None
352
+ assert out.shape == (num_tokens, hidden_size)
353
+ assert out.dtype == torch.bfloat16
354
assert not torch.isnan(out).any()
355
356
out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
vllm/model_executor/kernels/mhc/tilelang.py
+34
-15
@@ -29,7 +29,7 @@ def _tilelang_hc_prenorm_gemm(
29
n_thr: int = 512,
30
n_splits: int = 1,
31
) -> None:
32
- from vllm._tilelang_ops import (
32
+ from vllm.model_executor.kernels.mhc.tilelang_kernels import (
33
hc_prenorm_gemm_block_m_tilelang,
34
hc_prenorm_gemm_tilelang,
35
)
@@ -126,7 +126,7 @@ def mhc_pre_tilelang(
126
comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
127
layer_input: shape (..., hidden_size), dtype torch.bfloat16
128
"""
129
- from vllm._tilelang_ops import (
129
+ from vllm.model_executor.kernels.mhc.tilelang_kernels import (
130
compute_num_split,
131
mhc_pre_big_fuse_tilelang,
132
mhc_pre_big_fuse_with_norm_tilelang,
@@ -306,7 +306,9 @@ def mhc_post_tilelang(
306
post_layer_mix: torch.Tensor,
307
comb_res_mix: torch.Tensor,
308
) -> torch.Tensor:
309
- from vllm._tilelang_ops import mhc_post_tilelang as _mhc_post_kernel
309
+ from vllm.model_executor.kernels.mhc.tilelang_kernels import (
310
+ mhc_post_tilelang as _mhc_post_kernel,
311
+ )
312
313
out = torch.empty_like(residual)
314
_mhc_post_kernel(
@@ -353,7 +355,7 @@ def mhc_fused_post_pre_tilelang(
355
layer_input_cur: shape (..., hidden_size)
356
"""
357
356
- from vllm._tilelang_ops import (
358
+ from vllm.model_executor.kernels.mhc.tilelang_kernels import (
359
compute_num_split,
360
mhc_fused_tilelang,
361
mhc_post_tilelang,
@@ -608,21 +610,22 @@ def _mhc_post_tilelang_fake(
610
return torch.empty_like(residual)
611
612
611
-def _hc_head_fused_kernel_tilelang(
613
+def hc_head_fused_kernel_tilelang(
614
hs_flat: torch.Tensor,
615
fn: torch.Tensor,
616
hc_scale: torch.Tensor,
617
hc_base: torch.Tensor,
616
- out: torch.Tensor,
617
- hidden_size: int,
618
rms_eps: float,
619
hc_eps: float,
620
- hc_mult: int,
621
-) -> None:
622
- """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
623
- if hs_flat.shape[0] == 0:
624
- return
625
- from vllm._tilelang_ops import hc_head_fuse_tilelang
620
+) -> torch.Tensor:
621
+ """Apply the fused hc_head kernel and return the (T, H) bf16 result."""
622
+ num_tokens, hc_mult, hidden_size = hs_flat.shape
623
+ out = torch.empty(
624
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_flat.device
625
+ )
626
+ if num_tokens == 0:
627
+ return out
628
+ from vllm.model_executor.kernels.mhc.tilelang_kernels import hc_head_fuse_tilelang
629
630
hc_head_fuse_tilelang(
631
hs_flat,
@@ -635,6 +638,21 @@ def _hc_head_fused_kernel_tilelang(
638
hc_eps,
639
hc_mult,
640
)
641
+ return out
642
+
643
+
644
+def _hc_head_fused_kernel_tilelang_fake(
645
+ hs_flat: torch.Tensor,
646
+ fn: torch.Tensor,
647
+ hc_scale: torch.Tensor,
648
+ hc_base: torch.Tensor,
649
+ rms_eps: float,
650
+ hc_eps: float,
651
+) -> torch.Tensor:
652
+ num_tokens, _, hidden_size = hs_flat.shape
653
+ return torch.empty(
654
+ num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_flat.device
655
+ )
656
657
658
direct_register_custom_op(
@@ -659,6 +677,7 @@ direct_register_custom_op(
677
678
direct_register_custom_op(
679
op_name="hc_head_fused_kernel_tilelang",
662
- op_func=_hc_head_fused_kernel_tilelang,
663
- mutates_args=["out"],
680
+ op_func=hc_head_fused_kernel_tilelang,
681
+ mutates_args=[],
682
+ fake_impl=_hc_head_fused_kernel_tilelang_fake,
683
)
vllm/model_executor/kernels/mhc/tilelang_kernels.py
renamed
vllm/model_executor/layers/mhc.py
+9
-18
@@ -243,21 +243,13 @@ class HCHeadOp(CustomOp):
243
hc_mult, hidden_size = hidden_states.shape[-2:]
244
outer_shape = hidden_states.shape[:-2]
245
hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
246
- num_tokens = hs_flat.shape[0]
247
-
248
- out = torch.empty(
249
- num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
250
- )
251
- torch.ops.vllm.hc_head_fused_kernel_tilelang(
246
+ out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
247
hs_flat,
248
hc_fn,
249
hc_scale,
250
hc_base,
256
- out,
257
- hidden_size,
251
rms_norm_eps,
252
hc_eps,
260
- hc_mult,
253
)
254
return out.view(*outer_shape, hidden_size)
255
@@ -273,25 +265,24 @@ class HCHeadOp(CustomOp):
265
hc_mult, hidden_size = hidden_states.shape[-2:]
266
outer_shape = hidden_states.shape[:-2]
267
hs_flat = hidden_states.view(-1, hc_mult, hidden_size)
276
- num_tokens = hs_flat.shape[0]
277
-
278
- out = torch.empty(
279
- num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
280
- )
268
269
if HAS_TILELANG:
283
- torch.ops.vllm.hc_head_fused_kernel_tilelang(
270
+ out = torch.ops.vllm.hc_head_fused_kernel_tilelang(
271
hs_flat,
272
hc_fn,
273
hc_scale,
274
hc_base,
288
- out,
289
- hidden_size,
275
rms_norm_eps,
276
hc_eps,
292
- hc_mult,
277
)
278
else:
279
+ num_tokens = hs_flat.shape[0]
280
+ out = torch.empty(
281
+ num_tokens,
282
+ hidden_size,
283
+ dtype=torch.bfloat16,
284
+ device=hidden_states.device,
285
+ )
286
torch.ops.vllm.hc_head_triton(
287
hs_flat,
288
hc_fn,
vllm/models/deepseek_v4/nvidia/model.py
+20
-54
@@ -15,6 +15,12 @@ from vllm.distributed import (
15
get_tensor_model_parallel_rank,
16
get_tensor_model_parallel_world_size,
17
)
18
+from vllm.model_executor.kernels.mhc.tilelang import (
19
+ hc_head_fused_kernel_tilelang,
20
+ mhc_fused_post_pre_tilelang,
21
+ mhc_post_tilelang,
22
+ mhc_pre_tilelang,
23
+)
24
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
25
from vllm.model_executor.layers.fused_moe import FusedMoE
26
from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
@@ -28,12 +34,6 @@ from vllm.model_executor.layers.linear import (
34
RowParallelLinear,
35
)
36
from vllm.model_executor.layers.logits_processor import LogitsProcessor
31
-from vllm.model_executor.layers.mhc import (
32
- HCHeadOp,
33
- MHCFusedPostPreOp,
34
- MHCPostOp,
35
- MHCPreOp,
36
-)
37
from vllm.model_executor.layers.quantization import QuantizationConfig
38
from vllm.model_executor.layers.rotary_embedding import get_rope
39
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -794,10 +794,6 @@ class DeepseekV4DecoderLayer(nn.Module):
794
):
795
super().__init__()
796
797
- # Lazy import to avoid top-level tilelang dependency.
798
- # Registers both torch.ops.vllm.mhc_pre and mhc_post
799
- import vllm.model_executor.layers.mhc # noqa: F401
800
-
797
config = vllm_config.model_config.hf_config
798
self.hidden_size = config.hidden_size
799
@@ -860,42 +856,6 @@ class DeepseekV4DecoderLayer(nn.Module):
856
),
857
requires_grad=False,
858
)
863
- self.mhc_pre = MHCPreOp()
864
- self.mhc_post = MHCPostOp()
865
- self.mhc_fused_post_pre = MHCFusedPostPreOp()
866
-
867
- def hc_pre(
868
- self,
869
- x: torch.Tensor,
870
- hc_fn: torch.Tensor,
871
- hc_scale: torch.Tensor,
872
- hc_base: torch.Tensor,
873
- norm_weight: torch.Tensor | None = None,
874
- norm_eps: float = 1e-6,
875
- ):
876
- post_mix, res_mix, layer_input = self.mhc_pre(
877
- residual=x,
878
- fn=hc_fn,
879
- hc_scale=hc_scale,
880
- hc_base=hc_base,
881
- rms_eps=self.rms_norm_eps,
882
- hc_pre_eps=self.hc_eps,
883
- hc_sinkhorn_eps=self.hc_eps,
884
- hc_post_mult_value=self.hc_post_alpha,
885
- sinkhorn_repeat=self.hc_sinkhorn_iters,
886
- norm_weight=norm_weight,
887
- norm_eps=norm_eps,
888
- )
889
- return layer_input, post_mix, res_mix
890
-
891
- def hc_post(
892
- self,
893
- x: torch.Tensor,
894
- residual: torch.Tensor,
895
- post: torch.Tensor,
896
- comb: torch.Tensor,
897
- ):
898
- return self.mhc_post(x, residual, post, comb)
859
860
def forward(
861
self,
@@ -909,18 +869,23 @@ class DeepseekV4DecoderLayer(nn.Module):
869
attn_norm_weight = self.attn_norm.weight.data
870
attn_norm_eps = self.attn_norm.variance_epsilon
871
if residual is None:
912
- # Run standalone hc_pre on first layer
872
+ # Run standalone mhc_pre on first layer
873
residual = x
914
- x, post_mix, res_mix = self.hc_pre(
874
+ post_mix, res_mix, x = mhc_pre_tilelang(
875
x,
876
self.hc_attn_fn,
877
self.hc_attn_scale,
878
self.hc_attn_base,
879
+ self.rms_norm_eps,
880
+ self.hc_eps,
881
+ self.hc_eps,
882
+ self.hc_post_alpha,
883
+ self.hc_sinkhorn_iters,
884
norm_weight=attn_norm_weight,
885
norm_eps=attn_norm_eps,
886
)
887
else:
923
- residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
888
+ residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang(
889
x,
890
residual,
891
post_mix,
@@ -939,12 +904,12 @@ class DeepseekV4DecoderLayer(nn.Module):
904
norm_eps=attn_norm_eps,
905
)
906
942
- # attn_norm is fused into hc_pre / mhc_fused_post_pre above.
907
+ # attn_norm is fused into mhc_pre_tilelang / mhc_fused_post_pre above.
908
x = self.attn(positions, x, None)
909
910
ffn_norm_weight = self.ffn_norm.weight.data
911
ffn_norm_eps = self.ffn_norm.variance_epsilon
947
- residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
912
+ residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang(
913
x,
914
residual,
915
post_mix,
@@ -1047,7 +1012,6 @@ class DeepseekV4Model(nn.Module):
1012
torch.empty(1, dtype=torch.float32),
1013
requires_grad=False,
1014
)
1050
- self.hc_head_op = HCHeadOp()
1015
# Pre-hc_head residual stream buffer for the MTP draft. Stable
1016
# address (outside the cudagraph pool) so the copy_ in forward()
1017
# refreshes it correctly across captured shapes.
@@ -1117,7 +1081,9 @@ class DeepseekV4Model(nn.Module):
1081
residual,
1082
)
1083
if layer is not None:
1120
- hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
1084
+ hidden_states = mhc_post_tilelang(
1085
+ hidden_states, residual, post_mix, res_mix
1086
+ )
1087
1088
if not get_pp_group().is_last_rank:
1089
return IntermediateTensors({"hidden_states": hidden_states})
@@ -1126,7 +1092,7 @@ class DeepseekV4Model(nn.Module):
1092
num_tokens = hidden_states.shape[0]
1093
self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
1094
1129
- hidden_states = self.hc_head_op(
1095
+ hidden_states = hc_head_fused_kernel_tilelang(
1096
hidden_states,
1097
self.hc_head_fn,
1098
self.hc_head_scale,
vllm/models/deepseek_v4/nvidia/mtp.py
+6
-7
@@ -24,11 +24,14 @@ from vllm.distributed import (
24
get_tensor_model_parallel_world_size,
25
)
26
from vllm.logger import init_logger
27
+from vllm.model_executor.kernels.mhc.tilelang import (
28
+ hc_head_fused_kernel_tilelang,
29
+ mhc_post_tilelang,
30
+)
31
from vllm.model_executor.layers.fused_moe import FusedMoE
32
from vllm.model_executor.layers.layernorm import RMSNorm
33
from vllm.model_executor.layers.linear import ReplicatedLinear
34
from vllm.model_executor.layers.logits_processor import LogitsProcessor
31
-from vllm.model_executor.layers.mhc import HCHeadOp
35
from vllm.model_executor.layers.vocab_parallel_embedding import (
36
VocabParallelEmbedding,
37
)
@@ -122,8 +125,6 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
125
aux_stream_list=aux_stream_list,
126
)
127
125
- self.hc_head_op = HCHeadOp()
126
-
128
def forward(
129
self,
130
input_ids: torch.Tensor,
@@ -155,9 +156,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
156
hidden_states, residual, post_mix, res_mix = self.mtp_block(
157
positions=positions, x=hidden_states, input_ids=None
158
)
158
- hidden_states = self.mtp_block.hc_post(
159
- hidden_states, residual, post_mix, res_mix
160
- )
159
+ hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
160
# Return the flat pre-hc_head residual so it can be re-fed as the
161
# next spec step's `previous_hidden_states` when
162
# num_speculative_tokens > 1. hc_head is deferred to compute_logits.
@@ -237,7 +236,7 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
236
hidden_states = hidden_states.view(
237
-1, mtp_layer.hc_mult, mtp_layer.config.hidden_size
238
)
240
- hidden_states = mtp_layer.hc_head_op(
239
+ hidden_states = hc_head_fused_kernel_tilelang(
240
hidden_states,
241
mtp_layer.hc_head_fn,
242
mtp_layer.hc_head_scale,